diff --git a/docs/en/reference/utils/__init__.md b/docs/en/reference/utils/__init__.md
index 43b245f1..fc4c9a79 100644
--- a/docs/en/reference/utils/__init__.md
+++ b/docs/en/reference/utils/__init__.md
@@ -31,6 +31,10 @@ keywords: Ultralytics, Utils, utilitarian functions, colorstr, yaml_save, set_lo
+## ::: ultralytics.utils.Retry
+
+
+
## ::: ultralytics.utils.SettingsManager
diff --git a/tests/conftest.py b/tests/conftest.py
index 59955bd1..06c0d140 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -5,7 +5,7 @@ from pathlib import Path
import pytest
-TMP = Path(__file__).resolve().parent / 'tmp' # temp directory for test files
+TMP = Path(__file__).resolve().parent / "tmp" # temp directory for test files
def pytest_addoption(parser):
@@ -15,7 +15,7 @@ def pytest_addoption(parser):
Args:
parser (pytest.config.Parser): The pytest parser object.
"""
- parser.addoption('--slow', action='store_true', default=False, help='Run slow tests')
+ parser.addoption("--slow", action="store_true", default=False, help="Run slow tests")
def pytest_configure(config):
@@ -25,7 +25,7 @@ def pytest_configure(config):
Args:
config (pytest.config.Config): The pytest config object.
"""
- config.addinivalue_line('markers', 'slow: mark test as slow to run')
+ config.addinivalue_line("markers", "slow: mark test as slow to run")
def pytest_runtest_setup(item):
@@ -35,8 +35,8 @@ def pytest_runtest_setup(item):
Args:
item (pytest.Item): The test item object.
"""
- if 'slow' in item.keywords and not item.config.getoption('--slow'):
- pytest.skip('skip slow tests unless --slow is set')
+ if "slow" in item.keywords and not item.config.getoption("--slow"):
+ pytest.skip("skip slow tests unless --slow is set")
def pytest_collection_modifyitems(config, items):
@@ -47,9 +47,9 @@ def pytest_collection_modifyitems(config, items):
config (pytest.config.Config): The pytest config object.
items (list): List of test items to be executed.
"""
- if not config.getoption('--slow'):
+ if not config.getoption("--slow"):
# Remove the item entirely from the list of test items if it's marked as 'slow'
- items[:] = [item for item in items if 'slow' not in item.keywords]
+ items[:] = [item for item in items if "slow" not in item.keywords]
def pytest_sessionstart(session):
@@ -84,11 +84,11 @@ def pytest_terminal_summary(terminalreporter, exitstatus, config):
from ultralytics.utils import WEIGHTS_DIR
# Remove files
- models = [path for x in ['*.onnx', '*.torchscript'] for path in WEIGHTS_DIR.rglob(x)]
- for file in ['bus.jpg', 'yolov8n.onnx', 'yolov8n.torchscript'] + models:
+ models = [path for x in ["*.onnx", "*.torchscript"] for path in WEIGHTS_DIR.rglob(x)]
+ for file in ["bus.jpg", "yolov8n.onnx", "yolov8n.torchscript"] + models:
Path(file).unlink(missing_ok=True)
# Remove directories
- models = [path for x in ['*.mlpackage', '*_openvino_model'] for path in WEIGHTS_DIR.rglob(x)]
- for directory in [TMP.parents[1] / '.pytest_cache', TMP] + models:
+ models = [path for x in ["*.mlpackage", "*_openvino_model"] for path in WEIGHTS_DIR.rglob(x)]
+ for directory in [TMP.parents[1] / ".pytest_cache", TMP] + models:
shutil.rmtree(directory, ignore_errors=True)
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 994ce5a2..ea07d764 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -10,17 +10,19 @@ from ultralytics.utils.checks import cuda_device_count, cuda_is_available
CUDA_IS_AVAILABLE = cuda_is_available()
CUDA_DEVICE_COUNT = cuda_device_count()
TASK_ARGS = [
- ('detect', 'yolov8n', 'coco8.yaml'),
- ('segment', 'yolov8n-seg', 'coco8-seg.yaml'),
- ('classify', 'yolov8n-cls', 'imagenet10'),
- ('pose', 'yolov8n-pose', 'coco8-pose.yaml'),
- ('obb', 'yolov8n-obb', 'dota8.yaml'), ] # (task, model, data)
+ ("detect", "yolov8n", "coco8.yaml"),
+ ("segment", "yolov8n-seg", "coco8-seg.yaml"),
+ ("classify", "yolov8n-cls", "imagenet10"),
+ ("pose", "yolov8n-pose", "coco8-pose.yaml"),
+ ("obb", "yolov8n-obb", "dota8.yaml"),
+] # (task, model, data)
EXPORT_ARGS = [
- ('yolov8n', 'torchscript'),
- ('yolov8n-seg', 'torchscript'),
- ('yolov8n-cls', 'torchscript'),
- ('yolov8n-pose', 'torchscript'),
- ('yolov8n-obb', 'torchscript'), ] # (model, format)
+ ("yolov8n", "torchscript"),
+ ("yolov8n-seg", "torchscript"),
+ ("yolov8n-cls", "torchscript"),
+ ("yolov8n-pose", "torchscript"),
+ ("yolov8n-obb", "torchscript"),
+] # (model, format)
def run(cmd):
@@ -30,50 +32,50 @@ def run(cmd):
def test_special_modes():
"""Test various special command modes of YOLO."""
- run('yolo help')
- run('yolo checks')
- run('yolo version')
- run('yolo settings reset')
- run('yolo cfg')
+ run("yolo help")
+ run("yolo checks")
+ run("yolo version")
+ run("yolo settings reset")
+ run("yolo cfg")
-@pytest.mark.parametrize('task,model,data', TASK_ARGS)
+@pytest.mark.parametrize("task,model,data", TASK_ARGS)
def test_train(task, model, data):
"""Test YOLO training for a given task, model, and data."""
- run(f'yolo train {task} model={model}.yaml data={data} imgsz=32 epochs=1 cache=disk')
+ run(f"yolo train {task} model={model}.yaml data={data} imgsz=32 epochs=1 cache=disk")
-@pytest.mark.parametrize('task,model,data', TASK_ARGS)
+@pytest.mark.parametrize("task,model,data", TASK_ARGS)
def test_val(task, model, data):
"""Test YOLO validation for a given task, model, and data."""
- run(f'yolo val {task} model={WEIGHTS_DIR / model}.pt data={data} imgsz=32 save_txt save_json')
+ run(f"yolo val {task} model={WEIGHTS_DIR / model}.pt data={data} imgsz=32 save_txt save_json")
-@pytest.mark.parametrize('task,model,data', TASK_ARGS)
+@pytest.mark.parametrize("task,model,data", TASK_ARGS)
def test_predict(task, model, data):
"""Test YOLO prediction on sample assets for a given task and model."""
- run(f'yolo predict model={WEIGHTS_DIR / model}.pt source={ASSETS} imgsz=32 save save_crop save_txt')
+ run(f"yolo predict model={WEIGHTS_DIR / model}.pt source={ASSETS} imgsz=32 save save_crop save_txt")
-@pytest.mark.parametrize('model,format', EXPORT_ARGS)
+@pytest.mark.parametrize("model,format", EXPORT_ARGS)
def test_export(model, format):
"""Test exporting a YOLO model to different formats."""
- run(f'yolo export model={WEIGHTS_DIR / model}.pt format={format} imgsz=32')
+ run(f"yolo export model={WEIGHTS_DIR / model}.pt format={format} imgsz=32")
-def test_rtdetr(task='detect', model='yolov8n-rtdetr.yaml', data='coco8.yaml'):
+def test_rtdetr(task="detect", model="yolov8n-rtdetr.yaml", data="coco8.yaml"):
"""Test the RTDETR functionality with the Ultralytics framework."""
# Warning: MUST use imgsz=640
- run(f'yolo train {task} model={model} data={data} --imgsz= 640 epochs =1, cache = disk') # add coma, spaces to args
+ run(f"yolo train {task} model={model} data={data} --imgsz= 640 epochs =1, cache = disk") # add coma, spaces to args
run(f"yolo predict {task} model={model} source={ASSETS / 'bus.jpg'} imgsz=640 save save_crop save_txt")
-def test_fastsam(task='segment', model=WEIGHTS_DIR / 'FastSAM-s.pt', data='coco8-seg.yaml'):
+def test_fastsam(task="segment", model=WEIGHTS_DIR / "FastSAM-s.pt", data="coco8-seg.yaml"):
"""Test FastSAM segmentation functionality within Ultralytics."""
- source = ASSETS / 'bus.jpg'
+ source = ASSETS / "bus.jpg"
- run(f'yolo segment val {task} model={model} data={data} imgsz=32')
- run(f'yolo segment predict model={model} source={source} imgsz=32 save save_crop save_txt')
+ run(f"yolo segment val {task} model={model} data={data} imgsz=32")
+ run(f"yolo segment predict model={model} source={source} imgsz=32 save save_crop save_txt")
from ultralytics import FastSAM
from ultralytics.models.fastsam import FastSAMPrompt
@@ -83,26 +85,26 @@ def test_fastsam(task='segment', model=WEIGHTS_DIR / 'FastSAM-s.pt', data='coco8
sam_model = FastSAM(model) # or FastSAM-x.pt
# Run inference on an image
- everything_results = sam_model(source, device='cpu', retina_masks=True, imgsz=1024, conf=0.4, iou=0.9)
+ everything_results = sam_model(source, device="cpu", retina_masks=True, imgsz=1024, conf=0.4, iou=0.9)
# Remove small regions
new_masks, _ = Predictor.remove_small_regions(everything_results[0].masks.data, min_area=20)
# Everything prompt
- prompt_process = FastSAMPrompt(source, everything_results, device='cpu')
+ prompt_process = FastSAMPrompt(source, everything_results, device="cpu")
ann = prompt_process.everything_prompt()
# Bbox default shape [0,0,0,0] -> [x1,y1,x2,y2]
ann = prompt_process.box_prompt(bbox=[200, 200, 300, 300])
# Text prompt
- ann = prompt_process.text_prompt(text='a photo of a dog')
+ ann = prompt_process.text_prompt(text="a photo of a dog")
# Point prompt
# Points default [[0,0]] [[x1,y1],[x2,y2]]
# Point_label default [0] [1,0] 0:background, 1:foreground
ann = prompt_process.point_prompt(points=[[200, 200]], pointlabel=[1])
- prompt_process.plot(annotations=ann, output='./')
+ prompt_process.plot(annotations=ann, output="./")
def test_mobilesam():
@@ -110,10 +112,10 @@ def test_mobilesam():
from ultralytics import SAM
# Load the model
- model = SAM(WEIGHTS_DIR / 'mobile_sam.pt')
+ model = SAM(WEIGHTS_DIR / "mobile_sam.pt")
# Source
- source = ASSETS / 'zidane.jpg'
+ source = ASSETS / "zidane.jpg"
# Predict a segment based on a point prompt
model.predict(source, points=[900, 370], labels=[1])
@@ -127,10 +129,10 @@ def test_mobilesam():
# Slow Tests -----------------------------------------------------------------------------------------------------------
@pytest.mark.slow
-@pytest.mark.parametrize('task,model,data', TASK_ARGS)
-@pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason='CUDA is not available')
-@pytest.mark.skipif(CUDA_DEVICE_COUNT < 2, reason='DDP is not available')
+@pytest.mark.parametrize("task,model,data", TASK_ARGS)
+@pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason="CUDA is not available")
+@pytest.mark.skipif(CUDA_DEVICE_COUNT < 2, reason="DDP is not available")
def test_train_gpu(task, model, data):
"""Test YOLO training on GPU(s) for various tasks and models."""
- run(f'yolo train {task} model={model}.yaml data={data} imgsz=32 epochs=1 device=0') # single GPU
- run(f'yolo train {task} model={model}.pt data={data} imgsz=32 epochs=1 device=0,1') # multi GPU
+ run(f"yolo train {task} model={model}.yaml data={data} imgsz=32 epochs=1 device=0") # single GPU
+ run(f"yolo train {task} model={model}.pt data={data} imgsz=32 epochs=1 device=0,1") # multi GPU
diff --git a/tests/test_cuda.py b/tests/test_cuda.py
index eb49a08f..ed08f700 100644
--- a/tests/test_cuda.py
+++ b/tests/test_cuda.py
@@ -9,9 +9,9 @@ from ultralytics.utils import ASSETS, WEIGHTS_DIR, checks
CUDA_IS_AVAILABLE = checks.cuda_is_available()
CUDA_DEVICE_COUNT = checks.cuda_device_count()
-MODEL = WEIGHTS_DIR / 'path with spaces' / 'yolov8n.pt' # test spaces in path
-DATA = 'coco8.yaml'
-BUS = ASSETS / 'bus.jpg'
+MODEL = WEIGHTS_DIR / "path with spaces" / "yolov8n.pt" # test spaces in path
+DATA = "coco8.yaml"
+BUS = ASSETS / "bus.jpg"
def test_checks():
@@ -20,7 +20,7 @@ def test_checks():
assert torch.cuda.device_count() == CUDA_DEVICE_COUNT
-@pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason='CUDA is not available')
+@pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason="CUDA is not available")
def test_train():
"""Test model training on a minimal dataset."""
device = 0 if CUDA_DEVICE_COUNT == 1 else [0, 1]
@@ -28,32 +28,32 @@ def test_train():
@pytest.mark.slow
-@pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason='CUDA is not available')
+@pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason="CUDA is not available")
def test_predict_multiple_devices():
"""Validate model prediction on multiple devices."""
- model = YOLO('yolov8n.pt')
+ model = YOLO("yolov8n.pt")
model = model.cpu()
- assert str(model.device) == 'cpu'
+ assert str(model.device) == "cpu"
_ = model(BUS) # CPU inference
- assert str(model.device) == 'cpu'
+ assert str(model.device) == "cpu"
- model = model.to('cuda:0')
- assert str(model.device) == 'cuda:0'
+ model = model.to("cuda:0")
+ assert str(model.device) == "cuda:0"
_ = model(BUS) # CUDA inference
- assert str(model.device) == 'cuda:0'
+ assert str(model.device) == "cuda:0"
model = model.cpu()
- assert str(model.device) == 'cpu'
+ assert str(model.device) == "cpu"
_ = model(BUS) # CPU inference
- assert str(model.device) == 'cpu'
+ assert str(model.device) == "cpu"
model = model.cuda()
- assert str(model.device) == 'cuda:0'
+ assert str(model.device) == "cuda:0"
_ = model(BUS) # CUDA inference
- assert str(model.device) == 'cuda:0'
+ assert str(model.device) == "cuda:0"
-@pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason='CUDA is not available')
+@pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason="CUDA is not available")
def test_autobatch():
"""Check batch size for YOLO model using autobatch."""
from ultralytics.utils.autobatch import check_train_batch_size
@@ -62,24 +62,24 @@ def test_autobatch():
@pytest.mark.slow
-@pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason='CUDA is not available')
+@pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason="CUDA is not available")
def test_utils_benchmarks():
"""Profile YOLO models for performance benchmarks."""
from ultralytics.utils.benchmarks import ProfileModels
# Pre-export a dynamic engine model to use dynamic inference
- YOLO(MODEL).export(format='engine', imgsz=32, dynamic=True, batch=1)
+ YOLO(MODEL).export(format="engine", imgsz=32, dynamic=True, batch=1)
ProfileModels([MODEL], imgsz=32, half=False, min_time=1, num_timed_runs=3, num_warmup_runs=1).profile()
-@pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason='CUDA is not available')
+@pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason="CUDA is not available")
def test_predict_sam():
"""Test SAM model prediction with various prompts."""
from ultralytics import SAM
from ultralytics.models.sam import Predictor as SAMPredictor
# Load a model
- model = SAM(WEIGHTS_DIR / 'sam_b.pt')
+ model = SAM(WEIGHTS_DIR / "sam_b.pt")
# Display model information (optional)
model.info()
@@ -91,14 +91,14 @@ def test_predict_sam():
model(BUS, bboxes=[439, 437, 524, 709], device=0)
# Run inference with points prompt
- model(ASSETS / 'zidane.jpg', points=[900, 370], labels=[1], device=0)
+ model(ASSETS / "zidane.jpg", points=[900, 370], labels=[1], device=0)
# Create SAMPredictor
- overrides = dict(conf=0.25, task='segment', mode='predict', imgsz=1024, model=WEIGHTS_DIR / 'mobile_sam.pt')
+ overrides = dict(conf=0.25, task="segment", mode="predict", imgsz=1024, model=WEIGHTS_DIR / "mobile_sam.pt")
predictor = SAMPredictor(overrides=overrides)
# Set image
- predictor.set_image(ASSETS / 'zidane.jpg') # set with image file
+ predictor.set_image(ASSETS / "zidane.jpg") # set with image file
# predictor(bboxes=[439, 437, 524, 709])
# predictor(points=[900, 370], labels=[1])
diff --git a/tests/test_engine.py b/tests/test_engine.py
index ce328efc..bd528a8d 100644
--- a/tests/test_engine.py
+++ b/tests/test_engine.py
@@ -6,123 +6,123 @@ from ultralytics.engine.exporter import Exporter
from ultralytics.models.yolo import classify, detect, segment
from ultralytics.utils import ASSETS, DEFAULT_CFG, WEIGHTS_DIR
-CFG_DET = 'yolov8n.yaml'
-CFG_SEG = 'yolov8n-seg.yaml'
-CFG_CLS = 'yolov8n-cls.yaml' # or 'squeezenet1_0'
+CFG_DET = "yolov8n.yaml"
+CFG_SEG = "yolov8n-seg.yaml"
+CFG_CLS = "yolov8n-cls.yaml" # or 'squeezenet1_0'
CFG = get_cfg(DEFAULT_CFG)
-MODEL = WEIGHTS_DIR / 'yolov8n'
+MODEL = WEIGHTS_DIR / "yolov8n"
def test_func(*args): # noqa
"""Test function callback."""
- print('callback test passed')
+ print("callback test passed")
def test_export():
"""Test model exporting functionality."""
exporter = Exporter()
- exporter.add_callback('on_export_start', test_func)
- assert test_func in exporter.callbacks['on_export_start'], 'callback test failed'
+ exporter.add_callback("on_export_start", test_func)
+ assert test_func in exporter.callbacks["on_export_start"], "callback test failed"
f = exporter(model=YOLO(CFG_DET).model)
YOLO(f)(ASSETS) # exported model inference
def test_detect():
"""Test object detection functionality."""
- overrides = {'data': 'coco8.yaml', 'model': CFG_DET, 'imgsz': 32, 'epochs': 1, 'save': False}
- CFG.data = 'coco8.yaml'
+ overrides = {"data": "coco8.yaml", "model": CFG_DET, "imgsz": 32, "epochs": 1, "save": False}
+ CFG.data = "coco8.yaml"
CFG.imgsz = 32
# Trainer
trainer = detect.DetectionTrainer(overrides=overrides)
- trainer.add_callback('on_train_start', test_func)
- assert test_func in trainer.callbacks['on_train_start'], 'callback test failed'
+ trainer.add_callback("on_train_start", test_func)
+ assert test_func in trainer.callbacks["on_train_start"], "callback test failed"
trainer.train()
# Validator
val = detect.DetectionValidator(args=CFG)
- val.add_callback('on_val_start', test_func)
- assert test_func in val.callbacks['on_val_start'], 'callback test failed'
+ val.add_callback("on_val_start", test_func)
+ assert test_func in val.callbacks["on_val_start"], "callback test failed"
val(model=trainer.best) # validate best.pt
# Predictor
- pred = detect.DetectionPredictor(overrides={'imgsz': [64, 64]})
- pred.add_callback('on_predict_start', test_func)
- assert test_func in pred.callbacks['on_predict_start'], 'callback test failed'
- result = pred(source=ASSETS, model=f'{MODEL}.pt')
- assert len(result), 'predictor test failed'
+ pred = detect.DetectionPredictor(overrides={"imgsz": [64, 64]})
+ pred.add_callback("on_predict_start", test_func)
+ assert test_func in pred.callbacks["on_predict_start"], "callback test failed"
+ result = pred(source=ASSETS, model=f"{MODEL}.pt")
+ assert len(result), "predictor test failed"
- overrides['resume'] = trainer.last
+ overrides["resume"] = trainer.last
trainer = detect.DetectionTrainer(overrides=overrides)
try:
trainer.train()
except Exception as e:
- print(f'Expected exception caught: {e}')
+ print(f"Expected exception caught: {e}")
return
- Exception('Resume test failed!')
+ Exception("Resume test failed!")
def test_segment():
"""Test image segmentation functionality."""
- overrides = {'data': 'coco8-seg.yaml', 'model': CFG_SEG, 'imgsz': 32, 'epochs': 1, 'save': False}
- CFG.data = 'coco8-seg.yaml'
+ overrides = {"data": "coco8-seg.yaml", "model": CFG_SEG, "imgsz": 32, "epochs": 1, "save": False}
+ CFG.data = "coco8-seg.yaml"
CFG.imgsz = 32
# YOLO(CFG_SEG).train(**overrides) # works
# Trainer
trainer = segment.SegmentationTrainer(overrides=overrides)
- trainer.add_callback('on_train_start', test_func)
- assert test_func in trainer.callbacks['on_train_start'], 'callback test failed'
+ trainer.add_callback("on_train_start", test_func)
+ assert test_func in trainer.callbacks["on_train_start"], "callback test failed"
trainer.train()
# Validator
val = segment.SegmentationValidator(args=CFG)
- val.add_callback('on_val_start', test_func)
- assert test_func in val.callbacks['on_val_start'], 'callback test failed'
+ val.add_callback("on_val_start", test_func)
+ assert test_func in val.callbacks["on_val_start"], "callback test failed"
val(model=trainer.best) # validate best.pt
# Predictor
- pred = segment.SegmentationPredictor(overrides={'imgsz': [64, 64]})
- pred.add_callback('on_predict_start', test_func)
- assert test_func in pred.callbacks['on_predict_start'], 'callback test failed'
- result = pred(source=ASSETS, model=f'{MODEL}-seg.pt')
- assert len(result), 'predictor test failed'
+ pred = segment.SegmentationPredictor(overrides={"imgsz": [64, 64]})
+ pred.add_callback("on_predict_start", test_func)
+ assert test_func in pred.callbacks["on_predict_start"], "callback test failed"
+ result = pred(source=ASSETS, model=f"{MODEL}-seg.pt")
+ assert len(result), "predictor test failed"
# Test resume
- overrides['resume'] = trainer.last
+ overrides["resume"] = trainer.last
trainer = segment.SegmentationTrainer(overrides=overrides)
try:
trainer.train()
except Exception as e:
- print(f'Expected exception caught: {e}')
+ print(f"Expected exception caught: {e}")
return
- Exception('Resume test failed!')
+ Exception("Resume test failed!")
def test_classify():
"""Test image classification functionality."""
- overrides = {'data': 'imagenet10', 'model': CFG_CLS, 'imgsz': 32, 'epochs': 1, 'save': False}
- CFG.data = 'imagenet10'
+ overrides = {"data": "imagenet10", "model": CFG_CLS, "imgsz": 32, "epochs": 1, "save": False}
+ CFG.data = "imagenet10"
CFG.imgsz = 32
# YOLO(CFG_SEG).train(**overrides) # works
# Trainer
trainer = classify.ClassificationTrainer(overrides=overrides)
- trainer.add_callback('on_train_start', test_func)
- assert test_func in trainer.callbacks['on_train_start'], 'callback test failed'
+ trainer.add_callback("on_train_start", test_func)
+ assert test_func in trainer.callbacks["on_train_start"], "callback test failed"
trainer.train()
# Validator
val = classify.ClassificationValidator(args=CFG)
- val.add_callback('on_val_start', test_func)
- assert test_func in val.callbacks['on_val_start'], 'callback test failed'
+ val.add_callback("on_val_start", test_func)
+ assert test_func in val.callbacks["on_val_start"], "callback test failed"
val(model=trainer.best)
# Predictor
- pred = classify.ClassificationPredictor(overrides={'imgsz': [64, 64]})
- pred.add_callback('on_predict_start', test_func)
- assert test_func in pred.callbacks['on_predict_start'], 'callback test failed'
+ pred = classify.ClassificationPredictor(overrides={"imgsz": [64, 64]})
+ pred.add_callback("on_predict_start", test_func)
+ assert test_func in pred.callbacks["on_predict_start"], "callback test failed"
result = pred(source=ASSETS, model=trainer.best)
- assert len(result), 'predictor test failed'
+ assert len(result), "predictor test failed"
diff --git a/tests/test_explorer.py b/tests/test_explorer.py
index eed02ae5..12960dd2 100644
--- a/tests/test_explorer.py
+++ b/tests/test_explorer.py
@@ -12,7 +12,7 @@ def test_similarity():
exp.create_embeddings_table()
similar = exp.get_similar(idx=1)
assert len(similar) == 25
- similar = exp.get_similar(img=ASSETS / 'zidane.jpg')
+ similar = exp.get_similar(img=ASSETS / "zidane.jpg")
assert len(similar) == 25
similar = exp.get_similar(idx=[1, 2], limit=10)
assert len(similar) == 10
@@ -24,9 +24,9 @@ def test_similarity():
def test_det():
"""Test detection functionalities and ensure the embedding table has bounding boxes."""
- exp = Explorer(data='coco8.yaml', model='yolov8n.pt')
+ exp = Explorer(data="coco8.yaml", model="yolov8n.pt")
exp.create_embeddings_table(force=True)
- assert len(exp.table.head()['bboxes']) > 0
+ assert len(exp.table.head()["bboxes"]) > 0
similar = exp.get_similar(idx=[1, 2], limit=10)
assert len(similar) > 0
# This is a loose test, just checks errors not correctness
@@ -36,9 +36,9 @@ def test_det():
def test_seg():
"""Test segmentation functionalities and verify the embedding table includes masks."""
- exp = Explorer(data='coco8-seg.yaml', model='yolov8n-seg.pt')
+ exp = Explorer(data="coco8-seg.yaml", model="yolov8n-seg.pt")
exp.create_embeddings_table(force=True)
- assert len(exp.table.head()['masks']) > 0
+ assert len(exp.table.head()["masks"]) > 0
similar = exp.get_similar(idx=[1, 2], limit=10)
assert len(similar) > 0
similar = exp.plot_similar(idx=[1, 2], limit=10)
@@ -47,9 +47,9 @@ def test_seg():
def test_pose():
"""Test pose estimation functionalities and check the embedding table for keypoints."""
- exp = Explorer(data='coco8-pose.yaml', model='yolov8n-pose.pt')
+ exp = Explorer(data="coco8-pose.yaml", model="yolov8n-pose.pt")
exp.create_embeddings_table(force=True)
- assert len(exp.table.head()['keypoints']) > 0
+ assert len(exp.table.head()["keypoints"]) > 0
similar = exp.get_similar(idx=[1, 2], limit=10)
assert len(similar) > 0
similar = exp.plot_similar(idx=[1, 2], limit=10)
diff --git a/tests/test_integrations.py b/tests/test_integrations.py
index 0831f716..af8f507a 100644
--- a/tests/test_integrations.py
+++ b/tests/test_integrations.py
@@ -9,67 +9,67 @@ from ultralytics import YOLO, download
from ultralytics.utils import ASSETS, DATASETS_DIR, ROOT, SETTINGS, WEIGHTS_DIR
from ultralytics.utils.checks import check_requirements
-MODEL = WEIGHTS_DIR / 'path with spaces' / 'yolov8n.pt' # test spaces in path
-CFG = 'yolov8n.yaml'
-SOURCE = ASSETS / 'bus.jpg'
-TMP = (ROOT / '../tests/tmp').resolve() # temp directory for test files
+MODEL = WEIGHTS_DIR / "path with spaces" / "yolov8n.pt" # test spaces in path
+CFG = "yolov8n.yaml"
+SOURCE = ASSETS / "bus.jpg"
+TMP = (ROOT / "../tests/tmp").resolve() # temp directory for test files
-@pytest.mark.skipif(not check_requirements('ray', install=False), reason='ray[tune] not installed')
+@pytest.mark.skipif(not check_requirements("ray", install=False), reason="ray[tune] not installed")
def test_model_ray_tune():
"""Tune YOLO model with Ray optimization library."""
- YOLO('yolov8n-cls.yaml').tune(use_ray=True,
- data='imagenet10',
- grace_period=1,
- iterations=1,
- imgsz=32,
- epochs=1,
- plots=False,
- device='cpu')
+ YOLO("yolov8n-cls.yaml").tune(
+ use_ray=True, data="imagenet10", grace_period=1, iterations=1, imgsz=32, epochs=1, plots=False, device="cpu"
+ )
-@pytest.mark.skipif(not check_requirements('mlflow', install=False), reason='mlflow not installed')
+@pytest.mark.skipif(not check_requirements("mlflow", install=False), reason="mlflow not installed")
def test_mlflow():
"""Test training with MLflow tracking enabled."""
- SETTINGS['mlflow'] = True
- YOLO('yolov8n-cls.yaml').train(data='imagenet10', imgsz=32, epochs=3, plots=False, device='cpu')
+ SETTINGS["mlflow"] = True
+ YOLO("yolov8n-cls.yaml").train(data="imagenet10", imgsz=32, epochs=3, plots=False, device="cpu")
-@pytest.mark.skipif(not check_requirements('tritonclient', install=False), reason='tritonclient[all] not installed')
+@pytest.mark.skipif(not check_requirements("tritonclient", install=False), reason="tritonclient[all] not installed")
def test_triton():
"""Test NVIDIA Triton Server functionalities."""
- check_requirements('tritonclient[all]')
+ check_requirements("tritonclient[all]")
import subprocess
import time
from tritonclient.http import InferenceServerClient # noqa
# Create variables
- model_name = 'yolo'
- triton_repo_path = TMP / 'triton_repo'
+ model_name = "yolo"
+ triton_repo_path = TMP / "triton_repo"
triton_model_path = triton_repo_path / model_name
# Export model to ONNX
- f = YOLO(MODEL).export(format='onnx', dynamic=True)
+ f = YOLO(MODEL).export(format="onnx", dynamic=True)
# Prepare Triton repo
- (triton_model_path / '1').mkdir(parents=True, exist_ok=True)
- Path(f).rename(triton_model_path / '1' / 'model.onnx')
- (triton_model_path / 'config.pbtxt').touch()
+ (triton_model_path / "1").mkdir(parents=True, exist_ok=True)
+ Path(f).rename(triton_model_path / "1" / "model.onnx")
+ (triton_model_path / "config.pbtxt").touch()
# Define image https://catalog.ngc.nvidia.com/orgs/nvidia/containers/tritonserver
- tag = 'nvcr.io/nvidia/tritonserver:23.09-py3' # 6.4 GB
+ tag = "nvcr.io/nvidia/tritonserver:23.09-py3" # 6.4 GB
# Pull the image
- subprocess.call(f'docker pull {tag}', shell=True)
+ subprocess.call(f"docker pull {tag}", shell=True)
# Run the Triton server and capture the container ID
- container_id = subprocess.check_output(
- f'docker run -d --rm -v {triton_repo_path}:/models -p 8000:8000 {tag} tritonserver --model-repository=/models',
- shell=True).decode('utf-8').strip()
+ container_id = (
+ subprocess.check_output(
+ f"docker run -d --rm -v {triton_repo_path}:/models -p 8000:8000 {tag} tritonserver --model-repository=/models",
+ shell=True,
+ )
+ .decode("utf-8")
+ .strip()
+ )
# Wait for the Triton server to start
- triton_client = InferenceServerClient(url='localhost:8000', verbose=False, ssl=False)
+ triton_client = InferenceServerClient(url="localhost:8000", verbose=False, ssl=False)
# Wait until model is ready
for _ in range(10):
@@ -79,13 +79,13 @@ def test_triton():
time.sleep(1)
# Check Triton inference
- YOLO(f'http://localhost:8000/{model_name}', 'detect')(SOURCE) # exported model inference
+ YOLO(f"http://localhost:8000/{model_name}", "detect")(SOURCE) # exported model inference
# Kill and remove the container at the end of the test
- subprocess.call(f'docker kill {container_id}', shell=True)
+ subprocess.call(f"docker kill {container_id}", shell=True)
-@pytest.mark.skipif(not check_requirements('pycocotools', install=False), reason='pycocotools not installed')
+@pytest.mark.skipif(not check_requirements("pycocotools", install=False), reason="pycocotools not installed")
def test_pycocotools():
"""Validate model predictions using pycocotools."""
from ultralytics.models.yolo.detect import DetectionValidator
@@ -93,25 +93,25 @@ def test_pycocotools():
from ultralytics.models.yolo.segment import SegmentationValidator
# Download annotations after each dataset downloads first
- url = 'https://github.com/ultralytics/assets/releases/download/v8.1.0/'
+ url = "https://github.com/ultralytics/assets/releases/download/v8.1.0/"
- args = {'model': 'yolov8n.pt', 'data': 'coco8.yaml', 'save_json': True, 'imgsz': 64}
+ args = {"model": "yolov8n.pt", "data": "coco8.yaml", "save_json": True, "imgsz": 64}
validator = DetectionValidator(args=args)
validator()
validator.is_coco = True
- download(f'{url}instances_val2017.json', dir=DATASETS_DIR / 'coco8/annotations')
+ download(f"{url}instances_val2017.json", dir=DATASETS_DIR / "coco8/annotations")
_ = validator.eval_json(validator.stats)
- args = {'model': 'yolov8n-seg.pt', 'data': 'coco8-seg.yaml', 'save_json': True, 'imgsz': 64}
+ args = {"model": "yolov8n-seg.pt", "data": "coco8-seg.yaml", "save_json": True, "imgsz": 64}
validator = SegmentationValidator(args=args)
validator()
validator.is_coco = True
- download(f'{url}instances_val2017.json', dir=DATASETS_DIR / 'coco8-seg/annotations')
+ download(f"{url}instances_val2017.json", dir=DATASETS_DIR / "coco8-seg/annotations")
_ = validator.eval_json(validator.stats)
- args = {'model': 'yolov8n-pose.pt', 'data': 'coco8-pose.yaml', 'save_json': True, 'imgsz': 64}
+ args = {"model": "yolov8n-pose.pt", "data": "coco8-pose.yaml", "save_json": True, "imgsz": 64}
validator = PoseValidator(args=args)
validator()
validator.is_coco = True
- download(f'{url}person_keypoints_val2017.json', dir=DATASETS_DIR / 'coco8-pose/annotations')
+ download(f"{url}person_keypoints_val2017.json", dir=DATASETS_DIR / "coco8-pose/annotations")
_ = validator.eval_json(validator.stats)
diff --git a/tests/test_python.py b/tests/test_python.py
index 46d28762..61f43034 100644
--- a/tests/test_python.py
+++ b/tests/test_python.py
@@ -14,15 +14,27 @@ from torchvision.transforms import ToTensor
from ultralytics import RTDETR, YOLO
from ultralytics.cfg import TASK2DATA
from ultralytics.data.build import load_inference_source
-from ultralytics.utils import (ASSETS, DEFAULT_CFG, DEFAULT_CFG_PATH, LINUX, MACOS, ONLINE, ROOT, WEIGHTS_DIR, WINDOWS,
- checks, is_dir_writeable)
+from ultralytics.utils import (
+ ASSETS,
+ DEFAULT_CFG,
+ DEFAULT_CFG_PATH,
+ LINUX,
+ MACOS,
+ ONLINE,
+ ROOT,
+ WEIGHTS_DIR,
+ WINDOWS,
+ Retry,
+ checks,
+ is_dir_writeable,
+)
from ultralytics.utils.downloads import download
from ultralytics.utils.torch_utils import TORCH_1_9
-MODEL = WEIGHTS_DIR / 'path with spaces' / 'yolov8n.pt' # test spaces in path
-CFG = 'yolov8n.yaml'
-SOURCE = ASSETS / 'bus.jpg'
-TMP = (ROOT / '../tests/tmp').resolve() # temp directory for test files
+MODEL = WEIGHTS_DIR / "path with spaces" / "yolov8n.pt" # test spaces in path
+CFG = "yolov8n.yaml"
+SOURCE = ASSETS / "bus.jpg"
+TMP = (ROOT / "../tests/tmp").resolve() # temp directory for test files
IS_TMP_WRITEABLE = is_dir_writeable(TMP)
@@ -40,9 +52,9 @@ def test_model_methods():
model.info(verbose=True, detailed=True)
model = model.reset_weights()
model = model.load(MODEL)
- model.to('cpu')
+ model.to("cpu")
model.fuse()
- model.clear_callback('on_train_start')
+ model.clear_callback("on_train_start")
model.reset_callbacks()
# Model properties
@@ -61,23 +73,23 @@ def test_model_profile():
_ = model.predict(im, profile=True)
-@pytest.mark.skipif(not IS_TMP_WRITEABLE, reason='directory is not writeable')
+@pytest.mark.skipif(not IS_TMP_WRITEABLE, reason="directory is not writeable")
def test_predict_txt():
"""Test YOLO predictions with sources (file, dir, glob, recursive glob) specified in a text file."""
- txt_file = TMP / 'sources.txt'
- with open(txt_file, 'w') as f:
- for x in [ASSETS / 'bus.jpg', ASSETS, ASSETS / '*', ASSETS / '**/*.jpg']:
- f.write(f'{x}\n')
+ txt_file = TMP / "sources.txt"
+ with open(txt_file, "w") as f:
+ for x in [ASSETS / "bus.jpg", ASSETS, ASSETS / "*", ASSETS / "**/*.jpg"]:
+ f.write(f"{x}\n")
_ = YOLO(MODEL)(source=txt_file, imgsz=32)
def test_predict_img():
"""Test YOLO prediction on various types of image sources."""
model = YOLO(MODEL)
- seg_model = YOLO(WEIGHTS_DIR / 'yolov8n-seg.pt')
- cls_model = YOLO(WEIGHTS_DIR / 'yolov8n-cls.pt')
- pose_model = YOLO(WEIGHTS_DIR / 'yolov8n-pose.pt')
- obb_model = YOLO(WEIGHTS_DIR / 'yolov8n-obb.pt')
+ seg_model = YOLO(WEIGHTS_DIR / "yolov8n-seg.pt")
+ cls_model = YOLO(WEIGHTS_DIR / "yolov8n-cls.pt")
+ pose_model = YOLO(WEIGHTS_DIR / "yolov8n-pose.pt")
+ obb_model = YOLO(WEIGHTS_DIR / "yolov8n-obb.pt")
im = cv2.imread(str(SOURCE))
assert len(model(source=Image.open(SOURCE), save=True, verbose=True, imgsz=32)) == 1 # PIL
assert len(model(source=im, save=True, save_txt=True, imgsz=32)) == 1 # ndarray
@@ -87,10 +99,11 @@ def test_predict_img():
batch = [
str(SOURCE), # filename
Path(SOURCE), # Path
- 'https://ultralytics.com/images/zidane.jpg' if ONLINE else SOURCE, # URI
+ "https://ultralytics.com/images/zidane.jpg" if ONLINE else SOURCE, # URI
cv2.imread(str(SOURCE)), # OpenCV
Image.open(SOURCE), # PIL
- np.zeros((320, 640, 3))] # numpy
+ np.zeros((320, 640, 3)),
+ ] # numpy
assert len(model(batch, imgsz=32)) == len(batch) # multiple sources in a batch
# Test tensor inference
@@ -113,16 +126,16 @@ def test_predict_img():
def test_predict_grey_and_4ch():
"""Test YOLO prediction on SOURCE converted to greyscale and 4-channel images."""
im = Image.open(SOURCE)
- directory = TMP / 'im4'
+ directory = TMP / "im4"
directory.mkdir(parents=True, exist_ok=True)
- source_greyscale = directory / 'greyscale.jpg'
- source_rgba = directory / '4ch.png'
- source_non_utf = directory / 'non_UTF_测试文件_tést_image.jpg'
- source_spaces = directory / 'image with spaces.jpg'
+ source_greyscale = directory / "greyscale.jpg"
+ source_rgba = directory / "4ch.png"
+ source_non_utf = directory / "non_UTF_测试文件_tést_image.jpg"
+ source_spaces = directory / "image with spaces.jpg"
- im.convert('L').save(source_greyscale) # greyscale
- im.convert('RGBA').save(source_rgba) # 4-ch PNG with alpha
+ im.convert("L").save(source_greyscale) # greyscale
+ im.convert("RGBA").save(source_rgba) # 4-ch PNG with alpha
im.save(source_non_utf) # non-UTF characters in filename
im.save(source_spaces) # spaces in filename
@@ -136,7 +149,8 @@ def test_predict_grey_and_4ch():
@pytest.mark.slow
-@pytest.mark.skipif(not ONLINE, reason='environment is offline')
+@pytest.mark.skipif(not ONLINE, reason="environment is offline")
+@Retry(times=3, delay=10)
def test_youtube():
"""
Test YouTube inference.
@@ -144,11 +158,11 @@ def test_youtube():
Marked --slow to reduce YouTube API rate limits risk.
"""
model = YOLO(MODEL)
- model.predict('https://youtu.be/G17sBkb38XQ', imgsz=96, save=True)
+ model.predict("https://youtu.be/G17sBkb38XQ", imgsz=96, save=True)
-@pytest.mark.skipif(not ONLINE, reason='environment is offline')
-@pytest.mark.skipif(not IS_TMP_WRITEABLE, reason='directory is not writeable')
+@pytest.mark.skipif(not ONLINE, reason="environment is offline")
+@pytest.mark.skipif(not IS_TMP_WRITEABLE, reason="directory is not writeable")
def test_track_stream():
"""
Test streaming tracking (short 10 frame video) with non-default ByteTrack tracker.
@@ -157,56 +171,56 @@ def test_track_stream():
"""
import yaml
- video_url = 'https://ultralytics.com/assets/decelera_portrait_min.mov'
+ video_url = "https://ultralytics.com/assets/decelera_portrait_min.mov"
model = YOLO(MODEL)
- model.track(video_url, imgsz=160, tracker='bytetrack.yaml')
- model.track(video_url, imgsz=160, tracker='botsort.yaml', save_frames=True) # test frame saving also
+ model.track(video_url, imgsz=160, tracker="bytetrack.yaml")
+ model.track(video_url, imgsz=160, tracker="botsort.yaml", save_frames=True) # test frame saving also
# Test Global Motion Compensation (GMC) methods
- for gmc in 'orb', 'sift', 'ecc':
- with open(ROOT / 'cfg/trackers/botsort.yaml', encoding='utf-8') as f:
+ for gmc in "orb", "sift", "ecc":
+ with open(ROOT / "cfg/trackers/botsort.yaml", encoding="utf-8") as f:
data = yaml.safe_load(f)
- tracker = TMP / f'botsort-{gmc}.yaml'
- data['gmc_method'] = gmc
- with open(tracker, 'w', encoding='utf-8') as f:
+ tracker = TMP / f"botsort-{gmc}.yaml"
+ data["gmc_method"] = gmc
+ with open(tracker, "w", encoding="utf-8") as f:
yaml.safe_dump(data, f)
model.track(video_url, imgsz=160, tracker=tracker)
def test_val():
"""Test the validation mode of the YOLO model."""
- YOLO(MODEL).val(data='coco8.yaml', imgsz=32, save_hybrid=True)
+ YOLO(MODEL).val(data="coco8.yaml", imgsz=32, save_hybrid=True)
def test_train_scratch():
"""Test training the YOLO model from scratch."""
model = YOLO(CFG)
- model.train(data='coco8.yaml', epochs=2, imgsz=32, cache='disk', batch=-1, close_mosaic=1, name='model')
+ model.train(data="coco8.yaml", epochs=2, imgsz=32, cache="disk", batch=-1, close_mosaic=1, name="model")
model(SOURCE)
def test_train_pretrained():
"""Test training the YOLO model from a pre-trained state."""
- model = YOLO(WEIGHTS_DIR / 'yolov8n-seg.pt')
- model.train(data='coco8-seg.yaml', epochs=1, imgsz=32, cache='ram', copy_paste=0.5, mixup=0.5, name=0)
+ model = YOLO(WEIGHTS_DIR / "yolov8n-seg.pt")
+ model.train(data="coco8-seg.yaml", epochs=1, imgsz=32, cache="ram", copy_paste=0.5, mixup=0.5, name=0)
model(SOURCE)
def test_export_torchscript():
"""Test exporting the YOLO model to TorchScript format."""
- f = YOLO(MODEL).export(format='torchscript', optimize=False)
+ f = YOLO(MODEL).export(format="torchscript", optimize=False)
YOLO(f)(SOURCE) # exported model inference
def test_export_onnx():
"""Test exporting the YOLO model to ONNX format."""
- f = YOLO(MODEL).export(format='onnx', dynamic=True)
+ f = YOLO(MODEL).export(format="onnx", dynamic=True)
YOLO(f)(SOURCE) # exported model inference
def test_export_openvino():
"""Test exporting the YOLO model to OpenVINO format."""
- f = YOLO(MODEL).export(format='openvino')
+ f = YOLO(MODEL).export(format="openvino")
YOLO(f)(SOURCE) # exported model inference
@@ -214,10 +228,10 @@ def test_export_coreml():
"""Test exporting the YOLO model to CoreML format."""
if not WINDOWS: # RuntimeError: BlobWriter not loaded with coremltools 7.0 on windows
if MACOS:
- f = YOLO(MODEL).export(format='coreml')
+ f = YOLO(MODEL).export(format="coreml")
YOLO(f)(SOURCE) # model prediction only supported on macOS for nms=False models
else:
- YOLO(MODEL).export(format='coreml', nms=True)
+ YOLO(MODEL).export(format="coreml", nms=True)
def test_export_tflite(enabled=False):
@@ -228,7 +242,7 @@ def test_export_tflite(enabled=False):
"""
if enabled and LINUX:
model = YOLO(MODEL)
- f = model.export(format='tflite')
+ f = model.export(format="tflite")
YOLO(f)(SOURCE)
@@ -240,7 +254,7 @@ def test_export_pb(enabled=False):
"""
if enabled and LINUX:
model = YOLO(MODEL)
- f = model.export(format='pb')
+ f = model.export(format="pb")
YOLO(f)(SOURCE)
@@ -251,20 +265,20 @@ def test_export_paddle(enabled=False):
Note Paddle protobuf requirements conflicting with onnx protobuf requirements.
"""
if enabled:
- YOLO(MODEL).export(format='paddle')
+ YOLO(MODEL).export(format="paddle")
@pytest.mark.slow
def test_export_ncnn():
"""Test exporting the YOLO model to NCNN format."""
- f = YOLO(MODEL).export(format='ncnn')
+ f = YOLO(MODEL).export(format="ncnn")
YOLO(f)(SOURCE) # exported model inference
def test_all_model_yamls():
"""Test YOLO model creation for all available YAML configurations."""
- for m in (ROOT / 'cfg' / 'models').rglob('*.yaml'):
- if 'rtdetr' in m.name:
+ for m in (ROOT / "cfg" / "models").rglob("*.yaml"):
+ if "rtdetr" in m.name:
if TORCH_1_9: # torch<=1.8 issue - TypeError: __init__() got an unexpected keyword argument 'batch_first'
_ = RTDETR(m.name)(SOURCE, imgsz=640) # must be 640
else:
@@ -274,10 +288,10 @@ def test_all_model_yamls():
def test_workflow():
"""Test the complete workflow including training, validation, prediction, and exporting."""
model = YOLO(MODEL)
- model.train(data='coco8.yaml', epochs=1, imgsz=32, optimizer='SGD')
+ model.train(data="coco8.yaml", epochs=1, imgsz=32, optimizer="SGD")
model.val(imgsz=32)
model.predict(SOURCE, imgsz=32)
- model.export(format='onnx') # export a model to ONNX format
+ model.export(format="onnx") # export a model to ONNX format
def test_predict_callback_and_setup():
@@ -291,34 +305,34 @@ def test_predict_callback_and_setup():
predictor.results = zip(predictor.results, im0s, bs) # results is List[batch_size]
model = YOLO(MODEL)
- model.add_callback('on_predict_batch_end', on_predict_batch_end)
+ model.add_callback("on_predict_batch_end", on_predict_batch_end)
dataset = load_inference_source(source=SOURCE)
bs = dataset.bs # noqa access predictor properties
results = model.predict(dataset, stream=True, imgsz=160) # source already setup
for r, im0, bs in results:
- print('test_callback', im0.shape)
- print('test_callback', bs)
+ print("test_callback", im0.shape)
+ print("test_callback", bs)
boxes = r.boxes # Boxes object for bbox outputs
print(boxes)
def test_results():
"""Test various result formats for the YOLO model."""
- for m in 'yolov8n-pose.pt', 'yolov8n-seg.pt', 'yolov8n.pt', 'yolov8n-cls.pt':
+ for m in "yolov8n-pose.pt", "yolov8n-seg.pt", "yolov8n.pt", "yolov8n-cls.pt":
results = YOLO(WEIGHTS_DIR / m)([SOURCE, SOURCE], imgsz=160)
for r in results:
r = r.cpu().numpy()
- r = r.to(device='cpu', dtype=torch.float32)
- r.save_txt(txt_file=TMP / 'runs/tests/label.txt', save_conf=True)
- r.save_crop(save_dir=TMP / 'runs/tests/crops/')
+ r = r.to(device="cpu", dtype=torch.float32)
+ r.save_txt(txt_file=TMP / "runs/tests/label.txt", save_conf=True)
+ r.save_crop(save_dir=TMP / "runs/tests/crops/")
r.tojson(normalize=True)
r.plot(pil=True)
r.plot(conf=True, boxes=True)
print(r, len(r), r.path)
-@pytest.mark.skipif(not ONLINE, reason='environment is offline')
+@pytest.mark.skipif(not ONLINE, reason="environment is offline")
def test_data_utils():
"""Test utility functions in ultralytics/data/utils.py."""
from ultralytics.data.utils import HUBDatasetStats, autosplit
@@ -327,25 +341,25 @@ def test_data_utils():
# from ultralytics.utils.files import WorkingDirectory
# with WorkingDirectory(ROOT.parent / 'tests'):
- for task in 'detect', 'segment', 'pose', 'classify':
- file = Path(TASK2DATA[task]).with_suffix('.zip') # i.e. coco8.zip
- download(f'https://github.com/ultralytics/hub/raw/main/example_datasets/{file}', unzip=False, dir=TMP)
+ for task in "detect", "segment", "pose", "classify":
+ file = Path(TASK2DATA[task]).with_suffix(".zip") # i.e. coco8.zip
+ download(f"https://github.com/ultralytics/hub/raw/main/example_datasets/{file}", unzip=False, dir=TMP)
stats = HUBDatasetStats(TMP / file, task=task)
stats.get_json(save=True)
stats.process_images()
- autosplit(TMP / 'coco8')
- zip_directory(TMP / 'coco8/images/val') # zip
+ autosplit(TMP / "coco8")
+ zip_directory(TMP / "coco8/images/val") # zip
-@pytest.mark.skipif(not ONLINE, reason='environment is offline')
+@pytest.mark.skipif(not ONLINE, reason="environment is offline")
def test_data_converter():
"""Test dataset converters."""
from ultralytics.data.converter import coco80_to_coco91_class, convert_coco
- file = 'instances_val2017.json'
- download(f'https://github.com/ultralytics/yolov5/releases/download/v1.0/{file}', dir=TMP)
- convert_coco(labels_dir=TMP, save_dir=TMP / 'yolo_labels', use_segments=True, use_keypoints=False, cls91to80=True)
+ file = "instances_val2017.json"
+ download(f"https://github.com/ultralytics/yolov5/releases/download/v1.0/{file}", dir=TMP)
+ convert_coco(labels_dir=TMP, save_dir=TMP / "yolo_labels", use_segments=True, use_keypoints=False, cls91to80=True)
coco80_to_coco91_class()
@@ -353,10 +367,12 @@ def test_data_annotator():
"""Test automatic data annotation."""
from ultralytics.data.annotator import auto_annotate
- auto_annotate(ASSETS,
- det_model=WEIGHTS_DIR / 'yolov8n.pt',
- sam_model=WEIGHTS_DIR / 'mobile_sam.pt',
- output_dir=TMP / 'auto_annotate_labels')
+ auto_annotate(
+ ASSETS,
+ det_model=WEIGHTS_DIR / "yolov8n.pt",
+ sam_model=WEIGHTS_DIR / "mobile_sam.pt",
+ output_dir=TMP / "auto_annotate_labels",
+ )
def test_events():
@@ -366,7 +382,7 @@ def test_events():
events = Events()
events.enabled = True
cfg = copy(DEFAULT_CFG) # does not require deepcopy
- cfg.mode = 'test'
+ cfg.mode = "test"
events(cfg)
@@ -375,10 +391,10 @@ def test_cfg_init():
from ultralytics.cfg import check_dict_alignment, copy_default_cfg, smart_value
with contextlib.suppress(SyntaxError):
- check_dict_alignment({'a': 1}, {'b': 2})
+ check_dict_alignment({"a": 1}, {"b": 2})
copy_default_cfg()
- (Path.cwd() / DEFAULT_CFG_PATH.name.replace('.yaml', '_copy.yaml')).unlink(missing_ok=False)
- [smart_value(x) for x in ['none', 'true', 'false']]
+ (Path.cwd() / DEFAULT_CFG_PATH.name.replace(".yaml", "_copy.yaml")).unlink(missing_ok=False)
+ [smart_value(x) for x in ["none", "true", "false"]]
def test_utils_init():
@@ -393,12 +409,12 @@ def test_utils_init():
def test_utils_checks():
"""Test various utility checks."""
- checks.check_yolov5u_filename('yolov5n.pt')
+ checks.check_yolov5u_filename("yolov5n.pt")
checks.git_describe(ROOT)
checks.check_requirements() # check requirements.txt
checks.check_imgsz([600, 600], max_dim=1)
checks.check_imshow()
- checks.check_version('ultralytics', '8.0.0')
+ checks.check_version("ultralytics", "8.0.0")
checks.print_args()
# checks.check_imshow(warn=True)
@@ -407,7 +423,7 @@ def test_utils_benchmarks():
"""Test model benchmarking."""
from ultralytics.utils.benchmarks import ProfileModels
- ProfileModels(['yolov8n.yaml'], imgsz=32, min_time=1, num_timed_runs=3, num_warmup_runs=1).profile()
+ ProfileModels(["yolov8n.yaml"], imgsz=32, min_time=1, num_timed_runs=3, num_warmup_runs=1).profile()
def test_utils_torchutils():
@@ -423,18 +439,29 @@ def test_utils_torchutils():
time_sync()
-@pytest.mark.skipif(not ONLINE, reason='environment is offline')
+@pytest.mark.skipif(not ONLINE, reason="environment is offline")
def test_utils_downloads():
"""Test file download utilities."""
from ultralytics.utils.downloads import get_google_drive_file_info
- get_google_drive_file_info('https://drive.google.com/file/d/1cqT-cJgANNrhIHCrEufUYhQ4RqiWG_lJ/view?usp=drive_link')
+ get_google_drive_file_info("https://drive.google.com/file/d/1cqT-cJgANNrhIHCrEufUYhQ4RqiWG_lJ/view?usp=drive_link")
def test_utils_ops():
"""Test various operations utilities."""
- from ultralytics.utils.ops import (ltwh2xywh, ltwh2xyxy, make_divisible, xywh2ltwh, xywh2xyxy, xywhn2xyxy,
- xywhr2xyxyxyxy, xyxy2ltwh, xyxy2xywh, xyxy2xywhn, xyxyxyxy2xywhr)
+ from ultralytics.utils.ops import (
+ ltwh2xywh,
+ ltwh2xyxy,
+ make_divisible,
+ xywh2ltwh,
+ xywh2xyxy,
+ xywhn2xyxy,
+ xywhr2xyxyxyxy,
+ xyxy2ltwh,
+ xyxy2xywh,
+ xyxy2xywhn,
+ xyxyxyxy2xywhr,
+ )
make_divisible(17, torch.tensor([8]))
@@ -455,9 +482,9 @@ def test_utils_files():
file_age(SOURCE)
file_date(SOURCE)
- get_latest_run(ROOT / 'runs')
+ get_latest_run(ROOT / "runs")
- path = TMP / 'path/with spaces'
+ path = TMP / "path/with spaces"
path.mkdir(parents=True, exist_ok=True)
with spaces_in_path(path) as new_path:
print(new_path)
@@ -471,9 +498,9 @@ def test_utils_patches_torch_save():
mock = MagicMock(side_effect=RuntimeError)
- with patch('ultralytics.utils.patches._torch_save', new=mock):
+ with patch("ultralytics.utils.patches._torch_save", new=mock):
with pytest.raises(RuntimeError):
- torch_save(torch.zeros(1), TMP / 'test.pt')
+ torch_save(torch.zeros(1), TMP / "test.pt")
assert mock.call_count == 4, "torch_save was not attempted the expected number of times"
@@ -512,7 +539,7 @@ def test_nn_modules_block():
BottleneckCSP(c1, c2)(x)
-@pytest.mark.skipif(not ONLINE, reason='environment is offline')
+@pytest.mark.skipif(not ONLINE, reason="environment is offline")
def test_hub():
"""Test Ultralytics HUB functionalities."""
from ultralytics.hub import export_fmts_hub, logout
@@ -520,7 +547,7 @@ def test_hub():
export_fmts_hub()
logout()
- smart_request('GET', 'https://github.com', progress=True)
+ smart_request("GET", "https://github.com", progress=True)
@pytest.fixture
@@ -529,12 +556,13 @@ def image():
@pytest.mark.parametrize(
- 'auto_augment, erasing, force_color_jitter',
+ "auto_augment, erasing, force_color_jitter",
[
(None, 0.0, False),
- ('randaugment', 0.5, True),
- ('augmix', 0.2, False),
- ('autoaugment', 0.0, True), ],
+ ("randaugment", 0.5, True),
+ ("augmix", 0.2, False),
+ ("autoaugment", 0.0, True),
+ ],
)
def test_classify_transforms_train(image, auto_augment, erasing, force_color_jitter):
import torchvision.transforms as T
@@ -566,17 +594,17 @@ def test_classify_transforms_train(image, auto_augment, erasing, force_color_jit
@pytest.mark.slow
-@pytest.mark.skipif(not ONLINE, reason='environment is offline')
+@pytest.mark.skipif(not ONLINE, reason="environment is offline")
def test_model_tune():
"""Tune YOLO model for performance."""
- YOLO('yolov8n-pose.pt').tune(data='coco8-pose.yaml', plots=False, imgsz=32, epochs=1, iterations=2, device='cpu')
- YOLO('yolov8n-cls.pt').tune(data='imagenet10', plots=False, imgsz=32, epochs=1, iterations=2, device='cpu')
+ YOLO("yolov8n-pose.pt").tune(data="coco8-pose.yaml", plots=False, imgsz=32, epochs=1, iterations=2, device="cpu")
+ YOLO("yolov8n-cls.pt").tune(data="imagenet10", plots=False, imgsz=32, epochs=1, iterations=2, device="cpu")
def test_model_embeddings():
"""Test YOLO model embeddings."""
model_detect = YOLO(MODEL)
- model_segment = YOLO(WEIGHTS_DIR / 'yolov8n-seg.pt')
+ model_segment = YOLO(WEIGHTS_DIR / "yolov8n-seg.pt")
for batch in [SOURCE], [SOURCE, SOURCE]: # test batch size 1 and 2
assert len(model_detect.embed(source=batch, imgsz=32)) == len(batch)
diff --git a/ultralytics/utils/__init__.py b/ultralytics/utils/__init__.py
index d8cfccaa..2efa70bb 100644
--- a/ultralytics/utils/__init__.py
+++ b/ultralytics/utils/__init__.py
@@ -9,6 +9,7 @@ import re
import subprocess
import sys
import threading
+import time
import urllib
import uuid
from pathlib import Path
@@ -721,9 +722,19 @@ def remove_colorstr(input_string):
class TryExcept(contextlib.ContextDecorator):
"""
- YOLOv8 TryExcept class.
+ Ultralytics TryExcept class. Use as @TryExcept() decorator or 'with TryExcept():' context manager.
- Use as @TryExcept() decorator or 'with TryExcept():' context manager.
+ Examples:
+ As a decorator:
+ >>> @TryExcept(msg="Error occurred in func", verbose=True)
+ >>> def func():
+ >>> # Function logic here
+ >>> pass
+
+ As a context manager:
+ >>> with TryExcept(msg="Error occurred in block", verbose=True):
+ >>> # Code block here
+ >>> pass
"""
def __init__(self, msg="", verbose=True):
@@ -742,6 +753,64 @@ class TryExcept(contextlib.ContextDecorator):
return True
+class Retry(contextlib.ContextDecorator):
+ """
+ Retry class for function execution with exponential backoff.
+
+ Can be used as a decorator or a context manager to retry a function or block of code on exceptions, up to a
+ specified number of times with an exponentially increasing delay between retries.
+
+ Examples:
+ Example usage as a decorator:
+ >>> @Retry(times=3, delay=2)
+ >>> def test_func():
+ >>> # Replace with function logic that may raise exceptions
+ >>> return True
+
+ Example usage as a context manager:
+ >>> with Retry(times=3, delay=2):
+ >>> # Replace with code block that may raise exceptions
+ >>> pass
+ """
+
+ def __init__(self, times=3, delay=2):
+ """Initialize Retry class with specified number of retries and delay."""
+ self.times = times
+ self.delay = delay
+ self._attempts = 0
+
+ def __call__(self, func):
+ """Decorator implementation for Retry with exponential backoff."""
+
+ def wrapped_func(*args, **kwargs):
+ self._attempts = 0
+ while self._attempts < self.times:
+ try:
+ return func(*args, **kwargs)
+ except Exception as e:
+ self._attempts += 1
+ print(f"Retry {self._attempts}/{self.times} failed: {e}")
+ if self._attempts >= self.times:
+ raise e
+ time.sleep(self.delay * (2**self._attempts)) # exponential backoff delay
+
+ return wrapped_func
+
+ def __enter__(self):
+ """Enter the runtime context related to this object."""
+ self._attempts = 0
+
+ def __exit__(self, exc_type, exc_value, traceback):
+ """Exit the runtime context related to this object with exponential backoff."""
+ if exc_type is not None:
+ self._attempts += 1
+ if self._attempts < self.times:
+ print(f"Retry {self._attempts}/{self.times} failed: {exc_value}")
+ time.sleep(self.delay * (2**self._attempts)) # exponential backoff delay
+ return True # Suppresses the exception and retries
+ return False # Re-raises the exception if retries are exhausted
+
+
def threaded(func):
"""
Multi-threads a target function by default and returns the thread or function result.