File size: 3,463 Bytes
aa718fb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | 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
# output0 layout is [1, 4 + classes + 1, anchors].
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,
)
|