| from __future__ import annotations |
|
|
| import time |
|
|
| import aidlite |
| import numpy as np |
|
|
| from utils import obb_postprocess, preprocess_image |
|
|
|
|
| class Yolo26Obb: |
| """YOLO26 OBB (QNN context binary) via the AidLite SDK.""" |
|
|
| def __init__( |
| self, |
| model_path: str, |
| width: int = 1024, |
| height: int = 1024, |
| model_type: str = "qnn", |
| num_classes: int = 15, |
| ): |
| self.width = width |
| self.height = height |
| |
| self.output_shape = [1, 4 + num_classes + 1, 21504] |
|
|
| model = aidlite.Model.create_instance(model_path) |
| if model is None: |
| raise RuntimeError("Create model failed") |
| model.set_model_properties( |
| [[1, height, width, 3]], |
| aidlite.DataType.TYPE_FLOAT32, |
| [self.output_shape], |
| aidlite.DataType.TYPE_FLOAT32, |
| ) |
|
|
| config = aidlite.Config.create_instance() |
| if config is None: |
| raise RuntimeError("Create config failed") |
| config.implement_type = aidlite.ImplementType.TYPE_LOCAL |
| if model_type.lower() == "qnn": |
| config.framework_type = aidlite.FrameworkType.TYPE_QNN |
| elif model_type.lower() in ("snpe", "snpe2"): |
| config.framework_type = aidlite.FrameworkType.TYPE_SNPE2 |
| else: |
| raise ValueError(f"Unsupported model_type: {model_type}") |
| config.accelerate_type = aidlite.AccelerateType.TYPE_DSP |
| config.is_quantify_model = 1 |
|
|
| interpreter = aidlite.InterpreterBuilder.build_interpretper_from_model_and_config(model, config) |
| if interpreter is None: |
| raise RuntimeError("Build interpreter failed") |
| if interpreter.init() != 0: |
| raise RuntimeError("Interpreter init failed") |
| if interpreter.load_model() != 0: |
| raise RuntimeError("Interpreter load model failed") |
| self.interpreter = interpreter |
|
|
| def __del__(self): |
| interpreter = getattr(self, "interpreter", None) |
| if interpreter is not None: |
| interpreter.destory() |
|
|
| def __call__(self, image: np.ndarray, invoke_nums: int = 10, conf_thres: float = 0.25, iou_thres: float = 0.45): |
| input_tensor, ratio, pad_x, pad_y = preprocess_image(image, self.width) |
| if self.interpreter.set_input_tensor(0, input_tensor.data) != 0: |
| raise RuntimeError("set_input_tensor failed") |
|
|
| invoke_time = [] |
| for _ in range(invoke_nums): |
| t1 = time.time() |
| result = self.interpreter.invoke() |
| invoke_time.append((time.time() - t1) * 1000) |
| if result != 0: |
| raise RuntimeError("interpreter invoke failed") |
|
|
| print("====================================") |
| print( |
| f"QNN invoke {invoke_nums} times:\n" |
| f" --mean_invoke_time is {sum(invoke_time) / invoke_nums}\n" |
| f" --max_invoke_time is {max(invoke_time)}\n" |
| f" --min_invoke_time is {min(invoke_time)}\n" |
| f" --var_invoketime is {np.var(invoke_time)}" |
| ) |
| print("====================================") |
|
|
| output = self.interpreter.get_output_tensor(0).reshape(*self.output_shape) |
| return obb_postprocess( |
| output, |
| ratio, |
| pad_x, |
| pad_y, |
| image.shape, |
| conf_thres=conf_thres, |
| iou_thres=iou_thres, |
| ) |
|
|
