dgenerate-ultralytics-headless 8.3.214__py3-none-any.whl → 8.3.248__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {dgenerate_ultralytics_headless-8.3.214.dist-info → dgenerate_ultralytics_headless-8.3.248.dist-info}/METADATA +13 -14
- dgenerate_ultralytics_headless-8.3.248.dist-info/RECORD +298 -0
- tests/__init__.py +5 -7
- tests/conftest.py +8 -15
- tests/test_cli.py +1 -1
- tests/test_cuda.py +5 -8
- tests/test_engine.py +1 -1
- tests/test_exports.py +57 -12
- tests/test_integrations.py +4 -4
- tests/test_python.py +84 -53
- tests/test_solutions.py +160 -151
- ultralytics/__init__.py +1 -1
- ultralytics/cfg/__init__.py +56 -62
- ultralytics/cfg/datasets/Argoverse.yaml +7 -6
- ultralytics/cfg/datasets/DOTAv1.5.yaml +1 -1
- ultralytics/cfg/datasets/DOTAv1.yaml +1 -1
- ultralytics/cfg/datasets/ImageNet.yaml +1 -1
- ultralytics/cfg/datasets/VOC.yaml +15 -16
- ultralytics/cfg/datasets/african-wildlife.yaml +1 -1
- ultralytics/cfg/datasets/coco-pose.yaml +21 -0
- ultralytics/cfg/datasets/coco128-seg.yaml +1 -1
- ultralytics/cfg/datasets/coco8-pose.yaml +21 -0
- ultralytics/cfg/datasets/dog-pose.yaml +28 -0
- ultralytics/cfg/datasets/dota8-multispectral.yaml +1 -1
- ultralytics/cfg/datasets/dota8.yaml +2 -2
- ultralytics/cfg/datasets/hand-keypoints.yaml +26 -2
- ultralytics/cfg/datasets/kitti.yaml +27 -0
- ultralytics/cfg/datasets/lvis.yaml +5 -5
- ultralytics/cfg/datasets/open-images-v7.yaml +1 -1
- ultralytics/cfg/datasets/tiger-pose.yaml +16 -0
- ultralytics/cfg/datasets/xView.yaml +16 -16
- ultralytics/cfg/default.yaml +1 -1
- ultralytics/cfg/models/11/yolo11-pose.yaml +1 -1
- ultralytics/cfg/models/11/yoloe-11-seg.yaml +2 -2
- ultralytics/cfg/models/11/yoloe-11.yaml +2 -2
- ultralytics/cfg/models/rt-detr/rtdetr-l.yaml +1 -1
- ultralytics/cfg/models/rt-detr/rtdetr-resnet101.yaml +1 -1
- ultralytics/cfg/models/rt-detr/rtdetr-resnet50.yaml +1 -1
- ultralytics/cfg/models/rt-detr/rtdetr-x.yaml +1 -1
- ultralytics/cfg/models/v10/yolov10b.yaml +2 -2
- ultralytics/cfg/models/v10/yolov10l.yaml +2 -2
- ultralytics/cfg/models/v10/yolov10m.yaml +2 -2
- ultralytics/cfg/models/v10/yolov10n.yaml +2 -2
- ultralytics/cfg/models/v10/yolov10s.yaml +2 -2
- ultralytics/cfg/models/v10/yolov10x.yaml +2 -2
- ultralytics/cfg/models/v3/yolov3-tiny.yaml +1 -1
- ultralytics/cfg/models/v6/yolov6.yaml +1 -1
- ultralytics/cfg/models/v8/yoloe-v8-seg.yaml +9 -6
- ultralytics/cfg/models/v8/yoloe-v8.yaml +9 -6
- ultralytics/cfg/models/v8/yolov8-cls-resnet101.yaml +1 -1
- ultralytics/cfg/models/v8/yolov8-cls-resnet50.yaml +1 -1
- ultralytics/cfg/models/v8/yolov8-ghost-p2.yaml +2 -2
- ultralytics/cfg/models/v8/yolov8-ghost-p6.yaml +2 -2
- ultralytics/cfg/models/v8/yolov8-ghost.yaml +2 -2
- ultralytics/cfg/models/v8/yolov8-obb.yaml +1 -1
- ultralytics/cfg/models/v8/yolov8-p2.yaml +1 -1
- ultralytics/cfg/models/v8/yolov8-pose-p6.yaml +1 -1
- ultralytics/cfg/models/v8/yolov8-rtdetr.yaml +1 -1
- ultralytics/cfg/models/v8/yolov8-seg-p6.yaml +1 -1
- ultralytics/cfg/models/v8/yolov8-world.yaml +1 -1
- ultralytics/cfg/models/v8/yolov8-worldv2.yaml +6 -6
- ultralytics/cfg/models/v9/yolov9s.yaml +1 -1
- ultralytics/data/__init__.py +4 -4
- ultralytics/data/annotator.py +3 -4
- ultralytics/data/augment.py +285 -475
- ultralytics/data/base.py +18 -26
- ultralytics/data/build.py +147 -25
- ultralytics/data/converter.py +36 -46
- ultralytics/data/dataset.py +46 -74
- ultralytics/data/loaders.py +42 -49
- ultralytics/data/split.py +5 -6
- ultralytics/data/split_dota.py +8 -15
- ultralytics/data/utils.py +34 -43
- ultralytics/engine/exporter.py +319 -237
- ultralytics/engine/model.py +148 -188
- ultralytics/engine/predictor.py +29 -38
- ultralytics/engine/results.py +177 -311
- ultralytics/engine/trainer.py +83 -59
- ultralytics/engine/tuner.py +23 -34
- ultralytics/engine/validator.py +39 -22
- ultralytics/hub/__init__.py +16 -19
- ultralytics/hub/auth.py +6 -12
- ultralytics/hub/google/__init__.py +7 -10
- ultralytics/hub/session.py +15 -25
- ultralytics/hub/utils.py +5 -8
- ultralytics/models/__init__.py +1 -1
- ultralytics/models/fastsam/__init__.py +1 -1
- ultralytics/models/fastsam/model.py +8 -10
- ultralytics/models/fastsam/predict.py +17 -29
- ultralytics/models/fastsam/utils.py +1 -2
- ultralytics/models/fastsam/val.py +5 -7
- ultralytics/models/nas/__init__.py +1 -1
- ultralytics/models/nas/model.py +5 -8
- ultralytics/models/nas/predict.py +7 -9
- ultralytics/models/nas/val.py +1 -2
- ultralytics/models/rtdetr/__init__.py +1 -1
- ultralytics/models/rtdetr/model.py +5 -8
- ultralytics/models/rtdetr/predict.py +15 -19
- ultralytics/models/rtdetr/train.py +10 -13
- ultralytics/models/rtdetr/val.py +21 -23
- ultralytics/models/sam/__init__.py +15 -2
- ultralytics/models/sam/amg.py +14 -20
- ultralytics/models/sam/build.py +26 -19
- ultralytics/models/sam/build_sam3.py +377 -0
- ultralytics/models/sam/model.py +29 -32
- ultralytics/models/sam/modules/blocks.py +83 -144
- ultralytics/models/sam/modules/decoders.py +19 -37
- ultralytics/models/sam/modules/encoders.py +44 -101
- ultralytics/models/sam/modules/memory_attention.py +16 -30
- ultralytics/models/sam/modules/sam.py +200 -73
- ultralytics/models/sam/modules/tiny_encoder.py +64 -83
- ultralytics/models/sam/modules/transformer.py +18 -28
- ultralytics/models/sam/modules/utils.py +174 -50
- ultralytics/models/sam/predict.py +2248 -350
- ultralytics/models/sam/sam3/__init__.py +3 -0
- ultralytics/models/sam/sam3/decoder.py +546 -0
- ultralytics/models/sam/sam3/encoder.py +529 -0
- ultralytics/models/sam/sam3/geometry_encoders.py +415 -0
- ultralytics/models/sam/sam3/maskformer_segmentation.py +286 -0
- ultralytics/models/sam/sam3/model_misc.py +199 -0
- ultralytics/models/sam/sam3/necks.py +129 -0
- ultralytics/models/sam/sam3/sam3_image.py +339 -0
- ultralytics/models/sam/sam3/text_encoder_ve.py +307 -0
- ultralytics/models/sam/sam3/vitdet.py +547 -0
- ultralytics/models/sam/sam3/vl_combiner.py +160 -0
- ultralytics/models/utils/loss.py +14 -26
- ultralytics/models/utils/ops.py +13 -17
- ultralytics/models/yolo/__init__.py +1 -1
- ultralytics/models/yolo/classify/predict.py +9 -12
- ultralytics/models/yolo/classify/train.py +11 -32
- ultralytics/models/yolo/classify/val.py +29 -28
- ultralytics/models/yolo/detect/predict.py +7 -10
- ultralytics/models/yolo/detect/train.py +11 -20
- ultralytics/models/yolo/detect/val.py +70 -58
- ultralytics/models/yolo/model.py +36 -53
- ultralytics/models/yolo/obb/predict.py +5 -14
- ultralytics/models/yolo/obb/train.py +11 -14
- ultralytics/models/yolo/obb/val.py +39 -36
- ultralytics/models/yolo/pose/__init__.py +1 -1
- ultralytics/models/yolo/pose/predict.py +6 -21
- ultralytics/models/yolo/pose/train.py +10 -15
- ultralytics/models/yolo/pose/val.py +38 -57
- ultralytics/models/yolo/segment/predict.py +14 -18
- ultralytics/models/yolo/segment/train.py +3 -6
- ultralytics/models/yolo/segment/val.py +93 -45
- ultralytics/models/yolo/world/train.py +8 -14
- ultralytics/models/yolo/world/train_world.py +11 -34
- ultralytics/models/yolo/yoloe/__init__.py +7 -7
- ultralytics/models/yolo/yoloe/predict.py +16 -23
- ultralytics/models/yolo/yoloe/train.py +30 -43
- ultralytics/models/yolo/yoloe/train_seg.py +5 -10
- ultralytics/models/yolo/yoloe/val.py +15 -20
- ultralytics/nn/__init__.py +7 -7
- ultralytics/nn/autobackend.py +145 -77
- ultralytics/nn/modules/__init__.py +60 -60
- ultralytics/nn/modules/activation.py +4 -6
- ultralytics/nn/modules/block.py +132 -216
- ultralytics/nn/modules/conv.py +52 -97
- ultralytics/nn/modules/head.py +50 -103
- ultralytics/nn/modules/transformer.py +76 -88
- ultralytics/nn/modules/utils.py +16 -21
- ultralytics/nn/tasks.py +94 -154
- ultralytics/nn/text_model.py +40 -67
- ultralytics/solutions/__init__.py +12 -12
- ultralytics/solutions/ai_gym.py +11 -17
- ultralytics/solutions/analytics.py +15 -16
- ultralytics/solutions/config.py +5 -6
- ultralytics/solutions/distance_calculation.py +10 -13
- ultralytics/solutions/heatmap.py +7 -13
- ultralytics/solutions/instance_segmentation.py +5 -8
- ultralytics/solutions/object_blurrer.py +7 -10
- ultralytics/solutions/object_counter.py +12 -19
- ultralytics/solutions/object_cropper.py +8 -14
- ultralytics/solutions/parking_management.py +33 -31
- ultralytics/solutions/queue_management.py +10 -12
- ultralytics/solutions/region_counter.py +9 -12
- ultralytics/solutions/security_alarm.py +15 -20
- ultralytics/solutions/similarity_search.py +10 -15
- ultralytics/solutions/solutions.py +75 -74
- ultralytics/solutions/speed_estimation.py +7 -10
- ultralytics/solutions/streamlit_inference.py +2 -4
- ultralytics/solutions/templates/similarity-search.html +7 -18
- ultralytics/solutions/trackzone.py +7 -10
- ultralytics/solutions/vision_eye.py +5 -8
- ultralytics/trackers/__init__.py +1 -1
- ultralytics/trackers/basetrack.py +3 -5
- ultralytics/trackers/bot_sort.py +10 -27
- ultralytics/trackers/byte_tracker.py +14 -30
- ultralytics/trackers/track.py +3 -6
- ultralytics/trackers/utils/gmc.py +11 -22
- ultralytics/trackers/utils/kalman_filter.py +37 -48
- ultralytics/trackers/utils/matching.py +12 -15
- ultralytics/utils/__init__.py +116 -116
- ultralytics/utils/autobatch.py +2 -4
- ultralytics/utils/autodevice.py +17 -18
- ultralytics/utils/benchmarks.py +32 -46
- ultralytics/utils/callbacks/base.py +8 -10
- ultralytics/utils/callbacks/clearml.py +5 -13
- ultralytics/utils/callbacks/comet.py +32 -46
- ultralytics/utils/callbacks/dvc.py +13 -18
- ultralytics/utils/callbacks/mlflow.py +4 -5
- ultralytics/utils/callbacks/neptune.py +7 -15
- ultralytics/utils/callbacks/platform.py +314 -38
- ultralytics/utils/callbacks/raytune.py +3 -4
- ultralytics/utils/callbacks/tensorboard.py +23 -31
- ultralytics/utils/callbacks/wb.py +10 -13
- ultralytics/utils/checks.py +99 -76
- ultralytics/utils/cpu.py +3 -8
- ultralytics/utils/dist.py +8 -12
- ultralytics/utils/downloads.py +20 -30
- ultralytics/utils/errors.py +6 -14
- ultralytics/utils/events.py +2 -4
- ultralytics/utils/export/__init__.py +4 -236
- ultralytics/utils/export/engine.py +237 -0
- ultralytics/utils/export/imx.py +91 -55
- ultralytics/utils/export/tensorflow.py +231 -0
- ultralytics/utils/files.py +24 -28
- ultralytics/utils/git.py +9 -11
- ultralytics/utils/instance.py +30 -51
- ultralytics/utils/logger.py +212 -114
- ultralytics/utils/loss.py +14 -22
- ultralytics/utils/metrics.py +126 -155
- ultralytics/utils/nms.py +13 -16
- ultralytics/utils/ops.py +107 -165
- ultralytics/utils/patches.py +33 -21
- ultralytics/utils/plotting.py +72 -80
- ultralytics/utils/tal.py +25 -39
- ultralytics/utils/torch_utils.py +52 -78
- ultralytics/utils/tqdm.py +20 -20
- ultralytics/utils/triton.py +13 -19
- ultralytics/utils/tuner.py +17 -5
- dgenerate_ultralytics_headless-8.3.214.dist-info/RECORD +0 -283
- {dgenerate_ultralytics_headless-8.3.214.dist-info → dgenerate_ultralytics_headless-8.3.248.dist-info}/WHEEL +0 -0
- {dgenerate_ultralytics_headless-8.3.214.dist-info → dgenerate_ultralytics_headless-8.3.248.dist-info}/entry_points.txt +0 -0
- {dgenerate_ultralytics_headless-8.3.214.dist-info → dgenerate_ultralytics_headless-8.3.248.dist-info}/licenses/LICENSE +0 -0
- {dgenerate_ultralytics_headless-8.3.214.dist-info → dgenerate_ultralytics_headless-8.3.248.dist-info}/top_level.txt +0 -0
ultralytics/engine/exporter.py
CHANGED
|
@@ -20,6 +20,8 @@ MNN | `mnn` | yolo11n.mnn
|
|
|
20
20
|
NCNN | `ncnn` | yolo11n_ncnn_model/
|
|
21
21
|
IMX | `imx` | yolo11n_imx_model/
|
|
22
22
|
RKNN | `rknn` | yolo11n_rknn_model/
|
|
23
|
+
ExecuTorch | `executorch` | yolo11n_executorch_model/
|
|
24
|
+
Axelera | `axelera` | yolo11n_axelera_model/
|
|
23
25
|
|
|
24
26
|
Requirements:
|
|
25
27
|
$ pip install "ultralytics[export]"
|
|
@@ -48,6 +50,8 @@ Inference:
|
|
|
48
50
|
yolo11n_ncnn_model # NCNN
|
|
49
51
|
yolo11n_imx_model # IMX
|
|
50
52
|
yolo11n_rknn_model # RKNN
|
|
53
|
+
yolo11n_executorch_model # ExecuTorch
|
|
54
|
+
yolo11n_axelera_model # Axelera
|
|
51
55
|
|
|
52
56
|
TensorFlow.js:
|
|
53
57
|
$ cd .. && git clone https://github.com/zldrobit/tfjs-yolov5-example.git && cd tfjs-yolov5-example
|
|
@@ -62,7 +66,6 @@ import re
|
|
|
62
66
|
import shutil
|
|
63
67
|
import subprocess
|
|
64
68
|
import time
|
|
65
|
-
import warnings
|
|
66
69
|
from copy import deepcopy
|
|
67
70
|
from datetime import datetime
|
|
68
71
|
from pathlib import Path
|
|
@@ -82,13 +85,17 @@ from ultralytics.utils import (
|
|
|
82
85
|
ARM64,
|
|
83
86
|
DEFAULT_CFG,
|
|
84
87
|
IS_COLAB,
|
|
88
|
+
IS_DEBIAN_BOOKWORM,
|
|
89
|
+
IS_DEBIAN_TRIXIE,
|
|
90
|
+
IS_DOCKER,
|
|
85
91
|
IS_JETSON,
|
|
92
|
+
IS_RASPBERRYPI,
|
|
93
|
+
IS_UBUNTU,
|
|
86
94
|
LINUX,
|
|
87
95
|
LOGGER,
|
|
88
96
|
MACOS,
|
|
89
97
|
MACOS_VERSION,
|
|
90
98
|
RKNN_CHIPS,
|
|
91
|
-
ROOT,
|
|
92
99
|
SETTINGS,
|
|
93
100
|
TORCH_VERSION,
|
|
94
101
|
WINDOWS,
|
|
@@ -98,21 +105,38 @@ from ultralytics.utils import (
|
|
|
98
105
|
get_default_args,
|
|
99
106
|
)
|
|
100
107
|
from ultralytics.utils.checks import (
|
|
108
|
+
IS_PYTHON_3_10,
|
|
109
|
+
IS_PYTHON_MINIMUM_3_9,
|
|
110
|
+
check_apt_requirements,
|
|
101
111
|
check_imgsz,
|
|
102
|
-
check_is_path_safe,
|
|
103
112
|
check_requirements,
|
|
104
113
|
check_version,
|
|
105
114
|
is_intel,
|
|
106
115
|
is_sudo_available,
|
|
107
116
|
)
|
|
108
|
-
from ultralytics.utils.
|
|
109
|
-
|
|
110
|
-
|
|
117
|
+
from ultralytics.utils.export import (
|
|
118
|
+
keras2pb,
|
|
119
|
+
onnx2engine,
|
|
120
|
+
onnx2saved_model,
|
|
121
|
+
pb2tfjs,
|
|
122
|
+
tflite2edgetpu,
|
|
123
|
+
torch2imx,
|
|
124
|
+
torch2onnx,
|
|
125
|
+
)
|
|
126
|
+
from ultralytics.utils.files import file_size
|
|
111
127
|
from ultralytics.utils.metrics import batch_probiou
|
|
112
128
|
from ultralytics.utils.nms import TorchNMS
|
|
113
129
|
from ultralytics.utils.ops import Profile
|
|
114
130
|
from ultralytics.utils.patches import arange_patch
|
|
115
|
-
from ultralytics.utils.torch_utils import
|
|
131
|
+
from ultralytics.utils.torch_utils import (
|
|
132
|
+
TORCH_1_10,
|
|
133
|
+
TORCH_1_11,
|
|
134
|
+
TORCH_1_13,
|
|
135
|
+
TORCH_2_1,
|
|
136
|
+
TORCH_2_4,
|
|
137
|
+
TORCH_2_9,
|
|
138
|
+
select_device,
|
|
139
|
+
)
|
|
116
140
|
|
|
117
141
|
|
|
118
142
|
def export_formats():
|
|
@@ -148,18 +172,20 @@ def export_formats():
|
|
|
148
172
|
["NCNN", "ncnn", "_ncnn_model", True, True, ["batch", "half"]],
|
|
149
173
|
["IMX", "imx", "_imx_model", True, True, ["int8", "fraction", "nms"]],
|
|
150
174
|
["RKNN", "rknn", "_rknn_model", False, False, ["batch", "name"]],
|
|
175
|
+
["ExecuTorch", "executorch", "_executorch_model", True, False, ["batch"]],
|
|
176
|
+
["Axelera", "axelera", "_axelera_model", False, False, ["batch", "int8", "fraction"]],
|
|
151
177
|
]
|
|
152
178
|
return dict(zip(["Format", "Argument", "Suffix", "CPU", "GPU", "Arguments"], zip(*x)))
|
|
153
179
|
|
|
154
180
|
|
|
155
181
|
def best_onnx_opset(onnx, cuda=False) -> int:
|
|
156
182
|
"""Return max ONNX opset for this torch version with ONNX fallback."""
|
|
157
|
-
version = ".".join(TORCH_VERSION.split(".")[:2])
|
|
158
183
|
if TORCH_2_4: # _constants.ONNX_MAX_OPSET first defined in torch 1.13
|
|
159
184
|
opset = torch.onnx.utils._constants.ONNX_MAX_OPSET - 1 # use second-latest version for safety
|
|
160
185
|
if cuda:
|
|
161
186
|
opset -= 2 # fix CUDA ONNXRuntime NMS squeeze op errors
|
|
162
187
|
else:
|
|
188
|
+
version = ".".join(TORCH_VERSION.split(".")[:2])
|
|
163
189
|
opset = {
|
|
164
190
|
"1.8": 12,
|
|
165
191
|
"1.9": 12,
|
|
@@ -181,8 +207,7 @@ def best_onnx_opset(onnx, cuda=False) -> int:
|
|
|
181
207
|
|
|
182
208
|
|
|
183
209
|
def validate_args(format, passed_args, valid_args):
|
|
184
|
-
"""
|
|
185
|
-
Validate arguments based on the export format.
|
|
210
|
+
"""Validate arguments based on the export format.
|
|
186
211
|
|
|
187
212
|
Args:
|
|
188
213
|
format (str): The export format.
|
|
@@ -203,15 +228,6 @@ def validate_args(format, passed_args, valid_args):
|
|
|
203
228
|
assert arg in valid_args, f"ERROR ❌️ argument '{arg}' is not supported for format='{format}'"
|
|
204
229
|
|
|
205
230
|
|
|
206
|
-
def gd_outputs(gd):
|
|
207
|
-
"""Return TensorFlow GraphDef model output node names."""
|
|
208
|
-
name_list, input_list = [], []
|
|
209
|
-
for node in gd.node: # tensorflow.core.framework.node_def_pb2.NodeDef
|
|
210
|
-
name_list.append(node.name)
|
|
211
|
-
input_list.extend(node.input)
|
|
212
|
-
return sorted(f"{x}:0" for x in list(set(name_list) - set(input_list)) if not x.startswith("NoOp"))
|
|
213
|
-
|
|
214
|
-
|
|
215
231
|
def try_export(inner_func):
|
|
216
232
|
"""YOLO export decorator, i.e. @try_export."""
|
|
217
233
|
inner_args = get_default_args(inner_func)
|
|
@@ -236,8 +252,7 @@ def try_export(inner_func):
|
|
|
236
252
|
|
|
237
253
|
|
|
238
254
|
class Exporter:
|
|
239
|
-
"""
|
|
240
|
-
A class for exporting YOLO models to various formats.
|
|
255
|
+
"""A class for exporting YOLO models to various formats.
|
|
241
256
|
|
|
242
257
|
This class provides functionality to export YOLO models to different formats including ONNX, TensorRT, CoreML,
|
|
243
258
|
TensorFlow, and others. It handles format validation, device selection, model preparation, and the actual export
|
|
@@ -287,8 +302,7 @@ class Exporter:
|
|
|
287
302
|
"""
|
|
288
303
|
|
|
289
304
|
def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
|
|
290
|
-
"""
|
|
291
|
-
Initialize the Exporter class.
|
|
305
|
+
"""Initialize the Exporter class.
|
|
292
306
|
|
|
293
307
|
Args:
|
|
294
308
|
cfg (str, optional): Path to a configuration file.
|
|
@@ -300,7 +314,11 @@ class Exporter:
|
|
|
300
314
|
callbacks.add_integration_callbacks(self)
|
|
301
315
|
|
|
302
316
|
def __call__(self, model=None) -> str:
|
|
303
|
-
"""
|
|
317
|
+
"""Export a model and return the final exported path as a string.
|
|
318
|
+
|
|
319
|
+
Returns:
|
|
320
|
+
(str): Path to the exported file or directory (the last export artifact).
|
|
321
|
+
"""
|
|
304
322
|
t = time.time()
|
|
305
323
|
fmt = self.args.format.lower() # to lowercase
|
|
306
324
|
if fmt in {"tensorrt", "trt"}: # 'engine' aliases
|
|
@@ -322,9 +340,25 @@ class Exporter:
|
|
|
322
340
|
flags = [x == fmt for x in fmts]
|
|
323
341
|
if sum(flags) != 1:
|
|
324
342
|
raise ValueError(f"Invalid export format='{fmt}'. Valid formats are {fmts}")
|
|
325
|
-
(
|
|
326
|
-
|
|
327
|
-
|
|
343
|
+
(
|
|
344
|
+
jit,
|
|
345
|
+
onnx,
|
|
346
|
+
xml,
|
|
347
|
+
engine,
|
|
348
|
+
coreml,
|
|
349
|
+
saved_model,
|
|
350
|
+
pb,
|
|
351
|
+
tflite,
|
|
352
|
+
edgetpu,
|
|
353
|
+
tfjs,
|
|
354
|
+
paddle,
|
|
355
|
+
mnn,
|
|
356
|
+
ncnn,
|
|
357
|
+
imx,
|
|
358
|
+
rknn,
|
|
359
|
+
executorch,
|
|
360
|
+
axelera,
|
|
361
|
+
) = flags # export booleans
|
|
328
362
|
|
|
329
363
|
is_tf_format = any((saved_model, pb, tflite, edgetpu, tfjs))
|
|
330
364
|
|
|
@@ -334,9 +368,10 @@ class Exporter:
|
|
|
334
368
|
LOGGER.warning("TensorRT requires GPU export, automatically assigning device=0")
|
|
335
369
|
self.args.device = "0"
|
|
336
370
|
if engine and "dla" in str(self.args.device): # convert int/list to str first
|
|
337
|
-
|
|
371
|
+
device_str = str(self.args.device)
|
|
372
|
+
dla = device_str.rsplit(":", 1)[-1]
|
|
338
373
|
self.args.device = "0" # update device to "0"
|
|
339
|
-
assert dla in {"0", "1"}, f"Expected
|
|
374
|
+
assert dla in {"0", "1"}, f"Expected device 'dla:0' or 'dla:1', but got {device_str}."
|
|
340
375
|
if imx and self.args.device is None and torch.cuda.is_available():
|
|
341
376
|
LOGGER.warning("Exporting on CPU while CUDA is available, setting device=0 for faster export on GPU.")
|
|
342
377
|
self.args.device = "0" # update device to "0"
|
|
@@ -345,23 +380,37 @@ class Exporter:
|
|
|
345
380
|
# Argument compatibility checks
|
|
346
381
|
fmt_keys = fmts_dict["Arguments"][flags.index(True) + 1]
|
|
347
382
|
validate_args(fmt, self.args, fmt_keys)
|
|
383
|
+
if axelera:
|
|
384
|
+
if not IS_PYTHON_3_10:
|
|
385
|
+
raise SystemError("Axelera export only supported on Python 3.10.")
|
|
386
|
+
if not self.args.int8:
|
|
387
|
+
LOGGER.warning("Setting int8=True for Axelera mixed-precision export.")
|
|
388
|
+
self.args.int8 = True
|
|
389
|
+
if model.task not in {"detect"}:
|
|
390
|
+
raise ValueError("Axelera export only supported for detection models.")
|
|
391
|
+
if not self.args.data:
|
|
392
|
+
self.args.data = "coco128.yaml" # Axelera default to coco128.yaml
|
|
348
393
|
if imx:
|
|
349
394
|
if not self.args.int8:
|
|
350
395
|
LOGGER.warning("IMX export requires int8=True, setting int8=True.")
|
|
351
396
|
self.args.int8 = True
|
|
352
|
-
if not self.args.nms:
|
|
397
|
+
if not self.args.nms and model.task in {"detect", "pose", "segment"}:
|
|
353
398
|
LOGGER.warning("IMX export requires nms=True, setting nms=True.")
|
|
354
399
|
self.args.nms = True
|
|
355
|
-
if model.task not in {"detect", "pose"}:
|
|
356
|
-
raise ValueError(
|
|
400
|
+
if model.task not in {"detect", "pose", "classify", "segment"}:
|
|
401
|
+
raise ValueError(
|
|
402
|
+
"IMX export only supported for detection, pose estimation, classification, and segmentation models."
|
|
403
|
+
)
|
|
357
404
|
if not hasattr(model, "names"):
|
|
358
405
|
model.names = default_class_names()
|
|
359
406
|
model.names = check_class_names(model.names)
|
|
360
407
|
if self.args.half and self.args.int8:
|
|
361
408
|
LOGGER.warning("half=True and int8=True are mutually exclusive, setting half=False.")
|
|
362
409
|
self.args.half = False
|
|
363
|
-
if self.args.half and
|
|
364
|
-
LOGGER.warning(
|
|
410
|
+
if self.args.half and jit and self.device.type == "cpu":
|
|
411
|
+
LOGGER.warning(
|
|
412
|
+
"half=True only compatible with GPU export for TorchScript, i.e. use device=0, setting half=False."
|
|
413
|
+
)
|
|
365
414
|
self.args.half = False
|
|
366
415
|
self.imgsz = check_imgsz(self.args.imgsz, stride=model.stride, min_dim=2) # check image size
|
|
367
416
|
if self.args.optimize:
|
|
@@ -378,14 +427,12 @@ class Exporter:
|
|
|
378
427
|
assert self.args.name in RKNN_CHIPS, (
|
|
379
428
|
f"Invalid processor name '{self.args.name}' for Rockchip RKNN export. Valid names are {RKNN_CHIPS}."
|
|
380
429
|
)
|
|
381
|
-
if self.args.int8 and tflite:
|
|
382
|
-
assert not getattr(model, "end2end", False), "TFLite INT8 export not supported for end2end models."
|
|
383
430
|
if self.args.nms:
|
|
384
431
|
assert not isinstance(model, ClassificationModel), "'nms=True' is not valid for classification models."
|
|
385
432
|
assert not tflite or not ARM64 or not LINUX, "TFLite export with NMS unsupported on ARM64 Linux"
|
|
386
433
|
assert not is_tf_format or TORCH_1_13, "TensorFlow exports with NMS require torch>=1.13"
|
|
387
434
|
assert not onnx or TORCH_1_13, "ONNX export with NMS requires torch>=1.13"
|
|
388
|
-
if getattr(model, "end2end", False):
|
|
435
|
+
if getattr(model, "end2end", False) or isinstance(model.model[-1], RTDETRDecoder):
|
|
389
436
|
LOGGER.warning("'nms=True' is not available for end2end models. Forcing 'nms=False'.")
|
|
390
437
|
self.args.nms = False
|
|
391
438
|
self.args.conf = self.args.conf or 0.25 # set conf default value for nms export
|
|
@@ -445,6 +492,10 @@ class Exporter:
|
|
|
445
492
|
from ultralytics.utils.export.imx import FXModel
|
|
446
493
|
|
|
447
494
|
model = FXModel(model, self.imgsz)
|
|
495
|
+
if tflite or edgetpu:
|
|
496
|
+
from ultralytics.utils.export.tensorflow import tf_wrapper
|
|
497
|
+
|
|
498
|
+
model = tf_wrapper(model)
|
|
448
499
|
for m in model.modules():
|
|
449
500
|
if isinstance(m, Classify):
|
|
450
501
|
m.export = True
|
|
@@ -466,11 +517,6 @@ class Exporter:
|
|
|
466
517
|
if self.args.half and (onnx or jit) and self.device.type != "cpu":
|
|
467
518
|
im, model = im.half(), model.half() # to FP16
|
|
468
519
|
|
|
469
|
-
# Filter warnings
|
|
470
|
-
warnings.filterwarnings("ignore", category=torch.jit.TracerWarning) # suppress TracerWarning
|
|
471
|
-
warnings.filterwarnings("ignore", category=UserWarning) # suppress shape prim::Constant missing ONNX warning
|
|
472
|
-
warnings.filterwarnings("ignore", category=DeprecationWarning) # suppress CoreML np.bool deprecation warning
|
|
473
|
-
|
|
474
520
|
# Assign
|
|
475
521
|
self.im = im
|
|
476
522
|
self.model = model
|
|
@@ -502,6 +548,8 @@ class Exporter:
|
|
|
502
548
|
self.metadata["dla"] = dla # make sure `AutoBackend` uses correct dla device if it has one
|
|
503
549
|
if model.task == "pose":
|
|
504
550
|
self.metadata["kpt_shape"] = model.model[-1].kpt_shape
|
|
551
|
+
if hasattr(model, "kpt_names"):
|
|
552
|
+
self.metadata["kpt_names"] = model.kpt_names
|
|
505
553
|
|
|
506
554
|
LOGGER.info(
|
|
507
555
|
f"\n{colorstr('PyTorch:')} starting from '{file}' with input shape {tuple(im.shape)} BCHW and "
|
|
@@ -514,7 +562,7 @@ class Exporter:
|
|
|
514
562
|
f[0] = self.export_torchscript()
|
|
515
563
|
if engine: # TensorRT required before ONNX
|
|
516
564
|
f[1] = self.export_engine(dla=dla)
|
|
517
|
-
if onnx
|
|
565
|
+
if onnx: # ONNX
|
|
518
566
|
f[2] = self.export_onnx()
|
|
519
567
|
if xml: # OpenVINO
|
|
520
568
|
f[3] = self.export_openvino()
|
|
@@ -541,6 +589,10 @@ class Exporter:
|
|
|
541
589
|
f[13] = self.export_imx()
|
|
542
590
|
if rknn:
|
|
543
591
|
f[14] = self.export_rknn()
|
|
592
|
+
if executorch:
|
|
593
|
+
f[15] = self.export_executorch()
|
|
594
|
+
if axelera:
|
|
595
|
+
f[16] = self.export_axelera()
|
|
544
596
|
|
|
545
597
|
# Finish
|
|
546
598
|
f = [str(x) for x in f if x] # filter out '' and None
|
|
@@ -565,7 +617,7 @@ class Exporter:
|
|
|
565
617
|
)
|
|
566
618
|
|
|
567
619
|
self.run_callbacks("on_export_end")
|
|
568
|
-
return f #
|
|
620
|
+
return f # path to final export artifact
|
|
569
621
|
|
|
570
622
|
def get_int8_calibration_dataloader(self, prefix=""):
|
|
571
623
|
"""Build and return a dataloader for calibration of INT8 models."""
|
|
@@ -586,7 +638,9 @@ class Exporter:
|
|
|
586
638
|
f"The calibration dataset ({n} images) must have at least as many images as the batch size "
|
|
587
639
|
f"('batch={self.args.batch}')."
|
|
588
640
|
)
|
|
589
|
-
elif n <
|
|
641
|
+
elif self.args.format == "axelera" and n < 100:
|
|
642
|
+
LOGGER.warning(f"{prefix} >100 images required for Axelera calibration, found {n} images.")
|
|
643
|
+
elif self.args.format != "axelera" and n < 300:
|
|
590
644
|
LOGGER.warning(f"{prefix} >300 images recommended for INT8 calibration, found {n} images.")
|
|
591
645
|
return build_dataloader(dataset, batch=self.args.batch, workers=0, drop_last=True) # required for batch loading
|
|
592
646
|
|
|
@@ -610,11 +664,11 @@ class Exporter:
|
|
|
610
664
|
@try_export
|
|
611
665
|
def export_onnx(self, prefix=colorstr("ONNX:")):
|
|
612
666
|
"""Export YOLO model to ONNX format."""
|
|
613
|
-
requirements = ["onnx>=1.12.0"]
|
|
667
|
+
requirements = ["onnx>=1.12.0,<2.0.0"]
|
|
614
668
|
if self.args.simplify:
|
|
615
669
|
requirements += ["onnxslim>=0.1.71", "onnxruntime" + ("-gpu" if torch.cuda.is_available() else "")]
|
|
616
670
|
check_requirements(requirements)
|
|
617
|
-
import onnx
|
|
671
|
+
import onnx
|
|
618
672
|
|
|
619
673
|
opset = self.args.opset or best_onnx_opset(onnx, cuda="cuda" in self.device.type)
|
|
620
674
|
LOGGER.info(f"\n{prefix} starting export with onnx {onnx.__version__} opset {opset}...")
|
|
@@ -622,7 +676,7 @@ class Exporter:
|
|
|
622
676
|
assert TORCH_1_13, f"'nms=True' ONNX export requires torch>=1.13 (found torch=={TORCH_VERSION})"
|
|
623
677
|
|
|
624
678
|
f = str(self.file.with_suffix(".onnx"))
|
|
625
|
-
output_names = ["output0", "output1"] if
|
|
679
|
+
output_names = ["output0", "output1"] if self.model.task == "segment" else ["output0"]
|
|
626
680
|
dynamic = self.args.dynamic
|
|
627
681
|
if dynamic:
|
|
628
682
|
dynamic = {"images": {0: "batch", 2: "height", 3: "width"}} # shape(1,3,640,640)
|
|
@@ -671,6 +725,16 @@ class Exporter:
|
|
|
671
725
|
LOGGER.info(f"{prefix} limiting IR version {model_onnx.ir_version} to 10 for ONNXRuntime compatibility...")
|
|
672
726
|
model_onnx.ir_version = 10
|
|
673
727
|
|
|
728
|
+
# FP16 conversion for CPU export (GPU exports are already FP16 from model.half() during tracing)
|
|
729
|
+
if self.args.half and self.args.format == "onnx" and self.device.type == "cpu":
|
|
730
|
+
try:
|
|
731
|
+
from onnxruntime.transformers import float16
|
|
732
|
+
|
|
733
|
+
LOGGER.info(f"{prefix} converting to FP16...")
|
|
734
|
+
model_onnx = float16.convert_float_to_float16(model_onnx, keep_io_types=True)
|
|
735
|
+
except Exception as e:
|
|
736
|
+
LOGGER.warning(f"{prefix} FP16 conversion failure: {e}")
|
|
737
|
+
|
|
674
738
|
onnx.save(model_onnx, f)
|
|
675
739
|
return f
|
|
676
740
|
|
|
@@ -711,13 +775,6 @@ class Exporter:
|
|
|
711
775
|
check_requirements("nncf>=2.14.0")
|
|
712
776
|
import nncf
|
|
713
777
|
|
|
714
|
-
def transform_fn(data_item) -> np.ndarray:
|
|
715
|
-
"""Quantization transform function."""
|
|
716
|
-
data_item: torch.Tensor = data_item["img"] if isinstance(data_item, dict) else data_item
|
|
717
|
-
assert data_item.dtype == torch.uint8, "Input image must be uint8 for the quantization preprocessing"
|
|
718
|
-
im = data_item.numpy().astype(np.float32) / 255.0 # uint8 to fp16/32 and 0-255 to 0.0-1.0
|
|
719
|
-
return np.expand_dims(im, 0) if im.ndim == 3 else im
|
|
720
|
-
|
|
721
778
|
# Generate calibration data for integer quantization
|
|
722
779
|
ignored_scope = None
|
|
723
780
|
if isinstance(self.model.model[-1], Detect):
|
|
@@ -736,7 +793,7 @@ class Exporter:
|
|
|
736
793
|
|
|
737
794
|
quantized_ov_model = nncf.quantize(
|
|
738
795
|
model=ov_model,
|
|
739
|
-
calibration_dataset=nncf.Dataset(self.get_int8_calibration_dataloader(prefix),
|
|
796
|
+
calibration_dataset=nncf.Dataset(self.get_int8_calibration_dataloader(prefix), self._transform_fn),
|
|
740
797
|
preset=nncf.QuantizationPreset.MIXED,
|
|
741
798
|
ignored_scope=ignored_scope,
|
|
742
799
|
)
|
|
@@ -763,8 +820,8 @@ class Exporter:
|
|
|
763
820
|
"x2paddle",
|
|
764
821
|
)
|
|
765
822
|
)
|
|
766
|
-
import x2paddle
|
|
767
|
-
from x2paddle.convert import pytorch2paddle
|
|
823
|
+
import x2paddle
|
|
824
|
+
from x2paddle.convert import pytorch2paddle
|
|
768
825
|
|
|
769
826
|
LOGGER.info(f"\n{prefix} starting export with X2Paddle {x2paddle.__version__}...")
|
|
770
827
|
f = str(self.file).replace(self.file.suffix, f"_paddle_model{os.sep}")
|
|
@@ -776,10 +833,11 @@ class Exporter:
|
|
|
776
833
|
@try_export
|
|
777
834
|
def export_mnn(self, prefix=colorstr("MNN:")):
|
|
778
835
|
"""Export YOLO model to MNN format using MNN https://github.com/alibaba/MNN."""
|
|
836
|
+
assert TORCH_1_10, "MNN export requires torch>=1.10.0 to avoid segmentation faults"
|
|
779
837
|
f_onnx = self.export_onnx() # get onnx model first
|
|
780
838
|
|
|
781
839
|
check_requirements("MNN>=2.9.6")
|
|
782
|
-
import MNN
|
|
840
|
+
import MNN
|
|
783
841
|
from MNN.tools import mnnconvert
|
|
784
842
|
|
|
785
843
|
# Setup and checks
|
|
@@ -802,65 +860,31 @@ class Exporter:
|
|
|
802
860
|
def export_ncnn(self, prefix=colorstr("NCNN:")):
|
|
803
861
|
"""Export YOLO model to NCNN format using PNNX https://github.com/pnnx/pnnx."""
|
|
804
862
|
check_requirements("ncnn", cmds="--no-deps") # no deps to avoid installing opencv-python
|
|
805
|
-
|
|
863
|
+
check_requirements("pnnx")
|
|
864
|
+
import ncnn
|
|
865
|
+
import pnnx
|
|
806
866
|
|
|
807
|
-
LOGGER.info(f"\n{prefix} starting export with NCNN {ncnn.__version__}...")
|
|
867
|
+
LOGGER.info(f"\n{prefix} starting export with NCNN {ncnn.__version__} and PNNX {pnnx.__version__}...")
|
|
808
868
|
f = Path(str(self.file).replace(self.file.suffix, f"_ncnn_model{os.sep}"))
|
|
809
|
-
f_onnx = self.file.with_suffix(".onnx")
|
|
810
869
|
|
|
811
|
-
|
|
812
|
-
|
|
813
|
-
|
|
814
|
-
|
|
815
|
-
|
|
816
|
-
|
|
817
|
-
|
|
818
|
-
)
|
|
819
|
-
|
|
820
|
-
|
|
821
|
-
|
|
822
|
-
|
|
823
|
-
|
|
824
|
-
|
|
825
|
-
except Exception as e:
|
|
826
|
-
release = "20250930"
|
|
827
|
-
asset = f"pnnx-{release}-{system}.zip"
|
|
828
|
-
LOGGER.warning(f"{prefix} PNNX GitHub assets not found: {e}, using default {asset}")
|
|
829
|
-
unzip_dir = safe_download(f"https://github.com/pnnx/pnnx/releases/download/{release}/{asset}", delete=True)
|
|
830
|
-
if check_is_path_safe(Path.cwd(), unzip_dir): # avoid path traversal security vulnerability
|
|
831
|
-
shutil.move(src=unzip_dir / name, dst=pnnx) # move binary to ROOT
|
|
832
|
-
pnnx.chmod(0o777) # set read, write, and execute permissions for everyone
|
|
833
|
-
shutil.rmtree(unzip_dir) # delete unzip dir
|
|
834
|
-
|
|
835
|
-
ncnn_args = [
|
|
836
|
-
f"ncnnparam={f / 'model.ncnn.param'}",
|
|
837
|
-
f"ncnnbin={f / 'model.ncnn.bin'}",
|
|
838
|
-
f"ncnnpy={f / 'model_ncnn.py'}",
|
|
839
|
-
]
|
|
840
|
-
|
|
841
|
-
pnnx_args = [
|
|
842
|
-
f"pnnxparam={f / 'model.pnnx.param'}",
|
|
843
|
-
f"pnnxbin={f / 'model.pnnx.bin'}",
|
|
844
|
-
f"pnnxpy={f / 'model_pnnx.py'}",
|
|
845
|
-
f"pnnxonnx={f / 'model.pnnx.onnx'}",
|
|
846
|
-
]
|
|
847
|
-
|
|
848
|
-
cmd = [
|
|
849
|
-
str(pnnx),
|
|
850
|
-
str(f_onnx),
|
|
851
|
-
*ncnn_args,
|
|
852
|
-
*pnnx_args,
|
|
853
|
-
f"fp16={int(self.args.half)}",
|
|
854
|
-
f"device={self.device.type}",
|
|
855
|
-
f'inputshape="{[self.args.batch, 3, *self.imgsz]}"',
|
|
856
|
-
]
|
|
870
|
+
ncnn_args = dict(
|
|
871
|
+
ncnnparam=(f / "model.ncnn.param").as_posix(),
|
|
872
|
+
ncnnbin=(f / "model.ncnn.bin").as_posix(),
|
|
873
|
+
ncnnpy=(f / "model_ncnn.py").as_posix(),
|
|
874
|
+
)
|
|
875
|
+
|
|
876
|
+
pnnx_args = dict(
|
|
877
|
+
ptpath=(f / "model.pt").as_posix(),
|
|
878
|
+
pnnxparam=(f / "model.pnnx.param").as_posix(),
|
|
879
|
+
pnnxbin=(f / "model.pnnx.bin").as_posix(),
|
|
880
|
+
pnnxpy=(f / "model_pnnx.py").as_posix(),
|
|
881
|
+
pnnxonnx=(f / "model.pnnx.onnx").as_posix(),
|
|
882
|
+
)
|
|
883
|
+
|
|
857
884
|
f.mkdir(exist_ok=True) # make ncnn_model directory
|
|
858
|
-
|
|
859
|
-
subprocess.run(cmd, check=True)
|
|
885
|
+
pnnx.export(self.model, inputs=self.im, **ncnn_args, **pnnx_args, fp16=self.args.half, device=self.device.type)
|
|
860
886
|
|
|
861
|
-
|
|
862
|
-
pnnx_files = [x.rsplit("=", 1)[-1] for x in pnnx_args]
|
|
863
|
-
for f_debug in ("debug.bin", "debug.param", "debug2.bin", "debug2.param", *pnnx_files):
|
|
887
|
+
for f_debug in ("debug.bin", "debug.param", "debug2.bin", "debug2.param", *pnnx_args.values()):
|
|
864
888
|
Path(f_debug).unlink(missing_ok=True)
|
|
865
889
|
|
|
866
890
|
YAML.save(f / "metadata.yaml", self.metadata) # add metadata.yaml
|
|
@@ -870,8 +894,10 @@ class Exporter:
|
|
|
870
894
|
def export_coreml(self, prefix=colorstr("CoreML:")):
|
|
871
895
|
"""Export YOLO model to CoreML format."""
|
|
872
896
|
mlmodel = self.args.format.lower() == "mlmodel" # legacy *.mlmodel export format requested
|
|
873
|
-
check_requirements(
|
|
874
|
-
|
|
897
|
+
check_requirements(
|
|
898
|
+
["coremltools>=9.0", "numpy>=1.14.5,<=2.3.5"]
|
|
899
|
+
) # latest numpy 2.4.0rc1 breaks coremltools exports
|
|
900
|
+
import coremltools as ct
|
|
875
901
|
|
|
876
902
|
LOGGER.info(f"\n{prefix} starting export with coremltools {ct.__version__}...")
|
|
877
903
|
assert not WINDOWS, "CoreML export is not supported on Windows, please run on macOS or Linux."
|
|
@@ -917,7 +943,7 @@ class Exporter:
|
|
|
917
943
|
|
|
918
944
|
# Based on apple's documentation it is better to leave out the minimum_deployment target and let that get set
|
|
919
945
|
# Internally based on the model conversion and output type.
|
|
920
|
-
# Setting
|
|
946
|
+
# Setting minimum_deployment_target >= iOS16 will require setting compute_precision=ct.precision.FLOAT32.
|
|
921
947
|
# iOS16 adds in better support for FP16, but none of the CoreML NMS specifications handle FP16 as input.
|
|
922
948
|
ct_model = ct.convert(
|
|
923
949
|
ts,
|
|
@@ -967,12 +993,12 @@ class Exporter:
|
|
|
967
993
|
f_onnx = self.export_onnx() # run before TRT import https://github.com/ultralytics/ultralytics/issues/7016
|
|
968
994
|
|
|
969
995
|
try:
|
|
970
|
-
import tensorrt as trt
|
|
996
|
+
import tensorrt as trt
|
|
971
997
|
except ImportError:
|
|
972
998
|
if LINUX:
|
|
973
999
|
cuda_version = torch.version.cuda.split(".")[0]
|
|
974
1000
|
check_requirements(f"tensorrt-cu{cuda_version}>7.0.0,!=10.1.0")
|
|
975
|
-
import tensorrt as trt
|
|
1001
|
+
import tensorrt as trt
|
|
976
1002
|
check_version(trt.__version__, ">=7.0.0", hard=True)
|
|
977
1003
|
check_version(trt.__version__, "!=10.1.0", msg="https://github.com/ultralytics/ultralytics/pull/14239")
|
|
978
1004
|
|
|
@@ -1002,17 +1028,17 @@ class Exporter:
|
|
|
1002
1028
|
"""Export YOLO model to TensorFlow SavedModel format."""
|
|
1003
1029
|
cuda = torch.cuda.is_available()
|
|
1004
1030
|
try:
|
|
1005
|
-
import tensorflow as tf
|
|
1031
|
+
import tensorflow as tf
|
|
1006
1032
|
except ImportError:
|
|
1007
1033
|
check_requirements("tensorflow>=2.0.0,<=2.19.0")
|
|
1008
|
-
import tensorflow as tf
|
|
1034
|
+
import tensorflow as tf
|
|
1009
1035
|
check_requirements(
|
|
1010
1036
|
(
|
|
1011
1037
|
"tf_keras<=2.19.0", # required by 'onnx2tf' package
|
|
1012
1038
|
"sng4onnx>=1.0.1", # required by 'onnx2tf' package
|
|
1013
1039
|
"onnx_graphsurgeon>=0.3.26", # required by 'onnx2tf' package
|
|
1014
1040
|
"ai-edge-litert>=1.2.0" + (",<1.4.0" if MACOS else ""), # required by 'onnx2tf' package
|
|
1015
|
-
"onnx>=1.12.0",
|
|
1041
|
+
"onnx>=1.12.0,<2.0.0",
|
|
1016
1042
|
"onnx2tf>=1.26.3",
|
|
1017
1043
|
"onnxslim>=0.1.71",
|
|
1018
1044
|
"onnxruntime-gpu" if cuda else "onnxruntime",
|
|
@@ -1033,82 +1059,50 @@ class Exporter:
|
|
|
1033
1059
|
if f.is_dir():
|
|
1034
1060
|
shutil.rmtree(f) # delete output folder
|
|
1035
1061
|
|
|
1036
|
-
#
|
|
1037
|
-
|
|
1038
|
-
if
|
|
1039
|
-
|
|
1062
|
+
# Export to TF
|
|
1063
|
+
images = None
|
|
1064
|
+
if self.args.int8 and self.args.data:
|
|
1065
|
+
images = [batch["img"] for batch in self.get_int8_calibration_dataloader(prefix)]
|
|
1066
|
+
images = (
|
|
1067
|
+
torch.nn.functional.interpolate(torch.cat(images, 0).float(), size=self.imgsz)
|
|
1068
|
+
.permute(0, 2, 3, 1)
|
|
1069
|
+
.numpy()
|
|
1070
|
+
.astype(np.float32)
|
|
1071
|
+
)
|
|
1040
1072
|
|
|
1041
1073
|
# Export to ONNX
|
|
1042
|
-
if
|
|
1074
|
+
if isinstance(self.model.model[-1], RTDETRDecoder):
|
|
1043
1075
|
self.args.opset = self.args.opset or 19
|
|
1044
1076
|
assert 16 <= self.args.opset <= 19, "RTDETR export requires opset>=16;<=19"
|
|
1045
1077
|
self.args.simplify = True
|
|
1046
|
-
f_onnx = self.export_onnx()
|
|
1047
|
-
|
|
1048
|
-
|
|
1049
|
-
|
|
1050
|
-
|
|
1051
|
-
|
|
1052
|
-
|
|
1053
|
-
|
|
1054
|
-
images = [batch["img"] for batch in self.get_int8_calibration_dataloader(prefix)]
|
|
1055
|
-
images = torch.nn.functional.interpolate(torch.cat(images, 0).float(), size=self.imgsz).permute(
|
|
1056
|
-
0, 2, 3, 1
|
|
1057
|
-
)
|
|
1058
|
-
np.save(str(tmp_file), images.numpy().astype(np.float32)) # BHWC
|
|
1059
|
-
np_data = [["images", tmp_file, [[[[0, 0, 0]]]], [[[[255, 255, 255]]]]]]
|
|
1060
|
-
|
|
1061
|
-
import onnx2tf # scoped for after ONNX export for reduced conflict during import
|
|
1062
|
-
|
|
1063
|
-
LOGGER.info(f"{prefix} starting TFLite export with onnx2tf {onnx2tf.__version__}...")
|
|
1064
|
-
keras_model = onnx2tf.convert(
|
|
1065
|
-
input_onnx_file_path=f_onnx,
|
|
1066
|
-
output_folder_path=str(f),
|
|
1067
|
-
not_use_onnxsim=True,
|
|
1068
|
-
verbosity="error", # note INT8-FP16 activation bug https://github.com/ultralytics/ultralytics/issues/15873
|
|
1069
|
-
output_integer_quantized_tflite=self.args.int8,
|
|
1070
|
-
custom_input_op_name_np_data_path=np_data,
|
|
1071
|
-
enable_batchmatmul_unfold=True and not self.args.int8, # fix lower no. of detected objects on GPU delegate
|
|
1072
|
-
output_signaturedefs=True, # fix error with Attention block group convolution
|
|
1073
|
-
disable_group_convolution=self.args.format in {"tfjs", "edgetpu"}, # fix error with group convolution
|
|
1078
|
+
f_onnx = self.export_onnx() # ensure ONNX is available
|
|
1079
|
+
keras_model = onnx2saved_model(
|
|
1080
|
+
f_onnx,
|
|
1081
|
+
f,
|
|
1082
|
+
int8=self.args.int8,
|
|
1083
|
+
images=images,
|
|
1084
|
+
disable_group_convolution=self.args.format in {"tfjs", "edgetpu"},
|
|
1085
|
+
prefix=prefix,
|
|
1074
1086
|
)
|
|
1075
1087
|
YAML.save(f / "metadata.yaml", self.metadata) # add metadata.yaml
|
|
1076
|
-
|
|
1077
|
-
# Remove/rename TFLite models
|
|
1078
|
-
if self.args.int8:
|
|
1079
|
-
tmp_file.unlink(missing_ok=True)
|
|
1080
|
-
for file in f.rglob("*_dynamic_range_quant.tflite"):
|
|
1081
|
-
file.rename(file.with_name(file.stem.replace("_dynamic_range_quant", "_int8") + file.suffix))
|
|
1082
|
-
for file in f.rglob("*_integer_quant_with_int16_act.tflite"):
|
|
1083
|
-
file.unlink() # delete extra fp16 activation TFLite files
|
|
1084
|
-
|
|
1085
1088
|
# Add TFLite metadata
|
|
1086
1089
|
for file in f.rglob("*.tflite"):
|
|
1087
|
-
|
|
1090
|
+
file.unlink() if "quant_with_int16_act.tflite" in str(file) else self._add_tflite_metadata(file)
|
|
1088
1091
|
|
|
1089
1092
|
return str(f), keras_model # or keras_model = tf.saved_model.load(f, tags=None, options=None)
|
|
1090
1093
|
|
|
1091
1094
|
@try_export
|
|
1092
1095
|
def export_pb(self, keras_model, prefix=colorstr("TensorFlow GraphDef:")):
|
|
1093
1096
|
"""Export YOLO model to TensorFlow GraphDef *.pb format https://github.com/leimao/Frozen-Graph-TensorFlow."""
|
|
1094
|
-
import tensorflow as tf # noqa
|
|
1095
|
-
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2 # noqa
|
|
1096
|
-
|
|
1097
|
-
LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...")
|
|
1098
1097
|
f = self.file.with_suffix(".pb")
|
|
1099
|
-
|
|
1100
|
-
m = tf.function(lambda x: keras_model(x)) # full model
|
|
1101
|
-
m = m.get_concrete_function(tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype))
|
|
1102
|
-
frozen_func = convert_variables_to_constants_v2(m)
|
|
1103
|
-
frozen_func.graph.as_graph_def()
|
|
1104
|
-
tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=str(f.parent), name=f.name, as_text=False)
|
|
1098
|
+
keras2pb(keras_model, f, prefix)
|
|
1105
1099
|
return f
|
|
1106
1100
|
|
|
1107
1101
|
@try_export
|
|
1108
1102
|
def export_tflite(self, prefix=colorstr("TensorFlow Lite:")):
|
|
1109
1103
|
"""Export YOLO model to TensorFlow Lite format."""
|
|
1110
1104
|
# BUG https://github.com/ultralytics/ultralytics/issues/13436
|
|
1111
|
-
import tensorflow as tf
|
|
1105
|
+
import tensorflow as tf
|
|
1112
1106
|
|
|
1113
1107
|
LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...")
|
|
1114
1108
|
saved_model = Path(str(self.file).replace(self.file.suffix, "_saved_model"))
|
|
@@ -1120,6 +1114,110 @@ class Exporter:
|
|
|
1120
1114
|
f = saved_model / f"{self.file.stem}_float32.tflite"
|
|
1121
1115
|
return str(f)
|
|
1122
1116
|
|
|
1117
|
+
@try_export
|
|
1118
|
+
def export_axelera(self, prefix=colorstr("Axelera:")):
|
|
1119
|
+
"""YOLO Axelera export."""
|
|
1120
|
+
os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
|
|
1121
|
+
try:
|
|
1122
|
+
from axelera import compiler
|
|
1123
|
+
except ImportError:
|
|
1124
|
+
check_apt_requirements(
|
|
1125
|
+
["libllvm14", "libgirepository1.0-dev", "pkg-config", "libcairo2-dev", "build-essential", "cmake"]
|
|
1126
|
+
)
|
|
1127
|
+
|
|
1128
|
+
check_requirements(
|
|
1129
|
+
"axelera-voyager-sdk==1.5.2",
|
|
1130
|
+
cmds="--extra-index-url https://software.axelera.ai/artifactory/axelera-runtime-pypi "
|
|
1131
|
+
"--extra-index-url https://software.axelera.ai/artifactory/axelera-dev-pypi",
|
|
1132
|
+
)
|
|
1133
|
+
|
|
1134
|
+
from axelera import compiler
|
|
1135
|
+
from axelera.compiler import CompilerConfig
|
|
1136
|
+
|
|
1137
|
+
self.args.opset = 17 # hardcode opset for Axelera
|
|
1138
|
+
onnx_path = self.export_onnx()
|
|
1139
|
+
model_name = Path(onnx_path).stem
|
|
1140
|
+
export_path = Path(f"{model_name}_axelera_model")
|
|
1141
|
+
export_path.mkdir(exist_ok=True)
|
|
1142
|
+
|
|
1143
|
+
if "C2PSA" in self.model.__str__(): # YOLO11
|
|
1144
|
+
config = CompilerConfig(
|
|
1145
|
+
quantization_scheme="per_tensor_min_max",
|
|
1146
|
+
ignore_weight_buffers=False,
|
|
1147
|
+
resources_used=0.25,
|
|
1148
|
+
aipu_cores_used=1,
|
|
1149
|
+
multicore_mode="batch",
|
|
1150
|
+
output_axm_format=True,
|
|
1151
|
+
model_name=model_name,
|
|
1152
|
+
)
|
|
1153
|
+
else: # YOLOv8
|
|
1154
|
+
config = CompilerConfig(
|
|
1155
|
+
tiling_depth=6,
|
|
1156
|
+
split_buffer_promotion=True,
|
|
1157
|
+
resources_used=0.25,
|
|
1158
|
+
aipu_cores_used=1,
|
|
1159
|
+
multicore_mode="batch",
|
|
1160
|
+
output_axm_format=True,
|
|
1161
|
+
model_name=model_name,
|
|
1162
|
+
)
|
|
1163
|
+
|
|
1164
|
+
qmodel = compiler.quantize(
|
|
1165
|
+
model=onnx_path,
|
|
1166
|
+
calibration_dataset=self.get_int8_calibration_dataloader(prefix),
|
|
1167
|
+
config=config,
|
|
1168
|
+
transform_fn=self._transform_fn,
|
|
1169
|
+
)
|
|
1170
|
+
|
|
1171
|
+
compiler.compile(model=qmodel, config=config, output_dir=export_path)
|
|
1172
|
+
|
|
1173
|
+
axm_name = f"{model_name}.axm"
|
|
1174
|
+
axm_src = Path(axm_name)
|
|
1175
|
+
axm_dst = export_path / axm_name
|
|
1176
|
+
|
|
1177
|
+
if axm_src.exists():
|
|
1178
|
+
axm_src.replace(axm_dst)
|
|
1179
|
+
|
|
1180
|
+
YAML.save(export_path / "metadata.yaml", self.metadata)
|
|
1181
|
+
|
|
1182
|
+
return export_path
|
|
1183
|
+
|
|
1184
|
+
@try_export
|
|
1185
|
+
def export_executorch(self, prefix=colorstr("ExecuTorch:")):
|
|
1186
|
+
"""Exports a model to ExecuTorch (.pte) format into a dedicated directory and saves the required metadata,
|
|
1187
|
+
following Ultralytics conventions.
|
|
1188
|
+
"""
|
|
1189
|
+
LOGGER.info(f"\n{prefix} starting export with ExecuTorch...")
|
|
1190
|
+
assert TORCH_2_9, f"ExecuTorch export requires torch>=2.9.0 but torch=={TORCH_VERSION} is installed"
|
|
1191
|
+
|
|
1192
|
+
# BUG executorch build on arm64 Docker requires packaging>=22.0 https://github.com/pypa/setuptools/issues/4483
|
|
1193
|
+
if LINUX and ARM64 and IS_DOCKER:
|
|
1194
|
+
check_requirements("packaging>=22.0")
|
|
1195
|
+
|
|
1196
|
+
check_requirements("ruamel.yaml<0.19.0")
|
|
1197
|
+
check_requirements("executorch==1.0.1", "flatbuffers")
|
|
1198
|
+
# Pin numpy to avoid coremltools errors with numpy>=2.4.0, must be separate
|
|
1199
|
+
check_requirements("numpy<=2.3.5")
|
|
1200
|
+
|
|
1201
|
+
from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
|
|
1202
|
+
from executorch.exir import to_edge_transform_and_lower
|
|
1203
|
+
|
|
1204
|
+
file_directory = Path(str(self.file).replace(self.file.suffix, "_executorch_model"))
|
|
1205
|
+
file_directory.mkdir(parents=True, exist_ok=True)
|
|
1206
|
+
|
|
1207
|
+
file_pte = file_directory / self.file.with_suffix(".pte").name
|
|
1208
|
+
sample_inputs = (self.im,)
|
|
1209
|
+
|
|
1210
|
+
et_program = to_edge_transform_and_lower(
|
|
1211
|
+
torch.export.export(self.model, sample_inputs), partitioner=[XnnpackPartitioner()]
|
|
1212
|
+
).to_executorch()
|
|
1213
|
+
|
|
1214
|
+
with open(file_pte, "wb") as file:
|
|
1215
|
+
file.write(et_program.buffer)
|
|
1216
|
+
|
|
1217
|
+
YAML.save(file_directory / "metadata.yaml", self.metadata)
|
|
1218
|
+
|
|
1219
|
+
return str(file_directory)
|
|
1220
|
+
|
|
1123
1221
|
@try_export
|
|
1124
1222
|
def export_edgetpu(self, tflite_model="", prefix=colorstr("Edge TPU:")):
|
|
1125
1223
|
"""Export YOLO model to Edge TPU format https://coral.ai/docs/edgetpu/models-intro/."""
|
|
@@ -1128,30 +1226,19 @@ class Exporter:
|
|
|
1128
1226
|
assert LINUX, f"export only supported on Linux. See {help_url}"
|
|
1129
1227
|
if subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, shell=True).returncode != 0:
|
|
1130
1228
|
LOGGER.info(f"\n{prefix} export requires Edge TPU compiler. Attempting install from {help_url}")
|
|
1229
|
+
sudo = "sudo " if is_sudo_available() else ""
|
|
1131
1230
|
for c in (
|
|
1132
|
-
"
|
|
1133
|
-
|
|
1134
|
-
"sudo
|
|
1135
|
-
"sudo apt-get update",
|
|
1136
|
-
"sudo apt-get install edgetpu-compiler",
|
|
1231
|
+
f"{sudo}mkdir -p /etc/apt/keyrings",
|
|
1232
|
+
f"curl -fsSL https://packages.cloud.google.com/apt/doc/apt-key.gpg | {sudo}gpg --dearmor -o /etc/apt/keyrings/google.gpg",
|
|
1233
|
+
f'echo "deb [signed-by=/etc/apt/keyrings/google.gpg] https://packages.cloud.google.com/apt coral-edgetpu-stable main" | {sudo}tee /etc/apt/sources.list.d/coral-edgetpu.list',
|
|
1137
1234
|
):
|
|
1138
|
-
subprocess.run(c
|
|
1139
|
-
|
|
1235
|
+
subprocess.run(c, shell=True, check=True)
|
|
1236
|
+
check_apt_requirements(["edgetpu-compiler"])
|
|
1140
1237
|
|
|
1238
|
+
ver = subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().rsplit(maxsplit=1)[-1]
|
|
1141
1239
|
LOGGER.info(f"\n{prefix} starting export with Edge TPU compiler {ver}...")
|
|
1240
|
+
tflite2edgetpu(tflite_file=tflite_model, output_dir=tflite_model.parent, prefix=prefix)
|
|
1142
1241
|
f = str(tflite_model).replace(".tflite", "_edgetpu.tflite") # Edge TPU model
|
|
1143
|
-
|
|
1144
|
-
cmd = (
|
|
1145
|
-
"edgetpu_compiler "
|
|
1146
|
-
f'--out_dir "{Path(f).parent}" '
|
|
1147
|
-
"--show_operations "
|
|
1148
|
-
"--search_delegate "
|
|
1149
|
-
"--delegate_search_step 30 "
|
|
1150
|
-
"--timeout_sec 180 "
|
|
1151
|
-
f'"{tflite_model}"'
|
|
1152
|
-
)
|
|
1153
|
-
LOGGER.info(f"{prefix} running '{cmd}'")
|
|
1154
|
-
subprocess.run(cmd, shell=True)
|
|
1155
1242
|
self._add_tflite_metadata(f)
|
|
1156
1243
|
return f
|
|
1157
1244
|
|
|
@@ -1159,31 +1246,10 @@ class Exporter:
|
|
|
1159
1246
|
def export_tfjs(self, prefix=colorstr("TensorFlow.js:")):
|
|
1160
1247
|
"""Export YOLO model to TensorFlow.js format."""
|
|
1161
1248
|
check_requirements("tensorflowjs")
|
|
1162
|
-
import tensorflow as tf
|
|
1163
|
-
import tensorflowjs as tfjs # noqa
|
|
1164
1249
|
|
|
1165
|
-
LOGGER.info(f"\n{prefix} starting export with tensorflowjs {tfjs.__version__}...")
|
|
1166
1250
|
f = str(self.file).replace(self.file.suffix, "_web_model") # js dir
|
|
1167
1251
|
f_pb = str(self.file.with_suffix(".pb")) # *.pb path
|
|
1168
|
-
|
|
1169
|
-
gd = tf.Graph().as_graph_def() # TF GraphDef
|
|
1170
|
-
with open(f_pb, "rb") as file:
|
|
1171
|
-
gd.ParseFromString(file.read())
|
|
1172
|
-
outputs = ",".join(gd_outputs(gd))
|
|
1173
|
-
LOGGER.info(f"\n{prefix} output node names: {outputs}")
|
|
1174
|
-
|
|
1175
|
-
quantization = "--quantize_float16" if self.args.half else "--quantize_uint8" if self.args.int8 else ""
|
|
1176
|
-
with spaces_in_path(f_pb) as fpb_, spaces_in_path(f) as f_: # exporter can not handle spaces in path
|
|
1177
|
-
cmd = (
|
|
1178
|
-
"tensorflowjs_converter "
|
|
1179
|
-
f'--input_format=tf_frozen_model {quantization} --output_node_names={outputs} "{fpb_}" "{f_}"'
|
|
1180
|
-
)
|
|
1181
|
-
LOGGER.info(f"{prefix} running '{cmd}'")
|
|
1182
|
-
subprocess.run(cmd, shell=True)
|
|
1183
|
-
|
|
1184
|
-
if " " in f:
|
|
1185
|
-
LOGGER.warning(f"{prefix} your model may not work correctly with spaces in path '{f}'.")
|
|
1186
|
-
|
|
1252
|
+
pb2tfjs(pb_file=f_pb, output_dir=f, half=self.args.half, int8=self.args.int8, prefix=prefix)
|
|
1187
1253
|
# Add metadata
|
|
1188
1254
|
YAML.save(Path(f) / "metadata.yaml", self.metadata) # add metadata.yaml
|
|
1189
1255
|
return f
|
|
@@ -1219,16 +1285,24 @@ class Exporter:
|
|
|
1219
1285
|
def export_imx(self, prefix=colorstr("IMX:")):
|
|
1220
1286
|
"""Export YOLO model to IMX format."""
|
|
1221
1287
|
assert LINUX, (
|
|
1222
|
-
"
|
|
1223
|
-
"See https://developer.aitrios.sony-semicon.com/en/
|
|
1288
|
+
"Export only supported on Linux."
|
|
1289
|
+
"See https://developer.aitrios.sony-semicon.com/en/docs/raspberry-pi-ai-camera/imx500-converter?version=3.17.3&progLang="
|
|
1224
1290
|
)
|
|
1291
|
+
assert not ARM64, "IMX export is not supported on ARM64 architectures."
|
|
1292
|
+
assert IS_PYTHON_MINIMUM_3_9, "IMX export is only supported on Python 3.9 or above."
|
|
1293
|
+
|
|
1225
1294
|
if getattr(self.model, "end2end", False):
|
|
1226
1295
|
raise ValueError("IMX export is not supported for end2end models.")
|
|
1227
1296
|
check_requirements(
|
|
1228
|
-
(
|
|
1297
|
+
(
|
|
1298
|
+
"model-compression-toolkit>=2.4.1",
|
|
1299
|
+
"edge-mdt-cl<1.1.0",
|
|
1300
|
+
"edge-mdt-tpc>=1.2.0",
|
|
1301
|
+
"pydantic<=2.11.7",
|
|
1302
|
+
)
|
|
1229
1303
|
)
|
|
1230
|
-
|
|
1231
|
-
check_requirements("
|
|
1304
|
+
|
|
1305
|
+
check_requirements("imx500-converter[pt]>=3.17.3")
|
|
1232
1306
|
|
|
1233
1307
|
# Install Java>=17
|
|
1234
1308
|
try:
|
|
@@ -1237,8 +1311,12 @@ class Exporter:
|
|
|
1237
1311
|
java_version = int(version_match.group(1)) if version_match else 0
|
|
1238
1312
|
assert java_version >= 17, "Java version too old"
|
|
1239
1313
|
except (FileNotFoundError, subprocess.CalledProcessError, AssertionError):
|
|
1240
|
-
|
|
1241
|
-
|
|
1314
|
+
if IS_UBUNTU or IS_DEBIAN_TRIXIE:
|
|
1315
|
+
LOGGER.info(f"\n{prefix} installing Java 21 for Ubuntu...")
|
|
1316
|
+
check_apt_requirements(["openjdk-21-jre"])
|
|
1317
|
+
elif IS_RASPBERRYPI or IS_DEBIAN_BOOKWORM:
|
|
1318
|
+
LOGGER.info(f"\n{prefix} installing Java 17 for Raspberry Pi or Debian ...")
|
|
1319
|
+
check_apt_requirements(["openjdk-17-jre"])
|
|
1242
1320
|
|
|
1243
1321
|
return torch2imx(
|
|
1244
1322
|
self.model,
|
|
@@ -1260,7 +1338,7 @@ class Exporter:
|
|
|
1260
1338
|
|
|
1261
1339
|
def _pipeline_coreml(self, model, weights_dir=None, prefix=colorstr("CoreML Pipeline:")):
|
|
1262
1340
|
"""Create CoreML pipeline with NMS for YOLO detection models."""
|
|
1263
|
-
import coremltools as ct
|
|
1341
|
+
import coremltools as ct
|
|
1264
1342
|
|
|
1265
1343
|
LOGGER.info(f"{prefix} starting pipeline with coremltools {ct.__version__}...")
|
|
1266
1344
|
|
|
@@ -1353,6 +1431,14 @@ class Exporter:
|
|
|
1353
1431
|
LOGGER.info(f"{prefix} pipeline success")
|
|
1354
1432
|
return model
|
|
1355
1433
|
|
|
1434
|
+
@staticmethod
|
|
1435
|
+
def _transform_fn(data_item) -> np.ndarray:
|
|
1436
|
+
"""The transformation function for Axelera/OpenVINO quantization preprocessing."""
|
|
1437
|
+
data_item: torch.Tensor = data_item["img"] if isinstance(data_item, dict) else data_item
|
|
1438
|
+
assert data_item.dtype == torch.uint8, "Input image must be uint8 for the quantization preprocessing"
|
|
1439
|
+
im = data_item.numpy().astype(np.float32) / 255.0 # uint8 to fp16/32 and 0 - 255 to 0.0 - 1.0
|
|
1440
|
+
return im[None] if im.ndim == 3 else im
|
|
1441
|
+
|
|
1356
1442
|
def add_callback(self, event: str, callback):
|
|
1357
1443
|
"""Append the given callback to the specified event."""
|
|
1358
1444
|
self.callbacks[event].append(callback)
|
|
@@ -1367,8 +1453,7 @@ class IOSDetectModel(torch.nn.Module):
|
|
|
1367
1453
|
"""Wrap an Ultralytics YOLO model for Apple iOS CoreML export."""
|
|
1368
1454
|
|
|
1369
1455
|
def __init__(self, model, im, mlprogram=True):
|
|
1370
|
-
"""
|
|
1371
|
-
Initialize the IOSDetectModel class with a YOLO model and example image.
|
|
1456
|
+
"""Initialize the IOSDetectModel class with a YOLO model and example image.
|
|
1372
1457
|
|
|
1373
1458
|
Args:
|
|
1374
1459
|
model (torch.nn.Module): The YOLO model to wrap.
|
|
@@ -1402,8 +1487,7 @@ class NMSModel(torch.nn.Module):
|
|
|
1402
1487
|
"""Model wrapper with embedded NMS for Detect, Segment, Pose and OBB."""
|
|
1403
1488
|
|
|
1404
1489
|
def __init__(self, model, args):
|
|
1405
|
-
"""
|
|
1406
|
-
Initialize the NMSModel.
|
|
1490
|
+
"""Initialize the NMSModel.
|
|
1407
1491
|
|
|
1408
1492
|
Args:
|
|
1409
1493
|
model (torch.nn.Module): The model to wrap with NMS postprocessing.
|
|
@@ -1416,15 +1500,14 @@ class NMSModel(torch.nn.Module):
|
|
|
1416
1500
|
self.is_tf = self.args.format in frozenset({"saved_model", "tflite", "tfjs"})
|
|
1417
1501
|
|
|
1418
1502
|
def forward(self, x):
|
|
1419
|
-
"""
|
|
1420
|
-
Perform inference with NMS post-processing. Supports Detect, Segment, OBB and Pose.
|
|
1503
|
+
"""Perform inference with NMS post-processing. Supports Detect, Segment, OBB and Pose.
|
|
1421
1504
|
|
|
1422
1505
|
Args:
|
|
1423
1506
|
x (torch.Tensor): The preprocessed tensor with shape (N, 3, H, W).
|
|
1424
1507
|
|
|
1425
1508
|
Returns:
|
|
1426
|
-
(torch.Tensor): List of detections, each an (N, max_det, 4 + 2 + extra_shape) Tensor where N is the
|
|
1427
|
-
|
|
1509
|
+
(torch.Tensor): List of detections, each an (N, max_det, 4 + 2 + extra_shape) Tensor where N is the number
|
|
1510
|
+
of detections after NMS.
|
|
1428
1511
|
"""
|
|
1429
1512
|
from functools import partial
|
|
1430
1513
|
|
|
@@ -1455,17 +1538,16 @@ class NMSModel(torch.nn.Module):
|
|
|
1455
1538
|
box, score, cls, extra = box[mask], score[mask], cls[mask], extra[mask]
|
|
1456
1539
|
nmsbox = box.clone()
|
|
1457
1540
|
# `8` is the minimum value experimented to get correct NMS results for obb
|
|
1458
|
-
multiplier = 8 if self.obb else 1
|
|
1541
|
+
multiplier = 8 if self.obb else 1 / max(len(self.model.names), 1)
|
|
1459
1542
|
# Normalize boxes for NMS since large values for class offset causes issue with int8 quantization
|
|
1460
1543
|
if self.args.format == "tflite": # TFLite is already normalized
|
|
1461
1544
|
nmsbox *= multiplier
|
|
1462
1545
|
else:
|
|
1463
|
-
nmsbox = multiplier * nmsbox / torch.tensor(x.shape[2:], **kwargs).max()
|
|
1464
|
-
if not self.args.agnostic_nms: # class-
|
|
1546
|
+
nmsbox = multiplier * (nmsbox / torch.tensor(x.shape[2:], **kwargs).max())
|
|
1547
|
+
if not self.args.agnostic_nms: # class-wise NMS
|
|
1465
1548
|
end = 2 if self.obb else 4
|
|
1466
1549
|
# fully explicit expansion otherwise reshape error
|
|
1467
|
-
|
|
1468
|
-
cls_offset = cls.reshape(-1, 1).expand(nmsbox.shape[0], end)
|
|
1550
|
+
cls_offset = cls.view(cls.shape[0], 1).expand(cls.shape[0], end)
|
|
1469
1551
|
offbox = nmsbox[:, :end] + cls_offset * multiplier
|
|
1470
1552
|
nmsbox = torch.cat((offbox, nmsbox[:, end:]), dim=-1)
|
|
1471
1553
|
nms_fn = (
|