dgenerate-ultralytics-headless 8.3.137__py3-none-any.whl → 8.3.224__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 (215) hide show
  1. {dgenerate_ultralytics_headless-8.3.137.dist-info → dgenerate_ultralytics_headless-8.3.224.dist-info}/METADATA +41 -34
  2. dgenerate_ultralytics_headless-8.3.224.dist-info/RECORD +285 -0
  3. {dgenerate_ultralytics_headless-8.3.137.dist-info → dgenerate_ultralytics_headless-8.3.224.dist-info}/WHEEL +1 -1
  4. tests/__init__.py +7 -6
  5. tests/conftest.py +15 -39
  6. tests/test_cli.py +17 -17
  7. tests/test_cuda.py +17 -8
  8. tests/test_engine.py +36 -10
  9. tests/test_exports.py +98 -37
  10. tests/test_integrations.py +12 -15
  11. tests/test_python.py +126 -82
  12. tests/test_solutions.py +319 -135
  13. ultralytics/__init__.py +27 -9
  14. ultralytics/cfg/__init__.py +83 -87
  15. ultralytics/cfg/datasets/Argoverse.yaml +4 -4
  16. ultralytics/cfg/datasets/DOTAv1.5.yaml +2 -2
  17. ultralytics/cfg/datasets/DOTAv1.yaml +2 -2
  18. ultralytics/cfg/datasets/GlobalWheat2020.yaml +2 -2
  19. ultralytics/cfg/datasets/HomeObjects-3K.yaml +4 -5
  20. ultralytics/cfg/datasets/ImageNet.yaml +3 -3
  21. ultralytics/cfg/datasets/Objects365.yaml +24 -20
  22. ultralytics/cfg/datasets/SKU-110K.yaml +9 -9
  23. ultralytics/cfg/datasets/VOC.yaml +10 -13
  24. ultralytics/cfg/datasets/VisDrone.yaml +43 -33
  25. ultralytics/cfg/datasets/african-wildlife.yaml +5 -5
  26. ultralytics/cfg/datasets/brain-tumor.yaml +4 -5
  27. ultralytics/cfg/datasets/carparts-seg.yaml +5 -5
  28. ultralytics/cfg/datasets/coco-pose.yaml +26 -4
  29. ultralytics/cfg/datasets/coco.yaml +4 -4
  30. ultralytics/cfg/datasets/coco128-seg.yaml +2 -2
  31. ultralytics/cfg/datasets/coco128.yaml +2 -2
  32. ultralytics/cfg/datasets/coco8-grayscale.yaml +103 -0
  33. ultralytics/cfg/datasets/coco8-multispectral.yaml +2 -2
  34. ultralytics/cfg/datasets/coco8-pose.yaml +23 -2
  35. ultralytics/cfg/datasets/coco8-seg.yaml +2 -2
  36. ultralytics/cfg/datasets/coco8.yaml +2 -2
  37. ultralytics/cfg/datasets/construction-ppe.yaml +32 -0
  38. ultralytics/cfg/datasets/crack-seg.yaml +5 -5
  39. ultralytics/cfg/datasets/dog-pose.yaml +32 -4
  40. ultralytics/cfg/datasets/dota8-multispectral.yaml +2 -2
  41. ultralytics/cfg/datasets/dota8.yaml +2 -2
  42. ultralytics/cfg/datasets/hand-keypoints.yaml +29 -4
  43. ultralytics/cfg/datasets/lvis.yaml +9 -9
  44. ultralytics/cfg/datasets/medical-pills.yaml +4 -5
  45. ultralytics/cfg/datasets/open-images-v7.yaml +7 -10
  46. ultralytics/cfg/datasets/package-seg.yaml +5 -5
  47. ultralytics/cfg/datasets/signature.yaml +4 -4
  48. ultralytics/cfg/datasets/tiger-pose.yaml +20 -4
  49. ultralytics/cfg/datasets/xView.yaml +5 -5
  50. ultralytics/cfg/default.yaml +96 -93
  51. ultralytics/cfg/trackers/botsort.yaml +16 -17
  52. ultralytics/cfg/trackers/bytetrack.yaml +9 -11
  53. ultralytics/data/__init__.py +4 -4
  54. ultralytics/data/annotator.py +12 -12
  55. ultralytics/data/augment.py +531 -564
  56. ultralytics/data/base.py +76 -81
  57. ultralytics/data/build.py +206 -42
  58. ultralytics/data/converter.py +179 -78
  59. ultralytics/data/dataset.py +121 -121
  60. ultralytics/data/loaders.py +114 -91
  61. ultralytics/data/split.py +28 -15
  62. ultralytics/data/split_dota.py +67 -48
  63. ultralytics/data/utils.py +110 -89
  64. ultralytics/engine/exporter.py +422 -460
  65. ultralytics/engine/model.py +224 -252
  66. ultralytics/engine/predictor.py +94 -89
  67. ultralytics/engine/results.py +345 -595
  68. ultralytics/engine/trainer.py +231 -134
  69. ultralytics/engine/tuner.py +279 -73
  70. ultralytics/engine/validator.py +53 -46
  71. ultralytics/hub/__init__.py +26 -28
  72. ultralytics/hub/auth.py +30 -16
  73. ultralytics/hub/google/__init__.py +34 -36
  74. ultralytics/hub/session.py +53 -77
  75. ultralytics/hub/utils.py +23 -109
  76. ultralytics/models/__init__.py +1 -1
  77. ultralytics/models/fastsam/__init__.py +1 -1
  78. ultralytics/models/fastsam/model.py +36 -18
  79. ultralytics/models/fastsam/predict.py +33 -44
  80. ultralytics/models/fastsam/utils.py +4 -5
  81. ultralytics/models/fastsam/val.py +12 -14
  82. ultralytics/models/nas/__init__.py +1 -1
  83. ultralytics/models/nas/model.py +16 -20
  84. ultralytics/models/nas/predict.py +12 -14
  85. ultralytics/models/nas/val.py +4 -5
  86. ultralytics/models/rtdetr/__init__.py +1 -1
  87. ultralytics/models/rtdetr/model.py +9 -9
  88. ultralytics/models/rtdetr/predict.py +22 -17
  89. ultralytics/models/rtdetr/train.py +20 -16
  90. ultralytics/models/rtdetr/val.py +79 -59
  91. ultralytics/models/sam/__init__.py +8 -2
  92. ultralytics/models/sam/amg.py +53 -38
  93. ultralytics/models/sam/build.py +29 -31
  94. ultralytics/models/sam/model.py +33 -38
  95. ultralytics/models/sam/modules/blocks.py +159 -182
  96. ultralytics/models/sam/modules/decoders.py +38 -47
  97. ultralytics/models/sam/modules/encoders.py +114 -133
  98. ultralytics/models/sam/modules/memory_attention.py +38 -31
  99. ultralytics/models/sam/modules/sam.py +114 -93
  100. ultralytics/models/sam/modules/tiny_encoder.py +268 -291
  101. ultralytics/models/sam/modules/transformer.py +59 -66
  102. ultralytics/models/sam/modules/utils.py +55 -72
  103. ultralytics/models/sam/predict.py +745 -341
  104. ultralytics/models/utils/loss.py +118 -107
  105. ultralytics/models/utils/ops.py +118 -71
  106. ultralytics/models/yolo/__init__.py +1 -1
  107. ultralytics/models/yolo/classify/predict.py +28 -26
  108. ultralytics/models/yolo/classify/train.py +50 -81
  109. ultralytics/models/yolo/classify/val.py +68 -61
  110. ultralytics/models/yolo/detect/predict.py +12 -15
  111. ultralytics/models/yolo/detect/train.py +56 -46
  112. ultralytics/models/yolo/detect/val.py +279 -223
  113. ultralytics/models/yolo/model.py +167 -86
  114. ultralytics/models/yolo/obb/predict.py +7 -11
  115. ultralytics/models/yolo/obb/train.py +23 -25
  116. ultralytics/models/yolo/obb/val.py +107 -99
  117. ultralytics/models/yolo/pose/__init__.py +1 -1
  118. ultralytics/models/yolo/pose/predict.py +12 -14
  119. ultralytics/models/yolo/pose/train.py +31 -69
  120. ultralytics/models/yolo/pose/val.py +119 -254
  121. ultralytics/models/yolo/segment/predict.py +21 -25
  122. ultralytics/models/yolo/segment/train.py +12 -66
  123. ultralytics/models/yolo/segment/val.py +126 -305
  124. ultralytics/models/yolo/world/train.py +53 -45
  125. ultralytics/models/yolo/world/train_world.py +51 -32
  126. ultralytics/models/yolo/yoloe/__init__.py +7 -7
  127. ultralytics/models/yolo/yoloe/predict.py +30 -37
  128. ultralytics/models/yolo/yoloe/train.py +89 -71
  129. ultralytics/models/yolo/yoloe/train_seg.py +15 -17
  130. ultralytics/models/yolo/yoloe/val.py +56 -41
  131. ultralytics/nn/__init__.py +9 -11
  132. ultralytics/nn/autobackend.py +179 -107
  133. ultralytics/nn/modules/__init__.py +67 -67
  134. ultralytics/nn/modules/activation.py +8 -7
  135. ultralytics/nn/modules/block.py +302 -323
  136. ultralytics/nn/modules/conv.py +61 -104
  137. ultralytics/nn/modules/head.py +488 -186
  138. ultralytics/nn/modules/transformer.py +183 -123
  139. ultralytics/nn/modules/utils.py +15 -20
  140. ultralytics/nn/tasks.py +327 -203
  141. ultralytics/nn/text_model.py +81 -65
  142. ultralytics/py.typed +1 -0
  143. ultralytics/solutions/__init__.py +12 -12
  144. ultralytics/solutions/ai_gym.py +19 -27
  145. ultralytics/solutions/analytics.py +36 -26
  146. ultralytics/solutions/config.py +29 -28
  147. ultralytics/solutions/distance_calculation.py +23 -24
  148. ultralytics/solutions/heatmap.py +17 -19
  149. ultralytics/solutions/instance_segmentation.py +21 -19
  150. ultralytics/solutions/object_blurrer.py +16 -17
  151. ultralytics/solutions/object_counter.py +48 -53
  152. ultralytics/solutions/object_cropper.py +22 -16
  153. ultralytics/solutions/parking_management.py +61 -58
  154. ultralytics/solutions/queue_management.py +19 -19
  155. ultralytics/solutions/region_counter.py +63 -50
  156. ultralytics/solutions/security_alarm.py +22 -25
  157. ultralytics/solutions/similarity_search.py +107 -60
  158. ultralytics/solutions/solutions.py +343 -262
  159. ultralytics/solutions/speed_estimation.py +35 -31
  160. ultralytics/solutions/streamlit_inference.py +104 -40
  161. ultralytics/solutions/templates/similarity-search.html +31 -24
  162. ultralytics/solutions/trackzone.py +24 -24
  163. ultralytics/solutions/vision_eye.py +11 -12
  164. ultralytics/trackers/__init__.py +1 -1
  165. ultralytics/trackers/basetrack.py +18 -27
  166. ultralytics/trackers/bot_sort.py +48 -39
  167. ultralytics/trackers/byte_tracker.py +94 -94
  168. ultralytics/trackers/track.py +7 -16
  169. ultralytics/trackers/utils/gmc.py +37 -69
  170. ultralytics/trackers/utils/kalman_filter.py +68 -76
  171. ultralytics/trackers/utils/matching.py +13 -17
  172. ultralytics/utils/__init__.py +251 -275
  173. ultralytics/utils/autobatch.py +19 -7
  174. ultralytics/utils/autodevice.py +68 -38
  175. ultralytics/utils/benchmarks.py +169 -130
  176. ultralytics/utils/callbacks/base.py +12 -13
  177. ultralytics/utils/callbacks/clearml.py +14 -15
  178. ultralytics/utils/callbacks/comet.py +139 -66
  179. ultralytics/utils/callbacks/dvc.py +19 -27
  180. ultralytics/utils/callbacks/hub.py +8 -6
  181. ultralytics/utils/callbacks/mlflow.py +6 -10
  182. ultralytics/utils/callbacks/neptune.py +11 -19
  183. ultralytics/utils/callbacks/platform.py +73 -0
  184. ultralytics/utils/callbacks/raytune.py +3 -4
  185. ultralytics/utils/callbacks/tensorboard.py +9 -12
  186. ultralytics/utils/callbacks/wb.py +33 -30
  187. ultralytics/utils/checks.py +163 -114
  188. ultralytics/utils/cpu.py +89 -0
  189. ultralytics/utils/dist.py +24 -20
  190. ultralytics/utils/downloads.py +176 -146
  191. ultralytics/utils/errors.py +11 -13
  192. ultralytics/utils/events.py +113 -0
  193. ultralytics/utils/export/__init__.py +7 -0
  194. ultralytics/utils/{export.py → export/engine.py} +81 -63
  195. ultralytics/utils/export/imx.py +294 -0
  196. ultralytics/utils/export/tensorflow.py +217 -0
  197. ultralytics/utils/files.py +33 -36
  198. ultralytics/utils/git.py +137 -0
  199. ultralytics/utils/instance.py +105 -120
  200. ultralytics/utils/logger.py +404 -0
  201. ultralytics/utils/loss.py +99 -61
  202. ultralytics/utils/metrics.py +649 -478
  203. ultralytics/utils/nms.py +337 -0
  204. ultralytics/utils/ops.py +263 -451
  205. ultralytics/utils/patches.py +70 -31
  206. ultralytics/utils/plotting.py +253 -223
  207. ultralytics/utils/tal.py +48 -61
  208. ultralytics/utils/torch_utils.py +244 -251
  209. ultralytics/utils/tqdm.py +438 -0
  210. ultralytics/utils/triton.py +22 -23
  211. ultralytics/utils/tuner.py +11 -10
  212. dgenerate_ultralytics_headless-8.3.137.dist-info/RECORD +0 -272
  213. {dgenerate_ultralytics_headless-8.3.137.dist-info → dgenerate_ultralytics_headless-8.3.224.dist-info}/entry_points.txt +0 -0
  214. {dgenerate_ultralytics_headless-8.3.137.dist-info → dgenerate_ultralytics_headless-8.3.224.dist-info}/licenses/LICENSE +0 -0
  215. {dgenerate_ultralytics_headless-8.3.137.dist-info → dgenerate_ultralytics_headless-8.3.224.dist-info}/top_level.txt +0 -0
tests/conftest.py CHANGED
@@ -3,32 +3,18 @@
3
3
  import shutil
4
4
  from pathlib import Path
5
5
 
6
- from tests import TMP
7
-
8
6
 
9
7
  def pytest_addoption(parser):
10
- """
11
- Add custom command-line options to pytest.
12
-
13
- Args:
14
- parser (pytest.config.Parser): The pytest parser object for adding custom command-line options.
15
-
16
- Returns:
17
- (None)
18
- """
8
+ """Add custom command-line options to pytest."""
19
9
  parser.addoption("--slow", action="store_true", default=False, help="Run slow tests")
20
10
 
21
11
 
22
12
  def pytest_collection_modifyitems(config, items):
23
- """
24
- Modify the list of test items to exclude tests marked as slow if the --slow option is not specified.
13
+ """Modify the list of test items to exclude tests marked as slow if the --slow option is not specified.
25
14
 
26
15
  Args:
27
- config (pytest.config.Config): The pytest configuration object that provides access to command-line options.
16
+ config: The pytest configuration object that provides access to command-line options.
28
17
  items (list): The list of collected pytest item objects to be modified based on the presence of --slow option.
29
-
30
- Returns:
31
- (None): The function modifies the 'items' list in place.
32
18
  """
33
19
  if not config.getoption("--slow"):
34
20
  # Remove the item entirely from the list of test items if it's marked as 'slow'
@@ -36,48 +22,38 @@ def pytest_collection_modifyitems(config, items):
36
22
 
37
23
 
38
24
  def pytest_sessionstart(session):
39
- """
40
- Initialize session configurations for pytest.
25
+ """Initialize session configurations for pytest.
41
26
 
42
27
  This function is automatically called by pytest after the 'Session' object has been created but before performing
43
- test collection. It sets the initial seeds and prepares the temporary directory for the test session.
28
+ test collection. It sets the initial seeds for the test session.
44
29
 
45
30
  Args:
46
- session (pytest.Session): The pytest session object.
47
-
48
- Returns:
49
- (None)
31
+ session: The pytest session object.
50
32
  """
51
33
  from ultralytics.utils.torch_utils import init_seeds
52
34
 
53
35
  init_seeds()
54
- shutil.rmtree(TMP, ignore_errors=True) # delete any existing tests/tmp directory
55
- TMP.mkdir(parents=True, exist_ok=True) # create a new empty directory
56
36
 
57
37
 
58
38
  def pytest_terminal_summary(terminalreporter, exitstatus, config):
59
- """
60
- Cleanup operations after pytest session.
39
+ """Cleanup operations after pytest session.
61
40
 
62
- This function is automatically called by pytest at the end of the entire test session. It removes certain files
63
- and directories used during testing.
41
+ This function is automatically called by pytest at the end of the entire test session. It removes certain files and
42
+ directories used during testing.
64
43
 
65
44
  Args:
66
- terminalreporter (pytest.terminal.TerminalReporter): The terminal reporter object used for terminal output.
45
+ terminalreporter: The terminal reporter object used for terminal output.
67
46
  exitstatus (int): The exit status of the test run.
68
- config (pytest.config.Config): The pytest config object.
69
-
70
- Returns:
71
- (None)
47
+ config: The pytest config object.
72
48
  """
73
49
  from ultralytics.utils import WEIGHTS_DIR
74
50
 
75
51
  # Remove files
76
- models = [path for x in ["*.onnx", "*.torchscript"] for path in WEIGHTS_DIR.rglob(x)]
77
- for file in ["decelera_portrait_min.mov", "bus.jpg", "yolo11n.onnx", "yolo11n.torchscript"] + models:
52
+ models = [path for x in {"*.onnx", "*.torchscript"} for path in WEIGHTS_DIR.rglob(x)]
53
+ for file in ["decelera_portrait_min.mov", "bus.jpg", "yolo11n.onnx", "yolo11n.torchscript", *models]:
78
54
  Path(file).unlink(missing_ok=True)
79
55
 
80
56
  # Remove directories
81
- models = [path for x in ["*.mlpackage", "*_openvino_model"] for path in WEIGHTS_DIR.rglob(x)]
82
- for directory in [WEIGHTS_DIR / "path with spaces", TMP.parents[1] / ".pytest_cache", TMP] + models:
57
+ models = [path for x in {"*.mlpackage", "*_openvino_model"} for path in WEIGHTS_DIR.rglob(x)]
58
+ for directory in [WEIGHTS_DIR / "path with spaces", *models]:
83
59
  shutil.rmtree(directory, ignore_errors=True)
tests/test_cli.py CHANGED
@@ -1,18 +1,14 @@
1
1
  # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
2
2
 
3
3
  import subprocess
4
+ from pathlib import Path
4
5
 
5
6
  import pytest
6
7
  from PIL import Image
7
8
 
8
- from tests import CUDA_DEVICE_COUNT, CUDA_IS_AVAILABLE
9
- from ultralytics.cfg import TASK2DATA, TASK2MODEL, TASKS
9
+ from tests import CUDA_DEVICE_COUNT, CUDA_IS_AVAILABLE, MODELS, TASK_MODEL_DATA
10
10
  from ultralytics.utils import ARM64, ASSETS, LINUX, WEIGHTS_DIR, checks
11
- from ultralytics.utils.torch_utils import TORCH_1_9
12
-
13
- # Constants
14
- TASK_MODEL_DATA = [(task, WEIGHTS_DIR / TASK2MODEL[task], TASK2DATA[task]) for task in TASKS]
15
- MODELS = [WEIGHTS_DIR / TASK2MODEL[task] for task in TASKS]
11
+ from ultralytics.utils.torch_utils import TORCH_1_11, TORCH_2_9, WINDOWS
16
12
 
17
13
 
18
14
  def run(cmd: str) -> None:
@@ -38,13 +34,13 @@ def test_train(task: str, model: str, data: str) -> None:
38
34
  @pytest.mark.parametrize("task,model,data", TASK_MODEL_DATA)
39
35
  def test_val(task: str, model: str, data: str) -> None:
40
36
  """Test YOLO validation process for specified task, model, and data using a shell command."""
41
- run(f"yolo val {task} model={model} data={data} imgsz=32 save_txt save_json")
37
+ run(f"yolo val {task} model={model} data={data} imgsz=32 save_txt save_json visualize")
42
38
 
43
39
 
44
40
  @pytest.mark.parametrize("task,model,data", TASK_MODEL_DATA)
45
41
  def test_predict(task: str, model: str, data: str) -> None:
46
42
  """Test YOLO prediction on provided sample assets for specified task and model."""
47
- run(f"yolo predict model={model} source={ASSETS} imgsz=32 save save_crop save_txt")
43
+ run(f"yolo {task} predict model={model} source={ASSETS} imgsz=32 save save_crop save_txt visualize")
48
44
 
49
45
 
50
46
  @pytest.mark.parametrize("model", MODELS)
@@ -53,15 +49,12 @@ def test_export(model: str) -> None:
53
49
  run(f"yolo export model={model} format=torchscript imgsz=32")
54
50
 
55
51
 
56
- def test_rtdetr(task: str = "detect", model: str = "yolov8n-rtdetr.yaml", data: str = "coco8.yaml") -> None:
52
+ @pytest.mark.skipif(not TORCH_1_11, reason="RTDETR requires torch>=1.11")
53
+ def test_rtdetr(task: str = "detect", model: Path = WEIGHTS_DIR / "rtdetr-l.pt", data: str = "coco8.yaml") -> None:
57
54
  """Test the RTDETR functionality within Ultralytics for detection tasks using specified model and data."""
58
- # Warning: must use imgsz=640 (note also add coma, spaces, fraction=0.25 args to test single-image training)
59
- run(f"yolo train {task} model={model} data={data} --imgsz= 160 epochs =1, cache = disk fraction=0.25") # spaces
55
+ # Add comma, spaces, fraction=0.25 args to test single-image training
60
56
  run(f"yolo predict {task} model={model} source={ASSETS / 'bus.jpg'} imgsz=160 save save_crop save_txt")
61
- if TORCH_1_9:
62
- weights = WEIGHTS_DIR / "rtdetr-l.pt"
63
- run(f"yolo predict {task} model={weights} source={ASSETS / 'bus.jpg'} imgsz=160 save save_crop save_txt")
64
- run(f"yolo train {task} model={weights} epochs=1 imgsz=160 cache=disk data=coco8.yaml")
57
+ run(f"yolo train {task} model={model} data={data} --imgsz= 160 epochs =1, cache = disk fraction=0.25")
65
58
 
66
59
 
67
60
  @pytest.mark.skipif(checks.IS_PYTHON_3_12, reason="MobileSAM with CLIP is not supported in Python 3.12")
@@ -89,7 +82,7 @@ def test_fastsam(
89
82
  everything_results = sam_model(s, device="cpu", retina_masks=True, imgsz=320, conf=0.4, iou=0.9)
90
83
 
91
84
  # Remove small regions
92
- new_masks, _ = Predictor.remove_small_regions(everything_results[0].masks.data, min_area=20)
85
+ _new_masks, _ = Predictor.remove_small_regions(everything_results[0].masks.data, min_area=20)
93
86
 
94
87
  # Run inference with bboxes and points and texts prompt at the same time
95
88
  sam_model(source, bboxes=[439, 437, 524, 709], points=[[200, 200]], labels=[1], texts="a photo of a dog")
@@ -136,3 +129,10 @@ def test_train_gpu(task: str, model: str, data: str) -> None:
136
129
  def test_solutions(solution: str) -> None:
137
130
  """Test yolo solutions command-line modes."""
138
131
  run(f"yolo solutions {solution} verbose=False")
132
+
133
+
134
+ @pytest.mark.skipif(not checks.IS_PYTHON_MINIMUM_3_10 or not TORCH_2_9, reason="Requires Python>=3.10 and Torch>=2.9.0")
135
+ @pytest.mark.skipif(WINDOWS, reason="Skipping test on Windows")
136
+ def test_export_executorch() -> None:
137
+ """Test exporting a YOLO model to ExecuTorch format via CLI."""
138
+ run("yolo export model=yolo11n.pt format=executorch imgsz=32")
tests/test_cuda.py CHANGED
@@ -22,8 +22,12 @@ if CUDA_IS_AVAILABLE:
22
22
  else:
23
23
  gpu_info = GPUInfo()
24
24
  gpu_info.print_status()
25
- idle_gpus = gpu_info.select_idle_gpu(count=2, min_memory_mb=2048)
26
- if idle_gpus:
25
+ autodevice_fraction = __import__("os").environ.get("YOLO_AUTODEVICE_FRACTION_FREE", 0.3)
26
+ if idle_gpus := gpu_info.select_idle_gpu(
27
+ count=2,
28
+ min_memory_fraction=autodevice_fraction,
29
+ min_util_fraction=autodevice_fraction,
30
+ ):
27
31
  DEVICES = idle_gpus
28
32
 
29
33
 
@@ -41,7 +45,6 @@ def test_amp():
41
45
 
42
46
 
43
47
  @pytest.mark.slow
44
- @pytest.mark.skipif(IS_JETSON, reason="Temporary disable ONNX for Jetson")
45
48
  @pytest.mark.skipif(not DEVICES, reason="No CUDA devices available")
46
49
  @pytest.mark.parametrize(
47
50
  "task, dynamic, int8, half, batch, simplify, nms",
@@ -50,7 +53,9 @@ def test_amp():
50
53
  for task, dynamic, int8, half, batch, simplify, nms in product(
51
54
  TASKS, [True, False], [False], [False], [1, 2], [True, False], [True, False]
52
55
  )
53
- if not ((int8 and half) or (task == "classify" and nms) or (task == "obb" and nms and not TORCH_1_13))
56
+ if not (
57
+ (int8 and half) or (task == "classify" and nms) or (task == "obb" and nms and (not TORCH_1_13 or IS_JETSON))
58
+ )
54
59
  ],
55
60
  )
56
61
  def test_export_onnx_matrix(task, dynamic, int8, half, batch, simplify, nms):
@@ -65,13 +70,13 @@ def test_export_onnx_matrix(task, dynamic, int8, half, batch, simplify, nms):
65
70
  simplify=simplify,
66
71
  nms=nms,
67
72
  device=DEVICES[0],
73
+ # opset=20 if nms else None, # fix ONNX Runtime errors with NMS
68
74
  )
69
75
  YOLO(file)([SOURCE] * batch, imgsz=64 if dynamic else 32, device=DEVICES[0]) # exported model inference
70
76
  Path(file).unlink() # cleanup
71
77
 
72
78
 
73
79
  @pytest.mark.slow
74
- @pytest.mark.skipif(True, reason="CUDA export tests disabled pending additional Ultralytics GPU server availability")
75
80
  @pytest.mark.skipif(not DEVICES, reason="No CUDA devices available")
76
81
  @pytest.mark.parametrize(
77
82
  "task, dynamic, int8, half, batch",
@@ -99,7 +104,7 @@ def test_export_engine_matrix(task, dynamic, int8, half, batch):
99
104
  )
100
105
  YOLO(file)([SOURCE] * batch, imgsz=64 if dynamic else 32, device=DEVICES[0]) # exported model inference
101
106
  Path(file).unlink() # cleanup
102
- Path(file).with_suffix(".cache").unlink() if int8 else None # cleanup INT8 cache
107
+ Path(file).with_suffix(".cache").unlink(missing_ok=True) if int8 else None # cleanup INT8 cache
103
108
 
104
109
 
105
110
  @pytest.mark.skipif(not DEVICES, reason="No CUDA devices available")
@@ -108,12 +113,16 @@ def test_train():
108
113
  import os
109
114
 
110
115
  device = tuple(DEVICES) if len(DEVICES) > 1 else DEVICES[0]
111
- results = YOLO(MODEL).train(data="coco8.yaml", imgsz=64, epochs=1, device=device) # requires imgsz>=64
112
116
  # NVIDIA Jetson only has one GPU and therefore skipping checks
113
117
  if not IS_JETSON:
118
+ results = YOLO(MODEL).train(
119
+ data="coco8.yaml", imgsz=64, epochs=1, device=device, batch=15
120
+ ) # requires imgsz>=64
114
121
  visible = eval(os.environ["CUDA_VISIBLE_DEVICES"])
115
122
  assert visible == device, f"Passed GPUs '{device}', but used GPUs '{visible}'"
116
- assert results is (None if len(DEVICES) > 1 else not None) # DDP returns None, single-GPU returns metrics
123
+ assert (
124
+ (results is None) if len(DEVICES) > 1 else (results is not None)
125
+ ) # DDP returns None, single-GPU returns metrics
117
126
 
118
127
 
119
128
  @pytest.mark.slow
tests/test_engine.py CHANGED
@@ -3,6 +3,8 @@
3
3
  import sys
4
4
  from unittest import mock
5
5
 
6
+ import torch
7
+
6
8
  from tests import MODEL
7
9
  from ultralytics import YOLO
8
10
  from ultralytics.cfg import get_cfg
@@ -11,13 +13,13 @@ from ultralytics.models.yolo import classify, detect, segment
11
13
  from ultralytics.utils import ASSETS, DEFAULT_CFG, WEIGHTS_DIR
12
14
 
13
15
 
14
- def test_func(*args): # noqa
16
+ def test_func(*args):
15
17
  """Test function callback for evaluating YOLO model performance metrics."""
16
18
  print("callback test passed")
17
19
 
18
20
 
19
21
  def test_export():
20
- """Tests the model exporting function by adding a callback and asserting its execution."""
22
+ """Test model exporting functionality by adding a callback and verifying its execution."""
21
23
  exporter = Exporter()
22
24
  exporter.add_callback("on_export_start", test_func)
23
25
  assert test_func in exporter.callbacks["on_export_start"], "callback test failed"
@@ -48,11 +50,12 @@ def test_detect():
48
50
  pred = detect.DetectionPredictor(overrides={"imgsz": [64, 64]})
49
51
  pred.add_callback("on_predict_start", test_func)
50
52
  assert test_func in pred.callbacks["on_predict_start"], "callback test failed"
51
- # Confirm there is no issue with sys.argv being empty.
53
+ # Confirm there is no issue with sys.argv being empty
52
54
  with mock.patch.object(sys, "argv", []):
53
55
  result = pred(source=ASSETS, model=MODEL)
54
56
  assert len(result), "predictor test failed"
55
57
 
58
+ # Test resume functionality
56
59
  overrides["resume"] = trainer.last
57
60
  trainer = detect.DetectionTrainer(overrides=overrides)
58
61
  try:
@@ -61,16 +64,23 @@ def test_detect():
61
64
  print(f"Expected exception caught: {e}")
62
65
  return
63
66
 
64
- Exception("Resume test failed!")
67
+ raise Exception("Resume test failed!")
65
68
 
66
69
 
67
70
  def test_segment():
68
- """Tests image segmentation training, validation, and prediction pipelines using YOLO models."""
69
- overrides = {"data": "coco8-seg.yaml", "model": "yolo11n-seg.yaml", "imgsz": 32, "epochs": 1, "save": False}
71
+ """Test image segmentation training, validation, and prediction pipelines using YOLO models."""
72
+ overrides = {
73
+ "data": "coco8-seg.yaml",
74
+ "model": "yolo11n-seg.yaml",
75
+ "imgsz": 32,
76
+ "epochs": 1,
77
+ "save": False,
78
+ "mask_ratio": 1,
79
+ "overlap_mask": False,
80
+ }
70
81
  cfg = get_cfg(DEFAULT_CFG)
71
82
  cfg.data = "coco8-seg.yaml"
72
83
  cfg.imgsz = 32
73
- # YOLO(CFG_SEG).train(**overrides) # works
74
84
 
75
85
  # Trainer
76
86
  trainer = segment.SegmentationTrainer(overrides=overrides)
@@ -91,7 +101,7 @@ def test_segment():
91
101
  result = pred(source=ASSETS, model=WEIGHTS_DIR / "yolo11n-seg.pt")
92
102
  assert len(result), "predictor test failed"
93
103
 
94
- # Test resume
104
+ # Test resume functionality
95
105
  overrides["resume"] = trainer.last
96
106
  trainer = segment.SegmentationTrainer(overrides=overrides)
97
107
  try:
@@ -100,7 +110,7 @@ def test_segment():
100
110
  print(f"Expected exception caught: {e}")
101
111
  return
102
112
 
103
- Exception("Resume test failed!")
113
+ raise Exception("Resume test failed!")
104
114
 
105
115
 
106
116
  def test_classify():
@@ -109,7 +119,6 @@ def test_classify():
109
119
  cfg = get_cfg(DEFAULT_CFG)
110
120
  cfg.data = "imagenet10"
111
121
  cfg.imgsz = 32
112
- # YOLO(CFG_SEG).train(**overrides) # works
113
122
 
114
123
  # Trainer
115
124
  trainer = classify.ClassificationTrainer(overrides=overrides)
@@ -129,3 +138,20 @@ def test_classify():
129
138
  assert test_func in pred.callbacks["on_predict_start"], "callback test failed"
130
139
  result = pred(source=ASSETS, model=trainer.best)
131
140
  assert len(result), "predictor test failed"
141
+
142
+
143
+ def test_nan_recovery():
144
+ """Test NaN loss detection and recovery during training."""
145
+ nan_injected = [False]
146
+
147
+ def inject_nan(trainer):
148
+ """Inject NaN into loss during batch processing to test recovery mechanism."""
149
+ if trainer.epoch == 1 and trainer.tloss is not None and not nan_injected[0]:
150
+ trainer.tloss *= torch.tensor(float("nan"))
151
+ nan_injected[0] = True
152
+
153
+ overrides = {"data": "coco8.yaml", "model": "yolo11n.yaml", "imgsz": 32, "epochs": 3}
154
+ trainer = detect.DetectionTrainer(overrides=overrides)
155
+ trainer.add_callback("on_train_batch_end", inject_nan)
156
+ trainer.train()
157
+ assert nan_injected[0], "NaN injection failed"
tests/test_exports.py CHANGED
@@ -12,19 +12,12 @@ import pytest
12
12
  from tests import MODEL, SOURCE
13
13
  from ultralytics import YOLO
14
14
  from ultralytics.cfg import TASK2DATA, TASK2MODEL, TASKS
15
- from ultralytics.utils import (
16
- ARM64,
17
- IS_RASPBERRYPI,
18
- LINUX,
19
- MACOS,
20
- WINDOWS,
21
- checks,
22
- )
23
- from ultralytics.utils.torch_utils import TORCH_1_9, TORCH_1_13
15
+ from ultralytics.utils import ARM64, IS_RASPBERRYPI, LINUX, MACOS, WINDOWS, checks
16
+ from ultralytics.utils.torch_utils import TORCH_1_11, TORCH_1_13, TORCH_2_1, TORCH_2_9
24
17
 
25
18
 
26
19
  def test_export_torchscript():
27
- """Test YOLO model exporting to TorchScript format for compatibility and correctness."""
20
+ """Test YOLO model export to TorchScript format for compatibility and correctness."""
28
21
  file = YOLO(MODEL).export(format="torchscript", optimize=False, imgsz=32)
29
22
  YOLO(file)(SOURCE, imgsz=32) # exported model inference
30
23
 
@@ -35,15 +28,15 @@ def test_export_onnx():
35
28
  YOLO(file)(SOURCE, imgsz=32) # exported model inference
36
29
 
37
30
 
38
- @pytest.mark.skipif(not TORCH_1_13, reason="OpenVINO requires torch>=1.13")
31
+ @pytest.mark.skipif(not TORCH_2_1, reason="OpenVINO requires torch>=2.1")
39
32
  def test_export_openvino():
40
- """Test YOLO exports to OpenVINO format for model inference compatibility."""
33
+ """Test YOLO export to OpenVINO format for model inference compatibility."""
41
34
  file = YOLO(MODEL).export(format="openvino", imgsz=32)
42
35
  YOLO(file)(SOURCE, imgsz=32) # exported model inference
43
36
 
44
37
 
45
38
  @pytest.mark.slow
46
- @pytest.mark.skipif(not TORCH_1_13, reason="OpenVINO requires torch>=1.13")
39
+ @pytest.mark.skipif(not TORCH_2_1, reason="OpenVINO requires torch>=2.1")
47
40
  @pytest.mark.parametrize(
48
41
  "task, dynamic, int8, half, batch, nms",
49
42
  [ # generate all combinations except for exclusion cases
@@ -55,7 +48,7 @@ def test_export_openvino():
55
48
  ],
56
49
  )
57
50
  def test_export_openvino_matrix(task, dynamic, int8, half, batch, nms):
58
- """Test YOLO model exports to OpenVINO under various configuration matrix conditions."""
51
+ """Test YOLO model export to OpenVINO under various configuration matrix conditions."""
59
52
  file = YOLO(TASK2MODEL[task]).export(
60
53
  format="openvino",
61
54
  imgsz=32,
@@ -71,7 +64,7 @@ def test_export_openvino_matrix(task, dynamic, int8, half, batch, nms):
71
64
  # See https://github.com/ultralytics/ultralytics/actions/runs/8957949304/job/24601616830?pr=10423
72
65
  file = Path(file)
73
66
  file = file.rename(file.with_stem(f"{file.stem}-{uuid.uuid4()}"))
74
- YOLO(file)([SOURCE] * batch, imgsz=64 if dynamic else 32) # exported model inference
67
+ YOLO(file)([SOURCE] * batch, imgsz=64 if dynamic else 32, batch=batch) # exported model inference
75
68
  shutil.rmtree(file, ignore_errors=True) # retry in case of potential lingering multi-threaded file usage errors
76
69
 
77
70
 
@@ -83,11 +76,11 @@ def test_export_openvino_matrix(task, dynamic, int8, half, batch, nms):
83
76
  for task, dynamic, int8, half, batch, simplify, nms in product(
84
77
  TASKS, [True, False], [False], [False], [1, 2], [True, False], [True, False]
85
78
  )
86
- if not ((int8 and half) or (task == "classify" and nms) or (task == "obb" and nms and not TORCH_1_13))
79
+ if not ((int8 and half) or (task == "classify" and nms) or (nms and not TORCH_1_13))
87
80
  ],
88
81
  )
89
82
  def test_export_onnx_matrix(task, dynamic, int8, half, batch, simplify, nms):
90
- """Test YOLO exports to ONNX format with various configurations and parameters."""
83
+ """Test YOLO export to ONNX format with various configurations and parameters."""
91
84
  file = YOLO(TASK2MODEL[task]).export(
92
85
  format="onnx", imgsz=32, dynamic=dynamic, int8=int8, half=half, batch=batch, simplify=simplify, nms=nms
93
86
  )
@@ -100,12 +93,14 @@ def test_export_onnx_matrix(task, dynamic, int8, half, batch, simplify, nms):
100
93
  "task, dynamic, int8, half, batch, nms",
101
94
  [ # generate all combinations except for exclusion cases
102
95
  (task, dynamic, int8, half, batch, nms)
103
- for task, dynamic, int8, half, batch, nms in product(TASKS, [False], [False], [False], [1, 2], [True, False])
96
+ for task, dynamic, int8, half, batch, nms in product(
97
+ TASKS, [False, True], [False], [False, True], [1, 2], [True, False]
98
+ )
104
99
  if not (task == "classify" and nms)
105
100
  ],
106
101
  )
107
102
  def test_export_torchscript_matrix(task, dynamic, int8, half, batch, nms):
108
- """Tests YOLO model exports to TorchScript format under varied configurations."""
103
+ """Test YOLO model export to TorchScript format under varied configurations."""
109
104
  file = YOLO(TASK2MODEL[task]).export(
110
105
  format="torchscript", imgsz=32, dynamic=dynamic, int8=int8, half=half, batch=batch, nms=nms
111
106
  )
@@ -115,18 +110,23 @@ def test_export_torchscript_matrix(task, dynamic, int8, half, batch, nms):
115
110
 
116
111
  @pytest.mark.slow
117
112
  @pytest.mark.skipif(not MACOS, reason="CoreML inference only supported on macOS")
118
- @pytest.mark.skipif(not TORCH_1_9, reason="CoreML>=7.2 not supported with PyTorch<=1.8")
113
+ @pytest.mark.skipif(not TORCH_1_11, reason="CoreML export requires torch>=1.11")
119
114
  @pytest.mark.skipif(checks.IS_PYTHON_3_13, reason="CoreML not supported in Python 3.13")
120
115
  @pytest.mark.parametrize(
121
- "task, dynamic, int8, half, batch",
116
+ "task, dynamic, int8, half, nms, batch",
122
117
  [ # generate all combinations except for exclusion cases
123
- (task, dynamic, int8, half, batch)
124
- for task, dynamic, int8, half, batch in product(TASKS, [False], [True, False], [True, False], [1])
118
+ (task, dynamic, int8, half, nms, batch)
119
+ for task, dynamic, int8, half, nms, batch in product(
120
+ TASKS, [True, False], [True, False], [True, False], [True, False], [1]
121
+ )
125
122
  if not (int8 and half)
123
+ and not (task != "detect" and nms)
124
+ and not (dynamic and nms)
125
+ and not (task == "classify" and dynamic)
126
126
  ],
127
127
  )
128
- def test_export_coreml_matrix(task, dynamic, int8, half, batch):
129
- """Test YOLO exports to CoreML format with various parameter configurations."""
128
+ def test_export_coreml_matrix(task, dynamic, int8, half, nms, batch):
129
+ """Test YOLO export to CoreML format with various parameter configurations."""
130
130
  file = YOLO(TASK2MODEL[task]).export(
131
131
  format="coreml",
132
132
  imgsz=32,
@@ -134,8 +134,9 @@ def test_export_coreml_matrix(task, dynamic, int8, half, batch):
134
134
  int8=int8,
135
135
  half=half,
136
136
  batch=batch,
137
+ nms=nms,
137
138
  )
138
- YOLO(file)([SOURCE] * batch, imgsz=32) # exported model inference at batch=3
139
+ YOLO(file)([SOURCE] * batch, imgsz=32) # exported model inference
139
140
  shutil.rmtree(file) # cleanup
140
141
 
141
142
 
@@ -152,24 +153,24 @@ def test_export_coreml_matrix(task, dynamic, int8, half, batch):
152
153
  for task, dynamic, int8, half, batch, nms in product(
153
154
  TASKS, [False], [True, False], [True, False], [1], [True, False]
154
155
  )
155
- if not ((int8 and half) or (task == "classify" and nms) or (ARM64 and nms))
156
+ if not ((int8 and half) or (task == "classify" and nms) or (ARM64 and nms) or (nms and not TORCH_1_13))
156
157
  ],
157
158
  )
158
159
  def test_export_tflite_matrix(task, dynamic, int8, half, batch, nms):
159
- """Test YOLO exports to TFLite format considering various export configurations."""
160
+ """Test YOLO export to TFLite format considering various export configurations."""
160
161
  file = YOLO(TASK2MODEL[task]).export(
161
162
  format="tflite", imgsz=32, dynamic=dynamic, int8=int8, half=half, batch=batch, nms=nms
162
163
  )
163
- YOLO(file)([SOURCE] * batch, imgsz=32) # exported model inference at batch=3
164
+ YOLO(file)([SOURCE] * batch, imgsz=32) # exported model inference
164
165
  Path(file).unlink() # cleanup
165
166
 
166
167
 
167
- @pytest.mark.skipif(not TORCH_1_9, reason="CoreML>=7.2 not supported with PyTorch<=1.8")
168
+ @pytest.mark.skipif(not TORCH_1_11, reason="CoreML export requires torch>=1.11")
168
169
  @pytest.mark.skipif(WINDOWS, reason="CoreML not supported on Windows") # RuntimeError: BlobWriter not loaded
169
170
  @pytest.mark.skipif(LINUX and ARM64, reason="CoreML not supported on aarch64 Linux")
170
171
  @pytest.mark.skipif(checks.IS_PYTHON_3_13, reason="CoreML not supported in Python 3.13")
171
172
  def test_export_coreml():
172
- """Test YOLO exports to CoreML format and check for errors."""
173
+ """Test YOLO export to CoreML format and check for errors."""
173
174
  # Capture stdout and stderr
174
175
  stdout, stderr = io.StringIO(), io.StringIO()
175
176
  with redirect_stdout(stdout), redirect_stderr(stderr):
@@ -187,7 +188,7 @@ def test_export_coreml():
187
188
  @pytest.mark.skipif(not checks.IS_PYTHON_MINIMUM_3_10, reason="TFLite export requires Python>=3.10")
188
189
  @pytest.mark.skipif(not LINUX, reason="Test disabled as TF suffers from install conflicts on Windows and macOS")
189
190
  def test_export_tflite():
190
- """Test YOLO exports to TFLite format under specific OS and Python version conditions."""
191
+ """Test YOLO export to TFLite format under specific OS and Python version conditions."""
191
192
  model = YOLO(MODEL)
192
193
  file = model.export(format="tflite", imgsz=32)
193
194
  YOLO(file)(SOURCE, imgsz=32)
@@ -196,7 +197,7 @@ def test_export_tflite():
196
197
  @pytest.mark.skipif(True, reason="Test disabled")
197
198
  @pytest.mark.skipif(not LINUX, reason="TF suffers from install conflicts on Windows and macOS")
198
199
  def test_export_pb():
199
- """Test YOLO exports to TensorFlow's Protobuf (*.pb) format."""
200
+ """Test YOLO export to TensorFlow's Protobuf (*.pb) format."""
200
201
  model = YOLO(MODEL)
201
202
  file = model.export(format="pb", imgsz=32)
202
203
  YOLO(file)(SOURCE, imgsz=32)
@@ -204,28 +205,88 @@ def test_export_pb():
204
205
 
205
206
  @pytest.mark.skipif(True, reason="Test disabled as Paddle protobuf and ONNX protobuf requirements conflict.")
206
207
  def test_export_paddle():
207
- """Test YOLO exports to Paddle format, noting protobuf conflicts with ONNX."""
208
+ """Test YOLO export to Paddle format, noting protobuf conflicts with ONNX."""
208
209
  YOLO(MODEL).export(format="paddle", imgsz=32)
209
210
 
210
211
 
211
212
  @pytest.mark.slow
212
213
  def test_export_mnn():
213
- """Test YOLO exports to MNN format (WARNING: MNN test must precede NCNN test or CI error on Windows)."""
214
+ """Test YOLO export to MNN format (WARNING: MNN test must precede NCNN test or CI error on Windows)."""
214
215
  file = YOLO(MODEL).export(format="mnn", imgsz=32)
215
216
  YOLO(file)(SOURCE, imgsz=32) # exported model inference
216
217
 
217
218
 
219
+ @pytest.mark.slow
220
+ @pytest.mark.parametrize(
221
+ "task, int8, half, batch",
222
+ [ # generate all combinations except for exclusion cases
223
+ (task, int8, half, batch)
224
+ for task, int8, half, batch in product(TASKS, [True, False], [True, False], [1, 2])
225
+ if not (int8 and half)
226
+ ],
227
+ )
228
+ def test_export_mnn_matrix(task, int8, half, batch):
229
+ """Test YOLO export to MNN format considering various export configurations."""
230
+ file = YOLO(TASK2MODEL[task]).export(format="mnn", imgsz=32, int8=int8, half=half, batch=batch)
231
+ YOLO(file)([SOURCE] * batch, imgsz=32) # exported model inference
232
+ Path(file).unlink() # cleanup
233
+
234
+
218
235
  @pytest.mark.slow
219
236
  def test_export_ncnn():
220
- """Test YOLO exports to NCNN format."""
237
+ """Test YOLO export to NCNN format."""
221
238
  file = YOLO(MODEL).export(format="ncnn", imgsz=32)
222
239
  YOLO(file)(SOURCE, imgsz=32) # exported model inference
223
240
 
224
241
 
242
+ @pytest.mark.slow
243
+ @pytest.mark.parametrize("task, half, batch", list(product(TASKS, [True, False], [1])))
244
+ def test_export_ncnn_matrix(task, half, batch):
245
+ """Test YOLO export to NCNN format considering various export configurations."""
246
+ file = YOLO(TASK2MODEL[task]).export(format="ncnn", imgsz=32, half=half, batch=batch)
247
+ YOLO(file)([SOURCE] * batch, imgsz=32) # exported model inference
248
+ shutil.rmtree(file, ignore_errors=True) # retry in case of potential lingering multi-threaded file usage errors
249
+
250
+
225
251
  @pytest.mark.skipif(True, reason="Test disabled as keras and tensorflow version conflicts with TFlite export.")
226
252
  @pytest.mark.skipif(not LINUX or MACOS, reason="Skipping test on Windows and Macos")
227
253
  def test_export_imx():
228
- """Test YOLO exports to IMX format."""
254
+ """Test YOLO export to IMX format."""
229
255
  model = YOLO("yolov8n.pt")
230
256
  file = model.export(format="imx", imgsz=32)
231
257
  YOLO(file)(SOURCE, imgsz=32)
258
+
259
+
260
+ @pytest.mark.skipif(not checks.IS_PYTHON_MINIMUM_3_10 or not TORCH_2_9, reason="Requires Python>=3.10 and Torch>=2.9.0")
261
+ @pytest.mark.skipif(WINDOWS, reason="Skipping test on Windows")
262
+ def test_export_executorch():
263
+ """Test YOLO model export to ExecuTorch format."""
264
+ file = YOLO(MODEL).export(format="executorch", imgsz=32)
265
+ assert Path(file).exists(), f"ExecuTorch export failed, directory not found: {file}"
266
+ # Check that .pte file exists in the exported directory
267
+ pte_file = Path(file) / Path(MODEL).with_suffix(".pte").name
268
+ assert pte_file.exists(), f"ExecuTorch .pte file not found: {pte_file}"
269
+ # Check that metadata.yaml exists
270
+ metadata_file = Path(file) / "metadata.yaml"
271
+ assert metadata_file.exists(), f"ExecuTorch metadata.yaml not found: {metadata_file}"
272
+ # Note: Inference testing skipped as ExecuTorch requires special runtime setup
273
+ shutil.rmtree(file, ignore_errors=True) # cleanup
274
+
275
+
276
+ @pytest.mark.slow
277
+ @pytest.mark.skipif(not checks.IS_PYTHON_MINIMUM_3_10 or not TORCH_2_9, reason="Requires Python>=3.10 and Torch>=2.9.0")
278
+ @pytest.mark.skipif(WINDOWS, reason="Skipping test on Windows")
279
+ @pytest.mark.parametrize("task", TASKS)
280
+ def test_export_executorch_matrix(task):
281
+ """Test YOLO export to ExecuTorch format for various task types."""
282
+ file = YOLO(TASK2MODEL[task]).export(format="executorch", imgsz=32)
283
+ assert Path(file).exists(), f"ExecuTorch export failed for task '{task}', directory not found: {file}"
284
+ # Check that .pte file exists in the exported directory
285
+ model_name = Path(TASK2MODEL[task]).with_suffix(".pte").name
286
+ pte_file = Path(file) / model_name
287
+ assert pte_file.exists(), f"ExecuTorch .pte file not found for task '{task}': {pte_file}"
288
+ # Check that metadata.yaml exists
289
+ metadata_file = Path(file) / "metadata.yaml"
290
+ assert metadata_file.exists(), f"ExecuTorch metadata.yaml not found for task '{task}': {metadata_file}"
291
+ # Note: Inference testing skipped as ExecuTorch requires special runtime setup
292
+ shutil.rmtree(file, ignore_errors=True) # cleanup