HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

Support CoreML NMS export for Segment, Pose and OBB (#19173)

Browse source 
Signed-off-by: Mohammed Yasin <32206511+Y-T-G@users.noreply.github.com>
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: Laughing <61612323+Laughing-q@users.noreply.github.com>
This commit is contained in:
Mohammed Yasin 2025-02-17 15:01:12 +08:00 • committed by GitHub
parent d92ab8764b
commit 0ae4670da6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 20 additions and 167 deletions
Download patch file Download diff file Expand all files Collapse all files

13
tests/test_exports.py
View file

@ -116,14 +116,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_12, reason="CoreML not supported in Python 3.12")
@pytest.mark.parametrize(
"task, dynamic, int8, half, batch",
"task, dynamic, int8, half, batch, nms",
[ # 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, batch, nms)
for task, dynamic, int8, half, batch, nms in product(
TASKS, [False], [True, False], [True, False], [1], [True, False]
)
if not ((int8 and half) or (task == "classify" and nms))
],
)
def test_export_coreml_matrix(task, dynamic, int8, half, batch):
def test_export_coreml_matrix(task, dynamic, int8, half, batch, nms):
"""Test YOLO exports to CoreML format with various parameter configurations."""
file = YOLO(TASK2MODEL[task]).export(
format="coreml",
@ -132,6 +134,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 at batch=3
shutil.rmtree(file) # cleanup