dgenerate-ultralytics-headless 8.3.214__py3-none-any.whl → 8.4.7__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 (249) hide show
  1. {dgenerate_ultralytics_headless-8.3.214.dist-info → dgenerate_ultralytics_headless-8.4.7.dist-info}/METADATA +64 -74
  2. dgenerate_ultralytics_headless-8.4.7.dist-info/RECORD +311 -0
  3. {dgenerate_ultralytics_headless-8.3.214.dist-info → dgenerate_ultralytics_headless-8.4.7.dist-info}/WHEEL +1 -1
  4. tests/__init__.py +7 -9
  5. tests/conftest.py +8 -15
  6. tests/test_cli.py +1 -1
  7. tests/test_cuda.py +13 -10
  8. tests/test_engine.py +9 -9
  9. tests/test_exports.py +65 -13
  10. tests/test_integrations.py +13 -13
  11. tests/test_python.py +125 -69
  12. tests/test_solutions.py +161 -152
  13. ultralytics/__init__.py +1 -1
  14. ultralytics/cfg/__init__.py +86 -92
  15. ultralytics/cfg/datasets/Argoverse.yaml +7 -6
  16. ultralytics/cfg/datasets/DOTAv1.5.yaml +1 -1
  17. ultralytics/cfg/datasets/DOTAv1.yaml +1 -1
  18. ultralytics/cfg/datasets/ImageNet.yaml +1 -1
  19. ultralytics/cfg/datasets/TT100K.yaml +346 -0
  20. ultralytics/cfg/datasets/VOC.yaml +15 -16
  21. ultralytics/cfg/datasets/african-wildlife.yaml +1 -1
  22. ultralytics/cfg/datasets/coco-pose.yaml +21 -0
  23. ultralytics/cfg/datasets/coco12-formats.yaml +101 -0
  24. ultralytics/cfg/datasets/coco128-seg.yaml +1 -1
  25. ultralytics/cfg/datasets/coco8-pose.yaml +21 -0
  26. ultralytics/cfg/datasets/dog-pose.yaml +28 -0
  27. ultralytics/cfg/datasets/dota8-multispectral.yaml +1 -1
  28. ultralytics/cfg/datasets/dota8.yaml +2 -2
  29. ultralytics/cfg/datasets/hand-keypoints.yaml +26 -2
  30. ultralytics/cfg/datasets/kitti.yaml +27 -0
  31. ultralytics/cfg/datasets/lvis.yaml +5 -5
  32. ultralytics/cfg/datasets/open-images-v7.yaml +1 -1
  33. ultralytics/cfg/datasets/tiger-pose.yaml +16 -0
  34. ultralytics/cfg/datasets/xView.yaml +16 -16
  35. ultralytics/cfg/default.yaml +4 -2
  36. ultralytics/cfg/models/11/yolo11-pose.yaml +1 -1
  37. ultralytics/cfg/models/11/yoloe-11-seg.yaml +2 -2
  38. ultralytics/cfg/models/11/yoloe-11.yaml +2 -2
  39. ultralytics/cfg/models/26/yolo26-cls.yaml +33 -0
  40. ultralytics/cfg/models/26/yolo26-obb.yaml +52 -0
  41. ultralytics/cfg/models/26/yolo26-p2.yaml +60 -0
  42. ultralytics/cfg/models/26/yolo26-p6.yaml +62 -0
  43. ultralytics/cfg/models/26/yolo26-pose.yaml +53 -0
  44. ultralytics/cfg/models/26/yolo26-seg.yaml +52 -0
  45. ultralytics/cfg/models/26/yolo26.yaml +52 -0
  46. ultralytics/cfg/models/26/yoloe-26-seg.yaml +53 -0
  47. ultralytics/cfg/models/26/yoloe-26.yaml +53 -0
  48. ultralytics/cfg/models/rt-detr/rtdetr-l.yaml +1 -1
  49. ultralytics/cfg/models/rt-detr/rtdetr-resnet101.yaml +1 -1
  50. ultralytics/cfg/models/rt-detr/rtdetr-resnet50.yaml +1 -1
  51. ultralytics/cfg/models/rt-detr/rtdetr-x.yaml +1 -1
  52. ultralytics/cfg/models/v10/yolov10b.yaml +2 -2
  53. ultralytics/cfg/models/v10/yolov10l.yaml +2 -2
  54. ultralytics/cfg/models/v10/yolov10m.yaml +2 -2
  55. ultralytics/cfg/models/v10/yolov10n.yaml +2 -2
  56. ultralytics/cfg/models/v10/yolov10s.yaml +2 -2
  57. ultralytics/cfg/models/v10/yolov10x.yaml +2 -2
  58. ultralytics/cfg/models/v3/yolov3-tiny.yaml +1 -1
  59. ultralytics/cfg/models/v6/yolov6.yaml +1 -1
  60. ultralytics/cfg/models/v8/yoloe-v8-seg.yaml +9 -6
  61. ultralytics/cfg/models/v8/yoloe-v8.yaml +9 -6
  62. ultralytics/cfg/models/v8/yolov8-cls-resnet101.yaml +1 -1
  63. ultralytics/cfg/models/v8/yolov8-cls-resnet50.yaml +1 -1
  64. ultralytics/cfg/models/v8/yolov8-ghost-p2.yaml +2 -2
  65. ultralytics/cfg/models/v8/yolov8-ghost-p6.yaml +2 -2
  66. ultralytics/cfg/models/v8/yolov8-ghost.yaml +2 -2
  67. ultralytics/cfg/models/v8/yolov8-obb.yaml +1 -1
  68. ultralytics/cfg/models/v8/yolov8-p2.yaml +1 -1
  69. ultralytics/cfg/models/v8/yolov8-pose-p6.yaml +1 -1
  70. ultralytics/cfg/models/v8/yolov8-rtdetr.yaml +1 -1
  71. ultralytics/cfg/models/v8/yolov8-seg-p6.yaml +1 -1
  72. ultralytics/cfg/models/v8/yolov8-world.yaml +1 -1
  73. ultralytics/cfg/models/v8/yolov8-worldv2.yaml +6 -6
  74. ultralytics/cfg/models/v9/yolov9s.yaml +1 -1
  75. ultralytics/data/__init__.py +4 -4
  76. ultralytics/data/annotator.py +5 -6
  77. ultralytics/data/augment.py +300 -475
  78. ultralytics/data/base.py +18 -26
  79. ultralytics/data/build.py +147 -25
  80. ultralytics/data/converter.py +108 -87
  81. ultralytics/data/dataset.py +47 -75
  82. ultralytics/data/loaders.py +42 -49
  83. ultralytics/data/split.py +5 -6
  84. ultralytics/data/split_dota.py +8 -15
  85. ultralytics/data/utils.py +36 -45
  86. ultralytics/engine/exporter.py +351 -263
  87. ultralytics/engine/model.py +186 -225
  88. ultralytics/engine/predictor.py +45 -54
  89. ultralytics/engine/results.py +198 -325
  90. ultralytics/engine/trainer.py +165 -106
  91. ultralytics/engine/tuner.py +41 -43
  92. ultralytics/engine/validator.py +55 -38
  93. ultralytics/hub/__init__.py +16 -19
  94. ultralytics/hub/auth.py +6 -12
  95. ultralytics/hub/google/__init__.py +7 -10
  96. ultralytics/hub/session.py +15 -25
  97. ultralytics/hub/utils.py +5 -8
  98. ultralytics/models/__init__.py +1 -1
  99. ultralytics/models/fastsam/__init__.py +1 -1
  100. ultralytics/models/fastsam/model.py +8 -10
  101. ultralytics/models/fastsam/predict.py +18 -30
  102. ultralytics/models/fastsam/utils.py +1 -2
  103. ultralytics/models/fastsam/val.py +5 -7
  104. ultralytics/models/nas/__init__.py +1 -1
  105. ultralytics/models/nas/model.py +5 -8
  106. ultralytics/models/nas/predict.py +7 -9
  107. ultralytics/models/nas/val.py +1 -2
  108. ultralytics/models/rtdetr/__init__.py +1 -1
  109. ultralytics/models/rtdetr/model.py +5 -8
  110. ultralytics/models/rtdetr/predict.py +15 -19
  111. ultralytics/models/rtdetr/train.py +10 -13
  112. ultralytics/models/rtdetr/val.py +21 -23
  113. ultralytics/models/sam/__init__.py +15 -2
  114. ultralytics/models/sam/amg.py +14 -20
  115. ultralytics/models/sam/build.py +26 -19
  116. ultralytics/models/sam/build_sam3.py +377 -0
  117. ultralytics/models/sam/model.py +29 -32
  118. ultralytics/models/sam/modules/blocks.py +83 -144
  119. ultralytics/models/sam/modules/decoders.py +19 -37
  120. ultralytics/models/sam/modules/encoders.py +44 -101
  121. ultralytics/models/sam/modules/memory_attention.py +16 -30
  122. ultralytics/models/sam/modules/sam.py +200 -73
  123. ultralytics/models/sam/modules/tiny_encoder.py +64 -83
  124. ultralytics/models/sam/modules/transformer.py +18 -28
  125. ultralytics/models/sam/modules/utils.py +174 -50
  126. ultralytics/models/sam/predict.py +2248 -350
  127. ultralytics/models/sam/sam3/__init__.py +3 -0
  128. ultralytics/models/sam/sam3/decoder.py +546 -0
  129. ultralytics/models/sam/sam3/encoder.py +529 -0
  130. ultralytics/models/sam/sam3/geometry_encoders.py +415 -0
  131. ultralytics/models/sam/sam3/maskformer_segmentation.py +286 -0
  132. ultralytics/models/sam/sam3/model_misc.py +199 -0
  133. ultralytics/models/sam/sam3/necks.py +129 -0
  134. ultralytics/models/sam/sam3/sam3_image.py +339 -0
  135. ultralytics/models/sam/sam3/text_encoder_ve.py +307 -0
  136. ultralytics/models/sam/sam3/vitdet.py +547 -0
  137. ultralytics/models/sam/sam3/vl_combiner.py +160 -0
  138. ultralytics/models/utils/loss.py +14 -26
  139. ultralytics/models/utils/ops.py +13 -17
  140. ultralytics/models/yolo/__init__.py +1 -1
  141. ultralytics/models/yolo/classify/predict.py +10 -13
  142. ultralytics/models/yolo/classify/train.py +12 -33
  143. ultralytics/models/yolo/classify/val.py +30 -29
  144. ultralytics/models/yolo/detect/predict.py +9 -12
  145. ultralytics/models/yolo/detect/train.py +17 -23
  146. ultralytics/models/yolo/detect/val.py +77 -59
  147. ultralytics/models/yolo/model.py +43 -60
  148. ultralytics/models/yolo/obb/predict.py +7 -16
  149. ultralytics/models/yolo/obb/train.py +14 -17
  150. ultralytics/models/yolo/obb/val.py +40 -37
  151. ultralytics/models/yolo/pose/__init__.py +1 -1
  152. ultralytics/models/yolo/pose/predict.py +7 -22
  153. ultralytics/models/yolo/pose/train.py +13 -16
  154. ultralytics/models/yolo/pose/val.py +39 -58
  155. ultralytics/models/yolo/segment/predict.py +17 -21
  156. ultralytics/models/yolo/segment/train.py +7 -10
  157. ultralytics/models/yolo/segment/val.py +95 -47
  158. ultralytics/models/yolo/world/train.py +8 -14
  159. ultralytics/models/yolo/world/train_world.py +11 -34
  160. ultralytics/models/yolo/yoloe/__init__.py +7 -7
  161. ultralytics/models/yolo/yoloe/predict.py +16 -23
  162. ultralytics/models/yolo/yoloe/train.py +36 -44
  163. ultralytics/models/yolo/yoloe/train_seg.py +11 -11
  164. ultralytics/models/yolo/yoloe/val.py +15 -20
  165. ultralytics/nn/__init__.py +7 -7
  166. ultralytics/nn/autobackend.py +159 -85
  167. ultralytics/nn/modules/__init__.py +68 -60
  168. ultralytics/nn/modules/activation.py +4 -6
  169. ultralytics/nn/modules/block.py +260 -224
  170. ultralytics/nn/modules/conv.py +52 -97
  171. ultralytics/nn/modules/head.py +831 -299
  172. ultralytics/nn/modules/transformer.py +76 -88
  173. ultralytics/nn/modules/utils.py +16 -21
  174. ultralytics/nn/tasks.py +180 -195
  175. ultralytics/nn/text_model.py +45 -69
  176. ultralytics/optim/__init__.py +5 -0
  177. ultralytics/optim/muon.py +338 -0
  178. ultralytics/solutions/__init__.py +12 -12
  179. ultralytics/solutions/ai_gym.py +13 -19
  180. ultralytics/solutions/analytics.py +15 -16
  181. ultralytics/solutions/config.py +6 -7
  182. ultralytics/solutions/distance_calculation.py +10 -13
  183. ultralytics/solutions/heatmap.py +8 -14
  184. ultralytics/solutions/instance_segmentation.py +6 -9
  185. ultralytics/solutions/object_blurrer.py +7 -10
  186. ultralytics/solutions/object_counter.py +12 -19
  187. ultralytics/solutions/object_cropper.py +8 -14
  188. ultralytics/solutions/parking_management.py +34 -32
  189. ultralytics/solutions/queue_management.py +10 -12
  190. ultralytics/solutions/region_counter.py +9 -12
  191. ultralytics/solutions/security_alarm.py +15 -20
  192. ultralytics/solutions/similarity_search.py +10 -15
  193. ultralytics/solutions/solutions.py +77 -76
  194. ultralytics/solutions/speed_estimation.py +7 -10
  195. ultralytics/solutions/streamlit_inference.py +2 -4
  196. ultralytics/solutions/templates/similarity-search.html +7 -18
  197. ultralytics/solutions/trackzone.py +7 -10
  198. ultralytics/solutions/vision_eye.py +5 -8
  199. ultralytics/trackers/__init__.py +1 -1
  200. ultralytics/trackers/basetrack.py +3 -5
  201. ultralytics/trackers/bot_sort.py +10 -27
  202. ultralytics/trackers/byte_tracker.py +21 -37
  203. ultralytics/trackers/track.py +4 -7
  204. ultralytics/trackers/utils/gmc.py +11 -22
  205. ultralytics/trackers/utils/kalman_filter.py +37 -48
  206. ultralytics/trackers/utils/matching.py +12 -15
  207. ultralytics/utils/__init__.py +124 -124
  208. ultralytics/utils/autobatch.py +2 -4
  209. ultralytics/utils/autodevice.py +17 -18
  210. ultralytics/utils/benchmarks.py +57 -71
  211. ultralytics/utils/callbacks/base.py +8 -10
  212. ultralytics/utils/callbacks/clearml.py +5 -13
  213. ultralytics/utils/callbacks/comet.py +32 -46
  214. ultralytics/utils/callbacks/dvc.py +13 -18
  215. ultralytics/utils/callbacks/mlflow.py +4 -5
  216. ultralytics/utils/callbacks/neptune.py +7 -15
  217. ultralytics/utils/callbacks/platform.py +423 -38
  218. ultralytics/utils/callbacks/raytune.py +3 -4
  219. ultralytics/utils/callbacks/tensorboard.py +25 -31
  220. ultralytics/utils/callbacks/wb.py +16 -14
  221. ultralytics/utils/checks.py +127 -85
  222. ultralytics/utils/cpu.py +3 -8
  223. ultralytics/utils/dist.py +9 -12
  224. ultralytics/utils/downloads.py +25 -33
  225. ultralytics/utils/errors.py +6 -14
  226. ultralytics/utils/events.py +2 -4
  227. ultralytics/utils/export/__init__.py +4 -236
  228. ultralytics/utils/export/engine.py +246 -0
  229. ultralytics/utils/export/imx.py +117 -63
  230. ultralytics/utils/export/tensorflow.py +231 -0
  231. ultralytics/utils/files.py +26 -30
  232. ultralytics/utils/git.py +9 -11
  233. ultralytics/utils/instance.py +30 -51
  234. ultralytics/utils/logger.py +212 -114
  235. ultralytics/utils/loss.py +601 -215
  236. ultralytics/utils/metrics.py +128 -156
  237. ultralytics/utils/nms.py +13 -16
  238. ultralytics/utils/ops.py +117 -166
  239. ultralytics/utils/patches.py +75 -21
  240. ultralytics/utils/plotting.py +75 -80
  241. ultralytics/utils/tal.py +125 -59
  242. ultralytics/utils/torch_utils.py +53 -79
  243. ultralytics/utils/tqdm.py +24 -21
  244. ultralytics/utils/triton.py +13 -19
  245. ultralytics/utils/tuner.py +19 -10
  246. dgenerate_ultralytics_headless-8.3.214.dist-info/RECORD +0 -283
  247. {dgenerate_ultralytics_headless-8.3.214.dist-info → dgenerate_ultralytics_headless-8.4.7.dist-info}/entry_points.txt +0 -0
  248. {dgenerate_ultralytics_headless-8.3.214.dist-info → dgenerate_ultralytics_headless-8.4.7.dist-info}/licenses/LICENSE +0 -0
  249. {dgenerate_ultralytics_headless-8.3.214.dist-info → dgenerate_ultralytics_headless-8.4.7.dist-info}/top_level.txt +0 -0
@@ -4,36 +4,38 @@ Export a YOLO PyTorch model to other formats. TensorFlow exports authored by htt
4
4
 
5
5
  Format | `format=argument` | Model
6
6
  --- | --- | ---
7
- PyTorch | - | yolo11n.pt
8
- TorchScript | `torchscript` | yolo11n.torchscript
9
- ONNX | `onnx` | yolo11n.onnx
10
- OpenVINO | `openvino` | yolo11n_openvino_model/
11
- TensorRT | `engine` | yolo11n.engine
12
- CoreML | `coreml` | yolo11n.mlpackage
13
- TensorFlow SavedModel | `saved_model` | yolo11n_saved_model/
14
- TensorFlow GraphDef | `pb` | yolo11n.pb
15
- TensorFlow Lite | `tflite` | yolo11n.tflite
16
- TensorFlow Edge TPU | `edgetpu` | yolo11n_edgetpu.tflite
17
- TensorFlow.js | `tfjs` | yolo11n_web_model/
18
- PaddlePaddle | `paddle` | yolo11n_paddle_model/
19
- MNN | `mnn` | yolo11n.mnn
20
- NCNN | `ncnn` | yolo11n_ncnn_model/
21
- IMX | `imx` | yolo11n_imx_model/
22
- RKNN | `rknn` | yolo11n_rknn_model/
7
+ PyTorch | - | yolo26n.pt
8
+ TorchScript | `torchscript` | yolo26n.torchscript
9
+ ONNX | `onnx` | yolo26n.onnx
10
+ OpenVINO | `openvino` | yolo26n_openvino_model/
11
+ TensorRT | `engine` | yolo26n.engine
12
+ CoreML | `coreml` | yolo26n.mlpackage
13
+ TensorFlow SavedModel | `saved_model` | yolo26n_saved_model/
14
+ TensorFlow GraphDef | `pb` | yolo26n.pb
15
+ TensorFlow Lite | `tflite` | yolo26n.tflite
16
+ TensorFlow Edge TPU | `edgetpu` | yolo26n_edgetpu.tflite
17
+ TensorFlow.js | `tfjs` | yolo26n_web_model/
18
+ PaddlePaddle | `paddle` | yolo26n_paddle_model/
19
+ MNN | `mnn` | yolo26n.mnn
20
+ NCNN | `ncnn` | yolo26n_ncnn_model/
21
+ IMX | `imx` | yolo26n_imx_model/
22
+ RKNN | `rknn` | yolo26n_rknn_model/
23
+ ExecuTorch | `executorch` | yolo26n_executorch_model/
24
+ Axelera | `axelera` | yolo26n_axelera_model/
23
25
 
24
26
  Requirements:
25
27
  $ pip install "ultralytics[export]"
26
28
 
27
29
  Python:
28
30
  from ultralytics import YOLO
29
- model = YOLO('yolo11n.pt')
31
+ model = YOLO('yolo26n.pt')
30
32
  results = model.export(format='onnx')
31
33
 
32
34
  CLI:
33
- $ yolo mode=export model=yolo11n.pt format=onnx
35
+ $ yolo mode=export model=yolo26n.pt format=onnx
34
36
 
35
37
  Inference:
36
- $ yolo predict model=yolo11n.pt # PyTorch
38
+ $ yolo predict model=yolo26n.pt # PyTorch
37
39
  yolo11n.torchscript # TorchScript
38
40
  yolo11n.onnx # ONNX Runtime or OpenCV DNN with dnn=True
39
41
  yolo11n_openvino_model # OpenVINO
@@ -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
@@ -416,6 +463,9 @@ class Exporter:
416
463
  )
417
464
  if tfjs and (ARM64 and LINUX):
418
465
  raise SystemError("TF.js exports are not currently supported on ARM64 Linux")
466
+ if ncnn and hasattr(model.model[-1], "one2one_cv2"):
467
+ del model.model[-1].one2one_cv2 # Disable end2end branch for NCNN export as it does not support topk
468
+ LOGGER.warning("NCNN export does not support end2end models, disabling end2end branch.")
419
469
  # Recommend OpenVINO if export and Intel CPU
420
470
  if SETTINGS.get("openvino_msg"):
421
471
  if is_intel():
@@ -445,6 +495,10 @@ class Exporter:
445
495
  from ultralytics.utils.export.imx import FXModel
446
496
 
447
497
  model = FXModel(model, self.imgsz)
498
+ if tflite or edgetpu:
499
+ from ultralytics.utils.export.tensorflow import tf_wrapper
500
+
501
+ model = tf_wrapper(model)
448
502
  for m in model.modules():
449
503
  if isinstance(m, Classify):
450
504
  m.export = True
@@ -452,8 +506,11 @@ class Exporter:
452
506
  m.dynamic = self.args.dynamic
453
507
  m.export = True
454
508
  m.format = self.args.format
455
- m.max_det = self.args.max_det
509
+ # Clamp max_det to anchor count for small image sizes (required for TensorRT compatibility)
510
+ anchors = sum(int(self.imgsz[0] / s) * int(self.imgsz[1] / s) for s in model.stride.tolist())
511
+ m.max_det = min(self.args.max_det, anchors)
456
512
  m.xyxy = self.args.nms and not coreml
513
+ m.shape = None # reset cached shape for new export input size
457
514
  if hasattr(model, "pe") and hasattr(m, "fuse"): # for YOLOE models
458
515
  m.fuse(model.pe.to(self.device))
459
516
  elif isinstance(m, C2f) and not is_tf_format:
@@ -466,11 +523,6 @@ class Exporter:
466
523
  if self.args.half and (onnx or jit) and self.device.type != "cpu":
467
524
  im, model = im.half(), model.half() # to FP16
468
525
 
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
526
  # Assign
475
527
  self.im = im
476
528
  self.model = model
@@ -502,6 +554,10 @@ class Exporter:
502
554
  self.metadata["dla"] = dla # make sure `AutoBackend` uses correct dla device if it has one
503
555
  if model.task == "pose":
504
556
  self.metadata["kpt_shape"] = model.model[-1].kpt_shape
557
+ if hasattr(model, "kpt_names"):
558
+ self.metadata["kpt_names"] = model.kpt_names
559
+ if getattr(model.model[-1], "end2end", False):
560
+ self.metadata["end2end"] = True
505
561
 
506
562
  LOGGER.info(
507
563
  f"\n{colorstr('PyTorch:')} starting from '{file}' with input shape {tuple(im.shape)} BCHW and "
@@ -514,7 +570,7 @@ class Exporter:
514
570
  f[0] = self.export_torchscript()
515
571
  if engine: # TensorRT required before ONNX
516
572
  f[1] = self.export_engine(dla=dla)
517
- if onnx or ncnn: # ONNX
573
+ if onnx: # ONNX
518
574
  f[2] = self.export_onnx()
519
575
  if xml: # OpenVINO
520
576
  f[3] = self.export_openvino()
@@ -541,6 +597,10 @@ class Exporter:
541
597
  f[13] = self.export_imx()
542
598
  if rknn:
543
599
  f[14] = self.export_rknn()
600
+ if executorch:
601
+ f[15] = self.export_executorch()
602
+ if axelera:
603
+ f[16] = self.export_axelera()
544
604
 
545
605
  # Finish
546
606
  f = [str(x) for x in f if x] # filter out '' and None
@@ -554,18 +614,17 @@ class Exporter:
554
614
  f"work. Use export 'imgsz={max(self.imgsz)}' if val is required."
555
615
  )
556
616
  imgsz = self.imgsz[0] if square else str(self.imgsz)[1:-1].replace(" ", "")
557
- predict_data = f"data={data}" if model.task == "segment" and pb else ""
558
617
  q = "int8" if self.args.int8 else "half" if self.args.half else "" # quantization
559
618
  LOGGER.info(
560
619
  f"\nExport complete ({time.time() - t:.1f}s)"
561
620
  f"\nResults saved to {colorstr('bold', file.parent.resolve())}"
562
- f"\nPredict: yolo predict task={model.task} model={f} imgsz={imgsz} {q} {predict_data}"
621
+ f"\nPredict: yolo predict task={model.task} model={f} imgsz={imgsz} {q}"
563
622
  f"\nValidate: yolo val task={model.task} model={f} imgsz={imgsz} data={data} {q} {s}"
564
623
  f"\nVisualize: https://netron.app"
565
624
  )
566
625
 
567
626
  self.run_callbacks("on_export_end")
568
- return f # return list of exported files/dirs
627
+ return f # path to final export artifact
569
628
 
570
629
  def get_int8_calibration_dataloader(self, prefix=""):
571
630
  """Build and return a dataloader for calibration of INT8 models."""
@@ -586,7 +645,9 @@ class Exporter:
586
645
  f"The calibration dataset ({n} images) must have at least as many images as the batch size "
587
646
  f"('batch={self.args.batch}')."
588
647
  )
589
- elif n < 300:
648
+ elif self.args.format == "axelera" and n < 100:
649
+ LOGGER.warning(f"{prefix} >100 images required for Axelera calibration, found {n} images.")
650
+ elif self.args.format != "axelera" and n < 300:
590
651
  LOGGER.warning(f"{prefix} >300 images recommended for INT8 calibration, found {n} images.")
591
652
  return build_dataloader(dataset, batch=self.args.batch, workers=0, drop_last=True) # required for batch loading
592
653
 
@@ -610,11 +671,11 @@ class Exporter:
610
671
  @try_export
611
672
  def export_onnx(self, prefix=colorstr("ONNX:")):
612
673
  """Export YOLO model to ONNX format."""
613
- requirements = ["onnx>=1.12.0"]
674
+ requirements = ["onnx>=1.12.0,<2.0.0"]
614
675
  if self.args.simplify:
615
676
  requirements += ["onnxslim>=0.1.71", "onnxruntime" + ("-gpu" if torch.cuda.is_available() else "")]
616
677
  check_requirements(requirements)
617
- import onnx # noqa
678
+ import onnx
618
679
 
619
680
  opset = self.args.opset or best_onnx_opset(onnx, cuda="cuda" in self.device.type)
620
681
  LOGGER.info(f"\n{prefix} starting export with onnx {onnx.__version__} opset {opset}...")
@@ -622,7 +683,7 @@ class Exporter:
622
683
  assert TORCH_1_13, f"'nms=True' ONNX export requires torch>=1.13 (found torch=={TORCH_VERSION})"
623
684
 
624
685
  f = str(self.file.with_suffix(".onnx"))
625
- output_names = ["output0", "output1"] if isinstance(self.model, SegmentationModel) else ["output0"]
686
+ output_names = ["output0", "output1"] if self.model.task == "segment" else ["output0"]
626
687
  dynamic = self.args.dynamic
627
688
  if dynamic:
628
689
  dynamic = {"images": {0: "batch", 2: "height", 3: "width"}} # shape(1,3,640,640)
@@ -671,6 +732,16 @@ class Exporter:
671
732
  LOGGER.info(f"{prefix} limiting IR version {model_onnx.ir_version} to 10 for ONNXRuntime compatibility...")
672
733
  model_onnx.ir_version = 10
673
734
 
735
+ # FP16 conversion for CPU export (GPU exports are already FP16 from model.half() during tracing)
736
+ if self.args.half and self.args.format == "onnx" and self.device.type == "cpu":
737
+ try:
738
+ from onnxruntime.transformers import float16
739
+
740
+ LOGGER.info(f"{prefix} converting to FP16...")
741
+ model_onnx = float16.convert_float_to_float16(model_onnx, keep_io_types=True)
742
+ except Exception as e:
743
+ LOGGER.warning(f"{prefix} FP16 conversion failure: {e}")
744
+
674
745
  onnx.save(model_onnx, f)
675
746
  return f
676
747
 
@@ -711,13 +782,6 @@ class Exporter:
711
782
  check_requirements("nncf>=2.14.0")
712
783
  import nncf
713
784
 
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
785
  # Generate calibration data for integer quantization
722
786
  ignored_scope = None
723
787
  if isinstance(self.model.model[-1], Detect):
@@ -729,14 +793,13 @@ class Exporter:
729
793
  f".*{head_module_name}/.*/Sub*",
730
794
  f".*{head_module_name}/.*/Mul*",
731
795
  f".*{head_module_name}/.*/Div*",
732
- f".*{head_module_name}\\.dfl.*",
733
796
  ],
734
797
  types=["Sigmoid"],
735
798
  )
736
799
 
737
800
  quantized_ov_model = nncf.quantize(
738
801
  model=ov_model,
739
- calibration_dataset=nncf.Dataset(self.get_int8_calibration_dataloader(prefix), transform_fn),
802
+ calibration_dataset=nncf.Dataset(self.get_int8_calibration_dataloader(prefix), self._transform_fn),
740
803
  preset=nncf.QuantizationPreset.MIXED,
741
804
  ignored_scope=ignored_scope,
742
805
  )
@@ -755,16 +818,16 @@ class Exporter:
755
818
  assert not IS_JETSON, "Jetson Paddle exports not supported yet"
756
819
  check_requirements(
757
820
  (
758
- "paddlepaddle-gpu"
821
+ "paddlepaddle-gpu>=3.0.0,!=3.3.0" # exclude 3.3.0 https://github.com/PaddlePaddle/Paddle/issues/77340
759
822
  if torch.cuda.is_available()
760
823
  else "paddlepaddle==3.0.0" # pin 3.0.0 for ARM64
761
824
  if ARM64
762
- else "paddlepaddle>=3.0.0",
825
+ else "paddlepaddle>=3.0.0,!=3.3.0", # exclude 3.3.0 https://github.com/PaddlePaddle/Paddle/issues/77340
763
826
  "x2paddle",
764
827
  )
765
828
  )
766
- import x2paddle # noqa
767
- from x2paddle.convert import pytorch2paddle # noqa
829
+ import x2paddle
830
+ from x2paddle.convert import pytorch2paddle
768
831
 
769
832
  LOGGER.info(f"\n{prefix} starting export with X2Paddle {x2paddle.__version__}...")
770
833
  f = str(self.file).replace(self.file.suffix, f"_paddle_model{os.sep}")
@@ -776,10 +839,11 @@ class Exporter:
776
839
  @try_export
777
840
  def export_mnn(self, prefix=colorstr("MNN:")):
778
841
  """Export YOLO model to MNN format using MNN https://github.com/alibaba/MNN."""
842
+ assert TORCH_1_10, "MNN export requires torch>=1.10.0 to avoid segmentation faults"
779
843
  f_onnx = self.export_onnx() # get onnx model first
780
844
 
781
845
  check_requirements("MNN>=2.9.6")
782
- import MNN # noqa
846
+ import MNN
783
847
  from MNN.tools import mnnconvert
784
848
 
785
849
  # Setup and checks
@@ -802,65 +866,31 @@ class Exporter:
802
866
  def export_ncnn(self, prefix=colorstr("NCNN:")):
803
867
  """Export YOLO model to NCNN format using PNNX https://github.com/pnnx/pnnx."""
804
868
  check_requirements("ncnn", cmds="--no-deps") # no deps to avoid installing opencv-python
805
- import ncnn # noqa
869
+ check_requirements("pnnx")
870
+ import ncnn
871
+ import pnnx
806
872
 
807
- LOGGER.info(f"\n{prefix} starting export with NCNN {ncnn.__version__}...")
873
+ LOGGER.info(f"\n{prefix} starting export with NCNN {ncnn.__version__} and PNNX {pnnx.__version__}...")
808
874
  f = Path(str(self.file).replace(self.file.suffix, f"_ncnn_model{os.sep}"))
809
- f_onnx = self.file.with_suffix(".onnx")
810
875
 
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
- ]
876
+ ncnn_args = dict(
877
+ ncnnparam=(f / "model.ncnn.param").as_posix(),
878
+ ncnnbin=(f / "model.ncnn.bin").as_posix(),
879
+ ncnnpy=(f / "model_ncnn.py").as_posix(),
880
+ )
881
+
882
+ pnnx_args = dict(
883
+ ptpath=(f / "model.pt").as_posix(),
884
+ pnnxparam=(f / "model.pnnx.param").as_posix(),
885
+ pnnxbin=(f / "model.pnnx.bin").as_posix(),
886
+ pnnxpy=(f / "model_pnnx.py").as_posix(),
887
+ pnnxonnx=(f / "model.pnnx.onnx").as_posix(),
888
+ )
889
+
857
890
  f.mkdir(exist_ok=True) # make ncnn_model directory
858
- LOGGER.info(f"{prefix} running '{' '.join(cmd)}'")
859
- subprocess.run(cmd, check=True)
891
+ pnnx.export(self.model, inputs=self.im, **ncnn_args, **pnnx_args, fp16=self.args.half, device=self.device.type)
860
892
 
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):
893
+ for f_debug in ("debug.bin", "debug.param", "debug2.bin", "debug2.param", *pnnx_args.values()):
864
894
  Path(f_debug).unlink(missing_ok=True)
865
895
 
866
896
  YAML.save(f / "metadata.yaml", self.metadata) # add metadata.yaml
@@ -870,8 +900,10 @@ class Exporter:
870
900
  def export_coreml(self, prefix=colorstr("CoreML:")):
871
901
  """Export YOLO model to CoreML format."""
872
902
  mlmodel = self.args.format.lower() == "mlmodel" # legacy *.mlmodel export format requested
873
- check_requirements("coremltools>=8.0")
874
- import coremltools as ct # noqa
903
+ check_requirements(
904
+ ["coremltools>=9.0", "numpy>=1.14.5,<=2.3.5"]
905
+ ) # latest numpy 2.4.0rc1 breaks coremltools exports
906
+ import coremltools as ct
875
907
 
876
908
  LOGGER.info(f"\n{prefix} starting export with coremltools {ct.__version__}...")
877
909
  assert not WINDOWS, "CoreML export is not supported on Windows, please run on macOS or Linux."
@@ -897,7 +929,7 @@ class Exporter:
897
929
  model = IOSDetectModel(self.model, self.im, mlprogram=not mlmodel) if self.args.nms else self.model
898
930
  else:
899
931
  if self.args.nms:
900
- LOGGER.warning(f"{prefix} 'nms=True' is only available for Detect models like 'yolo11n.pt'.")
932
+ LOGGER.warning(f"{prefix} 'nms=True' is only available for Detect models like 'yolo26n.pt'.")
901
933
  # TODO CoreML Segment and Pose model pipelining
902
934
  model = self.model
903
935
  ts = torch.jit.trace(model.eval(), self.im, strict=False) # TorchScript model
@@ -917,7 +949,7 @@ class Exporter:
917
949
 
918
950
  # Based on apple's documentation it is better to leave out the minimum_deployment target and let that get set
919
951
  # Internally based on the model conversion and output type.
920
- # Setting minimum_depoloyment_target >= iOS16 will require setting compute_precision=ct.precision.FLOAT32.
952
+ # Setting minimum_deployment_target >= iOS16 will require setting compute_precision=ct.precision.FLOAT32.
921
953
  # iOS16 adds in better support for FP16, but none of the CoreML NMS specifications handle FP16 as input.
922
954
  ct_model = ct.convert(
923
955
  ts,
@@ -967,12 +999,12 @@ class Exporter:
967
999
  f_onnx = self.export_onnx() # run before TRT import https://github.com/ultralytics/ultralytics/issues/7016
968
1000
 
969
1001
  try:
970
- import tensorrt as trt # noqa
1002
+ import tensorrt as trt
971
1003
  except ImportError:
972
1004
  if LINUX:
973
1005
  cuda_version = torch.version.cuda.split(".")[0]
974
1006
  check_requirements(f"tensorrt-cu{cuda_version}>7.0.0,!=10.1.0")
975
- import tensorrt as trt # noqa
1007
+ import tensorrt as trt
976
1008
  check_version(trt.__version__, ">=7.0.0", hard=True)
977
1009
  check_version(trt.__version__, "!=10.1.0", msg="https://github.com/ultralytics/ultralytics/pull/14239")
978
1010
 
@@ -1002,17 +1034,17 @@ class Exporter:
1002
1034
  """Export YOLO model to TensorFlow SavedModel format."""
1003
1035
  cuda = torch.cuda.is_available()
1004
1036
  try:
1005
- import tensorflow as tf # noqa
1037
+ import tensorflow as tf
1006
1038
  except ImportError:
1007
1039
  check_requirements("tensorflow>=2.0.0,<=2.19.0")
1008
- import tensorflow as tf # noqa
1040
+ import tensorflow as tf
1009
1041
  check_requirements(
1010
1042
  (
1011
1043
  "tf_keras<=2.19.0", # required by 'onnx2tf' package
1012
1044
  "sng4onnx>=1.0.1", # required by 'onnx2tf' package
1013
1045
  "onnx_graphsurgeon>=0.3.26", # required by 'onnx2tf' package
1014
1046
  "ai-edge-litert>=1.2.0" + (",<1.4.0" if MACOS else ""), # required by 'onnx2tf' package
1015
- "onnx>=1.12.0",
1047
+ "onnx>=1.12.0,<2.0.0",
1016
1048
  "onnx2tf>=1.26.3",
1017
1049
  "onnxslim>=0.1.71",
1018
1050
  "onnxruntime-gpu" if cuda else "onnxruntime",
@@ -1033,82 +1065,50 @@ class Exporter:
1033
1065
  if f.is_dir():
1034
1066
  shutil.rmtree(f) # delete output folder
1035
1067
 
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)
1068
+ # Export to TF
1069
+ images = None
1070
+ if self.args.int8 and self.args.data:
1071
+ images = [batch["img"] for batch in self.get_int8_calibration_dataloader(prefix)]
1072
+ images = (
1073
+ torch.nn.functional.interpolate(torch.cat(images, 0).float(), size=self.imgsz)
1074
+ .permute(0, 2, 3, 1)
1075
+ .numpy()
1076
+ .astype(np.float32)
1077
+ )
1040
1078
 
1041
1079
  # Export to ONNX
1042
- if "rtdetr" in self.model.model[-1]._get_name().lower():
1080
+ if isinstance(self.model.model[-1], RTDETRDecoder):
1043
1081
  self.args.opset = self.args.opset or 19
1044
1082
  assert 16 <= self.args.opset <= 19, "RTDETR export requires opset>=16;<=19"
1045
1083
  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
1084
+ f_onnx = self.export_onnx() # ensure ONNX is available
1085
+ keras_model = onnx2saved_model(
1086
+ f_onnx,
1087
+ f,
1088
+ int8=self.args.int8,
1089
+ images=images,
1090
+ disable_group_convolution=self.args.format in {"tfjs", "edgetpu"},
1091
+ prefix=prefix,
1074
1092
  )
1075
1093
  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
1094
  # Add TFLite metadata
1086
1095
  for file in f.rglob("*.tflite"):
1087
- f.unlink() if "quant_with_int16_act.tflite" in str(f) else self._add_tflite_metadata(file)
1096
+ file.unlink() if "quant_with_int16_act.tflite" in str(file) else self._add_tflite_metadata(file)
1088
1097
 
1089
1098
  return str(f), keras_model # or keras_model = tf.saved_model.load(f, tags=None, options=None)
1090
1099
 
1091
1100
  @try_export
1092
1101
  def export_pb(self, keras_model, prefix=colorstr("TensorFlow GraphDef:")):
1093
1102
  """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
1103
  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)
1104
+ keras2pb(keras_model, f, prefix)
1105
1105
  return f
1106
1106
 
1107
1107
  @try_export
1108
1108
  def export_tflite(self, prefix=colorstr("TensorFlow Lite:")):
1109
1109
  """Export YOLO model to TensorFlow Lite format."""
1110
1110
  # BUG https://github.com/ultralytics/ultralytics/issues/13436
1111
- import tensorflow as tf # noqa
1111
+ import tensorflow as tf
1112
1112
 
1113
1113
  LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...")
1114
1114
  saved_model = Path(str(self.file).replace(self.file.suffix, "_saved_model"))
@@ -1120,6 +1120,110 @@ class Exporter:
1120
1120
  f = saved_model / f"{self.file.stem}_float32.tflite"
1121
1121
  return str(f)
1122
1122
 
1123
+ @try_export
1124
+ def export_axelera(self, prefix=colorstr("Axelera:")):
1125
+ """YOLO Axelera export."""
1126
+ os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
1127
+ try:
1128
+ from axelera import compiler
1129
+ except ImportError:
1130
+ check_apt_requirements(
1131
+ ["libllvm14", "libgirepository1.0-dev", "pkg-config", "libcairo2-dev", "build-essential", "cmake"]
1132
+ )
1133
+
1134
+ check_requirements(
1135
+ "axelera-voyager-sdk==1.5.2",
1136
+ cmds="--extra-index-url https://software.axelera.ai/artifactory/axelera-runtime-pypi "
1137
+ "--extra-index-url https://software.axelera.ai/artifactory/axelera-dev-pypi",
1138
+ )
1139
+
1140
+ from axelera import compiler
1141
+ from axelera.compiler import CompilerConfig
1142
+
1143
+ self.args.opset = 17 # hardcode opset for Axelera
1144
+ onnx_path = self.export_onnx()
1145
+ model_name = Path(onnx_path).stem
1146
+ export_path = Path(f"{model_name}_axelera_model")
1147
+ export_path.mkdir(exist_ok=True)
1148
+
1149
+ if "C2PSA" in self.model.__str__(): # YOLO11
1150
+ config = CompilerConfig(
1151
+ quantization_scheme="per_tensor_min_max",
1152
+ ignore_weight_buffers=False,
1153
+ resources_used=0.25,
1154
+ aipu_cores_used=1,
1155
+ multicore_mode="batch",
1156
+ output_axm_format=True,
1157
+ model_name=model_name,
1158
+ )
1159
+ else: # YOLOv8
1160
+ config = CompilerConfig(
1161
+ tiling_depth=6,
1162
+ split_buffer_promotion=True,
1163
+ resources_used=0.25,
1164
+ aipu_cores_used=1,
1165
+ multicore_mode="batch",
1166
+ output_axm_format=True,
1167
+ model_name=model_name,
1168
+ )
1169
+
1170
+ qmodel = compiler.quantize(
1171
+ model=onnx_path,
1172
+ calibration_dataset=self.get_int8_calibration_dataloader(prefix),
1173
+ config=config,
1174
+ transform_fn=self._transform_fn,
1175
+ )
1176
+
1177
+ compiler.compile(model=qmodel, config=config, output_dir=export_path)
1178
+
1179
+ axm_name = f"{model_name}.axm"
1180
+ axm_src = Path(axm_name)
1181
+ axm_dst = export_path / axm_name
1182
+
1183
+ if axm_src.exists():
1184
+ axm_src.replace(axm_dst)
1185
+
1186
+ YAML.save(export_path / "metadata.yaml", self.metadata)
1187
+
1188
+ return export_path
1189
+
1190
+ @try_export
1191
+ def export_executorch(self, prefix=colorstr("ExecuTorch:")):
1192
+ """Exports a model to ExecuTorch (.pte) format into a dedicated directory and saves the required metadata,
1193
+ following Ultralytics conventions.
1194
+ """
1195
+ LOGGER.info(f"\n{prefix} starting export with ExecuTorch...")
1196
+ assert TORCH_2_9, f"ExecuTorch export requires torch>=2.9.0 but torch=={TORCH_VERSION} is installed"
1197
+
1198
+ # BUG executorch build on arm64 Docker requires packaging>=22.0 https://github.com/pypa/setuptools/issues/4483
1199
+ if LINUX and ARM64 and IS_DOCKER:
1200
+ check_requirements("packaging>=22.0")
1201
+
1202
+ check_requirements("ruamel.yaml<0.19.0")
1203
+ check_requirements("executorch==1.0.1", "flatbuffers")
1204
+ # Pin numpy to avoid coremltools errors with numpy>=2.4.0, must be separate
1205
+ check_requirements("numpy<=2.3.5")
1206
+
1207
+ from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
1208
+ from executorch.exir import to_edge_transform_and_lower
1209
+
1210
+ file_directory = Path(str(self.file).replace(self.file.suffix, "_executorch_model"))
1211
+ file_directory.mkdir(parents=True, exist_ok=True)
1212
+
1213
+ file_pte = file_directory / self.file.with_suffix(".pte").name
1214
+ sample_inputs = (self.im,)
1215
+
1216
+ et_program = to_edge_transform_and_lower(
1217
+ torch.export.export(self.model, sample_inputs), partitioner=[XnnpackPartitioner()]
1218
+ ).to_executorch()
1219
+
1220
+ with open(file_pte, "wb") as file:
1221
+ file.write(et_program.buffer)
1222
+
1223
+ YAML.save(file_directory / "metadata.yaml", self.metadata)
1224
+
1225
+ return str(file_directory)
1226
+
1123
1227
  @try_export
1124
1228
  def export_edgetpu(self, tflite_model="", prefix=colorstr("Edge TPU:")):
1125
1229
  """Export YOLO model to Edge TPU format https://coral.ai/docs/edgetpu/models-intro/."""
@@ -1128,30 +1232,19 @@ class Exporter:
1128
1232
  assert LINUX, f"export only supported on Linux. See {help_url}"
1129
1233
  if subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, shell=True).returncode != 0:
1130
1234
  LOGGER.info(f"\n{prefix} export requires Edge TPU compiler. Attempting install from {help_url}")
1235
+ sudo = "sudo " if is_sudo_available() else ""
1131
1236
  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",
1237
+ f"{sudo}mkdir -p /etc/apt/keyrings",
1238
+ f"curl -fsSL https://packages.cloud.google.com/apt/doc/apt-key.gpg | {sudo}gpg --dearmor -o /etc/apt/keyrings/google.gpg",
1239
+ 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
1240
  ):
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]
1241
+ subprocess.run(c, shell=True, check=True)
1242
+ check_apt_requirements(["edgetpu-compiler"])
1140
1243
 
1244
+ ver = subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().rsplit(maxsplit=1)[-1]
1141
1245
  LOGGER.info(f"\n{prefix} starting export with Edge TPU compiler {ver}...")
1246
+ tflite2edgetpu(tflite_file=tflite_model, output_dir=tflite_model.parent, prefix=prefix)
1142
1247
  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
1248
  self._add_tflite_metadata(f)
1156
1249
  return f
1157
1250
 
@@ -1159,31 +1252,10 @@ class Exporter:
1159
1252
  def export_tfjs(self, prefix=colorstr("TensorFlow.js:")):
1160
1253
  """Export YOLO model to TensorFlow.js format."""
1161
1254
  check_requirements("tensorflowjs")
1162
- import tensorflow as tf
1163
- import tensorflowjs as tfjs # noqa
1164
1255
 
1165
- LOGGER.info(f"\n{prefix} starting export with tensorflowjs {tfjs.__version__}...")
1166
1256
  f = str(self.file).replace(self.file.suffix, "_web_model") # js dir
1167
1257
  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
-
1258
+ pb2tfjs(pb_file=f_pb, output_dir=f, half=self.args.half, int8=self.args.int8, prefix=prefix)
1187
1259
  # Add metadata
1188
1260
  YAML.save(Path(f) / "metadata.yaml", self.metadata) # add metadata.yaml
1189
1261
  return f
@@ -1219,16 +1291,24 @@ class Exporter:
1219
1291
  def export_imx(self, prefix=colorstr("IMX:")):
1220
1292
  """Export YOLO model to IMX format."""
1221
1293
  assert LINUX, (
1222
- "export only supported on Linux. "
1223
- "See https://developer.aitrios.sony-semicon.com/en/raspberrypi-ai-camera/documentation/imx500-converter"
1294
+ "Export only supported on Linux."
1295
+ "See https://developer.aitrios.sony-semicon.com/en/docs/raspberry-pi-ai-camera/imx500-converter?version=3.17.3&progLang="
1224
1296
  )
1297
+ assert not ARM64, "IMX export is not supported on ARM64 architectures."
1298
+ assert IS_PYTHON_MINIMUM_3_9, "IMX export is only supported on Python 3.9 or above."
1299
+
1225
1300
  if getattr(self.model, "end2end", False):
1226
1301
  raise ValueError("IMX export is not supported for end2end models.")
1227
1302
  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")
1303
+ (
1304
+ "model-compression-toolkit>=2.4.1",
1305
+ "edge-mdt-cl<1.1.0",
1306
+ "edge-mdt-tpc>=1.2.0",
1307
+ "pydantic<=2.11.7",
1308
+ )
1229
1309
  )
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
1310
+
1311
+ check_requirements("imx500-converter[pt]>=3.17.3")
1232
1312
 
1233
1313
  # Install Java>=17
1234
1314
  try:
@@ -1237,8 +1317,12 @@ class Exporter:
1237
1317
  java_version = int(version_match.group(1)) if version_match else 0
1238
1318
  assert java_version >= 17, "Java version too old"
1239
1319
  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)
1320
+ if IS_UBUNTU or IS_DEBIAN_TRIXIE:
1321
+ LOGGER.info(f"\n{prefix} installing Java 21 for Ubuntu...")
1322
+ check_apt_requirements(["openjdk-21-jre"])
1323
+ elif IS_RASPBERRYPI or IS_DEBIAN_BOOKWORM:
1324
+ LOGGER.info(f"\n{prefix} installing Java 17 for Raspberry Pi or Debian ...")
1325
+ check_apt_requirements(["openjdk-17-jre"])
1242
1326
 
1243
1327
  return torch2imx(
1244
1328
  self.model,
@@ -1260,7 +1344,7 @@ class Exporter:
1260
1344
 
1261
1345
  def _pipeline_coreml(self, model, weights_dir=None, prefix=colorstr("CoreML Pipeline:")):
1262
1346
  """Create CoreML pipeline with NMS for YOLO detection models."""
1263
- import coremltools as ct # noqa
1347
+ import coremltools as ct
1264
1348
 
1265
1349
  LOGGER.info(f"{prefix} starting pipeline with coremltools {ct.__version__}...")
1266
1350
 
@@ -1353,6 +1437,14 @@ class Exporter:
1353
1437
  LOGGER.info(f"{prefix} pipeline success")
1354
1438
  return model
1355
1439
 
1440
+ @staticmethod
1441
+ def _transform_fn(data_item) -> np.ndarray:
1442
+ """The transformation function for Axelera/OpenVINO quantization preprocessing."""
1443
+ data_item: torch.Tensor = data_item["img"] if isinstance(data_item, dict) else data_item
1444
+ assert data_item.dtype == torch.uint8, "Input image must be uint8 for the quantization preprocessing"
1445
+ im = data_item.numpy().astype(np.float32) / 255.0 # uint8 to fp16/32 and 0 - 255 to 0.0 - 1.0
1446
+ return im[None] if im.ndim == 3 else im
1447
+
1356
1448
  def add_callback(self, event: str, callback):
1357
1449
  """Append the given callback to the specified event."""
1358
1450
  self.callbacks[event].append(callback)
@@ -1367,8 +1459,7 @@ class IOSDetectModel(torch.nn.Module):
1367
1459
  """Wrap an Ultralytics YOLO model for Apple iOS CoreML export."""
1368
1460
 
1369
1461
  def __init__(self, model, im, mlprogram=True):
1370
- """
1371
- Initialize the IOSDetectModel class with a YOLO model and example image.
1462
+ """Initialize the IOSDetectModel class with a YOLO model and example image.
1372
1463
 
1373
1464
  Args:
1374
1465
  model (torch.nn.Module): The YOLO model to wrap.
@@ -1402,8 +1493,7 @@ class NMSModel(torch.nn.Module):
1402
1493
  """Model wrapper with embedded NMS for Detect, Segment, Pose and OBB."""
1403
1494
 
1404
1495
  def __init__(self, model, args):
1405
- """
1406
- Initialize the NMSModel.
1496
+ """Initialize the NMSModel.
1407
1497
 
1408
1498
  Args:
1409
1499
  model (torch.nn.Module): The model to wrap with NMS postprocessing.
@@ -1416,15 +1506,14 @@ class NMSModel(torch.nn.Module):
1416
1506
  self.is_tf = self.args.format in frozenset({"saved_model", "tflite", "tfjs"})
1417
1507
 
1418
1508
  def forward(self, x):
1419
- """
1420
- Perform inference with NMS post-processing. Supports Detect, Segment, OBB and Pose.
1509
+ """Perform inference with NMS post-processing. Supports Detect, Segment, OBB and Pose.
1421
1510
 
1422
1511
  Args:
1423
1512
  x (torch.Tensor): The preprocessed tensor with shape (N, 3, H, W).
1424
1513
 
1425
1514
  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.
1515
+ (torch.Tensor): List of detections, each an (N, max_det, 4 + 2 + extra_shape) Tensor where N is the number
1516
+ of detections after NMS.
1428
1517
  """
1429
1518
  from functools import partial
1430
1519
 
@@ -1455,17 +1544,16 @@ class NMSModel(torch.nn.Module):
1455
1544
  box, score, cls, extra = box[mask], score[mask], cls[mask], extra[mask]
1456
1545
  nmsbox = box.clone()
1457
1546
  # `8` is the minimum value experimented to get correct NMS results for obb
1458
- multiplier = 8 if self.obb else 1
1547
+ multiplier = 8 if self.obb else 1 / max(len(self.model.names), 1)
1459
1548
  # Normalize boxes for NMS since large values for class offset causes issue with int8 quantization
1460
1549
  if self.args.format == "tflite": # TFLite is already normalized
1461
1550
  nmsbox *= multiplier
1462
1551
  else:
1463
- nmsbox = multiplier * nmsbox / torch.tensor(x.shape[2:], **kwargs).max()
1464
- if not self.args.agnostic_nms: # class-specific NMS
1552
+ nmsbox = multiplier * (nmsbox / torch.tensor(x.shape[2:], **kwargs).max())
1553
+ if not self.args.agnostic_nms: # class-wise NMS
1465
1554
  end = 2 if self.obb else 4
1466
1555
  # 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)
1556
+ cls_offset = cls.view(cls.shape[0], 1).expand(cls.shape[0], end)
1469
1557
  offbox = nmsbox[:, :end] + cls_offset * multiplier
1470
1558
  nmsbox = torch.cat((offbox, nmsbox[:, end:]), dim=-1)
1471
1559
  nms_fn = (