microsoft/onnxruntime-extensions
Publicmirrored from https://github.com/microsoft/onnxruntime-extensionsAvailable
onnxruntime_extensions/tools/add_HuggingFace_CLIPImageProcessor_to_model.py
171lines · modecode
| 1 | # Copyright (c) Microsoft Corporation. All rights reserved. |
| 2 | # Licensed under the MIT License. |
| 3 | import argparse |
| 4 | import os |
| 5 | from pathlib import Path |
| 6 | import onnx |
| 7 | |
| 8 | from .pre_post_processing import * |
| 9 | |
| 10 | |
| 11 | class Dict2Class(object): |
| 12 | ''' |
| 13 | Convert dict to class |
| 14 | ''' |
| 15 | def __init__(self, my_dict): |
| 16 | for key in my_dict: |
| 17 | setattr(self, key, my_dict[key]) |
| 18 | |
| 19 | def image_processor(args: argparse.Namespace): |
| 20 | # support user providing encoded image bytes |
| 21 | steps = [ |
| 22 | ConvertImageToBGR(), # custom op to convert jpg/png to BGR (output is HWC) |
| 23 | ] # Normalization params are for RGB ordering |
| 24 | if args.do_convert_rgb: |
| 25 | steps.append(ReverseAxis(axis=2, dim_value=3, name="BGR_to_RGB")) |
| 26 | |
| 27 | if args.do_resize: |
| 28 | to_size = args.size |
| 29 | steps.append(Resize(to_size)) |
| 30 | |
| 31 | if args.do_center_crop: |
| 32 | to_size = args.crop_size |
| 33 | to_size = (to_size, to_size) |
| 34 | steps.append(CenterCrop(*to_size)) |
| 35 | |
| 36 | if args.do_rescale: |
| 37 | steps.append(ImageBytesToFloat(args.rescale_factor)) |
| 38 | |
| 39 | steps.append(ChannelsLastToChannelsFirst()) |
| 40 | |
| 41 | if args.do_normalize: |
| 42 | mean_std = list(zip(args.image_mean, args.image_std)) |
| 43 | layout = 'CHW' |
| 44 | steps.append(Normalize(mean_std, layout=layout)) |
| 45 | |
| 46 | steps.append(Unsqueeze([0])) # add batch dim |
| 47 | |
| 48 | return steps |
| 49 | |
| 50 | |
| 51 | def clip_image_processor(model_file: Path, output_file: Path, **kwargs): |
| 52 | """ |
| 53 | Used for models like stable-diffusion. should be compatible with |
| 54 | https://github.com/huggingface/transformers/blob/main/src/transformers/models/clip/image_processing_clip.py |
| 55 | |
| 56 | It's similar to 'CLIP' image processor, and aligns with the HuggingFace class name. |
| 57 | |
| 58 | A typical usage example is used in Stable diffusion. |
| 59 | |
| 60 | :param model_file: The input model file path. |
| 61 | :param output_file: The output file path, where the finalized model saved to. |
| 62 | :param Kwargs |
| 63 | onnx_opset: The opset version of onnx model, default(18). |
| 64 | do_convert_rgb: Convert image from BGR to RGB. default(True) |
| 65 | do_resize: Resize the image's (height, width) dimensions to the specified `size`. default(True) |
| 66 | size: The shortest edge of the image is resized to size. Default(224) |
| 67 | resample: An optional resampling filter. Default(cubic) |
| 68 | do_center_crop: Whether to center crop the image to the specified `crop_size`. Default(True) |
| 69 | crop_size: Size of the output image after applying `center_crop`. Default(224) |
| 70 | do_rescale: Whether to rescale the image by the specified scale (rescale_factor). Default(True) |
| 71 | rescale_factor: Scale factor to use if rescaling the image. Default(1/255) |
| 72 | do_normalize: Whether to normalize the image. Default(True) |
| 73 | image_mean: Mean values for image normalization. Default([0.485, 0.456, 0.406]) |
| 74 | image_std: Standard deviation values for image normalization. Default([0.229, 0.224, 0.225]) |
| 75 | |
| 76 | |
| 77 | """ |
| 78 | args = Dict2Class(kwargs) |
| 79 | # Load model |
| 80 | model = onnx.load(str(model_file.resolve(strict=True))) |
| 81 | inputs = [create_named_value("image", onnx.TensorProto.UINT8, ["num_bytes"])] |
| 82 | |
| 83 | pipeline = PrePostProcessor(inputs, args.opset) |
| 84 | |
| 85 | preprocessing = image_processor(args) |
| 86 | |
| 87 | pipeline.add_pre_processing(preprocessing) |
| 88 | |
| 89 | new_model = pipeline.run(model) |
| 90 | onnx.save_model(new_model, str(output_file.resolve())) |
| 91 | print(f"Updated model saved to {output_file}") |
| 92 | |
| 93 | |
| 94 | def main(): |
| 95 | parser = argparse.ArgumentParser( |
| 96 | os.path.basename(__file__), |
| 97 | description="""Add CLIPImageProcessor to a model. |
| 98 | |
| 99 | The updated model will be written in the same location as the original model, |
| 100 | with '.onnx' updated to '.with_clip_processor.onnx' |
| 101 | |
| 102 | Example usage: |
| 103 | object detection: |
| 104 | - python -m onnxruntime_extensions.tools.add_HuggingFace_CLIPImageProcessor_to_model model.onnx |
| 105 | """, |
| 106 | ) |
| 107 | |
| 108 | parser.add_argument( |
| 109 | "--opset", type=int, required=False, default=18, |
| 110 | help="ONNX opset to use. Minimum allowed is 16. Opset 18 is required for Resize with anti-aliasing.", |
| 111 | ) |
| 112 | |
| 113 | parser.add_argument( |
| 114 | "--do_resize", type=bool, required=False, default=True, |
| 115 | help="Whether to resize the image's (height, width) dimensions to the specified `size`. default(True)", |
| 116 | ) |
| 117 | parser.add_argument( |
| 118 | "--size", type=int, required=False, default=224, |
| 119 | help="The shortest edge of the image is resized to size. Default(224)", |
| 120 | ) |
| 121 | parser.add_argument( |
| 122 | "--resample", type=str, default="cubic", choices=["cubic", "nearest","linear"], |
| 123 | help="Whether to resize the image's (height, width) dimensions to the specified `size`. Default(cubic)", |
| 124 | ) |
| 125 | parser.add_argument( |
| 126 | "--do_center_crop", type=bool, default=True, |
| 127 | help="Whether to center crop the image to the specified `crop_size`. Default(True)", |
| 128 | ) |
| 129 | parser.add_argument( |
| 130 | "--crop_size", type=int, default=224, |
| 131 | help="Size of the output image after applying `center_crop`. Default(224)", |
| 132 | ) |
| 133 | parser.add_argument( |
| 134 | "--do_rescale", type=bool, default=True, |
| 135 | help="Whether to rescale the image by the specified scale (rescale_factor). Default(True)", |
| 136 | ) |
| 137 | parser.add_argument( |
| 138 | "--rescale_factor", type=float, default=1/255, |
| 139 | help="Scale factor to use if rescaling the image. Default(1/255)", |
| 140 | ) |
| 141 | parser.add_argument( |
| 142 | "--do_normalize", type=bool, default=True, |
| 143 | help="Whether to normalize the image. Default(True)", |
| 144 | ) |
| 145 | parser.add_argument( |
| 146 | "--image_mean", type=str, default="[0.48145466, 0.4578275, 0.40821073]", |
| 147 | help=" Mean to use if normalizing the image, default([0.48145466, 0.4578275, 0.40821073])", |
| 148 | ) |
| 149 | parser.add_argument( |
| 150 | "--image_std", type=str, default="[0.26862954, 0.26130258, 0.27577711]", |
| 151 | help="Image standard deviation., default([0.26862954, 0.26130258, 0.27577711]).", |
| 152 | ) |
| 153 | parser.add_argument( |
| 154 | "--do_convert_rgb", type=bool, default=True, |
| 155 | help="Convert image from BGR to RGB. Default(True)", |
| 156 | ) |
| 157 | parser.add_argument("model", type=Path, help="Provide path to ONNX model to update.") |
| 158 | |
| 159 | args = parser.parse_args() |
| 160 | |
| 161 | args.image_mean = [float(x) for x in args.image_mean.replace('[','').replace(']','').split(",")] |
| 162 | args.image_std = [float(x) for x in args.image_std.replace('[','').replace(']','').split(",")] |
| 163 | |
| 164 | model_path = args.model.resolve(strict=True) |
| 165 | new_model_path = model_path.with_suffix(".with_clip_processor.onnx") |
| 166 | |
| 167 | clip_image_processor(model_path, new_model_path, **vars(args)) |
| 168 | |
| 169 | |
| 170 | if __name__ == "__main__": |
| 171 | main() |