mirror of
https://github.com/azaion/detections.git
synced 2026-04-22 05:26:32 +00:00
Oleksandr Bezdieniezhnykh
fc57d677b4
- Updated various Cython files to explicitly cast types, enhancing type safety and readability. - Adjusted the `engine_name` property in `InferenceEngine` and its subclasses to be set directly in the constructor. - Modified the `request` method in `_SessionWithBase` to accept `*args` for better flexibility. - Ensured proper type casting for return values in methods across multiple classes, including `Inference`, `CoreMLEngine`, and `TensorRTEngine`. These changes aim to streamline the codebase and improve maintainability by enforcing consistent type usage.
86 lines
2.7 KiB
Cython
86 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}')
|
|
self.engine_name = <str>"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 <tuple>(self.img_height, self.img_width)
|
|
|
|
cdef int get_batch_size(self):
|
|
return <int>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, :, :]]
|