dgenerate-ultralytics-headless 8.3.134__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 (272) hide show
  1. dgenerate_ultralytics_headless-8.3.134.dist-info/METADATA +400 -0
  2. dgenerate_ultralytics_headless-8.3.134.dist-info/RECORD +272 -0
  3. dgenerate_ultralytics_headless-8.3.134.dist-info/WHEEL +5 -0
  4. dgenerate_ultralytics_headless-8.3.134.dist-info/entry_points.txt +3 -0
  5. dgenerate_ultralytics_headless-8.3.134.dist-info/licenses/LICENSE +661 -0
  6. dgenerate_ultralytics_headless-8.3.134.dist-info/top_level.txt +1 -0
  7. tests/__init__.py +22 -0
  8. tests/conftest.py +83 -0
  9. tests/test_cli.py +138 -0
  10. tests/test_cuda.py +215 -0
  11. tests/test_engine.py +131 -0
  12. tests/test_exports.py +236 -0
  13. tests/test_integrations.py +154 -0
  14. tests/test_python.py +694 -0
  15. tests/test_solutions.py +187 -0
  16. ultralytics/__init__.py +30 -0
  17. ultralytics/assets/bus.jpg +0 -0
  18. ultralytics/assets/zidane.jpg +0 -0
  19. ultralytics/cfg/__init__.py +1023 -0
  20. ultralytics/cfg/datasets/Argoverse.yaml +77 -0
  21. ultralytics/cfg/datasets/DOTAv1.5.yaml +37 -0
  22. ultralytics/cfg/datasets/DOTAv1.yaml +36 -0
  23. ultralytics/cfg/datasets/GlobalWheat2020.yaml +68 -0
  24. ultralytics/cfg/datasets/HomeObjects-3K.yaml +33 -0
  25. ultralytics/cfg/datasets/ImageNet.yaml +2025 -0
  26. ultralytics/cfg/datasets/Objects365.yaml +443 -0
  27. ultralytics/cfg/datasets/SKU-110K.yaml +58 -0
  28. ultralytics/cfg/datasets/VOC.yaml +106 -0
  29. ultralytics/cfg/datasets/VisDrone.yaml +77 -0
  30. ultralytics/cfg/datasets/african-wildlife.yaml +25 -0
  31. ultralytics/cfg/datasets/brain-tumor.yaml +23 -0
  32. ultralytics/cfg/datasets/carparts-seg.yaml +44 -0
  33. ultralytics/cfg/datasets/coco-pose.yaml +42 -0
  34. ultralytics/cfg/datasets/coco.yaml +118 -0
  35. ultralytics/cfg/datasets/coco128-seg.yaml +101 -0
  36. ultralytics/cfg/datasets/coco128.yaml +101 -0
  37. ultralytics/cfg/datasets/coco8-multispectral.yaml +104 -0
  38. ultralytics/cfg/datasets/coco8-pose.yaml +26 -0
  39. ultralytics/cfg/datasets/coco8-seg.yaml +101 -0
  40. ultralytics/cfg/datasets/coco8.yaml +101 -0
  41. ultralytics/cfg/datasets/crack-seg.yaml +22 -0
  42. ultralytics/cfg/datasets/dog-pose.yaml +24 -0
  43. ultralytics/cfg/datasets/dota8-multispectral.yaml +38 -0
  44. ultralytics/cfg/datasets/dota8.yaml +35 -0
  45. ultralytics/cfg/datasets/hand-keypoints.yaml +26 -0
  46. ultralytics/cfg/datasets/lvis.yaml +1240 -0
  47. ultralytics/cfg/datasets/medical-pills.yaml +22 -0
  48. ultralytics/cfg/datasets/open-images-v7.yaml +666 -0
  49. ultralytics/cfg/datasets/package-seg.yaml +22 -0
  50. ultralytics/cfg/datasets/signature.yaml +21 -0
  51. ultralytics/cfg/datasets/tiger-pose.yaml +25 -0
  52. ultralytics/cfg/datasets/xView.yaml +155 -0
  53. ultralytics/cfg/default.yaml +127 -0
  54. ultralytics/cfg/models/11/yolo11-cls-resnet18.yaml +17 -0
  55. ultralytics/cfg/models/11/yolo11-cls.yaml +33 -0
  56. ultralytics/cfg/models/11/yolo11-obb.yaml +50 -0
  57. ultralytics/cfg/models/11/yolo11-pose.yaml +51 -0
  58. ultralytics/cfg/models/11/yolo11-seg.yaml +50 -0
  59. ultralytics/cfg/models/11/yolo11.yaml +50 -0
  60. ultralytics/cfg/models/11/yoloe-11-seg.yaml +48 -0
  61. ultralytics/cfg/models/11/yoloe-11.yaml +48 -0
  62. ultralytics/cfg/models/12/yolo12-cls.yaml +32 -0
  63. ultralytics/cfg/models/12/yolo12-obb.yaml +48 -0
  64. ultralytics/cfg/models/12/yolo12-pose.yaml +49 -0
  65. ultralytics/cfg/models/12/yolo12-seg.yaml +48 -0
  66. ultralytics/cfg/models/12/yolo12.yaml +48 -0
  67. ultralytics/cfg/models/rt-detr/rtdetr-l.yaml +53 -0
  68. ultralytics/cfg/models/rt-detr/rtdetr-resnet101.yaml +45 -0
  69. ultralytics/cfg/models/rt-detr/rtdetr-resnet50.yaml +45 -0
  70. ultralytics/cfg/models/rt-detr/rtdetr-x.yaml +57 -0
  71. ultralytics/cfg/models/v10/yolov10b.yaml +45 -0
  72. ultralytics/cfg/models/v10/yolov10l.yaml +45 -0
  73. ultralytics/cfg/models/v10/yolov10m.yaml +45 -0
  74. ultralytics/cfg/models/v10/yolov10n.yaml +45 -0
  75. ultralytics/cfg/models/v10/yolov10s.yaml +45 -0
  76. ultralytics/cfg/models/v10/yolov10x.yaml +45 -0
  77. ultralytics/cfg/models/v3/yolov3-spp.yaml +49 -0
  78. ultralytics/cfg/models/v3/yolov3-tiny.yaml +40 -0
  79. ultralytics/cfg/models/v3/yolov3.yaml +49 -0
  80. ultralytics/cfg/models/v5/yolov5-p6.yaml +62 -0
  81. ultralytics/cfg/models/v5/yolov5.yaml +51 -0
  82. ultralytics/cfg/models/v6/yolov6.yaml +56 -0
  83. ultralytics/cfg/models/v8/yoloe-v8-seg.yaml +45 -0
  84. ultralytics/cfg/models/v8/yoloe-v8.yaml +45 -0
  85. ultralytics/cfg/models/v8/yolov8-cls-resnet101.yaml +28 -0
  86. ultralytics/cfg/models/v8/yolov8-cls-resnet50.yaml +28 -0
  87. ultralytics/cfg/models/v8/yolov8-cls.yaml +32 -0
  88. ultralytics/cfg/models/v8/yolov8-ghost-p2.yaml +58 -0
  89. ultralytics/cfg/models/v8/yolov8-ghost-p6.yaml +60 -0
  90. ultralytics/cfg/models/v8/yolov8-ghost.yaml +50 -0
  91. ultralytics/cfg/models/v8/yolov8-obb.yaml +49 -0
  92. ultralytics/cfg/models/v8/yolov8-p2.yaml +57 -0
  93. ultralytics/cfg/models/v8/yolov8-p6.yaml +59 -0
  94. ultralytics/cfg/models/v8/yolov8-pose-p6.yaml +60 -0
  95. ultralytics/cfg/models/v8/yolov8-pose.yaml +50 -0
  96. ultralytics/cfg/models/v8/yolov8-rtdetr.yaml +49 -0
  97. ultralytics/cfg/models/v8/yolov8-seg-p6.yaml +59 -0
  98. ultralytics/cfg/models/v8/yolov8-seg.yaml +49 -0
  99. ultralytics/cfg/models/v8/yolov8-world.yaml +51 -0
  100. ultralytics/cfg/models/v8/yolov8-worldv2.yaml +49 -0
  101. ultralytics/cfg/models/v8/yolov8.yaml +49 -0
  102. ultralytics/cfg/models/v9/yolov9c-seg.yaml +41 -0
  103. ultralytics/cfg/models/v9/yolov9c.yaml +41 -0
  104. ultralytics/cfg/models/v9/yolov9e-seg.yaml +64 -0
  105. ultralytics/cfg/models/v9/yolov9e.yaml +64 -0
  106. ultralytics/cfg/models/v9/yolov9m.yaml +41 -0
  107. ultralytics/cfg/models/v9/yolov9s.yaml +41 -0
  108. ultralytics/cfg/models/v9/yolov9t.yaml +41 -0
  109. ultralytics/cfg/trackers/botsort.yaml +22 -0
  110. ultralytics/cfg/trackers/bytetrack.yaml +14 -0
  111. ultralytics/data/__init__.py +26 -0
  112. ultralytics/data/annotator.py +66 -0
  113. ultralytics/data/augment.py +2945 -0
  114. ultralytics/data/base.py +438 -0
  115. ultralytics/data/build.py +258 -0
  116. ultralytics/data/converter.py +754 -0
  117. ultralytics/data/dataset.py +834 -0
  118. ultralytics/data/loaders.py +676 -0
  119. ultralytics/data/scripts/download_weights.sh +18 -0
  120. ultralytics/data/scripts/get_coco.sh +61 -0
  121. ultralytics/data/scripts/get_coco128.sh +18 -0
  122. ultralytics/data/scripts/get_imagenet.sh +52 -0
  123. ultralytics/data/split.py +125 -0
  124. ultralytics/data/split_dota.py +325 -0
  125. ultralytics/data/utils.py +777 -0
  126. ultralytics/engine/__init__.py +1 -0
  127. ultralytics/engine/exporter.py +1519 -0
  128. ultralytics/engine/model.py +1156 -0
  129. ultralytics/engine/predictor.py +502 -0
  130. ultralytics/engine/results.py +1840 -0
  131. ultralytics/engine/trainer.py +853 -0
  132. ultralytics/engine/tuner.py +243 -0
  133. ultralytics/engine/validator.py +377 -0
  134. ultralytics/hub/__init__.py +168 -0
  135. ultralytics/hub/auth.py +137 -0
  136. ultralytics/hub/google/__init__.py +176 -0
  137. ultralytics/hub/session.py +446 -0
  138. ultralytics/hub/utils.py +248 -0
  139. ultralytics/models/__init__.py +9 -0
  140. ultralytics/models/fastsam/__init__.py +7 -0
  141. ultralytics/models/fastsam/model.py +61 -0
  142. ultralytics/models/fastsam/predict.py +181 -0
  143. ultralytics/models/fastsam/utils.py +24 -0
  144. ultralytics/models/fastsam/val.py +40 -0
  145. ultralytics/models/nas/__init__.py +7 -0
  146. ultralytics/models/nas/model.py +102 -0
  147. ultralytics/models/nas/predict.py +58 -0
  148. ultralytics/models/nas/val.py +39 -0
  149. ultralytics/models/rtdetr/__init__.py +7 -0
  150. ultralytics/models/rtdetr/model.py +63 -0
  151. ultralytics/models/rtdetr/predict.py +84 -0
  152. ultralytics/models/rtdetr/train.py +85 -0
  153. ultralytics/models/rtdetr/val.py +191 -0
  154. ultralytics/models/sam/__init__.py +6 -0
  155. ultralytics/models/sam/amg.py +260 -0
  156. ultralytics/models/sam/build.py +358 -0
  157. ultralytics/models/sam/model.py +170 -0
  158. ultralytics/models/sam/modules/__init__.py +1 -0
  159. ultralytics/models/sam/modules/blocks.py +1129 -0
  160. ultralytics/models/sam/modules/decoders.py +515 -0
  161. ultralytics/models/sam/modules/encoders.py +854 -0
  162. ultralytics/models/sam/modules/memory_attention.py +299 -0
  163. ultralytics/models/sam/modules/sam.py +1006 -0
  164. ultralytics/models/sam/modules/tiny_encoder.py +1002 -0
  165. ultralytics/models/sam/modules/transformer.py +351 -0
  166. ultralytics/models/sam/modules/utils.py +394 -0
  167. ultralytics/models/sam/predict.py +1605 -0
  168. ultralytics/models/utils/__init__.py +1 -0
  169. ultralytics/models/utils/loss.py +455 -0
  170. ultralytics/models/utils/ops.py +268 -0
  171. ultralytics/models/yolo/__init__.py +7 -0
  172. ultralytics/models/yolo/classify/__init__.py +7 -0
  173. ultralytics/models/yolo/classify/predict.py +88 -0
  174. ultralytics/models/yolo/classify/train.py +233 -0
  175. ultralytics/models/yolo/classify/val.py +215 -0
  176. ultralytics/models/yolo/detect/__init__.py +7 -0
  177. ultralytics/models/yolo/detect/predict.py +124 -0
  178. ultralytics/models/yolo/detect/train.py +217 -0
  179. ultralytics/models/yolo/detect/val.py +451 -0
  180. ultralytics/models/yolo/model.py +354 -0
  181. ultralytics/models/yolo/obb/__init__.py +7 -0
  182. ultralytics/models/yolo/obb/predict.py +66 -0
  183. ultralytics/models/yolo/obb/train.py +81 -0
  184. ultralytics/models/yolo/obb/val.py +283 -0
  185. ultralytics/models/yolo/pose/__init__.py +7 -0
  186. ultralytics/models/yolo/pose/predict.py +79 -0
  187. ultralytics/models/yolo/pose/train.py +154 -0
  188. ultralytics/models/yolo/pose/val.py +394 -0
  189. ultralytics/models/yolo/segment/__init__.py +7 -0
  190. ultralytics/models/yolo/segment/predict.py +113 -0
  191. ultralytics/models/yolo/segment/train.py +123 -0
  192. ultralytics/models/yolo/segment/val.py +428 -0
  193. ultralytics/models/yolo/world/__init__.py +5 -0
  194. ultralytics/models/yolo/world/train.py +119 -0
  195. ultralytics/models/yolo/world/train_world.py +176 -0
  196. ultralytics/models/yolo/yoloe/__init__.py +22 -0
  197. ultralytics/models/yolo/yoloe/predict.py +169 -0
  198. ultralytics/models/yolo/yoloe/train.py +298 -0
  199. ultralytics/models/yolo/yoloe/train_seg.py +124 -0
  200. ultralytics/models/yolo/yoloe/val.py +191 -0
  201. ultralytics/nn/__init__.py +29 -0
  202. ultralytics/nn/autobackend.py +842 -0
  203. ultralytics/nn/modules/__init__.py +182 -0
  204. ultralytics/nn/modules/activation.py +53 -0
  205. ultralytics/nn/modules/block.py +1966 -0
  206. ultralytics/nn/modules/conv.py +712 -0
  207. ultralytics/nn/modules/head.py +880 -0
  208. ultralytics/nn/modules/transformer.py +713 -0
  209. ultralytics/nn/modules/utils.py +164 -0
  210. ultralytics/nn/tasks.py +1627 -0
  211. ultralytics/nn/text_model.py +351 -0
  212. ultralytics/solutions/__init__.py +41 -0
  213. ultralytics/solutions/ai_gym.py +116 -0
  214. ultralytics/solutions/analytics.py +252 -0
  215. ultralytics/solutions/config.py +106 -0
  216. ultralytics/solutions/distance_calculation.py +124 -0
  217. ultralytics/solutions/heatmap.py +127 -0
  218. ultralytics/solutions/instance_segmentation.py +84 -0
  219. ultralytics/solutions/object_blurrer.py +90 -0
  220. ultralytics/solutions/object_counter.py +195 -0
  221. ultralytics/solutions/object_cropper.py +84 -0
  222. ultralytics/solutions/parking_management.py +273 -0
  223. ultralytics/solutions/queue_management.py +93 -0
  224. ultralytics/solutions/region_counter.py +120 -0
  225. ultralytics/solutions/security_alarm.py +154 -0
  226. ultralytics/solutions/similarity_search.py +172 -0
  227. ultralytics/solutions/solutions.py +724 -0
  228. ultralytics/solutions/speed_estimation.py +110 -0
  229. ultralytics/solutions/streamlit_inference.py +196 -0
  230. ultralytics/solutions/templates/similarity-search.html +160 -0
  231. ultralytics/solutions/trackzone.py +88 -0
  232. ultralytics/solutions/vision_eye.py +68 -0
  233. ultralytics/trackers/__init__.py +7 -0
  234. ultralytics/trackers/basetrack.py +124 -0
  235. ultralytics/trackers/bot_sort.py +260 -0
  236. ultralytics/trackers/byte_tracker.py +480 -0
  237. ultralytics/trackers/track.py +125 -0
  238. ultralytics/trackers/utils/__init__.py +1 -0
  239. ultralytics/trackers/utils/gmc.py +376 -0
  240. ultralytics/trackers/utils/kalman_filter.py +493 -0
  241. ultralytics/trackers/utils/matching.py +157 -0
  242. ultralytics/utils/__init__.py +1435 -0
  243. ultralytics/utils/autobatch.py +106 -0
  244. ultralytics/utils/autodevice.py +174 -0
  245. ultralytics/utils/benchmarks.py +695 -0
  246. ultralytics/utils/callbacks/__init__.py +5 -0
  247. ultralytics/utils/callbacks/base.py +234 -0
  248. ultralytics/utils/callbacks/clearml.py +153 -0
  249. ultralytics/utils/callbacks/comet.py +552 -0
  250. ultralytics/utils/callbacks/dvc.py +205 -0
  251. ultralytics/utils/callbacks/hub.py +108 -0
  252. ultralytics/utils/callbacks/mlflow.py +138 -0
  253. ultralytics/utils/callbacks/neptune.py +140 -0
  254. ultralytics/utils/callbacks/raytune.py +43 -0
  255. ultralytics/utils/callbacks/tensorboard.py +132 -0
  256. ultralytics/utils/callbacks/wb.py +185 -0
  257. ultralytics/utils/checks.py +897 -0
  258. ultralytics/utils/dist.py +119 -0
  259. ultralytics/utils/downloads.py +499 -0
  260. ultralytics/utils/errors.py +43 -0
  261. ultralytics/utils/export.py +219 -0
  262. ultralytics/utils/files.py +221 -0
  263. ultralytics/utils/instance.py +499 -0
  264. ultralytics/utils/loss.py +813 -0
  265. ultralytics/utils/metrics.py +1356 -0
  266. ultralytics/utils/ops.py +885 -0
  267. ultralytics/utils/patches.py +143 -0
  268. ultralytics/utils/plotting.py +1011 -0
  269. ultralytics/utils/tal.py +416 -0
  270. ultralytics/utils/torch_utils.py +990 -0
  271. ultralytics/utils/triton.py +116 -0
  272. ultralytics/utils/tuner.py +159 -0
tests/test_python.py ADDED
@@ -0,0 +1,694 @@
1
+ # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
2
+
3
+ import contextlib
4
+ import csv
5
+ import urllib
6
+ from copy import copy
7
+ from pathlib import Path
8
+
9
+ import cv2
10
+ import numpy as np
11
+ import pytest
12
+ import torch
13
+ from PIL import Image
14
+
15
+ from tests import CFG, MODEL, SOURCE, SOURCES_LIST, TMP
16
+ from ultralytics import RTDETR, YOLO
17
+ from ultralytics.cfg import MODELS, TASK2DATA, TASKS
18
+ from ultralytics.data.build import load_inference_source
19
+ from ultralytics.utils import (
20
+ ARM64,
21
+ ASSETS,
22
+ DEFAULT_CFG,
23
+ DEFAULT_CFG_PATH,
24
+ LINUX,
25
+ LOGGER,
26
+ ONLINE,
27
+ ROOT,
28
+ WEIGHTS_DIR,
29
+ WINDOWS,
30
+ YAML,
31
+ checks,
32
+ is_dir_writeable,
33
+ is_github_action_running,
34
+ )
35
+ from ultralytics.utils.downloads import download
36
+ from ultralytics.utils.torch_utils import TORCH_1_9
37
+
38
+ IS_TMP_WRITEABLE = is_dir_writeable(TMP) # WARNING: must be run once tests start as TMP does not exist on tests/init
39
+
40
+
41
+ def test_model_forward():
42
+ """Test the forward pass of the YOLO model."""
43
+ model = YOLO(CFG)
44
+ model(source=None, imgsz=32, augment=True) # also test no source and augment
45
+
46
+
47
+ def test_model_methods():
48
+ """Test various methods and properties of the YOLO model to ensure correct functionality."""
49
+ model = YOLO(MODEL)
50
+
51
+ # Model methods
52
+ model.info(verbose=True, detailed=True)
53
+ model = model.reset_weights()
54
+ model = model.load(MODEL)
55
+ model.to("cpu")
56
+ model.fuse()
57
+ model.clear_callback("on_train_start")
58
+ model.reset_callbacks()
59
+
60
+ # Model properties
61
+ _ = model.names
62
+ _ = model.device
63
+ _ = model.transforms
64
+ _ = model.task_map
65
+
66
+
67
+ def test_model_profile():
68
+ """Test profiling of the YOLO model with `profile=True` to assess performance and resource usage."""
69
+ from ultralytics.nn.tasks import DetectionModel
70
+
71
+ model = DetectionModel() # build model
72
+ im = torch.randn(1, 3, 64, 64) # requires min imgsz=64
73
+ _ = model.predict(im, profile=True)
74
+
75
+
76
+ @pytest.mark.skipif(not IS_TMP_WRITEABLE, reason="directory is not writeable")
77
+ def test_predict_txt():
78
+ """Test YOLO predictions with file, directory, and pattern sources listed in a text file."""
79
+ file = TMP / "sources_multi_row.txt"
80
+ with open(file, "w") as f:
81
+ for src in SOURCES_LIST:
82
+ f.write(f"{src}\n")
83
+ results = YOLO(MODEL)(source=file, imgsz=32)
84
+ assert len(results) == 7 # 1 + 2 + 2 + 2 = 7 images
85
+
86
+
87
+ @pytest.mark.skipif(True, reason="disabled for testing")
88
+ @pytest.mark.skipif(not IS_TMP_WRITEABLE, reason="directory is not writeable")
89
+ def test_predict_csv_multi_row():
90
+ """Test YOLO predictions with sources listed in multiple rows of a CSV file."""
91
+ file = TMP / "sources_multi_row.csv"
92
+ with open(file, "w", newline="") as f:
93
+ writer = csv.writer(f)
94
+ writer.writerow(["source"])
95
+ writer.writerows([[src] for src in SOURCES_LIST])
96
+ results = YOLO(MODEL)(source=file, imgsz=32)
97
+ assert len(results) == 7 # 1 + 2 + 2 + 2 = 7 images
98
+
99
+
100
+ @pytest.mark.skipif(True, reason="disabled for testing")
101
+ @pytest.mark.skipif(not IS_TMP_WRITEABLE, reason="directory is not writeable")
102
+ def test_predict_csv_single_row():
103
+ """Test YOLO predictions with sources listed in a single row of a CSV file."""
104
+ file = TMP / "sources_single_row.csv"
105
+ with open(file, "w", newline="") as f:
106
+ writer = csv.writer(f)
107
+ writer.writerow(SOURCES_LIST)
108
+ results = YOLO(MODEL)(source=file, imgsz=32)
109
+ assert len(results) == 7 # 1 + 2 + 2 + 2 = 7 images
110
+
111
+
112
+ @pytest.mark.parametrize("model_name", MODELS)
113
+ def test_predict_img(model_name):
114
+ """Test YOLO model predictions on various image input types and sources, including online images."""
115
+ model = YOLO(WEIGHTS_DIR / model_name)
116
+ im = cv2.imread(str(SOURCE)) # uint8 numpy array
117
+ assert len(model(source=Image.open(SOURCE), save=True, verbose=True, imgsz=32)) == 1 # PIL
118
+ assert len(model(source=im, save=True, save_txt=True, imgsz=32)) == 1 # ndarray
119
+ assert len(model(torch.rand((2, 3, 32, 32)), imgsz=32)) == 2 # batch-size 2 Tensor, FP32 0.0-1.0 RGB order
120
+ assert len(model(source=[im, im], save=True, save_txt=True, imgsz=32)) == 2 # batch
121
+ assert len(list(model(source=[im, im], save=True, stream=True, imgsz=32))) == 2 # stream
122
+ assert len(model(torch.zeros(320, 640, 3).numpy().astype(np.uint8), imgsz=32)) == 1 # tensor to numpy
123
+ batch = [
124
+ str(SOURCE), # filename
125
+ Path(SOURCE), # Path
126
+ "https://github.com/ultralytics/assets/releases/download/v0.0.0/zidane.jpg" if ONLINE else SOURCE, # URI
127
+ cv2.imread(str(SOURCE)), # OpenCV
128
+ Image.open(SOURCE), # PIL
129
+ np.zeros((320, 640, 3), dtype=np.uint8), # numpy
130
+ ]
131
+ assert len(model(batch, imgsz=32, classes=0)) == len(batch) # multiple sources in a batch
132
+
133
+
134
+ @pytest.mark.parametrize("model", MODELS)
135
+ def test_predict_visualize(model):
136
+ """Test model prediction methods with 'visualize=True' to generate and display prediction visualizations."""
137
+ YOLO(WEIGHTS_DIR / model)(SOURCE, imgsz=32, visualize=True)
138
+
139
+
140
+ def test_predict_grey_and_4ch():
141
+ """Test YOLO prediction on SOURCE converted to greyscale and 4-channel images with various filenames."""
142
+ im = Image.open(SOURCE)
143
+ directory = TMP / "im4"
144
+ directory.mkdir(parents=True, exist_ok=True)
145
+
146
+ source_greyscale = directory / "greyscale.jpg"
147
+ source_rgba = directory / "4ch.png"
148
+ source_non_utf = directory / "non_UTF_测试文件_tést_image.jpg"
149
+ source_spaces = directory / "image with spaces.jpg"
150
+
151
+ im.convert("L").save(source_greyscale) # greyscale
152
+ im.convert("RGBA").save(source_rgba) # 4-ch PNG with alpha
153
+ im.save(source_non_utf) # non-UTF characters in filename
154
+ im.save(source_spaces) # spaces in filename
155
+
156
+ # Inference
157
+ model = YOLO(MODEL)
158
+ for f in source_rgba, source_greyscale, source_non_utf, source_spaces:
159
+ for source in Image.open(f), cv2.imread(str(f)), f:
160
+ results = model(source, save=True, verbose=True, imgsz=32)
161
+ assert len(results) == 1 # verify that an image was run
162
+ f.unlink() # cleanup
163
+
164
+
165
+ @pytest.mark.slow
166
+ @pytest.mark.skipif(not ONLINE, reason="environment is offline")
167
+ @pytest.mark.skipif(is_github_action_running(), reason="No auth https://github.com/JuanBindez/pytubefix/issues/166")
168
+ def test_youtube():
169
+ """Test YOLO model on a YouTube video stream, handling potential network-related errors."""
170
+ model = YOLO(MODEL)
171
+ try:
172
+ model.predict("https://youtu.be/G17sBkb38XQ", imgsz=96, save=True)
173
+ # Handle internet connection errors and 'urllib.error.HTTPError: HTTP Error 429: Too Many Requests'
174
+ except (urllib.error.HTTPError, ConnectionError) as e:
175
+ LOGGER.error(f"YouTube Test Error: {e}")
176
+
177
+
178
+ @pytest.mark.skipif(not ONLINE, reason="environment is offline")
179
+ @pytest.mark.skipif(not IS_TMP_WRITEABLE, reason="directory is not writeable")
180
+ def test_track_stream():
181
+ """
182
+ Test streaming tracking on a short 10 frame video using ByteTrack tracker and different GMC methods.
183
+
184
+ Note imgsz=160 required for tracking for higher confidence and better matches.
185
+ """
186
+ video_url = "https://github.com/ultralytics/assets/releases/download/v0.0.0/decelera_portrait_min.mov"
187
+ model = YOLO(MODEL)
188
+ model.track(video_url, imgsz=160, tracker="bytetrack.yaml")
189
+ model.track(video_url, imgsz=160, tracker="botsort.yaml", save_frames=True) # test frame saving also
190
+
191
+ # Test Global Motion Compensation (GMC) methods and ReID
192
+ for gmc, reidm in zip(["orb", "sift", "ecc"], ["auto", "auto", "yolo11n-cls.pt"]):
193
+ default_args = YAML.load(ROOT / "cfg/trackers/botsort.yaml")
194
+ custom_yaml = TMP / f"botsort-{gmc}.yaml"
195
+ YAML.save(custom_yaml, {**default_args, "gmc_method": gmc, "with_reid": True, "model": reidm})
196
+ model.track(video_url, imgsz=160, tracker=custom_yaml)
197
+
198
+
199
+ def test_val():
200
+ """Test the validation mode of the YOLO model."""
201
+ YOLO(MODEL).val(data="coco8.yaml", imgsz=32)
202
+
203
+
204
+ def test_train_scratch():
205
+ """Test training the YOLO model from scratch using the provided configuration."""
206
+ model = YOLO(CFG)
207
+ model.train(data="coco8.yaml", epochs=2, imgsz=32, cache="disk", batch=-1, close_mosaic=1, name="model")
208
+ model(SOURCE)
209
+
210
+
211
+ @pytest.mark.parametrize("scls", [False, True])
212
+ def test_train_pretrained(scls):
213
+ """Test training of the YOLO model starting from a pre-trained checkpoint."""
214
+ model = YOLO(WEIGHTS_DIR / "yolo11n-seg.pt")
215
+ model.train(
216
+ data="coco8-seg.yaml", epochs=1, imgsz=32, cache="ram", copy_paste=0.5, mixup=0.5, name=0, single_cls=scls
217
+ )
218
+ model(SOURCE)
219
+
220
+
221
+ def test_all_model_yamls():
222
+ """Test YOLO model creation for all available YAML configurations in the `cfg/models` directory."""
223
+ for m in (ROOT / "cfg" / "models").rglob("*.yaml"):
224
+ if "rtdetr" in m.name:
225
+ if TORCH_1_9: # torch<=1.8 issue - TypeError: __init__() got an unexpected keyword argument 'batch_first'
226
+ _ = RTDETR(m.name)(SOURCE, imgsz=640) # must be 640
227
+ else:
228
+ YOLO(m.name)
229
+
230
+
231
+ @pytest.mark.skipif(WINDOWS, reason="Windows slow CI export bug https://github.com/ultralytics/ultralytics/pull/16003")
232
+ def test_workflow():
233
+ """Test the complete workflow including training, validation, prediction, and exporting."""
234
+ model = YOLO(MODEL)
235
+ model.train(data="coco8.yaml", epochs=1, imgsz=32, optimizer="SGD")
236
+ model.val(imgsz=32)
237
+ model.predict(SOURCE, imgsz=32)
238
+ model.export(format="torchscript") # WARNING: Windows slow CI export bug
239
+
240
+
241
+ def test_predict_callback_and_setup():
242
+ """Test callback functionality during YOLO prediction setup and execution."""
243
+
244
+ def on_predict_batch_end(predictor):
245
+ """Callback function that handles operations at the end of a prediction batch."""
246
+ path, im0s, _ = predictor.batch
247
+ im0s = im0s if isinstance(im0s, list) else [im0s]
248
+ bs = [predictor.dataset.bs for _ in range(len(path))]
249
+ predictor.results = zip(predictor.results, im0s, bs) # results is List[batch_size]
250
+
251
+ model = YOLO(MODEL)
252
+ model.add_callback("on_predict_batch_end", on_predict_batch_end)
253
+
254
+ dataset = load_inference_source(source=SOURCE)
255
+ bs = dataset.bs # noqa access predictor properties
256
+ results = model.predict(dataset, stream=True, imgsz=160) # source already setup
257
+ for r, im0, bs in results:
258
+ print("test_callback", im0.shape)
259
+ print("test_callback", bs)
260
+ boxes = r.boxes # Boxes object for bbox outputs
261
+ print(boxes)
262
+
263
+
264
+ @pytest.mark.parametrize("model", MODELS)
265
+ def test_results(model):
266
+ """Test YOLO model results processing and output in various formats."""
267
+ results = YOLO(WEIGHTS_DIR / model)([SOURCE, SOURCE], imgsz=160)
268
+ for r in results:
269
+ r = r.cpu().numpy()
270
+ print(r, len(r), r.path) # print numpy attributes
271
+ r = r.to(device="cpu", dtype=torch.float32)
272
+ r.save_txt(txt_file=TMP / "runs/tests/label.txt", save_conf=True)
273
+ r.save_crop(save_dir=TMP / "runs/tests/crops/")
274
+ r.to_df(decimals=3) # Align to_ methods: https://docs.ultralytics.com/modes/predict/#working-with-results
275
+ r.to_csv()
276
+ r.to_xml()
277
+ r.to_html()
278
+ r.to_json(normalize=True)
279
+ r.to_sql()
280
+ r.plot(pil=True, save=True, filename=TMP / "results_plot_save.jpg")
281
+ r.plot(conf=True, boxes=True)
282
+ print(r, len(r), r.path) # print after methods
283
+
284
+
285
+ def test_labels_and_crops():
286
+ """Test output from prediction args for saving YOLO detection labels and crops."""
287
+ imgs = [SOURCE, ASSETS / "zidane.jpg"]
288
+ results = YOLO(WEIGHTS_DIR / "yolo11n.pt")(imgs, imgsz=160, save_txt=True, save_crop=True)
289
+ save_path = Path(results[0].save_dir)
290
+ for r in results:
291
+ im_name = Path(r.path).stem
292
+ cls_idxs = r.boxes.cls.int().tolist()
293
+ # Check correct detections
294
+ assert cls_idxs == ([0, 7, 0, 0] if r.path.endswith("bus.jpg") else [0, 0, 0]) # bus.jpg and zidane.jpg classes
295
+ # Check label path
296
+ labels = save_path / f"labels/{im_name}.txt"
297
+ assert labels.exists()
298
+ # Check detections match label count
299
+ assert len(r.boxes.data) == len([line for line in labels.read_text().splitlines() if line])
300
+ # Check crops path and files
301
+ crop_dirs = list((save_path / "crops").iterdir())
302
+ crop_files = [f for p in crop_dirs for f in p.glob("*")]
303
+ # Crop directories match detections
304
+ assert all(r.names.get(c) in {d.name for d in crop_dirs} for c in cls_idxs)
305
+ # Same number of crops as detections
306
+ assert len([f for f in crop_files if im_name in f.name]) == len(r.boxes.data)
307
+
308
+
309
+ @pytest.mark.skipif(not ONLINE, reason="environment is offline")
310
+ def test_data_utils():
311
+ """Test utility functions in ultralytics/data/utils.py, including dataset stats and auto-splitting."""
312
+ from ultralytics.data.split import autosplit
313
+ from ultralytics.data.utils import HUBDatasetStats
314
+ from ultralytics.utils.downloads import zip_directory
315
+
316
+ # from ultralytics.utils.files import WorkingDirectory
317
+ # with WorkingDirectory(ROOT.parent / 'tests'):
318
+
319
+ for task in TASKS:
320
+ file = Path(TASK2DATA[task]).with_suffix(".zip") # i.e. coco8.zip
321
+ download(f"https://github.com/ultralytics/hub/raw/main/example_datasets/{file}", unzip=False, dir=TMP)
322
+ stats = HUBDatasetStats(TMP / file, task=task)
323
+ stats.get_json(save=True)
324
+ stats.process_images()
325
+
326
+ autosplit(TMP / "coco8")
327
+ zip_directory(TMP / "coco8/images/val") # zip
328
+
329
+
330
+ @pytest.mark.skipif(not ONLINE, reason="environment is offline")
331
+ def test_data_converter():
332
+ """Test dataset conversion functions from COCO to YOLO format and class mappings."""
333
+ from ultralytics.data.converter import coco80_to_coco91_class, convert_coco
334
+
335
+ file = "instances_val2017.json"
336
+ download(f"https://github.com/ultralytics/assets/releases/download/v0.0.0/{file}", dir=TMP)
337
+ convert_coco(labels_dir=TMP, save_dir=TMP / "yolo_labels", use_segments=True, use_keypoints=False, cls91to80=True)
338
+ coco80_to_coco91_class()
339
+
340
+
341
+ def test_data_annotator():
342
+ """Test automatic annotation of data using detection and segmentation models."""
343
+ from ultralytics.data.annotator import auto_annotate
344
+
345
+ auto_annotate(
346
+ ASSETS,
347
+ det_model=WEIGHTS_DIR / "yolo11n.pt",
348
+ sam_model=WEIGHTS_DIR / "mobile_sam.pt",
349
+ output_dir=TMP / "auto_annotate_labels",
350
+ )
351
+
352
+
353
+ def test_events():
354
+ """Test event sending functionality."""
355
+ from ultralytics.hub.utils import Events
356
+
357
+ events = Events()
358
+ events.enabled = True
359
+ cfg = copy(DEFAULT_CFG) # does not require deepcopy
360
+ cfg.mode = "test"
361
+ events(cfg)
362
+
363
+
364
+ def test_cfg_init():
365
+ """Test configuration initialization utilities from the 'ultralytics.cfg' module."""
366
+ from ultralytics.cfg import check_dict_alignment, copy_default_cfg, smart_value
367
+
368
+ with contextlib.suppress(SyntaxError):
369
+ check_dict_alignment({"a": 1}, {"b": 2})
370
+ copy_default_cfg()
371
+ (Path.cwd() / DEFAULT_CFG_PATH.name.replace(".yaml", "_copy.yaml")).unlink(missing_ok=False)
372
+ [smart_value(x) for x in ["none", "true", "false"]]
373
+
374
+
375
+ def test_utils_init():
376
+ """Test initialization utilities in the Ultralytics library."""
377
+ from ultralytics.utils import get_git_branch, get_git_origin_url, get_ubuntu_version, is_github_action_running
378
+
379
+ get_ubuntu_version()
380
+ is_github_action_running()
381
+ get_git_origin_url()
382
+ get_git_branch()
383
+
384
+
385
+ def test_utils_checks():
386
+ """Test various utility checks for filenames, git status, requirements, image sizes, and versions."""
387
+ checks.check_yolov5u_filename("yolov5n.pt")
388
+ checks.git_describe(ROOT)
389
+ checks.check_requirements() # check requirements.txt
390
+ checks.check_imgsz([600, 600], max_dim=1)
391
+ checks.check_imshow(warn=True)
392
+ checks.check_version("ultralytics", "8.0.0")
393
+ checks.print_args()
394
+
395
+
396
+ @pytest.mark.skipif(WINDOWS, reason="Windows profiling is extremely slow (cause unknown)")
397
+ def test_utils_benchmarks():
398
+ """Benchmark model performance using 'ProfileModels' from 'ultralytics.utils.benchmarks'."""
399
+ from ultralytics.utils.benchmarks import ProfileModels
400
+
401
+ ProfileModels(["yolo11n.yaml"], imgsz=32, min_time=1, num_timed_runs=3, num_warmup_runs=1).run()
402
+
403
+
404
+ def test_utils_torchutils():
405
+ """Test Torch utility functions including profiling and FLOP calculations."""
406
+ from ultralytics.nn.modules.conv import Conv
407
+ from ultralytics.utils.torch_utils import get_flops_with_torch_profiler, profile_ops, time_sync
408
+
409
+ x = torch.randn(1, 64, 20, 20)
410
+ m = Conv(64, 64, k=1, s=2)
411
+
412
+ profile_ops(x, [m], n=3)
413
+ get_flops_with_torch_profiler(m)
414
+ time_sync()
415
+
416
+
417
+ def test_utils_ops():
418
+ """Test utility operations for coordinate transformations and normalizations."""
419
+ from ultralytics.utils.ops import (
420
+ ltwh2xywh,
421
+ ltwh2xyxy,
422
+ make_divisible,
423
+ xywh2ltwh,
424
+ xywh2xyxy,
425
+ xywhn2xyxy,
426
+ xywhr2xyxyxyxy,
427
+ xyxy2ltwh,
428
+ xyxy2xywh,
429
+ xyxy2xywhn,
430
+ xyxyxyxy2xywhr,
431
+ )
432
+
433
+ make_divisible(17, torch.tensor([8]))
434
+
435
+ boxes = torch.rand(10, 4) # xywh
436
+ torch.allclose(boxes, xyxy2xywh(xywh2xyxy(boxes)))
437
+ torch.allclose(boxes, xyxy2xywhn(xywhn2xyxy(boxes)))
438
+ torch.allclose(boxes, ltwh2xywh(xywh2ltwh(boxes)))
439
+ torch.allclose(boxes, xyxy2ltwh(ltwh2xyxy(boxes)))
440
+
441
+ boxes = torch.rand(10, 5) # xywhr for OBB
442
+ boxes[:, 4] = torch.randn(10) * 30
443
+ torch.allclose(boxes, xyxyxyxy2xywhr(xywhr2xyxyxyxy(boxes)), rtol=1e-3)
444
+
445
+
446
+ def test_utils_files():
447
+ """Test file handling utilities including file age, date, and paths with spaces."""
448
+ from ultralytics.utils.files import file_age, file_date, get_latest_run, spaces_in_path
449
+
450
+ file_age(SOURCE)
451
+ file_date(SOURCE)
452
+ get_latest_run(ROOT / "runs")
453
+
454
+ path = TMP / "path/with spaces"
455
+ path.mkdir(parents=True, exist_ok=True)
456
+ with spaces_in_path(path) as new_path:
457
+ print(new_path)
458
+
459
+
460
+ @pytest.mark.slow
461
+ def test_utils_patches_torch_save():
462
+ """Test torch_save backoff when _torch_save raises RuntimeError."""
463
+ from unittest.mock import MagicMock, patch
464
+
465
+ from ultralytics.utils.patches import torch_save
466
+
467
+ mock = MagicMock(side_effect=RuntimeError)
468
+
469
+ with patch("ultralytics.utils.patches._torch_save", new=mock):
470
+ with pytest.raises(RuntimeError):
471
+ torch_save(torch.zeros(1), TMP / "test.pt")
472
+
473
+ assert mock.call_count == 4, "torch_save was not attempted the expected number of times"
474
+
475
+
476
+ def test_nn_modules_conv():
477
+ """Test Convolutional Neural Network modules including CBAM, Conv2, and ConvTranspose."""
478
+ from ultralytics.nn.modules.conv import CBAM, Conv2, ConvTranspose, DWConvTranspose2d, Focus
479
+
480
+ c1, c2 = 8, 16 # input and output channels
481
+ x = torch.zeros(4, c1, 10, 10) # BCHW
482
+
483
+ # Run all modules not otherwise covered in tests
484
+ DWConvTranspose2d(c1, c2)(x)
485
+ ConvTranspose(c1, c2)(x)
486
+ Focus(c1, c2)(x)
487
+ CBAM(c1)(x)
488
+
489
+ # Fuse ops
490
+ m = Conv2(c1, c2)
491
+ m.fuse_convs()
492
+ m(x)
493
+
494
+
495
+ def test_nn_modules_block():
496
+ """Test various neural network block modules."""
497
+ from ultralytics.nn.modules.block import C1, C3TR, BottleneckCSP, C3Ghost, C3x
498
+
499
+ c1, c2 = 8, 16 # input and output channels
500
+ x = torch.zeros(4, c1, 10, 10) # BCHW
501
+
502
+ # Run all modules not otherwise covered in tests
503
+ C1(c1, c2)(x)
504
+ C3x(c1, c2)(x)
505
+ C3TR(c1, c2)(x)
506
+ C3Ghost(c1, c2)(x)
507
+ BottleneckCSP(c1, c2)(x)
508
+
509
+
510
+ @pytest.mark.skipif(not ONLINE, reason="environment is offline")
511
+ def test_hub():
512
+ """Test Ultralytics HUB functionalities."""
513
+ from ultralytics.hub import export_fmts_hub, logout
514
+ from ultralytics.hub.utils import smart_request
515
+
516
+ export_fmts_hub()
517
+ logout()
518
+ smart_request("GET", "https://github.com", progress=True)
519
+
520
+
521
+ @pytest.fixture
522
+ def image():
523
+ """Load and return an image from a predefined source."""
524
+ return cv2.imread(str(SOURCE))
525
+
526
+
527
+ @pytest.mark.parametrize(
528
+ "auto_augment, erasing, force_color_jitter",
529
+ [
530
+ (None, 0.0, False),
531
+ ("randaugment", 0.5, True),
532
+ ("augmix", 0.2, False),
533
+ ("autoaugment", 0.0, True),
534
+ ],
535
+ )
536
+ def test_classify_transforms_train(image, auto_augment, erasing, force_color_jitter):
537
+ """Test classification transforms during training with various augmentations."""
538
+ from ultralytics.data.augment import classify_augmentations
539
+
540
+ transform = classify_augmentations(
541
+ size=224,
542
+ mean=(0.5, 0.5, 0.5),
543
+ std=(0.5, 0.5, 0.5),
544
+ scale=(0.08, 1.0),
545
+ ratio=(3.0 / 4.0, 4.0 / 3.0),
546
+ hflip=0.5,
547
+ vflip=0.5,
548
+ auto_augment=auto_augment,
549
+ hsv_h=0.015,
550
+ hsv_s=0.4,
551
+ hsv_v=0.4,
552
+ force_color_jitter=force_color_jitter,
553
+ erasing=erasing,
554
+ )
555
+
556
+ transformed_image = transform(Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)))
557
+
558
+ assert transformed_image.shape == (3, 224, 224)
559
+ assert torch.is_tensor(transformed_image)
560
+ assert transformed_image.dtype == torch.float32
561
+
562
+
563
+ @pytest.mark.slow
564
+ @pytest.mark.skipif(not ONLINE, reason="environment is offline")
565
+ def test_model_tune():
566
+ """Tune YOLO model for performance improvement."""
567
+ YOLO("yolo11n-pose.pt").tune(data="coco8-pose.yaml", plots=False, imgsz=32, epochs=1, iterations=2, device="cpu")
568
+ YOLO("yolo11n-cls.pt").tune(data="imagenet10", plots=False, imgsz=32, epochs=1, iterations=2, device="cpu")
569
+
570
+
571
+ def test_model_embeddings():
572
+ """Test YOLO model embeddings extraction functionality."""
573
+ model_detect = YOLO(MODEL)
574
+ model_segment = YOLO(WEIGHTS_DIR / "yolo11n-seg.pt")
575
+
576
+ for batch in [SOURCE], [SOURCE, SOURCE]: # test batch size 1 and 2
577
+ assert len(model_detect.embed(source=batch, imgsz=32)) == len(batch)
578
+ assert len(model_segment.embed(source=batch, imgsz=32)) == len(batch)
579
+
580
+
581
+ @pytest.mark.skipif(checks.IS_PYTHON_3_12, reason="YOLOWorld with CLIP is not supported in Python 3.12")
582
+ @pytest.mark.skipif(
583
+ checks.IS_PYTHON_3_8 and LINUX and ARM64,
584
+ reason="YOLOWorld with CLIP is not supported in Python 3.8 and aarch64 Linux",
585
+ )
586
+ def test_yolo_world():
587
+ """Test YOLO world models with CLIP support."""
588
+ model = YOLO(WEIGHTS_DIR / "yolov8s-world.pt") # no YOLO11n-world model yet
589
+ model.set_classes(["tree", "window"])
590
+ model(SOURCE, conf=0.01)
591
+
592
+ model = YOLO(WEIGHTS_DIR / "yolov8s-worldv2.pt") # no YOLO11n-world model yet
593
+ # Training from a pretrained model. Eval is included at the final stage of training.
594
+ # Use dota8.yaml which has fewer categories to reduce the inference time of CLIP model
595
+ model.train(
596
+ data="dota8.yaml",
597
+ epochs=1,
598
+ imgsz=32,
599
+ cache="disk",
600
+ close_mosaic=1,
601
+ )
602
+
603
+ # test WorWorldTrainerFromScratch
604
+ from ultralytics.models.yolo.world.train_world import WorldTrainerFromScratch
605
+
606
+ model = YOLO("yolov8s-worldv2.yaml") # no YOLO11n-world model yet
607
+ model.train(
608
+ data={"train": {"yolo_data": ["dota8.yaml"]}, "val": {"yolo_data": ["dota8.yaml"]}},
609
+ epochs=1,
610
+ imgsz=32,
611
+ cache="disk",
612
+ close_mosaic=1,
613
+ trainer=WorldTrainerFromScratch,
614
+ )
615
+
616
+
617
+ @pytest.mark.skipif(checks.IS_PYTHON_3_12 or not TORCH_1_9, reason="YOLOE with CLIP is not supported in Python 3.12")
618
+ @pytest.mark.skipif(
619
+ checks.IS_PYTHON_3_8 and LINUX and ARM64,
620
+ reason="YOLOE with CLIP is not supported in Python 3.8 and aarch64 Linux",
621
+ )
622
+ def test_yoloe():
623
+ """Test YOLOE models with MobileClip support."""
624
+ # Predict
625
+ # text-prompts
626
+ model = YOLO(WEIGHTS_DIR / "yoloe-11s-seg.pt")
627
+ names = ["person", "bus"]
628
+ model.set_classes(names, model.get_text_pe(names))
629
+ model(SOURCE, conf=0.01)
630
+
631
+ import numpy as np
632
+
633
+ from ultralytics import YOLOE
634
+ from ultralytics.models.yolo.yoloe import YOLOEVPSegPredictor
635
+
636
+ # visual-prompts
637
+ visuals = dict(
638
+ bboxes=np.array(
639
+ [[221.52, 405.8, 344.98, 857.54], [120, 425, 160, 445]],
640
+ ),
641
+ cls=np.array([0, 1]),
642
+ )
643
+ model.predict(
644
+ SOURCE,
645
+ visual_prompts=visuals,
646
+ predictor=YOLOEVPSegPredictor,
647
+ )
648
+
649
+ # Val
650
+ model = YOLOE(WEIGHTS_DIR / "yoloe-11s-seg.pt")
651
+ # text prompts
652
+ model.val(data="coco128-seg.yaml", imgsz=32)
653
+ # visual prompts
654
+ model.val(data="coco128-seg.yaml", load_vp=True, imgsz=32)
655
+
656
+ # Train, fine-tune
657
+ from ultralytics.models.yolo.yoloe import YOLOEPESegTrainer
658
+
659
+ model = YOLOE("yoloe-11s-seg.pt")
660
+ model.train(
661
+ data="coco128-seg.yaml",
662
+ epochs=1,
663
+ close_mosaic=1,
664
+ trainer=YOLOEPESegTrainer,
665
+ imgsz=32,
666
+ )
667
+
668
+ # prompt-free
669
+ # predict
670
+ model = YOLOE(WEIGHTS_DIR / "yoloe-11s-seg-pf.pt")
671
+ model.predict(SOURCE)
672
+ # val
673
+ model = YOLOE("yoloe-11s-seg.pt") # or select yoloe-m/l-seg.pt for different sizes
674
+ model.val(data="coco128-seg.yaml", imgsz=32)
675
+
676
+
677
+ def test_yolov10():
678
+ """Test YOLOv10 model training, validation, and prediction functionality."""
679
+ model = YOLO("yolov10n.yaml")
680
+ # train/val/predict
681
+ model.train(data="coco8.yaml", epochs=1, imgsz=32, close_mosaic=1, cache="disk")
682
+ model.val(data="coco8.yaml", imgsz=32)
683
+ model.predict(imgsz=32, save_txt=True, save_crop=True, augment=True)
684
+ model(SOURCE)
685
+
686
+
687
+ def test_multichannel():
688
+ """Test YOLO model multi-channel training, validation, and prediction functionality."""
689
+ model = YOLO("yolo11n.pt")
690
+ model.train(data="coco8-multispectral.yaml", epochs=1, imgsz=32, close_mosaic=1, cache="disk")
691
+ model.val(data="coco8-multispectral.yaml")
692
+ im = np.zeros((32, 32, 10), dtype=np.uint8)
693
+ model.predict(source=im, imgsz=32, save_txt=True, save_crop=True, augment=True)
694
+ model.export(format="onnx")