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.
- dgenerate_ultralytics_headless-8.3.134.dist-info/METADATA +400 -0
- dgenerate_ultralytics_headless-8.3.134.dist-info/RECORD +272 -0
- dgenerate_ultralytics_headless-8.3.134.dist-info/WHEEL +5 -0
- dgenerate_ultralytics_headless-8.3.134.dist-info/entry_points.txt +3 -0
- dgenerate_ultralytics_headless-8.3.134.dist-info/licenses/LICENSE +661 -0
- dgenerate_ultralytics_headless-8.3.134.dist-info/top_level.txt +1 -0
- tests/__init__.py +22 -0
- tests/conftest.py +83 -0
- tests/test_cli.py +138 -0
- tests/test_cuda.py +215 -0
- tests/test_engine.py +131 -0
- tests/test_exports.py +236 -0
- tests/test_integrations.py +154 -0
- tests/test_python.py +694 -0
- tests/test_solutions.py +187 -0
- ultralytics/__init__.py +30 -0
- ultralytics/assets/bus.jpg +0 -0
- ultralytics/assets/zidane.jpg +0 -0
- ultralytics/cfg/__init__.py +1023 -0
- ultralytics/cfg/datasets/Argoverse.yaml +77 -0
- ultralytics/cfg/datasets/DOTAv1.5.yaml +37 -0
- ultralytics/cfg/datasets/DOTAv1.yaml +36 -0
- ultralytics/cfg/datasets/GlobalWheat2020.yaml +68 -0
- ultralytics/cfg/datasets/HomeObjects-3K.yaml +33 -0
- ultralytics/cfg/datasets/ImageNet.yaml +2025 -0
- ultralytics/cfg/datasets/Objects365.yaml +443 -0
- ultralytics/cfg/datasets/SKU-110K.yaml +58 -0
- ultralytics/cfg/datasets/VOC.yaml +106 -0
- ultralytics/cfg/datasets/VisDrone.yaml +77 -0
- ultralytics/cfg/datasets/african-wildlife.yaml +25 -0
- ultralytics/cfg/datasets/brain-tumor.yaml +23 -0
- ultralytics/cfg/datasets/carparts-seg.yaml +44 -0
- ultralytics/cfg/datasets/coco-pose.yaml +42 -0
- ultralytics/cfg/datasets/coco.yaml +118 -0
- ultralytics/cfg/datasets/coco128-seg.yaml +101 -0
- ultralytics/cfg/datasets/coco128.yaml +101 -0
- ultralytics/cfg/datasets/coco8-multispectral.yaml +104 -0
- ultralytics/cfg/datasets/coco8-pose.yaml +26 -0
- ultralytics/cfg/datasets/coco8-seg.yaml +101 -0
- ultralytics/cfg/datasets/coco8.yaml +101 -0
- ultralytics/cfg/datasets/crack-seg.yaml +22 -0
- ultralytics/cfg/datasets/dog-pose.yaml +24 -0
- ultralytics/cfg/datasets/dota8-multispectral.yaml +38 -0
- ultralytics/cfg/datasets/dota8.yaml +35 -0
- ultralytics/cfg/datasets/hand-keypoints.yaml +26 -0
- ultralytics/cfg/datasets/lvis.yaml +1240 -0
- ultralytics/cfg/datasets/medical-pills.yaml +22 -0
- ultralytics/cfg/datasets/open-images-v7.yaml +666 -0
- ultralytics/cfg/datasets/package-seg.yaml +22 -0
- ultralytics/cfg/datasets/signature.yaml +21 -0
- ultralytics/cfg/datasets/tiger-pose.yaml +25 -0
- ultralytics/cfg/datasets/xView.yaml +155 -0
- ultralytics/cfg/default.yaml +127 -0
- ultralytics/cfg/models/11/yolo11-cls-resnet18.yaml +17 -0
- ultralytics/cfg/models/11/yolo11-cls.yaml +33 -0
- ultralytics/cfg/models/11/yolo11-obb.yaml +50 -0
- ultralytics/cfg/models/11/yolo11-pose.yaml +51 -0
- ultralytics/cfg/models/11/yolo11-seg.yaml +50 -0
- ultralytics/cfg/models/11/yolo11.yaml +50 -0
- ultralytics/cfg/models/11/yoloe-11-seg.yaml +48 -0
- ultralytics/cfg/models/11/yoloe-11.yaml +48 -0
- ultralytics/cfg/models/12/yolo12-cls.yaml +32 -0
- ultralytics/cfg/models/12/yolo12-obb.yaml +48 -0
- ultralytics/cfg/models/12/yolo12-pose.yaml +49 -0
- ultralytics/cfg/models/12/yolo12-seg.yaml +48 -0
- ultralytics/cfg/models/12/yolo12.yaml +48 -0
- ultralytics/cfg/models/rt-detr/rtdetr-l.yaml +53 -0
- ultralytics/cfg/models/rt-detr/rtdetr-resnet101.yaml +45 -0
- ultralytics/cfg/models/rt-detr/rtdetr-resnet50.yaml +45 -0
- ultralytics/cfg/models/rt-detr/rtdetr-x.yaml +57 -0
- ultralytics/cfg/models/v10/yolov10b.yaml +45 -0
- ultralytics/cfg/models/v10/yolov10l.yaml +45 -0
- ultralytics/cfg/models/v10/yolov10m.yaml +45 -0
- ultralytics/cfg/models/v10/yolov10n.yaml +45 -0
- ultralytics/cfg/models/v10/yolov10s.yaml +45 -0
- ultralytics/cfg/models/v10/yolov10x.yaml +45 -0
- ultralytics/cfg/models/v3/yolov3-spp.yaml +49 -0
- ultralytics/cfg/models/v3/yolov3-tiny.yaml +40 -0
- ultralytics/cfg/models/v3/yolov3.yaml +49 -0
- ultralytics/cfg/models/v5/yolov5-p6.yaml +62 -0
- ultralytics/cfg/models/v5/yolov5.yaml +51 -0
- ultralytics/cfg/models/v6/yolov6.yaml +56 -0
- ultralytics/cfg/models/v8/yoloe-v8-seg.yaml +45 -0
- ultralytics/cfg/models/v8/yoloe-v8.yaml +45 -0
- ultralytics/cfg/models/v8/yolov8-cls-resnet101.yaml +28 -0
- ultralytics/cfg/models/v8/yolov8-cls-resnet50.yaml +28 -0
- ultralytics/cfg/models/v8/yolov8-cls.yaml +32 -0
- ultralytics/cfg/models/v8/yolov8-ghost-p2.yaml +58 -0
- ultralytics/cfg/models/v8/yolov8-ghost-p6.yaml +60 -0
- ultralytics/cfg/models/v8/yolov8-ghost.yaml +50 -0
- ultralytics/cfg/models/v8/yolov8-obb.yaml +49 -0
- ultralytics/cfg/models/v8/yolov8-p2.yaml +57 -0
- ultralytics/cfg/models/v8/yolov8-p6.yaml +59 -0
- ultralytics/cfg/models/v8/yolov8-pose-p6.yaml +60 -0
- ultralytics/cfg/models/v8/yolov8-pose.yaml +50 -0
- ultralytics/cfg/models/v8/yolov8-rtdetr.yaml +49 -0
- ultralytics/cfg/models/v8/yolov8-seg-p6.yaml +59 -0
- ultralytics/cfg/models/v8/yolov8-seg.yaml +49 -0
- ultralytics/cfg/models/v8/yolov8-world.yaml +51 -0
- ultralytics/cfg/models/v8/yolov8-worldv2.yaml +49 -0
- ultralytics/cfg/models/v8/yolov8.yaml +49 -0
- ultralytics/cfg/models/v9/yolov9c-seg.yaml +41 -0
- ultralytics/cfg/models/v9/yolov9c.yaml +41 -0
- ultralytics/cfg/models/v9/yolov9e-seg.yaml +64 -0
- ultralytics/cfg/models/v9/yolov9e.yaml +64 -0
- ultralytics/cfg/models/v9/yolov9m.yaml +41 -0
- ultralytics/cfg/models/v9/yolov9s.yaml +41 -0
- ultralytics/cfg/models/v9/yolov9t.yaml +41 -0
- ultralytics/cfg/trackers/botsort.yaml +22 -0
- ultralytics/cfg/trackers/bytetrack.yaml +14 -0
- ultralytics/data/__init__.py +26 -0
- ultralytics/data/annotator.py +66 -0
- ultralytics/data/augment.py +2945 -0
- ultralytics/data/base.py +438 -0
- ultralytics/data/build.py +258 -0
- ultralytics/data/converter.py +754 -0
- ultralytics/data/dataset.py +834 -0
- ultralytics/data/loaders.py +676 -0
- ultralytics/data/scripts/download_weights.sh +18 -0
- ultralytics/data/scripts/get_coco.sh +61 -0
- ultralytics/data/scripts/get_coco128.sh +18 -0
- ultralytics/data/scripts/get_imagenet.sh +52 -0
- ultralytics/data/split.py +125 -0
- ultralytics/data/split_dota.py +325 -0
- ultralytics/data/utils.py +777 -0
- ultralytics/engine/__init__.py +1 -0
- ultralytics/engine/exporter.py +1519 -0
- ultralytics/engine/model.py +1156 -0
- ultralytics/engine/predictor.py +502 -0
- ultralytics/engine/results.py +1840 -0
- ultralytics/engine/trainer.py +853 -0
- ultralytics/engine/tuner.py +243 -0
- ultralytics/engine/validator.py +377 -0
- ultralytics/hub/__init__.py +168 -0
- ultralytics/hub/auth.py +137 -0
- ultralytics/hub/google/__init__.py +176 -0
- ultralytics/hub/session.py +446 -0
- ultralytics/hub/utils.py +248 -0
- ultralytics/models/__init__.py +9 -0
- ultralytics/models/fastsam/__init__.py +7 -0
- ultralytics/models/fastsam/model.py +61 -0
- ultralytics/models/fastsam/predict.py +181 -0
- ultralytics/models/fastsam/utils.py +24 -0
- ultralytics/models/fastsam/val.py +40 -0
- ultralytics/models/nas/__init__.py +7 -0
- ultralytics/models/nas/model.py +102 -0
- ultralytics/models/nas/predict.py +58 -0
- ultralytics/models/nas/val.py +39 -0
- ultralytics/models/rtdetr/__init__.py +7 -0
- ultralytics/models/rtdetr/model.py +63 -0
- ultralytics/models/rtdetr/predict.py +84 -0
- ultralytics/models/rtdetr/train.py +85 -0
- ultralytics/models/rtdetr/val.py +191 -0
- ultralytics/models/sam/__init__.py +6 -0
- ultralytics/models/sam/amg.py +260 -0
- ultralytics/models/sam/build.py +358 -0
- ultralytics/models/sam/model.py +170 -0
- ultralytics/models/sam/modules/__init__.py +1 -0
- ultralytics/models/sam/modules/blocks.py +1129 -0
- ultralytics/models/sam/modules/decoders.py +515 -0
- ultralytics/models/sam/modules/encoders.py +854 -0
- ultralytics/models/sam/modules/memory_attention.py +299 -0
- ultralytics/models/sam/modules/sam.py +1006 -0
- ultralytics/models/sam/modules/tiny_encoder.py +1002 -0
- ultralytics/models/sam/modules/transformer.py +351 -0
- ultralytics/models/sam/modules/utils.py +394 -0
- ultralytics/models/sam/predict.py +1605 -0
- ultralytics/models/utils/__init__.py +1 -0
- ultralytics/models/utils/loss.py +455 -0
- ultralytics/models/utils/ops.py +268 -0
- ultralytics/models/yolo/__init__.py +7 -0
- ultralytics/models/yolo/classify/__init__.py +7 -0
- ultralytics/models/yolo/classify/predict.py +88 -0
- ultralytics/models/yolo/classify/train.py +233 -0
- ultralytics/models/yolo/classify/val.py +215 -0
- ultralytics/models/yolo/detect/__init__.py +7 -0
- ultralytics/models/yolo/detect/predict.py +124 -0
- ultralytics/models/yolo/detect/train.py +217 -0
- ultralytics/models/yolo/detect/val.py +451 -0
- ultralytics/models/yolo/model.py +354 -0
- ultralytics/models/yolo/obb/__init__.py +7 -0
- ultralytics/models/yolo/obb/predict.py +66 -0
- ultralytics/models/yolo/obb/train.py +81 -0
- ultralytics/models/yolo/obb/val.py +283 -0
- ultralytics/models/yolo/pose/__init__.py +7 -0
- ultralytics/models/yolo/pose/predict.py +79 -0
- ultralytics/models/yolo/pose/train.py +154 -0
- ultralytics/models/yolo/pose/val.py +394 -0
- ultralytics/models/yolo/segment/__init__.py +7 -0
- ultralytics/models/yolo/segment/predict.py +113 -0
- ultralytics/models/yolo/segment/train.py +123 -0
- ultralytics/models/yolo/segment/val.py +428 -0
- ultralytics/models/yolo/world/__init__.py +5 -0
- ultralytics/models/yolo/world/train.py +119 -0
- ultralytics/models/yolo/world/train_world.py +176 -0
- ultralytics/models/yolo/yoloe/__init__.py +22 -0
- ultralytics/models/yolo/yoloe/predict.py +169 -0
- ultralytics/models/yolo/yoloe/train.py +298 -0
- ultralytics/models/yolo/yoloe/train_seg.py +124 -0
- ultralytics/models/yolo/yoloe/val.py +191 -0
- ultralytics/nn/__init__.py +29 -0
- ultralytics/nn/autobackend.py +842 -0
- ultralytics/nn/modules/__init__.py +182 -0
- ultralytics/nn/modules/activation.py +53 -0
- ultralytics/nn/modules/block.py +1966 -0
- ultralytics/nn/modules/conv.py +712 -0
- ultralytics/nn/modules/head.py +880 -0
- ultralytics/nn/modules/transformer.py +713 -0
- ultralytics/nn/modules/utils.py +164 -0
- ultralytics/nn/tasks.py +1627 -0
- ultralytics/nn/text_model.py +351 -0
- ultralytics/solutions/__init__.py +41 -0
- ultralytics/solutions/ai_gym.py +116 -0
- ultralytics/solutions/analytics.py +252 -0
- ultralytics/solutions/config.py +106 -0
- ultralytics/solutions/distance_calculation.py +124 -0
- ultralytics/solutions/heatmap.py +127 -0
- ultralytics/solutions/instance_segmentation.py +84 -0
- ultralytics/solutions/object_blurrer.py +90 -0
- ultralytics/solutions/object_counter.py +195 -0
- ultralytics/solutions/object_cropper.py +84 -0
- ultralytics/solutions/parking_management.py +273 -0
- ultralytics/solutions/queue_management.py +93 -0
- ultralytics/solutions/region_counter.py +120 -0
- ultralytics/solutions/security_alarm.py +154 -0
- ultralytics/solutions/similarity_search.py +172 -0
- ultralytics/solutions/solutions.py +724 -0
- ultralytics/solutions/speed_estimation.py +110 -0
- ultralytics/solutions/streamlit_inference.py +196 -0
- ultralytics/solutions/templates/similarity-search.html +160 -0
- ultralytics/solutions/trackzone.py +88 -0
- ultralytics/solutions/vision_eye.py +68 -0
- ultralytics/trackers/__init__.py +7 -0
- ultralytics/trackers/basetrack.py +124 -0
- ultralytics/trackers/bot_sort.py +260 -0
- ultralytics/trackers/byte_tracker.py +480 -0
- ultralytics/trackers/track.py +125 -0
- ultralytics/trackers/utils/__init__.py +1 -0
- ultralytics/trackers/utils/gmc.py +376 -0
- ultralytics/trackers/utils/kalman_filter.py +493 -0
- ultralytics/trackers/utils/matching.py +157 -0
- ultralytics/utils/__init__.py +1435 -0
- ultralytics/utils/autobatch.py +106 -0
- ultralytics/utils/autodevice.py +174 -0
- ultralytics/utils/benchmarks.py +695 -0
- ultralytics/utils/callbacks/__init__.py +5 -0
- ultralytics/utils/callbacks/base.py +234 -0
- ultralytics/utils/callbacks/clearml.py +153 -0
- ultralytics/utils/callbacks/comet.py +552 -0
- ultralytics/utils/callbacks/dvc.py +205 -0
- ultralytics/utils/callbacks/hub.py +108 -0
- ultralytics/utils/callbacks/mlflow.py +138 -0
- ultralytics/utils/callbacks/neptune.py +140 -0
- ultralytics/utils/callbacks/raytune.py +43 -0
- ultralytics/utils/callbacks/tensorboard.py +132 -0
- ultralytics/utils/callbacks/wb.py +185 -0
- ultralytics/utils/checks.py +897 -0
- ultralytics/utils/dist.py +119 -0
- ultralytics/utils/downloads.py +499 -0
- ultralytics/utils/errors.py +43 -0
- ultralytics/utils/export.py +219 -0
- ultralytics/utils/files.py +221 -0
- ultralytics/utils/instance.py +499 -0
- ultralytics/utils/loss.py +813 -0
- ultralytics/utils/metrics.py +1356 -0
- ultralytics/utils/ops.py +885 -0
- ultralytics/utils/patches.py +143 -0
- ultralytics/utils/plotting.py +1011 -0
- ultralytics/utils/tal.py +416 -0
- ultralytics/utils/torch_utils.py +990 -0
- ultralytics/utils/triton.py +116 -0
- ultralytics/utils/tuner.py +159 -0
tests/__init__.py
ADDED
@@ -0,0 +1,22 @@
|
|
1
|
+
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
2
|
+
|
3
|
+
from ultralytics.utils import ASSETS, ROOT, WEIGHTS_DIR, checks
|
4
|
+
|
5
|
+
# Constants used in tests
|
6
|
+
MODEL = WEIGHTS_DIR / "path with spaces" / "yolo11n.pt" # test spaces in path
|
7
|
+
CFG = "yolo11n.yaml"
|
8
|
+
SOURCE = ASSETS / "bus.jpg"
|
9
|
+
SOURCES_LIST = [ASSETS / "bus.jpg", ASSETS, ASSETS / "*", ASSETS / "**/*.jpg"]
|
10
|
+
TMP = (ROOT / "../tests/tmp").resolve() # temp directory for test files
|
11
|
+
CUDA_IS_AVAILABLE = checks.cuda_is_available()
|
12
|
+
CUDA_DEVICE_COUNT = checks.cuda_device_count()
|
13
|
+
|
14
|
+
__all__ = (
|
15
|
+
"MODEL",
|
16
|
+
"CFG",
|
17
|
+
"SOURCE",
|
18
|
+
"SOURCES_LIST",
|
19
|
+
"TMP",
|
20
|
+
"CUDA_IS_AVAILABLE",
|
21
|
+
"CUDA_DEVICE_COUNT",
|
22
|
+
)
|
tests/conftest.py
ADDED
@@ -0,0 +1,83 @@
|
|
1
|
+
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
2
|
+
|
3
|
+
import shutil
|
4
|
+
from pathlib import Path
|
5
|
+
|
6
|
+
from tests import TMP
|
7
|
+
|
8
|
+
|
9
|
+
def pytest_addoption(parser):
|
10
|
+
"""
|
11
|
+
Add custom command-line options to pytest.
|
12
|
+
|
13
|
+
Args:
|
14
|
+
parser (pytest.config.Parser): The pytest parser object for adding custom command-line options.
|
15
|
+
|
16
|
+
Returns:
|
17
|
+
(None)
|
18
|
+
"""
|
19
|
+
parser.addoption("--slow", action="store_true", default=False, help="Run slow tests")
|
20
|
+
|
21
|
+
|
22
|
+
def pytest_collection_modifyitems(config, items):
|
23
|
+
"""
|
24
|
+
Modify the list of test items to exclude tests marked as slow if the --slow option is not specified.
|
25
|
+
|
26
|
+
Args:
|
27
|
+
config (pytest.config.Config): The pytest configuration object that provides access to command-line options.
|
28
|
+
items (list): The list of collected pytest item objects to be modified based on the presence of --slow option.
|
29
|
+
|
30
|
+
Returns:
|
31
|
+
(None): The function modifies the 'items' list in place.
|
32
|
+
"""
|
33
|
+
if not config.getoption("--slow"):
|
34
|
+
# Remove the item entirely from the list of test items if it's marked as 'slow'
|
35
|
+
items[:] = [item for item in items if "slow" not in item.keywords]
|
36
|
+
|
37
|
+
|
38
|
+
def pytest_sessionstart(session):
|
39
|
+
"""
|
40
|
+
Initialize session configurations for pytest.
|
41
|
+
|
42
|
+
This function is automatically called by pytest after the 'Session' object has been created but before performing
|
43
|
+
test collection. It sets the initial seeds and prepares the temporary directory for the test session.
|
44
|
+
|
45
|
+
Args:
|
46
|
+
session (pytest.Session): The pytest session object.
|
47
|
+
|
48
|
+
Returns:
|
49
|
+
(None)
|
50
|
+
"""
|
51
|
+
from ultralytics.utils.torch_utils import init_seeds
|
52
|
+
|
53
|
+
init_seeds()
|
54
|
+
shutil.rmtree(TMP, ignore_errors=True) # delete any existing tests/tmp directory
|
55
|
+
TMP.mkdir(parents=True, exist_ok=True) # create a new empty directory
|
56
|
+
|
57
|
+
|
58
|
+
def pytest_terminal_summary(terminalreporter, exitstatus, config):
|
59
|
+
"""
|
60
|
+
Cleanup operations after pytest session.
|
61
|
+
|
62
|
+
This function is automatically called by pytest at the end of the entire test session. It removes certain files
|
63
|
+
and directories used during testing.
|
64
|
+
|
65
|
+
Args:
|
66
|
+
terminalreporter (pytest.terminal.TerminalReporter): The terminal reporter object used for terminal output.
|
67
|
+
exitstatus (int): The exit status of the test run.
|
68
|
+
config (pytest.config.Config): The pytest config object.
|
69
|
+
|
70
|
+
Returns:
|
71
|
+
(None)
|
72
|
+
"""
|
73
|
+
from ultralytics.utils import WEIGHTS_DIR
|
74
|
+
|
75
|
+
# Remove files
|
76
|
+
models = [path for x in ["*.onnx", "*.torchscript"] for path in WEIGHTS_DIR.rglob(x)]
|
77
|
+
for file in ["decelera_portrait_min.mov", "bus.jpg", "yolo11n.onnx", "yolo11n.torchscript"] + models:
|
78
|
+
Path(file).unlink(missing_ok=True)
|
79
|
+
|
80
|
+
# Remove directories
|
81
|
+
models = [path for x in ["*.mlpackage", "*_openvino_model"] for path in WEIGHTS_DIR.rglob(x)]
|
82
|
+
for directory in [WEIGHTS_DIR / "path with spaces", TMP.parents[1] / ".pytest_cache", TMP] + models:
|
83
|
+
shutil.rmtree(directory, ignore_errors=True)
|
tests/test_cli.py
ADDED
@@ -0,0 +1,138 @@
|
|
1
|
+
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
2
|
+
|
3
|
+
import subprocess
|
4
|
+
|
5
|
+
import pytest
|
6
|
+
from PIL import Image
|
7
|
+
|
8
|
+
from tests import CUDA_DEVICE_COUNT, CUDA_IS_AVAILABLE
|
9
|
+
from ultralytics.cfg import TASK2DATA, TASK2MODEL, TASKS
|
10
|
+
from ultralytics.utils import ARM64, ASSETS, LINUX, WEIGHTS_DIR, checks
|
11
|
+
from ultralytics.utils.torch_utils import TORCH_1_9
|
12
|
+
|
13
|
+
# Constants
|
14
|
+
TASK_MODEL_DATA = [(task, WEIGHTS_DIR / TASK2MODEL[task], TASK2DATA[task]) for task in TASKS]
|
15
|
+
MODELS = [WEIGHTS_DIR / TASK2MODEL[task] for task in TASKS]
|
16
|
+
|
17
|
+
|
18
|
+
def run(cmd: str) -> None:
|
19
|
+
"""Execute a shell command using subprocess."""
|
20
|
+
subprocess.run(cmd.split(), check=True)
|
21
|
+
|
22
|
+
|
23
|
+
def test_special_modes() -> None:
|
24
|
+
"""Test various special command-line modes for YOLO functionality."""
|
25
|
+
run("yolo help")
|
26
|
+
run("yolo checks")
|
27
|
+
run("yolo version")
|
28
|
+
run("yolo settings reset")
|
29
|
+
run("yolo cfg")
|
30
|
+
|
31
|
+
|
32
|
+
@pytest.mark.parametrize("task,model,data", TASK_MODEL_DATA)
|
33
|
+
def test_train(task: str, model: str, data: str) -> None:
|
34
|
+
"""Test YOLO training for different tasks, models, and datasets."""
|
35
|
+
run(f"yolo train {task} model={model} data={data} imgsz=32 epochs=1 cache=disk")
|
36
|
+
|
37
|
+
|
38
|
+
@pytest.mark.parametrize("task,model,data", TASK_MODEL_DATA)
|
39
|
+
def test_val(task: str, model: str, data: str) -> None:
|
40
|
+
"""Test YOLO validation process for specified task, model, and data using a shell command."""
|
41
|
+
run(f"yolo val {task} model={model} data={data} imgsz=32 save_txt save_json")
|
42
|
+
|
43
|
+
|
44
|
+
@pytest.mark.parametrize("task,model,data", TASK_MODEL_DATA)
|
45
|
+
def test_predict(task: str, model: str, data: str) -> None:
|
46
|
+
"""Test YOLO prediction on provided sample assets for specified task and model."""
|
47
|
+
run(f"yolo predict model={model} source={ASSETS} imgsz=32 save save_crop save_txt")
|
48
|
+
|
49
|
+
|
50
|
+
@pytest.mark.parametrize("model", MODELS)
|
51
|
+
def test_export(model: str) -> None:
|
52
|
+
"""Test exporting a YOLO model to TorchScript format."""
|
53
|
+
run(f"yolo export model={model} format=torchscript imgsz=32")
|
54
|
+
|
55
|
+
|
56
|
+
def test_rtdetr(task: str = "detect", model: str = "yolov8n-rtdetr.yaml", data: str = "coco8.yaml") -> None:
|
57
|
+
"""Test the RTDETR functionality within Ultralytics for detection tasks using specified model and data."""
|
58
|
+
# Warning: must use imgsz=640 (note also add coma, spaces, fraction=0.25 args to test single-image training)
|
59
|
+
run(f"yolo train {task} model={model} data={data} --imgsz= 160 epochs =1, cache = disk fraction=0.25") # spaces
|
60
|
+
run(f"yolo predict {task} model={model} source={ASSETS / 'bus.jpg'} imgsz=160 save save_crop save_txt")
|
61
|
+
if TORCH_1_9:
|
62
|
+
weights = WEIGHTS_DIR / "rtdetr-l.pt"
|
63
|
+
run(f"yolo predict {task} model={weights} source={ASSETS / 'bus.jpg'} imgsz=160 save save_crop save_txt")
|
64
|
+
run(f"yolo train {task} model={weights} epochs=1 imgsz=160 cache=disk data=coco8.yaml")
|
65
|
+
|
66
|
+
|
67
|
+
@pytest.mark.skipif(checks.IS_PYTHON_3_12, reason="MobileSAM with CLIP is not supported in Python 3.12")
|
68
|
+
@pytest.mark.skipif(
|
69
|
+
checks.IS_PYTHON_3_8 and LINUX and ARM64,
|
70
|
+
reason="MobileSAM with CLIP is not supported in Python 3.8 and aarch64 Linux",
|
71
|
+
)
|
72
|
+
def test_fastsam(
|
73
|
+
task: str = "segment", model: str = WEIGHTS_DIR / "FastSAM-s.pt", data: str = "coco8-seg.yaml"
|
74
|
+
) -> None:
|
75
|
+
"""Test FastSAM model for segmenting objects in images using various prompts within Ultralytics."""
|
76
|
+
source = ASSETS / "bus.jpg"
|
77
|
+
|
78
|
+
run(f"yolo segment val {task} model={model} data={data} imgsz=32")
|
79
|
+
run(f"yolo segment predict model={model} source={source} imgsz=32 save save_crop save_txt")
|
80
|
+
|
81
|
+
from ultralytics import FastSAM
|
82
|
+
from ultralytics.models.sam import Predictor
|
83
|
+
|
84
|
+
# Create a FastSAM model
|
85
|
+
sam_model = FastSAM(model) # or FastSAM-x.pt
|
86
|
+
|
87
|
+
# Run inference on an image
|
88
|
+
for s in (source, Image.open(source)):
|
89
|
+
everything_results = sam_model(s, device="cpu", retina_masks=True, imgsz=320, conf=0.4, iou=0.9)
|
90
|
+
|
91
|
+
# Remove small regions
|
92
|
+
new_masks, _ = Predictor.remove_small_regions(everything_results[0].masks.data, min_area=20)
|
93
|
+
|
94
|
+
# Run inference with bboxes and points and texts prompt at the same time
|
95
|
+
sam_model(source, bboxes=[439, 437, 524, 709], points=[[200, 200]], labels=[1], texts="a photo of a dog")
|
96
|
+
|
97
|
+
|
98
|
+
def test_mobilesam() -> None:
|
99
|
+
"""Test MobileSAM segmentation with point prompts using Ultralytics."""
|
100
|
+
from ultralytics import SAM
|
101
|
+
|
102
|
+
# Load the model
|
103
|
+
model = SAM(WEIGHTS_DIR / "mobile_sam.pt")
|
104
|
+
|
105
|
+
# Source
|
106
|
+
source = ASSETS / "zidane.jpg"
|
107
|
+
|
108
|
+
# Predict a segment based on a 1D point prompt and 1D labels.
|
109
|
+
model.predict(source, points=[900, 370], labels=[1])
|
110
|
+
|
111
|
+
# Predict a segment based on 3D points and 2D labels (multiple points per object).
|
112
|
+
model.predict(source, points=[[[900, 370], [1000, 100]]], labels=[[1, 1]])
|
113
|
+
|
114
|
+
# Predict a segment based on a box prompt
|
115
|
+
model.predict(source, bboxes=[439, 437, 524, 709], save=True)
|
116
|
+
|
117
|
+
# Predict all
|
118
|
+
# model(source)
|
119
|
+
|
120
|
+
|
121
|
+
# Slow Tests -----------------------------------------------------------------------------------------------------------
|
122
|
+
@pytest.mark.slow
|
123
|
+
@pytest.mark.parametrize("task,model,data", TASK_MODEL_DATA)
|
124
|
+
@pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason="CUDA is not available")
|
125
|
+
@pytest.mark.skipif(CUDA_DEVICE_COUNT < 2, reason="DDP is not available")
|
126
|
+
def test_train_gpu(task: str, model: str, data: str) -> None:
|
127
|
+
"""Test YOLO training on GPU(s) for various tasks and models."""
|
128
|
+
run(f"yolo train {task} model={model} data={data} imgsz=32 epochs=1 device=0") # single GPU
|
129
|
+
run(f"yolo train {task} model={model} data={data} imgsz=32 epochs=1 device=0,1") # multi GPU
|
130
|
+
|
131
|
+
|
132
|
+
@pytest.mark.parametrize(
|
133
|
+
"solution",
|
134
|
+
["count", "blur", "workout", "heatmap", "isegment", "visioneye", "speed", "queue", "analytics", "trackzone"],
|
135
|
+
)
|
136
|
+
def test_solutions(solution: str) -> None:
|
137
|
+
"""Test yolo solutions command-line modes."""
|
138
|
+
run(f"yolo solutions {solution} verbose=False")
|
tests/test_cuda.py
ADDED
@@ -0,0 +1,215 @@
|
|
1
|
+
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
2
|
+
|
3
|
+
from itertools import product
|
4
|
+
from pathlib import Path
|
5
|
+
|
6
|
+
import pytest
|
7
|
+
import torch
|
8
|
+
|
9
|
+
from tests import CUDA_DEVICE_COUNT, CUDA_IS_AVAILABLE, MODEL, SOURCE
|
10
|
+
from ultralytics import YOLO
|
11
|
+
from ultralytics.cfg import TASK2DATA, TASK2MODEL, TASKS
|
12
|
+
from ultralytics.utils import ASSETS, IS_JETSON, WEIGHTS_DIR
|
13
|
+
from ultralytics.utils.autodevice import GPUInfo
|
14
|
+
from ultralytics.utils.checks import check_amp
|
15
|
+
from ultralytics.utils.torch_utils import TORCH_1_13
|
16
|
+
|
17
|
+
# Try to find idle devices if CUDA is available
|
18
|
+
DEVICES = []
|
19
|
+
if CUDA_IS_AVAILABLE:
|
20
|
+
if IS_JETSON:
|
21
|
+
DEVICES = [0] # NVIDIA Jetson only has one GPU and does not fully support pynvml library
|
22
|
+
else:
|
23
|
+
gpu_info = GPUInfo()
|
24
|
+
gpu_info.print_status()
|
25
|
+
idle_gpus = gpu_info.select_idle_gpu(count=2, min_memory_mb=2048)
|
26
|
+
if idle_gpus:
|
27
|
+
DEVICES = idle_gpus
|
28
|
+
|
29
|
+
|
30
|
+
def test_checks():
|
31
|
+
"""Validate CUDA settings against torch CUDA functions."""
|
32
|
+
assert torch.cuda.is_available() == CUDA_IS_AVAILABLE
|
33
|
+
assert torch.cuda.device_count() == CUDA_DEVICE_COUNT
|
34
|
+
|
35
|
+
|
36
|
+
@pytest.mark.skipif(not DEVICES, reason="No CUDA devices available")
|
37
|
+
def test_amp():
|
38
|
+
"""Test AMP training checks."""
|
39
|
+
model = YOLO("yolo11n.pt").model.to(f"cuda:{DEVICES[0]}")
|
40
|
+
assert check_amp(model)
|
41
|
+
|
42
|
+
|
43
|
+
@pytest.mark.slow
|
44
|
+
# @pytest.mark.skipif(IS_JETSON, reason="Temporary disable ONNX for Jetson")
|
45
|
+
@pytest.mark.skipif(not DEVICES, reason="No CUDA devices available")
|
46
|
+
@pytest.mark.parametrize(
|
47
|
+
"task, dynamic, int8, half, batch, simplify, nms",
|
48
|
+
[ # generate all combinations except for exclusion cases
|
49
|
+
(task, dynamic, int8, half, batch, simplify, nms)
|
50
|
+
for task, dynamic, int8, half, batch, simplify, nms in product(
|
51
|
+
TASKS, [True, False], [False], [False], [1, 2], [True, False], [True, False]
|
52
|
+
)
|
53
|
+
if not (
|
54
|
+
(int8 and half)
|
55
|
+
or (task == "classify" and nms)
|
56
|
+
or (task == "obb" and nms and (not TORCH_1_13 or IS_JETSON)) # obb nms fails on NVIDIA Jetson
|
57
|
+
or (simplify and dynamic) # onnxslim is slow when dynamic=True
|
58
|
+
)
|
59
|
+
],
|
60
|
+
)
|
61
|
+
def test_export_onnx_matrix(task, dynamic, int8, half, batch, simplify, nms):
|
62
|
+
"""Test YOLO exports to ONNX format with various configurations and parameters."""
|
63
|
+
file = YOLO(TASK2MODEL[task]).export(
|
64
|
+
format="onnx",
|
65
|
+
imgsz=32,
|
66
|
+
dynamic=dynamic,
|
67
|
+
int8=int8,
|
68
|
+
half=half,
|
69
|
+
batch=batch,
|
70
|
+
simplify=simplify,
|
71
|
+
nms=nms,
|
72
|
+
device=DEVICES[0],
|
73
|
+
)
|
74
|
+
YOLO(file)([SOURCE] * batch, imgsz=64 if dynamic else 32, device=DEVICES[0]) # exported model inference
|
75
|
+
Path(file).unlink() # cleanup
|
76
|
+
|
77
|
+
|
78
|
+
@pytest.mark.slow
|
79
|
+
@pytest.mark.skipif(True, reason="CUDA export tests disabled pending additional Ultralytics GPU server availability")
|
80
|
+
@pytest.mark.skipif(not DEVICES, reason="No CUDA devices available")
|
81
|
+
@pytest.mark.parametrize(
|
82
|
+
"task, dynamic, int8, half, batch",
|
83
|
+
[ # generate all combinations but exclude those where both int8 and half are True
|
84
|
+
(task, dynamic, int8, half, batch)
|
85
|
+
# Note: tests reduced below pending compute availability expansion as GPU CI runner utilization is high
|
86
|
+
# for task, dynamic, int8, half, batch in product(TASKS, [True, False], [True, False], [True, False], [1, 2])
|
87
|
+
for task, dynamic, int8, half, batch in product(TASKS, [True], [True], [False], [2])
|
88
|
+
if not (int8 and half) # exclude cases where both int8 and half are True
|
89
|
+
],
|
90
|
+
)
|
91
|
+
def test_export_engine_matrix(task, dynamic, int8, half, batch):
|
92
|
+
"""Test YOLO model export to TensorRT format for various configurations and run inference."""
|
93
|
+
file = YOLO(TASK2MODEL[task]).export(
|
94
|
+
format="engine",
|
95
|
+
imgsz=32,
|
96
|
+
dynamic=dynamic,
|
97
|
+
int8=int8,
|
98
|
+
half=half,
|
99
|
+
batch=batch,
|
100
|
+
data=TASK2DATA[task],
|
101
|
+
workspace=1, # reduce workspace GB for less resource utilization during testing
|
102
|
+
simplify=True,
|
103
|
+
device=DEVICES[0],
|
104
|
+
)
|
105
|
+
YOLO(file)([SOURCE] * batch, imgsz=64 if dynamic else 32, device=DEVICES[0]) # exported model inference
|
106
|
+
Path(file).unlink() # cleanup
|
107
|
+
Path(file).with_suffix(".cache").unlink() if int8 else None # cleanup INT8 cache
|
108
|
+
|
109
|
+
|
110
|
+
@pytest.mark.skipif(not DEVICES, reason="No CUDA devices available")
|
111
|
+
def test_train():
|
112
|
+
"""Test model training on a minimal dataset using available CUDA devices."""
|
113
|
+
import os
|
114
|
+
|
115
|
+
device = tuple(DEVICES) if len(DEVICES) > 1 else DEVICES[0]
|
116
|
+
results = YOLO(MODEL).train(data="coco8.yaml", imgsz=64, epochs=1, device=device) # requires imgsz>=64
|
117
|
+
# NVIDIA Jetson only has one GPU and therefore skipping checks
|
118
|
+
if not IS_JETSON:
|
119
|
+
visible = eval(os.environ["CUDA_VISIBLE_DEVICES"])
|
120
|
+
assert visible == device, f"Passed GPUs '{device}', but used GPUs '{visible}'"
|
121
|
+
assert results is (None if len(DEVICES) > 1 else not None) # DDP returns None, single-GPU returns metrics
|
122
|
+
|
123
|
+
|
124
|
+
@pytest.mark.slow
|
125
|
+
@pytest.mark.skipif(not DEVICES, reason="No CUDA devices available")
|
126
|
+
def test_predict_multiple_devices():
|
127
|
+
"""Validate model prediction consistency across CPU and CUDA devices."""
|
128
|
+
model = YOLO("yolo11n.pt")
|
129
|
+
|
130
|
+
# Test CPU
|
131
|
+
model = model.cpu()
|
132
|
+
assert str(model.device) == "cpu"
|
133
|
+
_ = model(SOURCE)
|
134
|
+
assert str(model.device) == "cpu"
|
135
|
+
|
136
|
+
# Test CUDA
|
137
|
+
cuda_device = f"cuda:{DEVICES[0]}"
|
138
|
+
model = model.to(cuda_device)
|
139
|
+
assert str(model.device) == cuda_device
|
140
|
+
_ = model(SOURCE)
|
141
|
+
assert str(model.device) == cuda_device
|
142
|
+
|
143
|
+
# Test CPU again
|
144
|
+
model = model.cpu()
|
145
|
+
assert str(model.device) == "cpu"
|
146
|
+
_ = model(SOURCE)
|
147
|
+
assert str(model.device) == "cpu"
|
148
|
+
|
149
|
+
# Test CUDA again
|
150
|
+
model = model.to(cuda_device)
|
151
|
+
assert str(model.device) == cuda_device
|
152
|
+
_ = model(SOURCE)
|
153
|
+
assert str(model.device) == cuda_device
|
154
|
+
|
155
|
+
|
156
|
+
@pytest.mark.skipif(not DEVICES, reason="No CUDA devices available")
|
157
|
+
def test_autobatch():
|
158
|
+
"""Check optimal batch size for YOLO model training using autobatch utility."""
|
159
|
+
from ultralytics.utils.autobatch import check_train_batch_size
|
160
|
+
|
161
|
+
check_train_batch_size(YOLO(MODEL).model.to(f"cuda:{DEVICES[0]}"), imgsz=128, amp=True)
|
162
|
+
|
163
|
+
|
164
|
+
@pytest.mark.slow
|
165
|
+
@pytest.mark.skipif(not DEVICES, reason="No CUDA devices available")
|
166
|
+
def test_utils_benchmarks():
|
167
|
+
"""Profile YOLO models for performance benchmarks."""
|
168
|
+
from ultralytics.utils.benchmarks import ProfileModels
|
169
|
+
|
170
|
+
# Pre-export a dynamic engine model to use dynamic inference
|
171
|
+
YOLO(MODEL).export(format="engine", imgsz=32, dynamic=True, batch=1, device=DEVICES[0])
|
172
|
+
ProfileModels(
|
173
|
+
[MODEL],
|
174
|
+
imgsz=32,
|
175
|
+
half=False,
|
176
|
+
min_time=1,
|
177
|
+
num_timed_runs=3,
|
178
|
+
num_warmup_runs=1,
|
179
|
+
device=DEVICES[0],
|
180
|
+
).run()
|
181
|
+
|
182
|
+
|
183
|
+
@pytest.mark.skipif(not DEVICES, reason="No CUDA devices available")
|
184
|
+
def test_predict_sam():
|
185
|
+
"""Test SAM model predictions using different prompts."""
|
186
|
+
from ultralytics import SAM
|
187
|
+
from ultralytics.models.sam import Predictor as SAMPredictor
|
188
|
+
|
189
|
+
model = SAM(WEIGHTS_DIR / "sam2.1_b.pt")
|
190
|
+
model.info()
|
191
|
+
|
192
|
+
# Run inference with various prompts
|
193
|
+
model(SOURCE, device=DEVICES[0])
|
194
|
+
model(SOURCE, bboxes=[439, 437, 524, 709], device=DEVICES[0])
|
195
|
+
model(ASSETS / "zidane.jpg", points=[900, 370], device=DEVICES[0])
|
196
|
+
model(ASSETS / "zidane.jpg", points=[900, 370], labels=[1], device=DEVICES[0])
|
197
|
+
model(ASSETS / "zidane.jpg", points=[[900, 370]], labels=[1], device=DEVICES[0])
|
198
|
+
model(ASSETS / "zidane.jpg", points=[[400, 370], [900, 370]], labels=[1, 1], device=DEVICES[0])
|
199
|
+
model(ASSETS / "zidane.jpg", points=[[[900, 370], [1000, 100]]], labels=[[1, 1]], device=DEVICES[0])
|
200
|
+
|
201
|
+
# Test predictor
|
202
|
+
predictor = SAMPredictor(
|
203
|
+
overrides=dict(
|
204
|
+
conf=0.25,
|
205
|
+
task="segment",
|
206
|
+
mode="predict",
|
207
|
+
imgsz=1024,
|
208
|
+
model=WEIGHTS_DIR / "mobile_sam.pt",
|
209
|
+
device=DEVICES[0],
|
210
|
+
)
|
211
|
+
)
|
212
|
+
predictor.set_image(ASSETS / "zidane.jpg")
|
213
|
+
# predictor(bboxes=[439, 437, 524, 709])
|
214
|
+
# predictor(points=[900, 370], labels=[1])
|
215
|
+
predictor.reset_image()
|
tests/test_engine.py
ADDED
@@ -0,0 +1,131 @@
|
|
1
|
+
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
2
|
+
|
3
|
+
import sys
|
4
|
+
from unittest import mock
|
5
|
+
|
6
|
+
from tests import MODEL
|
7
|
+
from ultralytics import YOLO
|
8
|
+
from ultralytics.cfg import get_cfg
|
9
|
+
from ultralytics.engine.exporter import Exporter
|
10
|
+
from ultralytics.models.yolo import classify, detect, segment
|
11
|
+
from ultralytics.utils import ASSETS, DEFAULT_CFG, WEIGHTS_DIR
|
12
|
+
|
13
|
+
|
14
|
+
def test_func(*args): # noqa
|
15
|
+
"""Test function callback for evaluating YOLO model performance metrics."""
|
16
|
+
print("callback test passed")
|
17
|
+
|
18
|
+
|
19
|
+
def test_export():
|
20
|
+
"""Tests the model exporting function by adding a callback and asserting its execution."""
|
21
|
+
exporter = Exporter()
|
22
|
+
exporter.add_callback("on_export_start", test_func)
|
23
|
+
assert test_func in exporter.callbacks["on_export_start"], "callback test failed"
|
24
|
+
f = exporter(model=YOLO("yolo11n.yaml").model)
|
25
|
+
YOLO(f)(ASSETS) # exported model inference
|
26
|
+
|
27
|
+
|
28
|
+
def test_detect():
|
29
|
+
"""Test YOLO object detection training, validation, and prediction functionality."""
|
30
|
+
overrides = {"data": "coco8.yaml", "model": "yolo11n.yaml", "imgsz": 32, "epochs": 1, "save": False}
|
31
|
+
cfg = get_cfg(DEFAULT_CFG)
|
32
|
+
cfg.data = "coco8.yaml"
|
33
|
+
cfg.imgsz = 32
|
34
|
+
|
35
|
+
# Trainer
|
36
|
+
trainer = detect.DetectionTrainer(overrides=overrides)
|
37
|
+
trainer.add_callback("on_train_start", test_func)
|
38
|
+
assert test_func in trainer.callbacks["on_train_start"], "callback test failed"
|
39
|
+
trainer.train()
|
40
|
+
|
41
|
+
# Validator
|
42
|
+
val = detect.DetectionValidator(args=cfg)
|
43
|
+
val.add_callback("on_val_start", test_func)
|
44
|
+
assert test_func in val.callbacks["on_val_start"], "callback test failed"
|
45
|
+
val(model=trainer.best) # validate best.pt
|
46
|
+
|
47
|
+
# Predictor
|
48
|
+
pred = detect.DetectionPredictor(overrides={"imgsz": [64, 64]})
|
49
|
+
pred.add_callback("on_predict_start", test_func)
|
50
|
+
assert test_func in pred.callbacks["on_predict_start"], "callback test failed"
|
51
|
+
# Confirm there is no issue with sys.argv being empty.
|
52
|
+
with mock.patch.object(sys, "argv", []):
|
53
|
+
result = pred(source=ASSETS, model=MODEL)
|
54
|
+
assert len(result), "predictor test failed"
|
55
|
+
|
56
|
+
overrides["resume"] = trainer.last
|
57
|
+
trainer = detect.DetectionTrainer(overrides=overrides)
|
58
|
+
try:
|
59
|
+
trainer.train()
|
60
|
+
except Exception as e:
|
61
|
+
print(f"Expected exception caught: {e}")
|
62
|
+
return
|
63
|
+
|
64
|
+
Exception("Resume test failed!")
|
65
|
+
|
66
|
+
|
67
|
+
def test_segment():
|
68
|
+
"""Tests image segmentation training, validation, and prediction pipelines using YOLO models."""
|
69
|
+
overrides = {"data": "coco8-seg.yaml", "model": "yolo11n-seg.yaml", "imgsz": 32, "epochs": 1, "save": False}
|
70
|
+
cfg = get_cfg(DEFAULT_CFG)
|
71
|
+
cfg.data = "coco8-seg.yaml"
|
72
|
+
cfg.imgsz = 32
|
73
|
+
# YOLO(CFG_SEG).train(**overrides) # works
|
74
|
+
|
75
|
+
# Trainer
|
76
|
+
trainer = segment.SegmentationTrainer(overrides=overrides)
|
77
|
+
trainer.add_callback("on_train_start", test_func)
|
78
|
+
assert test_func in trainer.callbacks["on_train_start"], "callback test failed"
|
79
|
+
trainer.train()
|
80
|
+
|
81
|
+
# Validator
|
82
|
+
val = segment.SegmentationValidator(args=cfg)
|
83
|
+
val.add_callback("on_val_start", test_func)
|
84
|
+
assert test_func in val.callbacks["on_val_start"], "callback test failed"
|
85
|
+
val(model=trainer.best) # validate best.pt
|
86
|
+
|
87
|
+
# Predictor
|
88
|
+
pred = segment.SegmentationPredictor(overrides={"imgsz": [64, 64]})
|
89
|
+
pred.add_callback("on_predict_start", test_func)
|
90
|
+
assert test_func in pred.callbacks["on_predict_start"], "callback test failed"
|
91
|
+
result = pred(source=ASSETS, model=WEIGHTS_DIR / "yolo11n-seg.pt")
|
92
|
+
assert len(result), "predictor test failed"
|
93
|
+
|
94
|
+
# Test resume
|
95
|
+
overrides["resume"] = trainer.last
|
96
|
+
trainer = segment.SegmentationTrainer(overrides=overrides)
|
97
|
+
try:
|
98
|
+
trainer.train()
|
99
|
+
except Exception as e:
|
100
|
+
print(f"Expected exception caught: {e}")
|
101
|
+
return
|
102
|
+
|
103
|
+
Exception("Resume test failed!")
|
104
|
+
|
105
|
+
|
106
|
+
def test_classify():
|
107
|
+
"""Test image classification including training, validation, and prediction phases."""
|
108
|
+
overrides = {"data": "imagenet10", "model": "yolo11n-cls.yaml", "imgsz": 32, "epochs": 1, "save": False}
|
109
|
+
cfg = get_cfg(DEFAULT_CFG)
|
110
|
+
cfg.data = "imagenet10"
|
111
|
+
cfg.imgsz = 32
|
112
|
+
# YOLO(CFG_SEG).train(**overrides) # works
|
113
|
+
|
114
|
+
# Trainer
|
115
|
+
trainer = classify.ClassificationTrainer(overrides=overrides)
|
116
|
+
trainer.add_callback("on_train_start", test_func)
|
117
|
+
assert test_func in trainer.callbacks["on_train_start"], "callback test failed"
|
118
|
+
trainer.train()
|
119
|
+
|
120
|
+
# Validator
|
121
|
+
val = classify.ClassificationValidator(args=cfg)
|
122
|
+
val.add_callback("on_val_start", test_func)
|
123
|
+
assert test_func in val.callbacks["on_val_start"], "callback test failed"
|
124
|
+
val(model=trainer.best)
|
125
|
+
|
126
|
+
# Predictor
|
127
|
+
pred = classify.ClassificationPredictor(overrides={"imgsz": [64, 64]})
|
128
|
+
pred.add_callback("on_predict_start", test_func)
|
129
|
+
assert test_func in pred.callbacks["on_predict_start"], "callback test failed"
|
130
|
+
result = pred(source=ASSETS, model=trainer.best)
|
131
|
+
assert len(result), "predictor test failed"
|