mirror of
https://github.com/azaion/detections.git
synced 2026-04-22 17:46:38 +00:00
Refactor inference and AI configuration handling
- 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.
This commit is contained in:
+39
-53
@@ -1,6 +1,7 @@
|
||||
from engines.inference_engine cimport InferenceEngine
|
||||
cimport constants_inf
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
import io
|
||||
import os
|
||||
import tempfile
|
||||
@@ -21,18 +22,12 @@ cdef class CoreMLEngine(InferenceEngine):
|
||||
model_path, compute_units=ct.ComputeUnit.ALL)
|
||||
spec = self.model.get_spec()
|
||||
|
||||
input_desc = spec.description.input[0]
|
||||
self.input_name = input_desc.name
|
||||
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
|
||||
|
||||
array_type = input_desc.type.multiArrayType
|
||||
self.input_shape = tuple(int(s) for s in array_type.shape)
|
||||
if len(self.input_shape) == 4:
|
||||
self.batch_size = self.input_shape[0] if self.input_shape[0] > 0 else batch_size
|
||||
|
||||
self._output_names = [o.name for o in spec.description.output]
|
||||
|
||||
constants_inf.log(<str>f'CoreML model: input={self.input_name} shape={self.input_shape}')
|
||||
constants_inf.log(<str>f'CoreML outputs: {self._output_names}')
|
||||
constants_inf.log(<str>f'CoreML model: {self.img_width}x{self.img_height}')
|
||||
|
||||
@property
|
||||
def engine_name(self):
|
||||
@@ -42,38 +37,6 @@ cdef class CoreMLEngine(InferenceEngine):
|
||||
def get_engine_filename():
|
||||
return "azaion_coreml.zip"
|
||||
|
||||
@staticmethod
|
||||
def convert_from_onnx(bytes onnx_bytes):
|
||||
import coremltools as ct
|
||||
|
||||
with tempfile.NamedTemporaryFile(suffix='.onnx', delete=False) as f:
|
||||
f.write(onnx_bytes)
|
||||
onnx_path = f.name
|
||||
|
||||
try:
|
||||
constants_inf.log(<str>'Converting ONNX to CoreML...')
|
||||
model = ct.convert(
|
||||
onnx_path,
|
||||
compute_units=ct.ComputeUnit.ALL,
|
||||
minimum_deployment_target=ct.target.macOS13,
|
||||
)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
pkg_path = os.path.join(tmpdir, "azaion.mlpackage")
|
||||
model.save(pkg_path)
|
||||
|
||||
buf = io.BytesIO()
|
||||
with zipfile.ZipFile(buf, 'w', zipfile.ZIP_DEFLATED) as zf:
|
||||
for root, dirs, files in os.walk(pkg_path):
|
||||
for fname in files:
|
||||
file_path = os.path.join(root, fname)
|
||||
arcname = os.path.relpath(file_path, tmpdir)
|
||||
zf.write(file_path, arcname)
|
||||
constants_inf.log(<str>'CoreML conversion done!')
|
||||
return buf.getvalue()
|
||||
finally:
|
||||
os.unlink(onnx_path)
|
||||
|
||||
@staticmethod
|
||||
def _extract_from_zip(model_bytes):
|
||||
tmpdir = tempfile.mkdtemp()
|
||||
@@ -86,17 +49,40 @@ cdef class CoreMLEngine(InferenceEngine):
|
||||
raise ValueError("No .mlpackage or .mlmodel found in zip")
|
||||
|
||||
cdef tuple get_input_shape(self):
|
||||
return self.input_shape[2], self.input_shape[3]
|
||||
return self.img_height, self.img_width
|
||||
|
||||
cdef int get_batch_size(self):
|
||||
return self.batch_size
|
||||
return 1
|
||||
|
||||
cdef run(self, input_data):
|
||||
prediction = self.model.predict({self.input_name: input_data})
|
||||
results = []
|
||||
for name in self._output_names:
|
||||
val = prediction[name]
|
||||
if not isinstance(val, np.ndarray):
|
||||
val = np.array(val)
|
||||
results.append(val)
|
||||
return results
|
||||
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, :, :]]
|
||||
|
||||
Reference in New Issue
Block a user