deep_learning/ultralytics
1
0
Fork 0
mirror of https://github.com/ultralytics/ultralytics.git synced 2025-09-15 15:48:41 +08:00
Code Issues Actions 2 Packages Projects Releases Wiki Activity

Support CoreML inference with NMS embedded models (#21948)

Browse Source 
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Jing Qiu <61612323+Laughing-q@users.noreply.github.com>
Co-authored-by: Laughing-q <1185102784@qq.com>
This commit is contained in:
Mohammed Yasin 2025-09-05 13:16:21 +08:00 committed by GitHub
parent cb0e6f08c4
commit d346c132a8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 16 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
tests/test_exports.py
View File

@ -120,14 +120,16 @@ def test_export_torchscript_matrix(task, dynamic, int8, half, batch, nms):
@pytest.mark.skipif(not TORCH_1_9, reason="CoreML>=7.2 not supported with PyTorch<=1.8")
@pytest.mark.skipif(checks.IS_PYTHON_3_13, reason="CoreML not supported in Python 3.13")
@pytest.mark.parametrize(
"task, dynamic, int8, half, batch",
"task, dynamic, int8, half, nms, batch",
[ # generate all combinations except for exclusion cases
(task, dynamic, int8, half, batch)
for task, dynamic, int8, half, batch in product(TASKS, [False], [True, False], [True, False], [1])
if not (int8 and half)
(task, dynamic, int8, half, nms, batch)
for task, dynamic, int8, half, nms, batch in product(
TASKS, [False], [True, False], [True, False], [True, False], [1]
)
if not (int8 and half) and not (task != "detect" and nms)
],
)
def test_export_coreml_matrix(task, dynamic, int8, half, batch):
def test_export_coreml_matrix(task, dynamic, int8, half, nms, batch):
"""Test YOLO export to CoreML format with various parameter configurations."""
file = YOLO(TASK2MODEL[task]).export(
format="coreml",
@ -136,6 +138,7 @@ def test_export_coreml_matrix(task, dynamic, int8, half, batch):
int8=int8,
half=half,
batch=batch,
nms=nms,
)
YOLO(file)([SOURCE] * batch, imgsz=32) # exported model inference
shutil.rmtree(file) # cleanup

19
ultralytics/nn/autobackend.py
View File

@ -724,17 +724,14 @@ class AutoBackend(nn.Module):
im_pil = Image.fromarray((im * 255).astype("uint8"))
# im = im.resize((192, 320), Image.BILINEAR)
y = self.model.predict({"image": im_pil}) # coordinates are xywh normalized
if "confidence" in y:
raise TypeError(
"Ultralytics only supports inference of non-pipelined CoreML models exported with "
f"'nms=False', but 'model={w}' has an NMS pipeline created by an 'nms=True' export."
)
# TODO: CoreML NMS inference handling
# from ultralytics.utils.ops import xywh2xyxy
# box = xywh2xyxy(y['coordinates'] * [[w, h, w, h]]) # xyxy pixels
# conf, cls = y['confidence'].max(1), y['confidence'].argmax(1).astype(np.float32)
# y = np.concatenate((box, conf.reshape(-1, 1), cls.reshape(-1, 1)), 1)
y = list(y.values())
if "confidence" in y: # NMS included
from ultralytics.utils.ops import xywh2xyxy
box = xywh2xyxy(y["coordinates"] * [[w, h, w, h]]) # xyxy pixels
cls = y["confidence"].argmax(1, keepdims=True)
y = np.concatenate((box, np.take_along_axis(y["confidence"], cls, axis=1), cls), 1)[None]
else:
y = list(y.values())
if len(y) == 2 and len(y[1].shape) != 4: # segmentation model
y = list(reversed(y)) # reversed for segmentation models (pred, proto)