|
| 1 | +from threading import Lock |
| 2 | +from typing import List, Optional, Tuple, Union |
| 3 | + |
| 4 | +import numpy as np |
| 5 | +import torch |
| 6 | +from inference_exp import ClassificationModel, ClassificationPrediction |
| 7 | +from inference_exp.configuration import DEFAULT_DEVICE |
| 8 | +from inference_exp.entities import ColorFormat |
| 9 | +from inference_exp.errors import ( |
| 10 | + CorruptedModelPackageError, |
| 11 | + EnvironmentConfigurationError, |
| 12 | + MissingDependencyError, |
| 13 | +) |
| 14 | +from inference_exp.models.base.types import PreprocessedInputs |
| 15 | +from inference_exp.models.common.model_packages import get_model_package_contents |
| 16 | +from inference_exp.models.common.onnx import ( |
| 17 | + run_session_with_batch_size_limit, |
| 18 | + set_execution_provider_defaults, |
| 19 | +) |
| 20 | +from inference_exp.models.common.roboflow.model_packages import ( |
| 21 | + InferenceConfig, |
| 22 | + ResizeMode, |
| 23 | + parse_class_names_file, |
| 24 | + parse_inference_config, |
| 25 | +) |
| 26 | +from inference_exp.models.common.roboflow.pre_processing import ( |
| 27 | + pre_process_network_input, |
| 28 | +) |
| 29 | +from inference_exp.utils.onnx_introspection import get_selected_onnx_execution_providers |
| 30 | + |
| 31 | +try: |
| 32 | + import onnxruntime |
| 33 | +except ImportError as import_error: |
| 34 | + raise MissingDependencyError( |
| 35 | + message=f"Could not import ResNet model with ONNX backend - this error means that some additional dependencies " |
| 36 | + f"are not installed in the environment. If you run the `inference-exp` library directly in your Python " |
| 37 | + f"program, make sure the following extras of the package are installed: \n" |
| 38 | + f"\t* `onnx-cpu` - when you wish to use library with CPU support only\n" |
| 39 | + f"\t* `onnx-cu12` - for running on GPU with Cuda 12 installed\n" |
| 40 | + f"\t* `onnx-cu118` - for running on GPU with Cuda 11.8 installed\n" |
| 41 | + f"\t* `onnx-jp6-cu126` - for running on Jetson with Jetpack 6\n" |
| 42 | + f"If you see this error using Roboflow infrastructure, make sure the service you use does support the model. " |
| 43 | + f"You can also contact Roboflow to get support.", |
| 44 | + help_url="https://todo", |
| 45 | + ) from import_error |
| 46 | + |
| 47 | + |
| 48 | +class YOLOv8ForClassificationOnnx(ClassificationModel[torch.Tensor, torch.Tensor]): |
| 49 | + |
| 50 | + @classmethod |
| 51 | + def from_pretrained( |
| 52 | + cls, |
| 53 | + model_name_or_path: str, |
| 54 | + onnx_execution_providers: Optional[List[Union[str, tuple]]] = None, |
| 55 | + default_onnx_trt_options: bool = True, |
| 56 | + device: torch.device = DEFAULT_DEVICE, |
| 57 | + **kwargs, |
| 58 | + ) -> "YOLOv8ForClassificationOnnx": |
| 59 | + if onnx_execution_providers is None: |
| 60 | + onnx_execution_providers = get_selected_onnx_execution_providers() |
| 61 | + if not onnx_execution_providers: |
| 62 | + raise EnvironmentConfigurationError( |
| 63 | + message=f"Could not initialize model - selected backend is ONNX which requires execution provider to " |
| 64 | + f"be specified - explicitly in `from_pretrained(...)` method or via env variable " |
| 65 | + f"`ONNXRUNTIME_EXECUTION_PROVIDERS`. If you run model locally - adjust your setup, otherwise " |
| 66 | + f"contact the platform support.", |
| 67 | + help_url="https://todo", |
| 68 | + ) |
| 69 | + onnx_execution_providers = set_execution_provider_defaults( |
| 70 | + providers=onnx_execution_providers, |
| 71 | + model_package_path=model_name_or_path, |
| 72 | + device=device, |
| 73 | + default_onnx_trt_options=default_onnx_trt_options, |
| 74 | + ) |
| 75 | + model_package_content = get_model_package_contents( |
| 76 | + model_package_dir=model_name_or_path, |
| 77 | + elements=[ |
| 78 | + "class_names.txt", |
| 79 | + "inference_config.json", |
| 80 | + "weights.onnx", |
| 81 | + ], |
| 82 | + ) |
| 83 | + class_names = parse_class_names_file( |
| 84 | + class_names_path=model_package_content["class_names.txt"] |
| 85 | + ) |
| 86 | + inference_config = parse_inference_config( |
| 87 | + config_path=model_package_content["inference_config.json"], |
| 88 | + allowed_resize_modes={ |
| 89 | + ResizeMode.STRETCH_TO, |
| 90 | + ResizeMode.LETTERBOX, |
| 91 | + ResizeMode.CENTER_CROP, |
| 92 | + ResizeMode.LETTERBOX_REFLECT_EDGES, |
| 93 | + }, |
| 94 | + ) |
| 95 | + if inference_config.post_processing.type != "softmax": |
| 96 | + raise CorruptedModelPackageError( |
| 97 | + message="Expected Softmax to be the post-processing", |
| 98 | + help_url="https://todo", |
| 99 | + ) |
| 100 | + session = onnxruntime.InferenceSession( |
| 101 | + path_or_bytes=model_package_content["weights.onnx"], |
| 102 | + providers=onnx_execution_providers, |
| 103 | + ) |
| 104 | + input_shape = session.get_inputs()[0].shape |
| 105 | + input_batch_size = input_shape[0] |
| 106 | + if isinstance(input_batch_size, str): |
| 107 | + input_batch_size = None |
| 108 | + input_name = session.get_inputs()[0].name |
| 109 | + return cls( |
| 110 | + session=session, |
| 111 | + input_name=input_name, |
| 112 | + inference_config=inference_config, |
| 113 | + class_names=class_names, |
| 114 | + device=device, |
| 115 | + input_batch_size=input_batch_size, |
| 116 | + ) |
| 117 | + |
| 118 | + def __init__( |
| 119 | + self, |
| 120 | + session: onnxruntime.InferenceSession, |
| 121 | + input_name: str, |
| 122 | + inference_config: InferenceConfig, |
| 123 | + class_names: List[str], |
| 124 | + device: torch.device, |
| 125 | + input_batch_size: Optional[int], |
| 126 | + ): |
| 127 | + self._session = session |
| 128 | + self._input_name = input_name |
| 129 | + self._inference_config = inference_config |
| 130 | + self._class_names = class_names |
| 131 | + self._device = device |
| 132 | + self._input_batch_size = input_batch_size |
| 133 | + self._session_thread_lock = Lock() |
| 134 | + |
| 135 | + @property |
| 136 | + def class_names(self) -> List[str]: |
| 137 | + return self._class_names |
| 138 | + |
| 139 | + def pre_process( |
| 140 | + self, |
| 141 | + images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]], |
| 142 | + input_color_format: Optional[ColorFormat] = None, |
| 143 | + image_size: Optional[Tuple[int, int]] = None, |
| 144 | + **kwargs, |
| 145 | + ) -> torch.Tensor: |
| 146 | + return pre_process_network_input( |
| 147 | + images=images, |
| 148 | + image_pre_processing=self._inference_config.image_pre_processing, |
| 149 | + network_input=self._inference_config.network_input, |
| 150 | + target_device=self._device, |
| 151 | + input_color_format=input_color_format, |
| 152 | + image_size_wh=image_size, |
| 153 | + )[0] |
| 154 | + |
| 155 | + def forward( |
| 156 | + self, pre_processed_images: PreprocessedInputs, **kwargs |
| 157 | + ) -> torch.Tensor: |
| 158 | + with self._session_thread_lock: |
| 159 | + return run_session_with_batch_size_limit( |
| 160 | + session=self._session, |
| 161 | + inputs={self._input_name: pre_processed_images}, |
| 162 | + min_batch_size=self._input_batch_size, |
| 163 | + max_batch_size=self._input_batch_size, |
| 164 | + )[0] |
| 165 | + |
| 166 | + def post_process( |
| 167 | + self, |
| 168 | + model_results: torch.Tensor, |
| 169 | + **kwargs, |
| 170 | + ) -> ClassificationPrediction: |
| 171 | + if self._inference_config.post_processing.fused: |
| 172 | + confidence = model_results |
| 173 | + else: |
| 174 | + confidence = torch.nn.functional.softmax(model_results, dim=-1) |
| 175 | + return ClassificationPrediction( |
| 176 | + class_id=confidence.argmax(dim=-1), |
| 177 | + confidence=confidence, |
| 178 | + ) |
0 commit comments