mirror of
https://github.com/azaion/detections.git
synced 2026-04-22 10:26:32 +00:00
Oleksandr Bezdieniezhnykh
27f4aceb52
- Updated the `Inference` class to replace the `get_onnx_engine_bytes` method with `download_model`, allowing for dynamic model loading based on a specified filename. - Modified the `convert_and_upload_model` method to accept `source_bytes` instead of `onnx_engine_bytes`, enhancing flexibility in model conversion. - Introduced a new property `engine_name` to the `Inference` class for better access to engine details. - Adjusted the `AIRecognitionConfig` structure to include a new method pointer `from_dict`, improving configuration handling. - Updated various test cases to reflect changes in model paths and timeout settings, ensuring consistency and reliability in testing.
89 lines
2.7 KiB
Cython
89 lines
2.7 KiB
Cython
from engines.inference_engine cimport InferenceEngine
|
|
cimport constants_inf
|
|
import numpy as np
|
|
from PIL import Image
|
|
import io
|
|
import os
|
|
import tempfile
|
|
import zipfile
|
|
|
|
|
|
cdef class CoreMLEngine(InferenceEngine):
|
|
|
|
def __init__(self, model_bytes: bytes, batch_size: int = 1, **kwargs):
|
|
super().__init__(model_bytes, batch_size)
|
|
import coremltools as ct
|
|
|
|
model_path = kwargs.get('model_path')
|
|
if model_path is None:
|
|
model_path = self._extract_from_zip(model_bytes)
|
|
|
|
self.model = ct.models.MLModel(
|
|
model_path, compute_units=ct.ComputeUnit.ALL)
|
|
spec = self.model.get_spec()
|
|
|
|
img_input = spec.description.input[0]
|
|
self.img_width = int(img_input.type.imageType.width)
|
|
self.img_height = int(img_input.type.imageType.height)
|
|
self.batch_size = 1
|
|
|
|
constants_inf.log(<str>f'CoreML model: {self.img_width}x{self.img_height}')
|
|
|
|
@property
|
|
def engine_name(self):
|
|
return "coreml"
|
|
|
|
@staticmethod
|
|
def get_engine_filename():
|
|
return "azaion_coreml.zip"
|
|
|
|
@staticmethod
|
|
def _extract_from_zip(model_bytes):
|
|
tmpdir = tempfile.mkdtemp()
|
|
buf = io.BytesIO(model_bytes)
|
|
with zipfile.ZipFile(buf, 'r') as zf:
|
|
zf.extractall(tmpdir)
|
|
for item in os.listdir(tmpdir):
|
|
if item.endswith('.mlpackage') or item.endswith('.mlmodel'):
|
|
return os.path.join(tmpdir, item)
|
|
raise ValueError("No .mlpackage or .mlmodel found in zip")
|
|
|
|
cdef tuple get_input_shape(self):
|
|
return self.img_height, self.img_width
|
|
|
|
cdef int get_batch_size(self):
|
|
return 1
|
|
|
|
cdef run(self, input_data):
|
|
cdef int w = self.img_width
|
|
cdef int h = self.img_height
|
|
|
|
blob = input_data[0]
|
|
img_array = np.clip(blob * 255.0, 0, 255).astype(np.uint8)
|
|
img_array = np.transpose(img_array, (1, 2, 0))
|
|
pil_img = Image.fromarray(img_array, 'RGB')
|
|
|
|
pred = self.model.predict({
|
|
'image': pil_img,
|
|
'iouThreshold': 0.45,
|
|
'confidenceThreshold': 0.25,
|
|
})
|
|
|
|
coords = pred.get('coordinates', np.empty((0, 4), dtype=np.float32))
|
|
confs = pred.get('confidence', np.empty((0, 80), dtype=np.float32))
|
|
|
|
if coords.size == 0:
|
|
return [np.zeros((1, 0, 6), dtype=np.float32)]
|
|
|
|
cx, cy, bw, bh = coords[:, 0], coords[:, 1], coords[:, 2], coords[:, 3]
|
|
x1 = (cx - bw / 2) * w
|
|
y1 = (cy - bh / 2) * h
|
|
x2 = (cx + bw / 2) * w
|
|
y2 = (cy + bh / 2) * h
|
|
|
|
class_ids = np.argmax(confs, axis=1).astype(np.float32)
|
|
conf_values = np.max(confs, axis=1)
|
|
|
|
dets = np.stack([x1, y1, x2, y2, conf_values, class_ids], axis=1)
|
|
return [dets[np.newaxis, :, :]]
|