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.
Files changed (236) hide show
  1. {dgenerate_ultralytics_headless-8.3.214.dist-info → dgenerate_ultralytics_headless-8.3.248.dist-info}/METADATA +13 -14
  2. dgenerate_ultralytics_headless-8.3.248.dist-info/RECORD +298 -0
  3. tests/__init__.py +5 -7
  4. tests/conftest.py +8 -15
  5. tests/test_cli.py +1 -1
  6. tests/test_cuda.py +5 -8
  7. tests/test_engine.py +1 -1
  8. tests/test_exports.py +57 -12
  9. tests/test_integrations.py +4 -4
  10. tests/test_python.py +84 -53
  11. tests/test_solutions.py +160 -151
  12. ultralytics/__init__.py +1 -1
  13. ultralytics/cfg/__init__.py +56 -62
  14. ultralytics/cfg/datasets/Argoverse.yaml +7 -6
  15. ultralytics/cfg/datasets/DOTAv1.5.yaml +1 -1
  16. ultralytics/cfg/datasets/DOTAv1.yaml +1 -1
  17. ultralytics/cfg/datasets/ImageNet.yaml +1 -1
  18. ultralytics/cfg/datasets/VOC.yaml +15 -16
  19. ultralytics/cfg/datasets/african-wildlife.yaml +1 -1
  20. ultralytics/cfg/datasets/coco-pose.yaml +21 -0
  21. ultralytics/cfg/datasets/coco128-seg.yaml +1 -1
  22. ultralytics/cfg/datasets/coco8-pose.yaml +21 -0
  23. ultralytics/cfg/datasets/dog-pose.yaml +28 -0
  24. ultralytics/cfg/datasets/dota8-multispectral.yaml +1 -1
  25. ultralytics/cfg/datasets/dota8.yaml +2 -2
  26. ultralytics/cfg/datasets/hand-keypoints.yaml +26 -2
  27. ultralytics/cfg/datasets/kitti.yaml +27 -0
  28. ultralytics/cfg/datasets/lvis.yaml +5 -5
  29. ultralytics/cfg/datasets/open-images-v7.yaml +1 -1
  30. ultralytics/cfg/datasets/tiger-pose.yaml +16 -0
  31. ultralytics/cfg/datasets/xView.yaml +16 -16
  32. ultralytics/cfg/default.yaml +1 -1
  33. ultralytics/cfg/models/11/yolo11-pose.yaml +1 -1
  34. ultralytics/cfg/models/11/yoloe-11-seg.yaml +2 -2
  35. ultralytics/cfg/models/11/yoloe-11.yaml +2 -2
  36. ultralytics/cfg/models/rt-detr/rtdetr-l.yaml +1 -1
  37. ultralytics/cfg/models/rt-detr/rtdetr-resnet101.yaml +1 -1
  38. ultralytics/cfg/models/rt-detr/rtdetr-resnet50.yaml +1 -1
  39. ultralytics/cfg/models/rt-detr/rtdetr-x.yaml +1 -1
  40. ultralytics/cfg/models/v10/yolov10b.yaml +2 -2
  41. ultralytics/cfg/models/v10/yolov10l.yaml +2 -2
  42. ultralytics/cfg/models/v10/yolov10m.yaml +2 -2
  43. ultralytics/cfg/models/v10/yolov10n.yaml +2 -2
  44. ultralytics/cfg/models/v10/yolov10s.yaml +2 -2
  45. ultralytics/cfg/models/v10/yolov10x.yaml +2 -2
  46. ultralytics/cfg/models/v3/yolov3-tiny.yaml +1 -1
  47. ultralytics/cfg/models/v6/yolov6.yaml +1 -1
  48. ultralytics/cfg/models/v8/yoloe-v8-seg.yaml +9 -6
  49. ultralytics/cfg/models/v8/yoloe-v8.yaml +9 -6
  50. ultralytics/cfg/models/v8/yolov8-cls-resnet101.yaml +1 -1
  51. ultralytics/cfg/models/v8/yolov8-cls-resnet50.yaml +1 -1
  52. ultralytics/cfg/models/v8/yolov8-ghost-p2.yaml +2 -2
  53. ultralytics/cfg/models/v8/yolov8-ghost-p6.yaml +2 -2
  54. ultralytics/cfg/models/v8/yolov8-ghost.yaml +2 -2
  55. ultralytics/cfg/models/v8/yolov8-obb.yaml +1 -1
  56. ultralytics/cfg/models/v8/yolov8-p2.yaml +1 -1
  57. ultralytics/cfg/models/v8/yolov8-pose-p6.yaml +1 -1
  58. ultralytics/cfg/models/v8/yolov8-rtdetr.yaml +1 -1
  59. ultralytics/cfg/models/v8/yolov8-seg-p6.yaml +1 -1
  60. ultralytics/cfg/models/v8/yolov8-world.yaml +1 -1
  61. ultralytics/cfg/models/v8/yolov8-worldv2.yaml +6 -6
  62. ultralytics/cfg/models/v9/yolov9s.yaml +1 -1
  63. ultralytics/data/__init__.py +4 -4
  64. ultralytics/data/annotator.py +3 -4
  65. ultralytics/data/augment.py +285 -475
  66. ultralytics/data/base.py +18 -26
  67. ultralytics/data/build.py +147 -25
  68. ultralytics/data/converter.py +36 -46
  69. ultralytics/data/dataset.py +46 -74
  70. ultralytics/data/loaders.py +42 -49
  71. ultralytics/data/split.py +5 -6
  72. ultralytics/data/split_dota.py +8 -15
  73. ultralytics/data/utils.py +34 -43
  74. ultralytics/engine/exporter.py +319 -237
  75. ultralytics/engine/model.py +148 -188
  76. ultralytics/engine/predictor.py +29 -38
  77. ultralytics/engine/results.py +177 -311
  78. ultralytics/engine/trainer.py +83 -59
  79. ultralytics/engine/tuner.py +23 -34
  80. ultralytics/engine/validator.py +39 -22
  81. ultralytics/hub/__init__.py +16 -19
  82. ultralytics/hub/auth.py +6 -12
  83. ultralytics/hub/google/__init__.py +7 -10
  84. ultralytics/hub/session.py +15 -25
  85. ultralytics/hub/utils.py +5 -8
  86. ultralytics/models/__init__.py +1 -1
  87. ultralytics/models/fastsam/__init__.py +1 -1
  88. ultralytics/models/fastsam/model.py +8 -10
  89. ultralytics/models/fastsam/predict.py +17 -29
  90. ultralytics/models/fastsam/utils.py +1 -2
  91. ultralytics/models/fastsam/val.py +5 -7
  92. ultralytics/models/nas/__init__.py +1 -1
  93. ultralytics/models/nas/model.py +5 -8
  94. ultralytics/models/nas/predict.py +7 -9
  95. ultralytics/models/nas/val.py +1 -2
  96. ultralytics/models/rtdetr/__init__.py +1 -1
  97. ultralytics/models/rtdetr/model.py +5 -8
  98. ultralytics/models/rtdetr/predict.py +15 -19
  99. ultralytics/models/rtdetr/train.py +10 -13
  100. ultralytics/models/rtdetr/val.py +21 -23
  101. ultralytics/models/sam/__init__.py +15 -2
  102. ultralytics/models/sam/amg.py +14 -20
  103. ultralytics/models/sam/build.py +26 -19
  104. ultralytics/models/sam/build_sam3.py +377 -0
  105. ultralytics/models/sam/model.py +29 -32
  106. ultralytics/models/sam/modules/blocks.py +83 -144
  107. ultralytics/models/sam/modules/decoders.py +19 -37
  108. ultralytics/models/sam/modules/encoders.py +44 -101
  109. ultralytics/models/sam/modules/memory_attention.py +16 -30
  110. ultralytics/models/sam/modules/sam.py +200 -73
  111. ultralytics/models/sam/modules/tiny_encoder.py +64 -83
  112. ultralytics/models/sam/modules/transformer.py +18 -28
  113. ultralytics/models/sam/modules/utils.py +174 -50
  114. ultralytics/models/sam/predict.py +2248 -350
  115. ultralytics/models/sam/sam3/__init__.py +3 -0
  116. ultralytics/models/sam/sam3/decoder.py +546 -0
  117. ultralytics/models/sam/sam3/encoder.py +529 -0
  118. ultralytics/models/sam/sam3/geometry_encoders.py +415 -0
  119. ultralytics/models/sam/sam3/maskformer_segmentation.py +286 -0
  120. ultralytics/models/sam/sam3/model_misc.py +199 -0
  121. ultralytics/models/sam/sam3/necks.py +129 -0
  122. ultralytics/models/sam/sam3/sam3_image.py +339 -0
  123. ultralytics/models/sam/sam3/text_encoder_ve.py +307 -0
  124. ultralytics/models/sam/sam3/vitdet.py +547 -0
  125. ultralytics/models/sam/sam3/vl_combiner.py +160 -0
  126. ultralytics/models/utils/loss.py +14 -26
  127. ultralytics/models/utils/ops.py +13 -17
  128. ultralytics/models/yolo/__init__.py +1 -1
  129. ultralytics/models/yolo/classify/predict.py +9 -12
  130. ultralytics/models/yolo/classify/train.py +11 -32
  131. ultralytics/models/yolo/classify/val.py +29 -28
  132. ultralytics/models/yolo/detect/predict.py +7 -10
  133. ultralytics/models/yolo/detect/train.py +11 -20
  134. ultralytics/models/yolo/detect/val.py +70 -58
  135. ultralytics/models/yolo/model.py +36 -53
  136. ultralytics/models/yolo/obb/predict.py +5 -14
  137. ultralytics/models/yolo/obb/train.py +11 -14
  138. ultralytics/models/yolo/obb/val.py +39 -36
  139. ultralytics/models/yolo/pose/__init__.py +1 -1
  140. ultralytics/models/yolo/pose/predict.py +6 -21
  141. ultralytics/models/yolo/pose/train.py +10 -15
  142. ultralytics/models/yolo/pose/val.py +38 -57
  143. ultralytics/models/yolo/segment/predict.py +14 -18
  144. ultralytics/models/yolo/segment/train.py +3 -6
  145. ultralytics/models/yolo/segment/val.py +93 -45
  146. ultralytics/models/yolo/world/train.py +8 -14
  147. ultralytics/models/yolo/world/train_world.py +11 -34
  148. ultralytics/models/yolo/yoloe/__init__.py +7 -7
  149. ultralytics/models/yolo/yoloe/predict.py +16 -23
  150. ultralytics/models/yolo/yoloe/train.py +30 -43
  151. ultralytics/models/yolo/yoloe/train_seg.py +5 -10
  152. ultralytics/models/yolo/yoloe/val.py +15 -20
  153. ultralytics/nn/__init__.py +7 -7
  154. ultralytics/nn/autobackend.py +145 -77
  155. ultralytics/nn/modules/__init__.py +60 -60
  156. ultralytics/nn/modules/activation.py +4 -6
  157. ultralytics/nn/modules/block.py +132 -216
  158. ultralytics/nn/modules/conv.py +52 -97
  159. ultralytics/nn/modules/head.py +50 -103
  160. ultralytics/nn/modules/transformer.py +76 -88
  161. ultralytics/nn/modules/utils.py +16 -21
  162. ultralytics/nn/tasks.py +94 -154
  163. ultralytics/nn/text_model.py +40 -67
  164. ultralytics/solutions/__init__.py +12 -12
  165. ultralytics/solutions/ai_gym.py +11 -17
  166. ultralytics/solutions/analytics.py +15 -16
  167. ultralytics/solutions/config.py +5 -6
  168. ultralytics/solutions/distance_calculation.py +10 -13
  169. ultralytics/solutions/heatmap.py +7 -13
  170. ultralytics/solutions/instance_segmentation.py +5 -8
  171. ultralytics/solutions/object_blurrer.py +7 -10
  172. ultralytics/solutions/object_counter.py +12 -19
  173. ultralytics/solutions/object_cropper.py +8 -14
  174. ultralytics/solutions/parking_management.py +33 -31
  175. ultralytics/solutions/queue_management.py +10 -12
  176. ultralytics/solutions/region_counter.py +9 -12
  177. ultralytics/solutions/security_alarm.py +15 -20
  178. ultralytics/solutions/similarity_search.py +10 -15
  179. ultralytics/solutions/solutions.py +75 -74
  180. ultralytics/solutions/speed_estimation.py +7 -10
  181. ultralytics/solutions/streamlit_inference.py +2 -4
  182. ultralytics/solutions/templates/similarity-search.html +7 -18
  183. ultralytics/solutions/trackzone.py +7 -10
  184. ultralytics/solutions/vision_eye.py +5 -8
  185. ultralytics/trackers/__init__.py +1 -1
  186. ultralytics/trackers/basetrack.py +3 -5
  187. ultralytics/trackers/bot_sort.py +10 -27
  188. ultralytics/trackers/byte_tracker.py +14 -30
  189. ultralytics/trackers/track.py +3 -6
  190. ultralytics/trackers/utils/gmc.py +11 -22
  191. ultralytics/trackers/utils/kalman_filter.py +37 -48
  192. ultralytics/trackers/utils/matching.py +12 -15
  193. ultralytics/utils/__init__.py +116 -116
  194. ultralytics/utils/autobatch.py +2 -4
  195. ultralytics/utils/autodevice.py +17 -18
  196. ultralytics/utils/benchmarks.py +32 -46
  197. ultralytics/utils/callbacks/base.py +8 -10
  198. ultralytics/utils/callbacks/clearml.py +5 -13
  199. ultralytics/utils/callbacks/comet.py +32 -46
  200. ultralytics/utils/callbacks/dvc.py +13 -18
  201. ultralytics/utils/callbacks/mlflow.py +4 -5
  202. ultralytics/utils/callbacks/neptune.py +7 -15
  203. ultralytics/utils/callbacks/platform.py +314 -38
  204. ultralytics/utils/callbacks/raytune.py +3 -4
  205. ultralytics/utils/callbacks/tensorboard.py +23 -31
  206. ultralytics/utils/callbacks/wb.py +10 -13
  207. ultralytics/utils/checks.py +99 -76
  208. ultralytics/utils/cpu.py +3 -8
  209. ultralytics/utils/dist.py +8 -12
  210. ultralytics/utils/downloads.py +20 -30
  211. ultralytics/utils/errors.py +6 -14
  212. ultralytics/utils/events.py +2 -4
  213. ultralytics/utils/export/__init__.py +4 -236
  214. ultralytics/utils/export/engine.py +237 -0
  215. ultralytics/utils/export/imx.py +91 -55
  216. ultralytics/utils/export/tensorflow.py +231 -0
  217. ultralytics/utils/files.py +24 -28
  218. ultralytics/utils/git.py +9 -11
  219. ultralytics/utils/instance.py +30 -51
  220. ultralytics/utils/logger.py +212 -114
  221. ultralytics/utils/loss.py +14 -22
  222. ultralytics/utils/metrics.py +126 -155
  223. ultralytics/utils/nms.py +13 -16
  224. ultralytics/utils/ops.py +107 -165
  225. ultralytics/utils/patches.py +33 -21
  226. ultralytics/utils/plotting.py +72 -80
  227. ultralytics/utils/tal.py +25 -39
  228. ultralytics/utils/torch_utils.py +52 -78
  229. ultralytics/utils/tqdm.py +20 -20
  230. ultralytics/utils/triton.py +13 -19
  231. ultralytics/utils/tuner.py +17 -5
  232. dgenerate_ultralytics_headless-8.3.214.dist-info/RECORD +0 -283
  233. {dgenerate_ultralytics_headless-8.3.214.dist-info → dgenerate_ultralytics_headless-8.3.248.dist-info}/WHEEL +0 -0
  234. {dgenerate_ultralytics_headless-8.3.214.dist-info → dgenerate_ultralytics_headless-8.3.248.dist-info}/entry_points.txt +0 -0
  235. {dgenerate_ultralytics_headless-8.3.214.dist-info → dgenerate_ultralytics_headless-8.3.248.dist-info}/licenses/LICENSE +0 -0
  236. {dgenerate_ultralytics_headless-8.3.214.dist-info → dgenerate_ultralytics_headless-8.3.248.dist-info}/top_level.txt +0 -0
@@ -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.downloads import attempt_download_asset, get_github_assets, safe_download
109
- from ultralytics.utils.export import onnx2engine, torch2imx, torch2onnx
110
- from ultralytics.utils.files import file_size, spaces_in_path
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 TORCH_1_11, TORCH_1_13, TORCH_2_1, TORCH_2_4, select_device
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
- """Return list of exported files/dirs after running callbacks."""
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
- (jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle, mnn, ncnn, imx, rknn) = (
326
- flags # export booleans
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
- dla = self.args.device.rsplit(":", 1)[-1]
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 self.args.device='dla:0' or 'dla:1, but got {self.args.device}."
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("IMX export only supported for detection and pose estimation models.")
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 (onnx or jit) and self.device.type == "cpu":
364
- LOGGER.warning("half=True only compatible with GPU export, i.e. use device=0, setting half=False.")
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 or ncnn: # 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 # return list of exported files/dirs
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 < 300:
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 # noqa
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 isinstance(self.model, SegmentationModel) else ["output0"]
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), transform_fn),
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 # noqa
767
- from x2paddle.convert import pytorch2paddle # noqa
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 # noqa
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
- import ncnn # noqa
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
- name = Path("pnnx.exe" if WINDOWS else "pnnx") # PNNX filename
812
- pnnx = name if name.is_file() else (ROOT / name)
813
- if not pnnx.is_file():
814
- LOGGER.warning(
815
- f"{prefix} PNNX not found. Attempting to download binary file from "
816
- "https://github.com/pnnx/pnnx/.\nNote PNNX Binary file must be placed in current working directory "
817
- f"or in {ROOT}. See PNNX repo for full installation instructions."
818
- )
819
- system = "macos" if MACOS else "windows" if WINDOWS else "linux-aarch64" if ARM64 else "linux"
820
- try:
821
- release, assets = get_github_assets(repo="pnnx/pnnx")
822
- asset = [x for x in assets if f"{system}.zip" in x][0]
823
- assert isinstance(asset, str), "Unable to retrieve PNNX repo assets" # i.e. pnnx-20250930-macos.zip
824
- LOGGER.info(f"{prefix} successfully found latest PNNX asset file {asset}")
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
- LOGGER.info(f"{prefix} running '{' '.join(cmd)}'")
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
- # Remove debug files
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("coremltools>=8.0")
874
- import coremltools as ct # noqa
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 minimum_depoloyment_target >= iOS16 will require setting compute_precision=ct.precision.FLOAT32.
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 # noqa
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 # noqa
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 # noqa
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 # noqa
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
- # Pre-download calibration file to fix https://github.com/PINTO0309/onnx2tf/issues/545
1037
- onnx2tf_file = Path("calibration_image_sample_data_20x128x128x3_float32.npy")
1038
- if not onnx2tf_file.exists():
1039
- attempt_download_asset(f"{onnx2tf_file}.zip", unzip=True, delete=True)
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 "rtdetr" in self.model.model[-1]._get_name().lower():
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
- # Export to TF
1049
- np_data = None
1050
- if self.args.int8:
1051
- tmp_file = f / "tmp_tflite_int8_calibration_images.npy" # int8 calibration images file
1052
- if self.args.data:
1053
- f.mkdir()
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
- f.unlink() if "quant_with_int16_act.tflite" in str(f) else self._add_tflite_metadata(file)
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 # noqa
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
- "curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -",
1133
- 'echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | '
1134
- "sudo tee /etc/apt/sources.list.d/coral-edgetpu.list",
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 if is_sudo_available() else c.replace("sudo ", ""), shell=True, check=True)
1139
- ver = subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().rsplit(maxsplit=1)[-1]
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
- "export only supported on Linux. "
1223
- "See https://developer.aitrios.sony-semicon.com/en/raspberrypi-ai-camera/documentation/imx500-converter"
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
- ("model-compression-toolkit>=2.4.1", "sony-custom-layers>=0.3.0", "edge-mdt-tpc>=1.1.0", "pydantic<=2.11.7")
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
- check_requirements("imx500-converter[pt]>=3.16.1") # Separate requirements for imx500-converter
1231
- check_requirements("mct-quantizers>=1.6.0") # Separate for compatibility with model-compression-toolkit
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
- cmd = (["sudo"] if is_sudo_available() else []) + ["apt", "install", "-y", "openjdk-21-jre"]
1241
- subprocess.run(cmd, check=True)
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 # noqa
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
- number of detections after NMS.
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-specific NMS
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
- # large max_wh causes issues when quantizing
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 = (