test_python.py 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615
  1. # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
  2. import contextlib
  3. import csv
  4. import urllib
  5. from copy import copy
  6. from pathlib import Path
  7. import cv2
  8. import numpy as np
  9. import pytest
  10. import torch
  11. import yaml
  12. from PIL import Image
  13. from tests import CFG, MODEL, SOURCE, SOURCES_LIST, TMP
  14. from ultralytics import RTDETR, YOLO
  15. from ultralytics.cfg import MODELS, TASK2DATA, TASKS
  16. from ultralytics.data.build import load_inference_source
  17. from ultralytics.utils import (
  18. ASSETS,
  19. DEFAULT_CFG,
  20. DEFAULT_CFG_PATH,
  21. LOGGER,
  22. ONLINE,
  23. ROOT,
  24. WEIGHTS_DIR,
  25. WINDOWS,
  26. checks,
  27. is_dir_writeable,
  28. is_github_action_running,
  29. )
  30. from ultralytics.utils.downloads import download
  31. from ultralytics.utils.torch_utils import TORCH_1_9
  32. IS_TMP_WRITEABLE = is_dir_writeable(TMP) # WARNING: must be run once tests start as TMP does not exist on tests/init
  33. def test_model_forward():
  34. """Test the forward pass of the YOLO model."""
  35. model = YOLO(CFG)
  36. model(source=None, imgsz=32, augment=True) # also test no source and augment
  37. def test_model_methods():
  38. """Test various methods and properties of the YOLO model to ensure correct functionality."""
  39. model = YOLO(MODEL)
  40. # Model methods
  41. model.info(verbose=True, detailed=True)
  42. model = model.reset_weights()
  43. model = model.load(MODEL)
  44. model.to("cpu")
  45. model.fuse()
  46. model.clear_callback("on_train_start")
  47. model.reset_callbacks()
  48. # Model properties
  49. _ = model.names
  50. _ = model.device
  51. _ = model.transforms
  52. _ = model.task_map
  53. def test_model_profile():
  54. """Test profiling of the YOLO model with `profile=True` to assess performance and resource usage."""
  55. from ultralytics.nn.tasks import DetectionModel
  56. model = DetectionModel() # build model
  57. im = torch.randn(1, 3, 64, 64) # requires min imgsz=64
  58. _ = model.predict(im, profile=True)
  59. @pytest.mark.skipif(not IS_TMP_WRITEABLE, reason="directory is not writeable")
  60. def test_predict_txt():
  61. """Tests YOLO predictions with file, directory, and pattern sources listed in a text file."""
  62. file = TMP / "sources_multi_row.txt"
  63. with open(file, "w") as f:
  64. for src in SOURCES_LIST:
  65. f.write(f"{src}\n")
  66. results = YOLO(MODEL)(source=file, imgsz=32)
  67. assert len(results) == 7 # 1 + 2 + 2 + 2 = 7 images
  68. @pytest.mark.skipif(True, reason="disabled for testing")
  69. @pytest.mark.skipif(not IS_TMP_WRITEABLE, reason="directory is not writeable")
  70. def test_predict_csv_multi_row():
  71. """Tests YOLO predictions with sources listed in multiple rows of a CSV file."""
  72. file = TMP / "sources_multi_row.csv"
  73. with open(file, "w", newline="") as f:
  74. writer = csv.writer(f)
  75. writer.writerow(["source"])
  76. writer.writerows([[src] for src in SOURCES_LIST])
  77. results = YOLO(MODEL)(source=file, imgsz=32)
  78. assert len(results) == 7 # 1 + 2 + 2 + 2 = 7 images
  79. @pytest.mark.skipif(True, reason="disabled for testing")
  80. @pytest.mark.skipif(not IS_TMP_WRITEABLE, reason="directory is not writeable")
  81. def test_predict_csv_single_row():
  82. """Tests YOLO predictions with sources listed in a single row of a CSV file."""
  83. file = TMP / "sources_single_row.csv"
  84. with open(file, "w", newline="") as f:
  85. writer = csv.writer(f)
  86. writer.writerow(SOURCES_LIST)
  87. results = YOLO(MODEL)(source=file, imgsz=32)
  88. assert len(results) == 7 # 1 + 2 + 2 + 2 = 7 images
  89. @pytest.mark.parametrize("model_name", MODELS)
  90. def test_predict_img(model_name):
  91. """Test YOLO model predictions on various image input types and sources, including online images."""
  92. model = YOLO(WEIGHTS_DIR / model_name)
  93. im = cv2.imread(str(SOURCE)) # uint8 numpy array
  94. assert len(model(source=Image.open(SOURCE), save=True, verbose=True, imgsz=32)) == 1 # PIL
  95. assert len(model(source=im, save=True, save_txt=True, imgsz=32)) == 1 # ndarray
  96. assert len(model(torch.rand((2, 3, 32, 32)), imgsz=32)) == 2 # batch-size 2 Tensor, FP32 0.0-1.0 RGB order
  97. assert len(model(source=[im, im], save=True, save_txt=True, imgsz=32)) == 2 # batch
  98. assert len(list(model(source=[im, im], save=True, stream=True, imgsz=32))) == 2 # stream
  99. assert len(model(torch.zeros(320, 640, 3).numpy().astype(np.uint8), imgsz=32)) == 1 # tensor to numpy
  100. batch = [
  101. str(SOURCE), # filename
  102. Path(SOURCE), # Path
  103. "https://github.com/ultralytics/assets/releases/download/v0.0.0/zidane.jpg" if ONLINE else SOURCE, # URI
  104. cv2.imread(str(SOURCE)), # OpenCV
  105. Image.open(SOURCE), # PIL
  106. np.zeros((320, 640, 3), dtype=np.uint8), # numpy
  107. ]
  108. assert len(model(batch, imgsz=32)) == len(batch) # multiple sources in a batch
  109. @pytest.mark.parametrize("model", MODELS)
  110. def test_predict_visualize(model):
  111. """Test model prediction methods with 'visualize=True' to generate and display prediction visualizations."""
  112. YOLO(WEIGHTS_DIR / model)(SOURCE, imgsz=32, visualize=True)
  113. def test_predict_grey_and_4ch():
  114. """Test YOLO prediction on SOURCE converted to greyscale and 4-channel images with various filenames."""
  115. im = Image.open(SOURCE)
  116. directory = TMP / "im4"
  117. directory.mkdir(parents=True, exist_ok=True)
  118. source_greyscale = directory / "greyscale.jpg"
  119. source_rgba = directory / "4ch.png"
  120. source_non_utf = directory / "non_UTF_测试文件_tést_image.jpg"
  121. source_spaces = directory / "image with spaces.jpg"
  122. im.convert("L").save(source_greyscale) # greyscale
  123. im.convert("RGBA").save(source_rgba) # 4-ch PNG with alpha
  124. im.save(source_non_utf) # non-UTF characters in filename
  125. im.save(source_spaces) # spaces in filename
  126. # Inference
  127. model = YOLO(MODEL)
  128. for f in source_rgba, source_greyscale, source_non_utf, source_spaces:
  129. for source in Image.open(f), cv2.imread(str(f)), f:
  130. results = model(source, save=True, verbose=True, imgsz=32)
  131. assert len(results) == 1 # verify that an image was run
  132. f.unlink() # cleanup
  133. @pytest.mark.slow
  134. @pytest.mark.skipif(not ONLINE, reason="environment is offline")
  135. @pytest.mark.skipif(is_github_action_running(), reason="No auth https://github.com/JuanBindez/pytubefix/issues/166")
  136. def test_youtube():
  137. """Test YOLO model on a YouTube video stream, handling potential network-related errors."""
  138. model = YOLO(MODEL)
  139. try:
  140. model.predict("https://youtu.be/G17sBkb38XQ", imgsz=96, save=True)
  141. # Handle internet connection errors and 'urllib.error.HTTPError: HTTP Error 429: Too Many Requests'
  142. except (urllib.error.HTTPError, ConnectionError) as e:
  143. LOGGER.warning(f"WARNING: YouTube Test Error: {e}")
  144. @pytest.mark.skipif(not ONLINE, reason="environment is offline")
  145. @pytest.mark.skipif(not IS_TMP_WRITEABLE, reason="directory is not writeable")
  146. def test_track_stream():
  147. """
  148. Tests streaming tracking on a short 10 frame video using ByteTrack tracker and different GMC methods.
  149. Note imgsz=160 required for tracking for higher confidence and better matches.
  150. """
  151. video_url = "https://github.com/ultralytics/assets/releases/download/v0.0.0/decelera_portrait_min.mov"
  152. model = YOLO(MODEL)
  153. model.track(video_url, imgsz=160, tracker="bytetrack.yaml")
  154. model.track(video_url, imgsz=160, tracker="botsort.yaml", save_frames=True) # test frame saving also
  155. # Test Global Motion Compensation (GMC) methods
  156. for gmc in "orb", "sift", "ecc":
  157. with open(ROOT / "cfg/trackers/botsort.yaml", encoding="utf-8") as f:
  158. data = yaml.safe_load(f)
  159. tracker = TMP / f"botsort-{gmc}.yaml"
  160. data["gmc_method"] = gmc
  161. with open(tracker, "w", encoding="utf-8") as f:
  162. yaml.safe_dump(data, f)
  163. model.track(video_url, imgsz=160, tracker=tracker)
  164. def test_val():
  165. """Test the validation mode of the YOLO model."""
  166. YOLO(MODEL).val(data="coco8.yaml", imgsz=32, save_hybrid=True)
  167. def test_train_scratch():
  168. """Test training the YOLO model from scratch using the provided configuration."""
  169. model = YOLO(CFG)
  170. model.train(data="coco8.yaml", epochs=2, imgsz=32, cache="disk", batch=-1, close_mosaic=1, name="model")
  171. model(SOURCE)
  172. def test_train_pretrained():
  173. """Test training of the YOLO model starting from a pre-trained checkpoint."""
  174. model = YOLO(WEIGHTS_DIR / "yolo11n-seg.pt")
  175. model.train(data="coco8-seg.yaml", epochs=1, imgsz=32, cache="ram", copy_paste=0.5, mixup=0.5, name=0)
  176. model(SOURCE)
  177. def test_all_model_yamls():
  178. """Test YOLO model creation for all available YAML configurations in the `cfg/models` directory."""
  179. for m in (ROOT / "cfg" / "models").rglob("*.yaml"):
  180. if "rtdetr" in m.name:
  181. if TORCH_1_9: # torch<=1.8 issue - TypeError: __init__() got an unexpected keyword argument 'batch_first'
  182. _ = RTDETR(m.name)(SOURCE, imgsz=640) # must be 640
  183. else:
  184. YOLO(m.name)
  185. @pytest.mark.skipif(WINDOWS, reason="Windows slow CI export bug https://github.com/ultralytics/ultralytics/pull/16003")
  186. def test_workflow():
  187. """Test the complete workflow including training, validation, prediction, and exporting."""
  188. model = YOLO(MODEL)
  189. model.train(data="coco8.yaml", epochs=1, imgsz=32, optimizer="SGD")
  190. model.val(imgsz=32)
  191. model.predict(SOURCE, imgsz=32)
  192. model.export(format="torchscript") # WARNING: Windows slow CI export bug
  193. def test_predict_callback_and_setup():
  194. """Test callback functionality during YOLO prediction setup and execution."""
  195. def on_predict_batch_end(predictor):
  196. """Callback function that handles operations at the end of a prediction batch."""
  197. path, im0s, _ = predictor.batch
  198. im0s = im0s if isinstance(im0s, list) else [im0s]
  199. bs = [predictor.dataset.bs for _ in range(len(path))]
  200. predictor.results = zip(predictor.results, im0s, bs) # results is List[batch_size]
  201. model = YOLO(MODEL)
  202. model.add_callback("on_predict_batch_end", on_predict_batch_end)
  203. dataset = load_inference_source(source=SOURCE)
  204. bs = dataset.bs # noqa access predictor properties
  205. results = model.predict(dataset, stream=True, imgsz=160) # source already setup
  206. for r, im0, bs in results:
  207. print("test_callback", im0.shape)
  208. print("test_callback", bs)
  209. boxes = r.boxes # Boxes object for bbox outputs
  210. print(boxes)
  211. @pytest.mark.parametrize("model", MODELS)
  212. def test_results(model):
  213. """Ensure YOLO model predictions can be processed and printed in various formats."""
  214. results = YOLO(WEIGHTS_DIR / model)([SOURCE, SOURCE], imgsz=160)
  215. for r in results:
  216. r = r.cpu().numpy()
  217. print(r, len(r), r.path) # print numpy attributes
  218. r = r.to(device="cpu", dtype=torch.float32)
  219. r.save_txt(txt_file=TMP / "runs/tests/label.txt", save_conf=True)
  220. r.save_crop(save_dir=TMP / "runs/tests/crops/")
  221. r.to_json(normalize=True)
  222. r.to_df(decimals=3)
  223. r.to_csv()
  224. r.to_xml()
  225. r.plot(pil=True)
  226. r.plot(conf=True, boxes=True)
  227. print(r, len(r), r.path) # print after methods
  228. def test_labels_and_crops():
  229. """Test output from prediction args for saving YOLO detection labels and crops; ensures accurate saving."""
  230. imgs = [SOURCE, ASSETS / "zidane.jpg"]
  231. results = YOLO(WEIGHTS_DIR / "yolo11n.pt")(imgs, imgsz=160, save_txt=True, save_crop=True)
  232. save_path = Path(results[0].save_dir)
  233. for r in results:
  234. im_name = Path(r.path).stem
  235. cls_idxs = r.boxes.cls.int().tolist()
  236. # Check correct detections
  237. assert cls_idxs == ([0, 7, 0, 0] if r.path.endswith("bus.jpg") else [0, 0, 0]) # bus.jpg and zidane.jpg classes
  238. # Check label path
  239. labels = save_path / f"labels/{im_name}.txt"
  240. assert labels.exists()
  241. # Check detections match label count
  242. assert len(r.boxes.data) == len([line for line in labels.read_text().splitlines() if line])
  243. # Check crops path and files
  244. crop_dirs = list((save_path / "crops").iterdir())
  245. crop_files = [f for p in crop_dirs for f in p.glob("*")]
  246. # Crop directories match detections
  247. assert all(r.names.get(c) in {d.name for d in crop_dirs} for c in cls_idxs)
  248. # Same number of crops as detections
  249. assert len([f for f in crop_files if im_name in f.name]) == len(r.boxes.data)
  250. @pytest.mark.skipif(not ONLINE, reason="environment is offline")
  251. def test_data_utils():
  252. """Test utility functions in ultralytics/data/utils.py, including dataset stats and auto-splitting."""
  253. from ultralytics.data.utils import HUBDatasetStats, autosplit
  254. from ultralytics.utils.downloads import zip_directory
  255. # from ultralytics.utils.files import WorkingDirectory
  256. # with WorkingDirectory(ROOT.parent / 'tests'):
  257. for task in TASKS:
  258. file = Path(TASK2DATA[task]).with_suffix(".zip") # i.e. coco8.zip
  259. download(f"https://github.com/ultralytics/hub/raw/main/example_datasets/{file}", unzip=False, dir=TMP)
  260. stats = HUBDatasetStats(TMP / file, task=task)
  261. stats.get_json(save=True)
  262. stats.process_images()
  263. autosplit(TMP / "coco8")
  264. zip_directory(TMP / "coco8/images/val") # zip
  265. @pytest.mark.skipif(not ONLINE, reason="environment is offline")
  266. def test_data_converter():
  267. """Test dataset conversion functions from COCO to YOLO format and class mappings."""
  268. from ultralytics.data.converter import coco80_to_coco91_class, convert_coco
  269. file = "instances_val2017.json"
  270. download(f"https://github.com/ultralytics/assets/releases/download/v0.0.0/{file}", dir=TMP)
  271. convert_coco(labels_dir=TMP, save_dir=TMP / "yolo_labels", use_segments=True, use_keypoints=False, cls91to80=True)
  272. coco80_to_coco91_class()
  273. def test_data_annotator():
  274. """Automatically annotate data using specified detection and segmentation models."""
  275. from ultralytics.data.annotator import auto_annotate
  276. auto_annotate(
  277. ASSETS,
  278. det_model=WEIGHTS_DIR / "yolo11n.pt",
  279. sam_model=WEIGHTS_DIR / "mobile_sam.pt",
  280. output_dir=TMP / "auto_annotate_labels",
  281. )
  282. def test_events():
  283. """Test event sending functionality."""
  284. from ultralytics.hub.utils import Events
  285. events = Events()
  286. events.enabled = True
  287. cfg = copy(DEFAULT_CFG) # does not require deepcopy
  288. cfg.mode = "test"
  289. events(cfg)
  290. def test_cfg_init():
  291. """Test configuration initialization utilities from the 'ultralytics.cfg' module."""
  292. from ultralytics.cfg import check_dict_alignment, copy_default_cfg, smart_value
  293. with contextlib.suppress(SyntaxError):
  294. check_dict_alignment({"a": 1}, {"b": 2})
  295. copy_default_cfg()
  296. (Path.cwd() / DEFAULT_CFG_PATH.name.replace(".yaml", "_copy.yaml")).unlink(missing_ok=False)
  297. [smart_value(x) for x in ["none", "true", "false"]]
  298. def test_utils_init():
  299. """Test initialization utilities in the Ultralytics library."""
  300. from ultralytics.utils import get_git_branch, get_git_origin_url, get_ubuntu_version, is_github_action_running
  301. get_ubuntu_version()
  302. is_github_action_running()
  303. get_git_origin_url()
  304. get_git_branch()
  305. def test_utils_checks():
  306. """Test various utility checks for filenames, git status, requirements, image sizes, and versions."""
  307. checks.check_yolov5u_filename("yolov5n.pt")
  308. checks.git_describe(ROOT)
  309. checks.check_requirements() # check requirements.txt
  310. checks.check_imgsz([600, 600], max_dim=1)
  311. checks.check_imshow(warn=True)
  312. checks.check_version("ultralytics", "8.0.0")
  313. checks.print_args()
  314. @pytest.mark.skipif(WINDOWS, reason="Windows profiling is extremely slow (cause unknown)")
  315. def test_utils_benchmarks():
  316. """Benchmark model performance using 'ProfileModels' from 'ultralytics.utils.benchmarks'."""
  317. from ultralytics.utils.benchmarks import ProfileModels
  318. ProfileModels(["yolo11n.yaml"], imgsz=32, min_time=1, num_timed_runs=3, num_warmup_runs=1).profile()
  319. def test_utils_torchutils():
  320. """Test Torch utility functions including profiling and FLOP calculations."""
  321. from ultralytics.nn.modules.conv import Conv
  322. from ultralytics.utils.torch_utils import get_flops_with_torch_profiler, profile, time_sync
  323. x = torch.randn(1, 64, 20, 20)
  324. m = Conv(64, 64, k=1, s=2)
  325. profile(x, [m], n=3)
  326. get_flops_with_torch_profiler(m)
  327. time_sync()
  328. def test_utils_ops():
  329. """Test utility operations functions for coordinate transformation and normalization."""
  330. from ultralytics.utils.ops import (
  331. ltwh2xywh,
  332. ltwh2xyxy,
  333. make_divisible,
  334. xywh2ltwh,
  335. xywh2xyxy,
  336. xywhn2xyxy,
  337. xywhr2xyxyxyxy,
  338. xyxy2ltwh,
  339. xyxy2xywh,
  340. xyxy2xywhn,
  341. xyxyxyxy2xywhr,
  342. )
  343. make_divisible(17, torch.tensor([8]))
  344. boxes = torch.rand(10, 4) # xywh
  345. torch.allclose(boxes, xyxy2xywh(xywh2xyxy(boxes)))
  346. torch.allclose(boxes, xyxy2xywhn(xywhn2xyxy(boxes)))
  347. torch.allclose(boxes, ltwh2xywh(xywh2ltwh(boxes)))
  348. torch.allclose(boxes, xyxy2ltwh(ltwh2xyxy(boxes)))
  349. boxes = torch.rand(10, 5) # xywhr for OBB
  350. boxes[:, 4] = torch.randn(10) * 30
  351. torch.allclose(boxes, xyxyxyxy2xywhr(xywhr2xyxyxyxy(boxes)), rtol=1e-3)
  352. def test_utils_files():
  353. """Test file handling utilities including file age, date, and paths with spaces."""
  354. from ultralytics.utils.files import file_age, file_date, get_latest_run, spaces_in_path
  355. file_age(SOURCE)
  356. file_date(SOURCE)
  357. get_latest_run(ROOT / "runs")
  358. path = TMP / "path/with spaces"
  359. path.mkdir(parents=True, exist_ok=True)
  360. with spaces_in_path(path) as new_path:
  361. print(new_path)
  362. @pytest.mark.slow
  363. def test_utils_patches_torch_save():
  364. """Test torch_save backoff when _torch_save raises RuntimeError to ensure robustness."""
  365. from unittest.mock import MagicMock, patch
  366. from ultralytics.utils.patches import torch_save
  367. mock = MagicMock(side_effect=RuntimeError)
  368. with patch("ultralytics.utils.patches._torch_save", new=mock):
  369. with pytest.raises(RuntimeError):
  370. torch_save(torch.zeros(1), TMP / "test.pt")
  371. assert mock.call_count == 4, "torch_save was not attempted the expected number of times"
  372. def test_nn_modules_conv():
  373. """Test Convolutional Neural Network modules including CBAM, Conv2, and ConvTranspose."""
  374. from ultralytics.nn.modules.conv import CBAM, Conv2, ConvTranspose, DWConvTranspose2d, Focus
  375. c1, c2 = 8, 16 # input and output channels
  376. x = torch.zeros(4, c1, 10, 10) # BCHW
  377. # Run all modules not otherwise covered in tests
  378. DWConvTranspose2d(c1, c2)(x)
  379. ConvTranspose(c1, c2)(x)
  380. Focus(c1, c2)(x)
  381. CBAM(c1)(x)
  382. # Fuse ops
  383. m = Conv2(c1, c2)
  384. m.fuse_convs()
  385. m(x)
  386. def test_nn_modules_block():
  387. """Test various blocks in neural network modules including C1, C3TR, BottleneckCSP, C3Ghost, and C3x."""
  388. from ultralytics.nn.modules.block import C1, C3TR, BottleneckCSP, C3Ghost, C3x
  389. c1, c2 = 8, 16 # input and output channels
  390. x = torch.zeros(4, c1, 10, 10) # BCHW
  391. # Run all modules not otherwise covered in tests
  392. C1(c1, c2)(x)
  393. C3x(c1, c2)(x)
  394. C3TR(c1, c2)(x)
  395. C3Ghost(c1, c2)(x)
  396. BottleneckCSP(c1, c2)(x)
  397. @pytest.mark.skipif(not ONLINE, reason="environment is offline")
  398. def test_hub():
  399. """Test Ultralytics HUB functionalities (e.g. export formats, logout)."""
  400. from ultralytics.hub import export_fmts_hub, logout
  401. from ultralytics.hub.utils import smart_request
  402. export_fmts_hub()
  403. logout()
  404. smart_request("GET", "https://github.com", progress=True)
  405. @pytest.fixture
  406. def image():
  407. """Load and return an image from a predefined source using OpenCV."""
  408. return cv2.imread(str(SOURCE))
  409. @pytest.mark.parametrize(
  410. "auto_augment, erasing, force_color_jitter",
  411. [
  412. (None, 0.0, False),
  413. ("randaugment", 0.5, True),
  414. ("augmix", 0.2, False),
  415. ("autoaugment", 0.0, True),
  416. ],
  417. )
  418. def test_classify_transforms_train(image, auto_augment, erasing, force_color_jitter):
  419. """Tests classification transforms during training with various augmentations to ensure proper functionality."""
  420. from ultralytics.data.augment import classify_augmentations
  421. transform = classify_augmentations(
  422. size=224,
  423. mean=(0.5, 0.5, 0.5),
  424. std=(0.5, 0.5, 0.5),
  425. scale=(0.08, 1.0),
  426. ratio=(3.0 / 4.0, 4.0 / 3.0),
  427. hflip=0.5,
  428. vflip=0.5,
  429. auto_augment=auto_augment,
  430. hsv_h=0.015,
  431. hsv_s=0.4,
  432. hsv_v=0.4,
  433. force_color_jitter=force_color_jitter,
  434. erasing=erasing,
  435. )
  436. transformed_image = transform(Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)))
  437. assert transformed_image.shape == (3, 224, 224)
  438. assert torch.is_tensor(transformed_image)
  439. assert transformed_image.dtype == torch.float32
  440. @pytest.mark.slow
  441. @pytest.mark.skipif(not ONLINE, reason="environment is offline")
  442. def test_model_tune():
  443. """Tune YOLO model for performance improvement."""
  444. YOLO("yolo11n-pose.pt").tune(data="coco8-pose.yaml", plots=False, imgsz=32, epochs=1, iterations=2, device="cpu")
  445. YOLO("yolo11n-cls.pt").tune(data="imagenet10", plots=False, imgsz=32, epochs=1, iterations=2, device="cpu")
  446. def test_model_embeddings():
  447. """Test YOLO model embeddings."""
  448. model_detect = YOLO(MODEL)
  449. model_segment = YOLO(WEIGHTS_DIR / "yolo11n-seg.pt")
  450. for batch in [SOURCE], [SOURCE, SOURCE]: # test batch size 1 and 2
  451. assert len(model_detect.embed(source=batch, imgsz=32)) == len(batch)
  452. assert len(model_segment.embed(source=batch, imgsz=32)) == len(batch)
  453. @pytest.mark.skipif(checks.IS_PYTHON_3_12, reason="YOLOWorld with CLIP is not supported in Python 3.12")
  454. def test_yolo_world():
  455. """Tests YOLO world models with CLIP support, including detection and training scenarios."""
  456. model = YOLO(WEIGHTS_DIR / "yolov8s-world.pt") # no YOLO11n-world model yet
  457. model.set_classes(["tree", "window"])
  458. model(SOURCE, conf=0.01)
  459. model = YOLO(WEIGHTS_DIR / "yolov8s-worldv2.pt") # no YOLO11n-world model yet
  460. # Training from a pretrained model. Eval is included at the final stage of training.
  461. # Use dota8.yaml which has fewer categories to reduce the inference time of CLIP model
  462. model.train(
  463. data="dota8.yaml",
  464. epochs=1,
  465. imgsz=32,
  466. cache="disk",
  467. close_mosaic=1,
  468. )
  469. # test WorWorldTrainerFromScratch
  470. from ultralytics.models.yolo.world.train_world import WorldTrainerFromScratch
  471. model = YOLO("yolov8s-worldv2.yaml") # no YOLO11n-world model yet
  472. model.train(
  473. data={"train": {"yolo_data": ["dota8.yaml"]}, "val": {"yolo_data": ["dota8.yaml"]}},
  474. epochs=1,
  475. imgsz=32,
  476. cache="disk",
  477. close_mosaic=1,
  478. trainer=WorldTrainerFromScratch,
  479. )
  480. def test_yolov10():
  481. """Test YOLOv10 model training, validation, and prediction steps with minimal configurations."""
  482. model = YOLO("yolov10n.yaml")
  483. # train/val/predict
  484. model.train(data="coco8.yaml", epochs=1, imgsz=32, close_mosaic=1, cache="disk")
  485. model.val(data="coco8.yaml", imgsz=32)
  486. model.predict(imgsz=32, save_txt=True, save_crop=True, augment=True)
  487. model(SOURCE)