Skip to content

GitLab

Commit e63cf68a authored by chenzk's avatar chenzk
Browse files

v1.0

parents
Pipeline #2842 canceled with stages
Hide whitespace changes
Inline Side-by-side
Showing with 2888 additions and 0 deletions
model.properties 0 → 100644
View file @ e63cf68a
# 模型编码
modelCode=1667
# 模型名称
modelName=yolov13_pytorch
# 模型描述
modelDescription=YOLOv13引入超图理论，其中YOLOv13-N相比YOLO11-N提升了3.0%的mAP，相比YOLOv12-N提升了1.5%的mAP。
# 应用场景
appScenario=推理,训练,目标检测,制造,电商,医疗,能源,教育
# 框架类型
frameType=pytorch
prediction.py 0 → 100644
View file @ e63cf68a
from ultralytics import YOLO
model = YOLO('yolov13n.pt') # Replace with the desired model scale
model.predict()
pyproject.toml 0 → 100644
View file @ e63cf68a
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
# Overview:
# This pyproject.toml file manages the build, packaging, and distribution of the Ultralytics library.
# It defines essential project metadata, dependencies, and settings used to develop and deploy the library.
# Key Sections:
# - [build-system]: Specifies the build requirements and backend (e.g., setuptools, wheel).
# - [project]: Includes details like name, version, description, authors, dependencies and more.
# - [project.optional-dependencies]: Provides additional, optional packages for extended features.
# - [tool.*]: Configures settings for various tools (pytest, yapf, etc.) used in the project.
# Installation:
# The Ultralytics library can be installed using the command: 'pip install ultralytics'
# For development purposes, you can install the package in editable mode with: 'pip install -e .'
# This approach allows for real-time code modifications without the need for re-installation.
# Documentation:
# For comprehensive documentation and usage instructions, visit: https://docs.ultralytics.com
[build-system]
requires = ["setuptools>=70.0.0", "wheel"]
build-backend = "setuptools.build_meta"
# Project settings -----------------------------------------------------------------------------------------------------
[project]
name = "ultralytics"
dynamic = ["version"]
description = "Ultralytics YOLO 🚀 for SOTA object detection, multi-object tracking, instance segmentation, pose estimation and image classification."
readme = "README.md"
requires-python = ">=3.8"
license = { "text" = "AGPL-3.0" }
keywords = ["machine-learning", "deep-learning", "computer-vision", "ML", "DL", "AI", "YOLO", "YOLOv3", "YOLOv5", "YOLOv8", "YOLOv9", "YOLOv10", "YOLO11", "HUB", "Ultralytics"]
authors = [
{ name = "Glenn Jocher", email = "glenn.jocher@ultralytics.com" },
{ name = "Jing Qiu", email = "jing.qiu@ultralytics.com" },
]
maintainers = [
{ name = "Ultralytics", email = "hello@ultralytics.com" },
]
classifiers = [
"Development Status :: 4 - Beta",
"Intended Audience :: Developers",
"Intended Audience :: Education",
"Intended Audience :: Science/Research",
"License :: OSI Approved :: GNU Affero General Public License v3 or later (AGPLv3+)",
"Programming Language :: Python :: 3",
"Programming Language :: Python :: 3.8",
"Programming Language :: Python :: 3.9",
"Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.11",
"Programming Language :: Python :: 3.12",
"Topic :: Software Development",
"Topic :: Scientific/Engineering",
"Topic :: Scientific/Engineering :: Artificial Intelligence",
"Topic :: Scientific/Engineering :: Image Recognition",
"Operating System :: POSIX :: Linux",
"Operating System :: MacOS",
"Operating System :: Microsoft :: Windows",
]
# Required dependencies ------------------------------------------------------------------------------------------------
dependencies = [
"numpy>=1.23.0",
"numpy<2.0.0; sys_platform == 'darwin'", # macOS OpenVINO errors https://github.com/ultralytics/ultralytics/pull/17221
"matplotlib>=3.3.0",
"opencv-python>=4.6.0",
"pillow>=7.1.2",
"pyyaml>=5.3.1",
"requests>=2.23.0",
"scipy>=1.4.1",
"torch>=1.8.0",
"torch>=1.8.0,!=2.4.0; sys_platform == 'win32'", # Windows CPU errors w/ 2.4.0 https://github.com/ultralytics/ultralytics/issues/15049
"torchvision>=0.9.0",
"tqdm>=4.64.0", # progress bars
"psutil", # system utilization
"py-cpuinfo", # display CPU info
"pandas>=1.1.4",
"seaborn>=0.11.0", # plotting
"ultralytics-thop>=2.0.0", # FLOPs computation https://github.com/ultralytics/thop
]
# Optional dependencies ------------------------------------------------------------------------------------------------
[project.optional-dependencies]
dev = [
"ipython",
"pytest",
"pytest-cov",
"coverage[toml]",
"mkdocs>=1.6.0",
"mkdocs-material>=9.5.9",
"mkdocstrings[python]",
"mkdocs-redirects", # 301 redirects
"mkdocs-ultralytics-plugin>=0.1.8", # for meta descriptions and images, dates and authors
"mkdocs-macros-plugin>=1.0.5" # duplicating content (i.e. export tables) in multiple places
]
export = [
"onnx>=1.12.0", # ONNX export
"coremltools>=7.0; platform_system != 'Windows' and python_version <= '3.11'", # CoreML supported on macOS and Linux
"scikit-learn>=1.3.2; platform_system != 'Windows' and python_version <= '3.11'", # CoreML k-means quantization
"openvino>=2024.0.0", # OpenVINO export
"tensorflow>=2.0.0", # TF bug https://github.com/ultralytics/ultralytics/issues/5161
"tensorflowjs>=3.9.0", # TF.js export, automatically installs tensorflow
"tensorstore>=0.1.63; platform_machine == 'aarch64' and python_version >= '3.9'", # for TF Raspberry Pi exports
"keras", # not installed automatically by tensorflow>=2.16
"flatbuffers>=23.5.26,<100; platform_machine == 'aarch64'", # update old 'flatbuffers' included inside tensorflow package
"numpy==1.23.5; platform_machine == 'aarch64'", # fix error: `np.bool` was a deprecated alias for the builtin `bool` when using TensorRT models on NVIDIA Jetson
"h5py!=3.11.0; platform_machine == 'aarch64'", # fix h5py build issues due to missing aarch64 wheels in 3.11 release
]
solutions = [
"shapely>=2.0.0", # shapely for point and polygon data matching
"streamlit", # for live inference on web browser i.e `yolo streamlit-predict`
]
logging = [
"comet", # https://docs.ultralytics.com/integrations/comet/
"tensorboard>=2.13.0",
"dvclive>=2.12.0",
]
extra = [
"hub-sdk>=0.0.12", # Ultralytics HUB
"ipython", # interactive notebook
"albumentations>=1.4.6", # training augmentations
"pycocotools>=2.0.7", # COCO mAP
]
[project.urls]
"Homepage" = "https://ultralytics.com"
"Source" = "https://github.com/ultralytics/ultralytics"
"Documentation" = "https://docs.ultralytics.com"
"Bug Reports" = "https://github.com/ultralytics/ultralytics/issues"
"Changelog" = "https://github.com/ultralytics/ultralytics/releases"
[project.scripts]
yolo = "ultralytics.cfg:entrypoint"
ultralytics = "ultralytics.cfg:entrypoint"
# Tools settings -------------------------------------------------------------------------------------------------------
[tool.setuptools] # configuration specific to the `setuptools` build backend.
packages = { find = { where = ["."], include = ["ultralytics", "ultralytics.*"] } }
package-data = { "ultralytics" = ["**/*.yaml", "../tests/*.py"], "ultralytics.assets" = ["*.jpg"] }
[tool.setuptools.dynamic]
version = { attr = "ultralytics.__version__" }
[tool.pytest.ini_options]
addopts = "--doctest-modules --durations=30 --color=yes"
markers = [
"slow: skip slow tests unless --slow is set",
]
norecursedirs = [".git", "dist", "build"]
[tool.coverage.run]
source = ["ultralytics/"]
data_file = "tests/.coverage"
omit = ["ultralytics/utils/callbacks/*"]
[tool.isort]
line_length = 120
multi_line_output = 0
[tool.yapf]
based_on_style = "pep8"
spaces_before_comment = 2
column_limit = 120
coalesce_brackets = true
spaces_around_power_operator = true
space_between_ending_comma_and_closing_bracket = true
split_before_closing_bracket = false
split_before_first_argument = false
[tool.ruff]
line-length = 120
[tool.ruff.format]
docstring-code-format = true
[tool.docformatter]
wrap-summaries = 120
wrap-descriptions = 120
pre-summary-newline = true
close-quotes-on-newline = true
in-place = true
[tool.codespell]
ignore-words-list = "crate,nd,ned,strack,dota,ane,segway,fo,gool,winn,commend,bloc,nam,afterall"
skip = '*.pt,*.pth,*.torchscript,*.onnx,*.tflite,*.pb,*.bin,*.param,*.mlmodel,*.engine,*.npy,*.data*,*.csv,*pnnx*,*venv*,*translat*,__pycache__*,*.ico,*.jpg,*.png,*.mp4,*.mov,/runs,/.git,./docs/??/*.md,./docs/mkdocs_??.yml'
requirements.txt 0 → 100644
View file @ e63cf68a
# torch==2.2.2
# torchvision==0.17.2
# flash_attn-2.7.3+cu11torch2.2cxx11abiFALSE-cp311-cp311-linux_x86_64.whl
timm==1.0.14
albumentations==2.0.4
onnx==1.14.0
# onnxruntime==1.15.1
pycocotools==2.0.7
PyYAML==6.0.1
scipy # 1.13.0
onnxslim==0.1.31
# onnxruntime-gpu==1.18.0
gradio==4.44.1
opencv-python # 4.9.0.80
psutil==5.9.8
py-cpuinfo==9.0.0
huggingface-hub==0.33.2
safetensors==0.4.3
numpy==1.26.4
supervision==0.22.0
tests/__init__.py 0 → 100644
View file @ e63cf68a
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
from ultralytics.utils import ASSETS, ROOT, WEIGHTS_DIR, checks
# Constants used in tests
MODEL = WEIGHTS_DIR / "path with spaces" / "yolo11n.pt" # test spaces in path
CFG = "yolo11n.yaml"
SOURCE = ASSETS / "bus.jpg"
SOURCES_LIST = [ASSETS / "bus.jpg", ASSETS, ASSETS / "*", ASSETS / "**/*.jpg"]
TMP = (ROOT / "../tests/tmp").resolve() # temp directory for test files
CUDA_IS_AVAILABLE = checks.cuda_is_available()
CUDA_DEVICE_COUNT = checks.cuda_device_count()
__all__ = (
"MODEL",
"CFG",
"SOURCE",
"SOURCES_LIST",
"TMP",
"CUDA_IS_AVAILABLE",
"CUDA_DEVICE_COUNT",
)
tests/conftest.py 0 → 100644
View file @ e63cf68a
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
import shutil
from pathlib import Path
from tests import TMP
def pytest_addoption(parser):
"""
Add custom command-line options to pytest.
Args:
parser (pytest.config.Parser): The pytest parser object for adding custom command-line options.
Returns:
(None)
"""
parser.addoption("--slow", action="store_true", default=False, help="Run slow tests")
def pytest_collection_modifyitems(config, items):
"""
Modify the list of test items to exclude tests marked as slow if the --slow option is not specified.
Args:
config (pytest.config.Config): The pytest configuration object that provides access to command-line options.
items (list): The list of collected pytest item objects to be modified based on the presence of --slow option.
Returns:
(None) The function modifies the 'items' list in place, and does not return a value.
"""
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]
def pytest_sessionstart(session):
"""
Initialize session configurations for pytest.
This function is automatically called by pytest after the 'Session' object has been created but before performing
test collection. It sets the initial seeds and prepares the temporary directory for the test session.
Args:
session (pytest.Session): The pytest session object.
Returns:
(None)
"""
from ultralytics.utils.torch_utils import init_seeds
init_seeds()
shutil.rmtree(TMP, ignore_errors=True) # delete any existing tests/tmp directory
TMP.mkdir(parents=True, exist_ok=True) # create a new empty directory
def pytest_terminal_summary(terminalreporter, exitstatus, config):
"""
Cleanup operations after pytest session.
This function is automatically called by pytest at the end of the entire test session. It removes certain files
and directories used during testing.
Args:
terminalreporter (pytest.terminal.TerminalReporter): The terminal reporter object used for terminal output.
exitstatus (int): The exit status of the test run.
config (pytest.config.Config): The pytest config object.
Returns:
(None)
"""
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 ["decelera_portrait_min.mov", "bus.jpg", "yolo11n.onnx", "yolo11n.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 [WEIGHTS_DIR / "path with spaces", TMP.parents[1] / ".pytest_cache", TMP] + models:
shutil.rmtree(directory, ignore_errors=True)
tests/test_cli.py 0 → 100644
View file @ e63cf68a
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
import subprocess
import pytest
from PIL import Image
from tests import CUDA_DEVICE_COUNT, CUDA_IS_AVAILABLE
from ultralytics.cfg import TASK2DATA, TASK2MODEL, TASKS
from ultralytics.utils import ASSETS, WEIGHTS_DIR, checks
from ultralytics.utils.torch_utils import TORCH_1_9
# Constants
TASK_MODEL_DATA = [(task, WEIGHTS_DIR / TASK2MODEL[task], TASK2DATA[task]) for task in TASKS]
MODELS = [WEIGHTS_DIR / TASK2MODEL[task] for task in TASKS]
def run(cmd):
"""Execute a shell command using subprocess."""
subprocess.run(cmd.split(), check=True)
def test_special_modes():
"""Test various special command-line modes for YOLO functionality."""
run("yolo help")
run("yolo checks")
run("yolo version")
run("yolo settings reset")
run("yolo cfg")
@pytest.mark.parametrize("task,model,data", TASK_MODEL_DATA)
def test_train(task, model, data):
"""Test YOLO training for different tasks, models, and datasets."""
run(f"yolo train {task} model={model} data={data} imgsz=32 epochs=1 cache=disk")
@pytest.mark.parametrize("task,model,data", TASK_MODEL_DATA)
def test_val(task, model, data):
"""Test YOLO validation process for specified task, model, and data using a shell command."""
run(f"yolo val {task} model={model} data={data} imgsz=32 save_txt save_json")
@pytest.mark.parametrize("task,model,data", TASK_MODEL_DATA)
def test_predict(task, model, data):
"""Test YOLO prediction on provided sample assets for specified task and model."""
run(f"yolo predict model={model} source={ASSETS} imgsz=32 save save_crop save_txt")
@pytest.mark.parametrize("model", MODELS)
def test_export(model):
"""Test exporting a YOLO model to TorchScript format."""
run(f"yolo export model={model} format=torchscript imgsz=32")
def test_rtdetr(task="detect", model="yolov8n-rtdetr.yaml", data="coco8.yaml"):
"""Test the RTDETR functionality within Ultralytics for detection tasks using specified model and data."""
# Warning: must use imgsz=640 (note also add coma, spaces, fraction=0.25 args to test single-image training)
run(f"yolo train {task} model={model} data={data} --imgsz= 160 epochs =1, cache = disk fraction=0.25")
run(f"yolo predict {task} model={model} source={ASSETS / 'bus.jpg'} imgsz=160 save save_crop save_txt")
if TORCH_1_9:
weights = WEIGHTS_DIR / "rtdetr-l.pt"
run(f"yolo predict {task} model={weights} source={ASSETS / 'bus.jpg'} imgsz=160 save save_crop save_txt")
@pytest.mark.skipif(checks.IS_PYTHON_3_12, reason="MobileSAM with CLIP is not supported in Python 3.12")
def test_fastsam(task="segment", model=WEIGHTS_DIR / "FastSAM-s.pt", data="coco8-seg.yaml"):
"""Test FastSAM model for segmenting objects in images using various prompts within Ultralytics."""
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")
from ultralytics import FastSAM
from ultralytics.models.sam import Predictor
# Create a FastSAM model
sam_model = FastSAM(model) # or FastSAM-x.pt
# Run inference on an image
for s in (source, Image.open(source)):
everything_results = sam_model(s, device="cpu", retina_masks=True, imgsz=320, conf=0.4, iou=0.9)
# Remove small regions
new_masks, _ = Predictor.remove_small_regions(everything_results[0].masks.data, min_area=20)
# Run inference with bboxes and points and texts prompt at the same time
sam_model(source, bboxes=[439, 437, 524, 709], points=[[200, 200]], labels=[1], texts="a photo of a dog")
def test_mobilesam():
"""Test MobileSAM segmentation with point prompts using Ultralytics."""
from ultralytics import SAM
# Load the model
model = SAM(WEIGHTS_DIR / "mobile_sam.pt")
# Source
source = ASSETS / "zidane.jpg"
# Predict a segment based on a 1D point prompt and 1D labels.
model.predict(source, points=[900, 370], labels=[1])
# Predict a segment based on 3D points and 2D labels (multiple points per object).
model.predict(source, points=[[[900, 370], [1000, 100]]], labels=[[1, 1]])
# Predict a segment based on a box prompt
model.predict(source, bboxes=[439, 437, 524, 709], save=True)
# Predict all
# model(source)
# Slow Tests -----------------------------------------------------------------------------------------------------------
@pytest.mark.slow
@pytest.mark.parametrize("task,model,data", TASK_MODEL_DATA)
@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} data={data} imgsz=32 epochs=1 device=0") # single GPU
run(f"yolo train {task} model={model} data={data} imgsz=32 epochs=1 device=0,1") # multi GPU
tests/test_cuda.py 0 → 100644
View file @ e63cf68a
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
from itertools import product
from pathlib import Path
import pytest
import torch
from tests import CUDA_DEVICE_COUNT, CUDA_IS_AVAILABLE, MODEL, SOURCE
from ultralytics import YOLO
from ultralytics.cfg import TASK2DATA, TASK2MODEL, TASKS
from ultralytics.utils import ASSETS, WEIGHTS_DIR
from ultralytics.utils.checks import check_amp
def test_checks():
"""Validate CUDA settings against torch CUDA functions."""
assert torch.cuda.is_available() == CUDA_IS_AVAILABLE
assert torch.cuda.device_count() == CUDA_DEVICE_COUNT
@pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason="CUDA is not available")
def test_amp():
"""Test AMP training checks."""
model = YOLO("yolo11n.pt").model.cuda()
assert check_amp(model)
@pytest.mark.slow
@pytest.mark.skipif(True, reason="CUDA export tests disabled pending additional Ultralytics GPU server availability")
@pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason="CUDA is not available")
@pytest.mark.parametrize(
"task, dynamic, int8, half, batch",
[ # generate all combinations but exclude those where both int8 and half are True
(task, dynamic, int8, half, batch)
# Note: tests reduced below pending compute availability expansion as GPU CI runner utilization is high
# for task, dynamic, int8, half, batch in product(TASKS, [True, False], [True, False], [True, False], [1, 2])
for task, dynamic, int8, half, batch in product(TASKS, [True], [True], [False], [2])
if not (int8 and half) # exclude cases where both int8 and half are True
],
)
def test_export_engine_matrix(task, dynamic, int8, half, batch):
"""Test YOLO model export to TensorRT format for various configurations and run inference."""
file = YOLO(TASK2MODEL[task]).export(
format="engine",
imgsz=32,
dynamic=dynamic,
int8=int8,
half=half,
batch=batch,
data=TASK2DATA[task],
workspace=1, # reduce workspace GB for less resource utilization during testing
simplify=True, # use 'onnxslim'
)
YOLO(file)([SOURCE] * batch, imgsz=64 if dynamic else 32) # exported model inference
Path(file).unlink() # cleanup
Path(file).with_suffix(".cache").unlink() if int8 else None # cleanup INT8 cache
@pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason="CUDA is not available")
def test_train():
"""Test model training on a minimal dataset using available CUDA devices."""
device = 0 if CUDA_DEVICE_COUNT == 1 else [0, 1]
YOLO(MODEL).train(data="coco8.yaml", imgsz=64, epochs=1, device=device) # requires imgsz>=64
@pytest.mark.slow
@pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason="CUDA is not available")
def test_predict_multiple_devices():
"""Validate model prediction consistency across CPU and CUDA devices."""
model = YOLO("yolo11n.pt")
model = model.cpu()
assert str(model.device) == "cpu"
_ = model(SOURCE) # CPU inference
assert str(model.device) == "cpu"
model = model.to("cuda:0")
assert str(model.device) == "cuda:0"
_ = model(SOURCE) # CUDA inference
assert str(model.device) == "cuda:0"
model = model.cpu()
assert str(model.device) == "cpu"
_ = model(SOURCE) # CPU inference
assert str(model.device) == "cpu"
model = model.cuda()
assert str(model.device) == "cuda:0"
_ = model(SOURCE) # CUDA inference
assert str(model.device) == "cuda:0"
@pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason="CUDA is not available")
def test_autobatch():
"""Check optimal batch size for YOLO model training using autobatch utility."""
from ultralytics.utils.autobatch import check_train_batch_size
check_train_batch_size(YOLO(MODEL).model.cuda(), imgsz=128, amp=True)
@pytest.mark.slow
@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)
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")
def test_predict_sam():
"""Test SAM model predictions using different prompts, including bounding boxes and point annotations."""
from ultralytics import SAM
from ultralytics.models.sam import Predictor as SAMPredictor
# Load a model
model = SAM(WEIGHTS_DIR / "sam2.1_b.pt")
# Display model information (optional)
model.info()
# Run inference
model(SOURCE, device=0)
# Run inference with bboxes prompt
model(SOURCE, bboxes=[439, 437, 524, 709], device=0)
# Run inference with no labels
model(ASSETS / "zidane.jpg", points=[900, 370], device=0)
# Run inference with 1D points and 1D labels
model(ASSETS / "zidane.jpg", points=[900, 370], labels=[1], device=0)
# Run inference with 2D points and 1D labels
model(ASSETS / "zidane.jpg", points=[[900, 370]], labels=[1], device=0)
# Run inference with multiple 2D points and 1D labels
model(ASSETS / "zidane.jpg", points=[[400, 370], [900, 370]], labels=[1, 1], device=0)
# Run inference with 3D points and 2D labels (multiple points per object)
model(ASSETS / "zidane.jpg", points=[[[900, 370], [1000, 100]]], labels=[[1, 1]], device=0)
# Create SAMPredictor
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(bboxes=[439, 437, 524, 709])
# predictor(points=[900, 370], labels=[1])
# Reset image
predictor.reset_image()
tests/test_engine.py 0 → 100644
View file @ e63cf68a
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
import sys
from unittest import mock
from tests import MODEL
from ultralytics import YOLO
from ultralytics.cfg import get_cfg
from ultralytics.engine.exporter import Exporter
from ultralytics.models.yolo import classify, detect, segment
from ultralytics.utils import ASSETS, DEFAULT_CFG, WEIGHTS_DIR
def test_func(*args): # noqa
"""Test function callback for evaluating YOLO model performance metrics."""
print("callback test passed")
def test_export():
"""Tests the model exporting function by adding a callback and asserting its execution."""
exporter = Exporter()
exporter.add_callback("on_export_start", test_func)
assert test_func in exporter.callbacks["on_export_start"], "callback test failed"
f = exporter(model=YOLO("yolo11n.yaml").model)
YOLO(f)(ASSETS) # exported model inference
def test_detect():
"""Test YOLO object detection training, validation, and prediction functionality."""
overrides = {"data": "coco8.yaml", "model": "yolo11n.yaml", "imgsz": 32, "epochs": 1, "save": False}
cfg = get_cfg(DEFAULT_CFG)
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.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(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"
# Confirm there is no issue with sys.argv being empty.
with mock.patch.object(sys, "argv", []):
result = pred(source=ASSETS, model=MODEL)
assert len(result), "predictor test failed"
overrides["resume"] = trainer.last
trainer = detect.DetectionTrainer(overrides=overrides)
try:
trainer.train()
except Exception as e:
print(f"Expected exception caught: {e}")
return
Exception("Resume test failed!")
def test_segment():
"""Tests image segmentation training, validation, and prediction pipelines using YOLO models."""
overrides = {"data": "coco8-seg.yaml", "model": "yolo11n-seg.yaml", "imgsz": 32, "epochs": 1, "save": False}
cfg = get_cfg(DEFAULT_CFG)
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.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(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=WEIGHTS_DIR / "yolo11n-seg.pt")
assert len(result), "predictor test failed"
# Test resume
overrides["resume"] = trainer.last
trainer = segment.SegmentationTrainer(overrides=overrides)
try:
trainer.train()
except Exception as e:
print(f"Expected exception caught: {e}")
return
Exception("Resume test failed!")
def test_classify():
"""Test image classification including training, validation, and prediction phases."""
overrides = {"data": "imagenet10", "model": "yolo11n-cls.yaml", "imgsz": 32, "epochs": 1, "save": False}
cfg = get_cfg(DEFAULT_CFG)
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.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(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"
result = pred(source=ASSETS, model=trainer.best)
assert len(result), "predictor test failed"
tests/test_exports.py 0 → 100644
View file @ e63cf68a
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
import shutil
import uuid
from itertools import product
from pathlib import Path
import pytest
from tests import MODEL, SOURCE
from ultralytics import YOLO
from ultralytics.cfg import TASK2DATA, TASK2MODEL, TASKS
from ultralytics.utils import (
IS_RASPBERRYPI,
LINUX,
MACOS,
WINDOWS,
checks,
)
from ultralytics.utils.torch_utils import TORCH_1_9, TORCH_1_13
def test_export_torchscript():
"""Test YOLO model exporting to TorchScript format for compatibility and correctness."""
file = YOLO(MODEL).export(format="torchscript", optimize=False, imgsz=32)
YOLO(file)(SOURCE, imgsz=32) # exported model inference
def test_export_onnx():
"""Test YOLO model export to ONNX format with dynamic axes."""
file = YOLO(MODEL).export(format="onnx", dynamic=True, imgsz=32)
YOLO(file)(SOURCE, imgsz=32) # exported model inference
@pytest.mark.skipif(not TORCH_1_13, reason="OpenVINO requires torch>=1.13")
def test_export_openvino():
"""Test YOLO exports to OpenVINO format for model inference compatibility."""
file = YOLO(MODEL).export(format="openvino", imgsz=32)
YOLO(file)(SOURCE, imgsz=32) # exported model inference
@pytest.mark.slow
@pytest.mark.skipif(not TORCH_1_13, reason="OpenVINO requires torch>=1.13")
@pytest.mark.parametrize(
"task, dynamic, int8, half, batch",
[ # generate all combinations but exclude those where both int8 and half are True
(task, dynamic, int8, half, batch)
for task, dynamic, int8, half, batch in product(TASKS, [True, False], [True, False], [True, False], [1, 2])
if not (int8 and half) # exclude cases where both int8 and half are True
],
)
def test_export_openvino_matrix(task, dynamic, int8, half, batch):
"""Test YOLO model exports to OpenVINO under various configuration matrix conditions."""
file = YOLO(TASK2MODEL[task]).export(
format="openvino",
imgsz=32,
dynamic=dynamic,
int8=int8,
half=half,
batch=batch,
data=TASK2DATA[task],
)
if WINDOWS:
# Use unique filenames due to Windows file permissions bug possibly due to latent threaded use
# See https://github.com/ultralytics/ultralytics/actions/runs/8957949304/job/24601616830?pr=10423
file = Path(file)
file = file.rename(file.with_stem(f"{file.stem}-{uuid.uuid4()}"))
YOLO(file)([SOURCE] * batch, imgsz=64 if dynamic else 32) # exported model inference
shutil.rmtree(file, ignore_errors=True) # retry in case of potential lingering multi-threaded file usage errors
@pytest.mark.slow
@pytest.mark.parametrize(
"task, dynamic, int8, half, batch, simplify", product(TASKS, [True, False], [False], [False], [1, 2], [True, False])
)
def test_export_onnx_matrix(task, dynamic, int8, half, batch, simplify):
"""Test YOLO exports to ONNX format with various configurations and parameters."""
file = YOLO(TASK2MODEL[task]).export(
format="onnx",
imgsz=32,
dynamic=dynamic,
int8=int8,
half=half,
batch=batch,
simplify=simplify,
)
YOLO(file)([SOURCE] * batch, imgsz=64 if dynamic else 32) # exported model inference
Path(file).unlink() # cleanup
@pytest.mark.slow
@pytest.mark.parametrize("task, dynamic, int8, half, batch", product(TASKS, [False], [False], [False], [1, 2]))
def test_export_torchscript_matrix(task, dynamic, int8, half, batch):
"""Tests YOLO model exports to TorchScript format under varied configurations."""
file = YOLO(TASK2MODEL[task]).export(
format="torchscript",
imgsz=32,
dynamic=dynamic,
int8=int8,
half=half,
batch=batch,
)
YOLO(file)([SOURCE] * 3, imgsz=64 if dynamic else 32) # exported model inference at batch=3
Path(file).unlink() # cleanup
@pytest.mark.slow
@pytest.mark.skipif(not MACOS, reason="CoreML inference only supported on macOS")
@pytest.mark.skipif(not TORCH_1_9, reason="CoreML>=7.2 not supported with PyTorch<=1.8")
@pytest.mark.skipif(checks.IS_PYTHON_3_12, reason="CoreML not supported in Python 3.12")
@pytest.mark.parametrize(
"task, dynamic, int8, half, batch",
[ # generate all combinations but exclude those where both int8 and half are True
(task, dynamic, int8, half, batch)
for task, dynamic, int8, half, batch in product(TASKS, [False], [True, False], [True, False], [1])
if not (int8 and half) # exclude cases where both int8 and half are True
],
)
def test_export_coreml_matrix(task, dynamic, int8, half, batch):
"""Test YOLO exports to CoreML format with various parameter configurations."""
file = YOLO(TASK2MODEL[task]).export(
format="coreml",
imgsz=32,
dynamic=dynamic,
int8=int8,
half=half,
batch=batch,
)
YOLO(file)([SOURCE] * batch, imgsz=32) # exported model inference at batch=3
shutil.rmtree(file) # cleanup
@pytest.mark.slow
@pytest.mark.skipif(not checks.IS_PYTHON_MINIMUM_3_10, reason="TFLite export requires Python>=3.10")
@pytest.mark.skipif(not LINUX, reason="Test disabled as TF suffers from install conflicts on Windows and macOS")
@pytest.mark.parametrize(
"task, dynamic, int8, half, batch",
[ # generate all combinations but exclude those where both int8 and half are True
(task, dynamic, int8, half, batch)
for task, dynamic, int8, half, batch in product(TASKS, [False], [True, False], [True, False], [1])
if not (int8 and half) # exclude cases where both int8 and half are True
],
)
def test_export_tflite_matrix(task, dynamic, int8, half, batch):
"""Test YOLO exports to TFLite format considering various export configurations."""
file = YOLO(TASK2MODEL[task]).export(
format="tflite",
imgsz=32,
dynamic=dynamic,
int8=int8,
half=half,
batch=batch,
)
YOLO(file)([SOURCE] * batch, imgsz=32) # exported model inference at batch=3
Path(file).unlink() # cleanup
@pytest.mark.skipif(not TORCH_1_9, reason="CoreML>=7.2 not supported with PyTorch<=1.8")
@pytest.mark.skipif(WINDOWS, reason="CoreML not supported on Windows") # RuntimeError: BlobWriter not loaded
@pytest.mark.skipif(IS_RASPBERRYPI, reason="CoreML not supported on Raspberry Pi")
@pytest.mark.skipif(checks.IS_PYTHON_3_12, reason="CoreML not supported in Python 3.12")
def test_export_coreml():
"""Test YOLO exports to CoreML format, optimized for macOS only."""
if MACOS:
file = YOLO(MODEL).export(format="coreml", imgsz=32)
YOLO(file)(SOURCE, imgsz=32) # model prediction only supported on macOS for nms=False models
else:
YOLO(MODEL).export(format="coreml", nms=True, imgsz=32)
@pytest.mark.skipif(not checks.IS_PYTHON_MINIMUM_3_10, reason="TFLite export requires Python>=3.10")
@pytest.mark.skipif(not LINUX, reason="Test disabled as TF suffers from install conflicts on Windows and macOS")
def test_export_tflite():
"""Test YOLO exports to TFLite format under specific OS and Python version conditions."""
model = YOLO(MODEL)
file = model.export(format="tflite", imgsz=32)
YOLO(file)(SOURCE, imgsz=32)
@pytest.mark.skipif(True, reason="Test disabled")
@pytest.mark.skipif(not LINUX, reason="TF suffers from install conflicts on Windows and macOS")
def test_export_pb():
"""Test YOLO exports to TensorFlow's Protobuf (*.pb) format."""
model = YOLO(MODEL)
file = model.export(format="pb", imgsz=32)
YOLO(file)(SOURCE, imgsz=32)
@pytest.mark.skipif(True, reason="Test disabled as Paddle protobuf and ONNX protobuf requirements conflict.")
def test_export_paddle():
"""Test YOLO exports to Paddle format, noting protobuf conflicts with ONNX."""
YOLO(MODEL).export(format="paddle", imgsz=32)
@pytest.mark.slow
@pytest.mark.skipif(IS_RASPBERRYPI, reason="MNN not supported on Raspberry Pi")
def test_export_mnn():
"""Test YOLO exports to MNN format (WARNING: MNN test must precede NCNN test or CI error on Windows)."""
file = YOLO(MODEL).export(format="mnn", imgsz=32)
YOLO(file)(SOURCE, imgsz=32) # exported model inference
@pytest.mark.slow
def test_export_ncnn():
"""Test YOLO exports to NCNN format."""
file = YOLO(MODEL).export(format="ncnn", imgsz=32)
YOLO(file)(SOURCE, imgsz=32) # exported model inference
@pytest.mark.skipif(True, reason="Test disabled as keras and tensorflow version conflicts with tflite export.")
@pytest.mark.skipif(not LINUX or MACOS, reason="Skipping test on Windows and Macos")
def test_export_imx():
"""Test YOLOv8n exports to IMX format."""
model = YOLO("yolov8n.pt")
file = model.export(format="imx", imgsz=32)
YOLO(file)(SOURCE, imgsz=32)
tests/test_integrations.py 0 → 100644
View file @ e63cf68a
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
import contextlib
import os
import subprocess
import time
from pathlib import Path
import pytest
from tests import MODEL, SOURCE, TMP
from ultralytics import YOLO, download
from ultralytics.utils import DATASETS_DIR, SETTINGS
from ultralytics.utils.checks import check_requirements
@pytest.mark.skipif(not check_requirements("ray", install=False), reason="ray[tune] not installed")
def test_model_ray_tune():
"""Tune YOLO model using Ray for hyperparameter optimization."""
YOLO("yolo11n-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")
def test_mlflow():
"""Test training with MLflow tracking enabled (see https://mlflow.org/ for details)."""
SETTINGS["mlflow"] = True
YOLO("yolo11n-cls.yaml").train(data="imagenet10", imgsz=32, epochs=3, plots=False, device="cpu")
SETTINGS["mlflow"] = False
@pytest.mark.skipif(True, reason="Test failing in scheduled CI https://github.com/ultralytics/ultralytics/pull/8868")
@pytest.mark.skipif(not check_requirements("mlflow", install=False), reason="mlflow not installed")
def test_mlflow_keep_run_active():
"""Ensure MLflow run status matches MLFLOW_KEEP_RUN_ACTIVE environment variable settings."""
import mlflow
SETTINGS["mlflow"] = True
run_name = "Test Run"
os.environ["MLFLOW_RUN"] = run_name
# Test with MLFLOW_KEEP_RUN_ACTIVE=True
os.environ["MLFLOW_KEEP_RUN_ACTIVE"] = "True"
YOLO("yolo11n-cls.yaml").train(data="imagenet10", imgsz=32, epochs=1, plots=False, device="cpu")
status = mlflow.active_run().info.status
assert status == "RUNNING", "MLflow run should be active when MLFLOW_KEEP_RUN_ACTIVE=True"
run_id = mlflow.active_run().info.run_id
# Test with MLFLOW_KEEP_RUN_ACTIVE=False
os.environ["MLFLOW_KEEP_RUN_ACTIVE"] = "False"
YOLO("yolo11n-cls.yaml").train(data="imagenet10", imgsz=32, epochs=1, plots=False, device="cpu")
status = mlflow.get_run(run_id=run_id).info.status
assert status == "FINISHED", "MLflow run should be ended when MLFLOW_KEEP_RUN_ACTIVE=False"
# Test with MLFLOW_KEEP_RUN_ACTIVE not set
os.environ.pop("MLFLOW_KEEP_RUN_ACTIVE", None)
YOLO("yolo11n-cls.yaml").train(data="imagenet10", imgsz=32, epochs=1, plots=False, device="cpu")
status = mlflow.get_run(run_id=run_id).info.status
assert status == "FINISHED", "MLflow run should be ended by default when MLFLOW_KEEP_RUN_ACTIVE is not set"
SETTINGS["mlflow"] = False
@pytest.mark.skipif(not check_requirements("tritonclient", install=False), reason="tritonclient[all] not installed")
def test_triton():
"""
Test NVIDIA Triton Server functionalities with YOLO model.
See https://catalog.ngc.nvidia.com/orgs/nvidia/containers/tritonserver.
"""
check_requirements("tritonclient[all]")
from tritonclient.http import InferenceServerClient # noqa
# Create variables
model_name = "yolo"
triton_repo = TMP / "triton_repo" # Triton repo path
triton_model = triton_repo / model_name # Triton model path
# Export model to ONNX
f = YOLO(MODEL).export(format="onnx", dynamic=True)
# Prepare Triton repo
(triton_model / "1").mkdir(parents=True, exist_ok=True)
Path(f).rename(triton_model / "1" / "model.onnx")
(triton_model / "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
# Pull the image
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}:/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)
# Wait until model is ready
for _ in range(10):
with contextlib.suppress(Exception):
assert triton_client.is_model_ready(model_name)
break
time.sleep(1)
# Check Triton 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)
@pytest.mark.skipif(not check_requirements("pycocotools", install=False), reason="pycocotools not installed")
def test_pycocotools():
"""Validate YOLO model predictions on COCO dataset using pycocotools."""
from ultralytics.models.yolo.detect import DetectionValidator
from ultralytics.models.yolo.pose import PoseValidator
from ultralytics.models.yolo.segment import SegmentationValidator
# Download annotations after each dataset downloads first
url = "https://github.com/ultralytics/assets/releases/download/v0.0.0/"
args = {"model": "yolo11n.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")
_ = validator.eval_json(validator.stats)
args = {"model": "yolo11n-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")
_ = validator.eval_json(validator.stats)
args = {"model": "yolo11n-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")
_ = validator.eval_json(validator.stats)
tests/test_python.py 0 → 100644
View file @ e63cf68a
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
import contextlib
import csv
import urllib
from copy import copy
from pathlib import Path
import cv2
import numpy as np
import pytest
import torch
import yaml
from PIL import Image
from tests import CFG, MODEL, SOURCE, SOURCES_LIST, TMP
from ultralytics import RTDETR, YOLO
from ultralytics.cfg import MODELS, TASK2DATA, TASKS
from ultralytics.data.build import load_inference_source
from ultralytics.utils import (
ASSETS,
DEFAULT_CFG,
DEFAULT_CFG_PATH,
LOGGER,
ONLINE,
ROOT,
WEIGHTS_DIR,
WINDOWS,
checks,
is_dir_writeable,
is_github_action_running,
)
from ultralytics.utils.downloads import download
from ultralytics.utils.torch_utils import TORCH_1_9
IS_TMP_WRITEABLE = is_dir_writeable(TMP) # WARNING: must be run once tests start as TMP does not exist on tests/init
def test_model_forward():
"""Test the forward pass of the YOLO model."""
model = YOLO(CFG)
model(source=None, imgsz=32, augment=True) # also test no source and augment
def test_model_methods():
"""Test various methods and properties of the YOLO model to ensure correct functionality."""
model = YOLO(MODEL)
# Model methods
model.info(verbose=True, detailed=True)
model = model.reset_weights()
model = model.load(MODEL)
model.to("cpu")
model.fuse()
model.clear_callback("on_train_start")
model.reset_callbacks()
# Model properties
_ = model.names
_ = model.device
_ = model.transforms
_ = model.task_map
def test_model_profile():
"""Test profiling of the YOLO model with `profile=True` to assess performance and resource usage."""
from ultralytics.nn.tasks import DetectionModel
model = DetectionModel() # build model
im = torch.randn(1, 3, 64, 64) # requires min imgsz=64
_ = model.predict(im, profile=True)
@pytest.mark.skipif(not IS_TMP_WRITEABLE, reason="directory is not writeable")
def test_predict_txt():
"""Tests YOLO predictions with file, directory, and pattern sources listed in a text file."""
file = TMP / "sources_multi_row.txt"
with open(file, "w") as f:
for src in SOURCES_LIST:
f.write(f"{src}\n")
results = YOLO(MODEL)(source=file, imgsz=32)
assert len(results) == 7 # 1 + 2 + 2 + 2 = 7 images
@pytest.mark.skipif(True, reason="disabled for testing")
@pytest.mark.skipif(not IS_TMP_WRITEABLE, reason="directory is not writeable")
def test_predict_csv_multi_row():
"""Tests YOLO predictions with sources listed in multiple rows of a CSV file."""
file = TMP / "sources_multi_row.csv"
with open(file, "w", newline="") as f:
writer = csv.writer(f)
writer.writerow(["source"])
writer.writerows([[src] for src in SOURCES_LIST])
results = YOLO(MODEL)(source=file, imgsz=32)
assert len(results) == 7 # 1 + 2 + 2 + 2 = 7 images
@pytest.mark.skipif(True, reason="disabled for testing")
@pytest.mark.skipif(not IS_TMP_WRITEABLE, reason="directory is not writeable")
def test_predict_csv_single_row():
"""Tests YOLO predictions with sources listed in a single row of a CSV file."""
file = TMP / "sources_single_row.csv"
with open(file, "w", newline="") as f:
writer = csv.writer(f)
writer.writerow(SOURCES_LIST)
results = YOLO(MODEL)(source=file, imgsz=32)
assert len(results) == 7 # 1 + 2 + 2 + 2 = 7 images
@pytest.mark.parametrize("model_name", MODELS)
def test_predict_img(model_name):
"""Test YOLO model predictions on various image input types and sources, including online images."""
model = YOLO(WEIGHTS_DIR / model_name)
im = cv2.imread(str(SOURCE)) # uint8 numpy array
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
assert len(model(torch.rand((2, 3, 32, 32)), imgsz=32)) == 2 # batch-size 2 Tensor, FP32 0.0-1.0 RGB order
assert len(model(source=[im, im], save=True, save_txt=True, imgsz=32)) == 2 # batch
assert len(list(model(source=[im, im], save=True, stream=True, imgsz=32))) == 2 # stream
assert len(model(torch.zeros(320, 640, 3).numpy().astype(np.uint8), imgsz=32)) == 1 # tensor to numpy
batch = [
str(SOURCE), # filename
Path(SOURCE), # Path
"https://github.com/ultralytics/assets/releases/download/v0.0.0/zidane.jpg" if ONLINE else SOURCE, # URI
cv2.imread(str(SOURCE)), # OpenCV
Image.open(SOURCE), # PIL
np.zeros((320, 640, 3), dtype=np.uint8), # numpy
]
assert len(model(batch, imgsz=32)) == len(batch) # multiple sources in a batch
@pytest.mark.parametrize("model", MODELS)
def test_predict_visualize(model):
"""Test model prediction methods with 'visualize=True' to generate and display prediction visualizations."""
YOLO(WEIGHTS_DIR / model)(SOURCE, imgsz=32, visualize=True)
def test_predict_grey_and_4ch():
"""Test YOLO prediction on SOURCE converted to greyscale and 4-channel images with various filenames."""
im = Image.open(SOURCE)
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"
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
# Inference
model = YOLO(MODEL)
for f in source_rgba, source_greyscale, source_non_utf, source_spaces:
for source in Image.open(f), cv2.imread(str(f)), f:
results = model(source, save=True, verbose=True, imgsz=32)
assert len(results) == 1 # verify that an image was run
f.unlink() # cleanup
@pytest.mark.slow
@pytest.mark.skipif(not ONLINE, reason="environment is offline")
@pytest.mark.skipif(is_github_action_running(), reason="No auth https://github.com/JuanBindez/pytubefix/issues/166")
def test_youtube():
"""Test YOLO model on a YouTube video stream, handling potential network-related errors."""
model = YOLO(MODEL)
try:
model.predict("https://youtu.be/G17sBkb38XQ", imgsz=96, save=True)
# Handle internet connection errors and 'urllib.error.HTTPError: HTTP Error 429: Too Many Requests'
except (urllib.error.HTTPError, ConnectionError) as e:
LOGGER.warning(f"WARNING: YouTube Test Error: {e}")
@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():
"""
Tests streaming tracking on a short 10 frame video using ByteTrack tracker and different GMC methods.
Note imgsz=160 required for tracking for higher confidence and better matches.
"""
video_url = "https://github.com/ultralytics/assets/releases/download/v0.0.0/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
# Test Global Motion Compensation (GMC) methods
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:
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)
def test_train_scratch():
"""Test training the YOLO model from scratch using the provided configuration."""
model = YOLO(CFG)
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 of the YOLO model starting from a pre-trained checkpoint."""
model = YOLO(WEIGHTS_DIR / "yolo11n-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_all_model_yamls():
"""Test YOLO model creation for all available YAML configurations in the `cfg/models` directory."""
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:
YOLO(m.name)
@pytest.mark.skipif(WINDOWS, reason="Windows slow CI export bug https://github.com/ultralytics/ultralytics/pull/16003")
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.val(imgsz=32)
model.predict(SOURCE, imgsz=32)
model.export(format="torchscript") # WARNING: Windows slow CI export bug
def test_predict_callback_and_setup():
"""Test callback functionality during YOLO prediction setup and execution."""
def on_predict_batch_end(predictor):
"""Callback function that handles operations at the end of a prediction batch."""
path, im0s, _ = predictor.batch
im0s = im0s if isinstance(im0s, list) else [im0s]
bs = [predictor.dataset.bs for _ in range(len(path))]
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)
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)
boxes = r.boxes # Boxes object for bbox outputs
print(boxes)
@pytest.mark.parametrize("model", MODELS)
def test_results(model):
"""Ensure YOLO model predictions can be processed and printed in various formats."""
results = YOLO(WEIGHTS_DIR / model)([SOURCE, SOURCE], imgsz=160)
for r in results:
r = r.cpu().numpy()
print(r, len(r), r.path) # print numpy attributes
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.to_json(normalize=True)
r.to_df(decimals=3)
r.to_csv()
r.to_xml()
r.plot(pil=True)
r.plot(conf=True, boxes=True)
print(r, len(r), r.path) # print after methods
def test_labels_and_crops():
"""Test output from prediction args for saving YOLO detection labels and crops; ensures accurate saving."""
imgs = [SOURCE, ASSETS / "zidane.jpg"]
results = YOLO(WEIGHTS_DIR / "yolo11n.pt")(imgs, imgsz=160, save_txt=True, save_crop=True)
save_path = Path(results[0].save_dir)
for r in results:
im_name = Path(r.path).stem
cls_idxs = r.boxes.cls.int().tolist()
# Check correct detections
assert cls_idxs == ([0, 7, 0, 0] if r.path.endswith("bus.jpg") else [0, 0, 0]) # bus.jpg and zidane.jpg classes
# Check label path
labels = save_path / f"labels/{im_name}.txt"
assert labels.exists()
# Check detections match label count
assert len(r.boxes.data) == len([line for line in labels.read_text().splitlines() if line])
# Check crops path and files
crop_dirs = list((save_path / "crops").iterdir())
crop_files = [f for p in crop_dirs for f in p.glob("*")]
# Crop directories match detections
assert all(r.names.get(c) in {d.name for d in crop_dirs} for c in cls_idxs)
# Same number of crops as detections
assert len([f for f in crop_files if im_name in f.name]) == len(r.boxes.data)
@pytest.mark.skipif(not ONLINE, reason="environment is offline")
def test_data_utils():
"""Test utility functions in ultralytics/data/utils.py, including dataset stats and auto-splitting."""
from ultralytics.data.utils import HUBDatasetStats, autosplit
from ultralytics.utils.downloads import zip_directory
# from ultralytics.utils.files import WorkingDirectory
# with WorkingDirectory(ROOT.parent / 'tests'):
for task in TASKS:
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
@pytest.mark.skipif(not ONLINE, reason="environment is offline")
def test_data_converter():
"""Test dataset conversion functions from COCO to YOLO format and class mappings."""
from ultralytics.data.converter import coco80_to_coco91_class, convert_coco
file = "instances_val2017.json"
download(f"https://github.com/ultralytics/assets/releases/download/v0.0.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()
def test_data_annotator():
"""Automatically annotate data using specified detection and segmentation models."""
from ultralytics.data.annotator import auto_annotate
auto_annotate(
ASSETS,
det_model=WEIGHTS_DIR / "yolo11n.pt",
sam_model=WEIGHTS_DIR / "mobile_sam.pt",
output_dir=TMP / "auto_annotate_labels",
)
def test_events():
"""Test event sending functionality."""
from ultralytics.hub.utils import Events
events = Events()
events.enabled = True
cfg = copy(DEFAULT_CFG) # does not require deepcopy
cfg.mode = "test"
events(cfg)
def test_cfg_init():
"""Test configuration initialization utilities from the 'ultralytics.cfg' module."""
from ultralytics.cfg import check_dict_alignment, copy_default_cfg, smart_value
with contextlib.suppress(SyntaxError):
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"]]
def test_utils_init():
"""Test initialization utilities in the Ultralytics library."""
from ultralytics.utils import get_git_branch, get_git_origin_url, get_ubuntu_version, is_github_action_running
get_ubuntu_version()
is_github_action_running()
get_git_origin_url()
get_git_branch()
def test_utils_checks():
"""Test various utility checks for filenames, git status, requirements, image sizes, and versions."""
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(warn=True)
checks.check_version("ultralytics", "8.0.0")
checks.print_args()
@pytest.mark.skipif(WINDOWS, reason="Windows profiling is extremely slow (cause unknown)")
def test_utils_benchmarks():
"""Benchmark model performance using 'ProfileModels' from 'ultralytics.utils.benchmarks'."""
from ultralytics.utils.benchmarks import ProfileModels
ProfileModels(["yolo11n.yaml"], imgsz=32, min_time=1, num_timed_runs=3, num_warmup_runs=1).profile()
def test_utils_torchutils():
"""Test Torch utility functions including profiling and FLOP calculations."""
from ultralytics.nn.modules.conv import Conv
from ultralytics.utils.torch_utils import get_flops_with_torch_profiler, profile, time_sync
x = torch.randn(1, 64, 20, 20)
m = Conv(64, 64, k=1, s=2)
profile(x, [m], n=3)
get_flops_with_torch_profiler(m)
time_sync()
def test_utils_ops():
"""Test utility operations functions for coordinate transformation and normalization."""
from ultralytics.utils.ops import (
ltwh2xywh,
ltwh2xyxy,
make_divisible,
xywh2ltwh,
xywh2xyxy,
xywhn2xyxy,
xywhr2xyxyxyxy,
xyxy2ltwh,
xyxy2xywh,
xyxy2xywhn,
xyxyxyxy2xywhr,
)
make_divisible(17, torch.tensor([8]))
boxes = torch.rand(10, 4) # xywh
torch.allclose(boxes, xyxy2xywh(xywh2xyxy(boxes)))
torch.allclose(boxes, xyxy2xywhn(xywhn2xyxy(boxes)))
torch.allclose(boxes, ltwh2xywh(xywh2ltwh(boxes)))
torch.allclose(boxes, xyxy2ltwh(ltwh2xyxy(boxes)))
boxes = torch.rand(10, 5) # xywhr for OBB
boxes[:, 4] = torch.randn(10) * 30
torch.allclose(boxes, xyxyxyxy2xywhr(xywhr2xyxyxyxy(boxes)), rtol=1e-3)
def test_utils_files():
"""Test file handling utilities including file age, date, and paths with spaces."""
from ultralytics.utils.files import file_age, file_date, get_latest_run, spaces_in_path
file_age(SOURCE)
file_date(SOURCE)
get_latest_run(ROOT / "runs")
path = TMP / "path/with spaces"
path.mkdir(parents=True, exist_ok=True)
with spaces_in_path(path) as new_path:
print(new_path)
@pytest.mark.slow
def test_utils_patches_torch_save():
"""Test torch_save backoff when _torch_save raises RuntimeError to ensure robustness."""
from unittest.mock import MagicMock, patch
from ultralytics.utils.patches import torch_save
mock = MagicMock(side_effect=RuntimeError)
with patch("ultralytics.utils.patches._torch_save", new=mock):
with pytest.raises(RuntimeError):
torch_save(torch.zeros(1), TMP / "test.pt")
assert mock.call_count == 4, "torch_save was not attempted the expected number of times"
def test_nn_modules_conv():
"""Test Convolutional Neural Network modules including CBAM, Conv2, and ConvTranspose."""
from ultralytics.nn.modules.conv import CBAM, Conv2, ConvTranspose, DWConvTranspose2d, Focus
c1, c2 = 8, 16 # input and output channels
x = torch.zeros(4, c1, 10, 10) # BCHW
# Run all modules not otherwise covered in tests
DWConvTranspose2d(c1, c2)(x)
ConvTranspose(c1, c2)(x)
Focus(c1, c2)(x)
CBAM(c1)(x)
# Fuse ops
m = Conv2(c1, c2)
m.fuse_convs()
m(x)
def test_nn_modules_block():
"""Test various blocks in neural network modules including C1, C3TR, BottleneckCSP, C3Ghost, and C3x."""
from ultralytics.nn.modules.block import C1, C3TR, BottleneckCSP, C3Ghost, C3x
c1, c2 = 8, 16 # input and output channels
x = torch.zeros(4, c1, 10, 10) # BCHW
# Run all modules not otherwise covered in tests
C1(c1, c2)(x)
C3x(c1, c2)(x)
C3TR(c1, c2)(x)
C3Ghost(c1, c2)(x)
BottleneckCSP(c1, c2)(x)
@pytest.mark.skipif(not ONLINE, reason="environment is offline")
def test_hub():
"""Test Ultralytics HUB functionalities (e.g. export formats, logout)."""
from ultralytics.hub import export_fmts_hub, logout
from ultralytics.hub.utils import smart_request
export_fmts_hub()
logout()
smart_request("GET", "https://github.com", progress=True)
@pytest.fixture
def image():
"""Load and return an image from a predefined source using OpenCV."""
return cv2.imread(str(SOURCE))
@pytest.mark.parametrize(
"auto_augment, erasing, force_color_jitter",
[
(None, 0.0, False),
("randaugment", 0.5, True),
("augmix", 0.2, False),
("autoaugment", 0.0, True),
],
)
def test_classify_transforms_train(image, auto_augment, erasing, force_color_jitter):
"""Tests classification transforms during training with various augmentations to ensure proper functionality."""
from ultralytics.data.augment import classify_augmentations
transform = classify_augmentations(
size=224,
mean=(0.5, 0.5, 0.5),
std=(0.5, 0.5, 0.5),
scale=(0.08, 1.0),
ratio=(3.0 / 4.0, 4.0 / 3.0),
hflip=0.5,
vflip=0.5,
auto_augment=auto_augment,
hsv_h=0.015,
hsv_s=0.4,
hsv_v=0.4,
force_color_jitter=force_color_jitter,
erasing=erasing,
)
transformed_image = transform(Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)))
assert transformed_image.shape == (3, 224, 224)
assert torch.is_tensor(transformed_image)
assert transformed_image.dtype == torch.float32
@pytest.mark.slow
@pytest.mark.skipif(not ONLINE, reason="environment is offline")
def test_model_tune():
"""Tune YOLO model for performance improvement."""
YOLO("yolo11n-pose.pt").tune(data="coco8-pose.yaml", plots=False, imgsz=32, epochs=1, iterations=2, device="cpu")
YOLO("yolo11n-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 / "yolo11n-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)
assert len(model_segment.embed(source=batch, imgsz=32)) == len(batch)
@pytest.mark.skipif(checks.IS_PYTHON_3_12, reason="YOLOWorld with CLIP is not supported in Python 3.12")
def test_yolo_world():
"""Tests YOLO world models with CLIP support, including detection and training scenarios."""
model = YOLO(WEIGHTS_DIR / "yolov8s-world.pt") # no YOLO11n-world model yet
model.set_classes(["tree", "window"])
model(SOURCE, conf=0.01)
model = YOLO(WEIGHTS_DIR / "yolov8s-worldv2.pt") # no YOLO11n-world model yet
# Training from a pretrained model. Eval is included at the final stage of training.
# Use dota8.yaml which has fewer categories to reduce the inference time of CLIP model
model.train(
data="dota8.yaml",
epochs=1,
imgsz=32,
cache="disk",
close_mosaic=1,
)
# test WorWorldTrainerFromScratch
from ultralytics.models.yolo.world.train_world import WorldTrainerFromScratch
model = YOLO("yolov8s-worldv2.yaml") # no YOLO11n-world model yet
model.train(
data={"train": {"yolo_data": ["dota8.yaml"]}, "val": {"yolo_data": ["dota8.yaml"]}},
epochs=1,
imgsz=32,
cache="disk",
close_mosaic=1,
trainer=WorldTrainerFromScratch,
)
def test_yolov10():
"""Test YOLOv10 model training, validation, and prediction steps with minimal configurations."""
model = YOLO("yolov10n.yaml")
# train/val/predict
model.train(data="coco8.yaml", epochs=1, imgsz=32, close_mosaic=1, cache="disk")
model.val(data="coco8.yaml", imgsz=32)
model.predict(imgsz=32, save_txt=True, save_crop=True, augment=True)
model(SOURCE)
tests/test_solutions.py 0 → 100644
View file @ e63cf68a
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
import cv2
import pytest
from tests import TMP
from ultralytics import YOLO, solutions
from ultralytics.utils import ASSETS_URL, WEIGHTS_DIR
from ultralytics.utils.downloads import safe_download
DEMO_VIDEO = "solutions_ci_demo.mp4"
POSE_VIDEO = "solution_ci_pose_demo.mp4"
@pytest.mark.slow
def test_major_solutions():
"""Test the object counting, heatmap, speed estimation, trackzone and queue management solution."""
safe_download(url=f"{ASSETS_URL}/{DEMO_VIDEO}", dir=TMP)
cap = cv2.VideoCapture(str(TMP / DEMO_VIDEO))
assert cap.isOpened(), "Error reading video file"
region_points = [(20, 400), (1080, 400), (1080, 360), (20, 360)]
counter = solutions.ObjectCounter(region=region_points, model="yolo11n.pt", show=False) # Test object counter
heatmap = solutions.Heatmap(colormap=cv2.COLORMAP_PARULA, model="yolo11n.pt", show=False) # Test heatmaps
heatmap_count = solutions.Heatmap(
colormap=cv2.COLORMAP_PARULA, model="yolo11n.pt", show=False, region=region_points
) # Test heatmaps with object counting
speed = solutions.SpeedEstimator(region=region_points, model="yolo11n.pt", show=False) # Test queue manager
queue = solutions.QueueManager(region=region_points, model="yolo11n.pt", show=False) # Test speed estimation
line_analytics = solutions.Analytics(analytics_type="line", model="yolo11n.pt", show=False) # line analytics
pie_analytics = solutions.Analytics(analytics_type="pie", model="yolo11n.pt", show=False) # line analytics
bar_analytics = solutions.Analytics(analytics_type="bar", model="yolo11n.pt", show=False) # line analytics
area_analytics = solutions.Analytics(analytics_type="area", model="yolo11n.pt", show=False) # line analytics
trackzone = solutions.TrackZone(region=region_points, model="yolo11n.pt", show=False) # Test trackzone
frame_count = 0 # Required for analytics
while cap.isOpened():
success, im0 = cap.read()
if not success:
break
frame_count += 1
original_im0 = im0.copy()
_ = counter.count(original_im0.copy())
_ = heatmap.generate_heatmap(original_im0.copy())
_ = heatmap_count.generate_heatmap(original_im0.copy())
_ = speed.estimate_speed(original_im0.copy())
_ = queue.process_queue(original_im0.copy())
_ = line_analytics.process_data(original_im0.copy(), frame_count)
_ = pie_analytics.process_data(original_im0.copy(), frame_count)
_ = bar_analytics.process_data(original_im0.copy(), frame_count)
_ = area_analytics.process_data(original_im0.copy(), frame_count)
_ = trackzone.trackzone(original_im0.copy())
cap.release()
# Test workouts monitoring
safe_download(url=f"{ASSETS_URL}/{POSE_VIDEO}", dir=TMP)
cap = cv2.VideoCapture(str(TMP / POSE_VIDEO))
assert cap.isOpened(), "Error reading video file"
gym = solutions.AIGym(kpts=[5, 11, 13], show=False)
while cap.isOpened():
success, im0 = cap.read()
if not success:
break
_ = gym.monitor(im0)
cap.release()
@pytest.mark.slow
def test_instance_segmentation():
"""Test the instance segmentation solution."""
from ultralytics.utils.plotting import Annotator, colors
model = YOLO(WEIGHTS_DIR / "yolo11n-seg.pt")
names = model.names
cap = cv2.VideoCapture(TMP / DEMO_VIDEO)
assert cap.isOpened(), "Error reading video file"
while cap.isOpened():
success, im0 = cap.read()
if not success:
break
results = model.predict(im0)
annotator = Annotator(im0, line_width=2)
if results[0].masks is not None:
clss = results[0].boxes.cls.cpu().tolist()
masks = results[0].masks.xy
for mask, cls in zip(masks, clss):
color = colors(int(cls), True)
annotator.seg_bbox(mask=mask, mask_color=color, label=names[int(cls)])
cap.release()
cv2.destroyAllWindows()
@pytest.mark.slow
def test_streamlit_predict():
"""Test streamlit predict live inference solution."""
solutions.Inference().inference()
train.py 0 → 100644
View file @ e63cf68a
from ultralytics import YOLO
model = YOLO('yolov13n.yaml')
# Train the model
results = model.train(
# resume=True,
# model="runs/detect/train/weights/last.pt",
data='coco.yaml',
epochs=600,
batch=256,
imgsz=640,
scale=0.5, # S:0.9; L:0.9; X:0.9
mosaic=1.0,
mixup=0.0, # S:0.05; L:0.15; X:0.2
copy_paste=0.1, # S:0.15; L:0.5; X:0.6
device="0,1,2,3",
)
# Evaluate model performance on the validation set
metrics = model.val('coco.yaml')
# Perform object detection on an image
results = model("./ultralytics/assets/bus.jpg")
results[0].show()
ultralytics/__init__.py 0 → 100644
View file @ e63cf68a
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
__version__ = "8.3.63"
import os
# Set ENV variables (place before imports)
if not os.environ.get("OMP_NUM_THREADS"):
os.environ["OMP_NUM_THREADS"] = "1" # default for reduced CPU utilization during training
from ultralytics.models import NAS, RTDETR, SAM, YOLO, FastSAM, YOLOWorld
from ultralytics.utils import ASSETS, SETTINGS
from ultralytics.utils.checks import check_yolo as checks
from ultralytics.utils.downloads import download
settings = SETTINGS
__all__ = (
"__version__",
"ASSETS",
"YOLO",
"YOLOWorld",
"NAS",
"SAM",
"FastSAM",
"RTDETR",
"checks",
"download",
"settings",
)
ultralytics/cfg/__init__.py 0 → 100644
View file @ e63cf68a
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
import shutil
import subprocess
import sys
from pathlib import Path
from types import SimpleNamespace
from typing import Dict, List, Union
import cv2
from ultralytics.utils import (
ASSETS,
DEFAULT_CFG,
DEFAULT_CFG_DICT,
DEFAULT_CFG_PATH,
DEFAULT_SOL_DICT,
IS_VSCODE,
LOGGER,
RANK,
ROOT,
RUNS_DIR,
SETTINGS,
SETTINGS_FILE,
TESTS_RUNNING,
IterableSimpleNamespace,
__version__,
checks,
colorstr,
deprecation_warn,
vscode_msg,
yaml_load,
yaml_print,
)
# Define valid solutions
SOLUTION_MAP = {
"count": ("ObjectCounter", "count"),
"heatmap": ("Heatmap", "generate_heatmap"),
"queue": ("QueueManager", "process_queue"),
"speed": ("SpeedEstimator", "estimate_speed"),
"workout": ("AIGym", "monitor"),
"analytics": ("Analytics", "process_data"),
"trackzone": ("TrackZone", "trackzone"),
"inference": ("Inference", "inference"),
"help": None,
}
# Define valid tasks and modes
MODES = {"train", "val", "predict", "export", "track", "benchmark"}
TASKS = {"detect", "segment", "classify", "pose", "obb"}
TASK2DATA = {
"detect": "coco8.yaml",
"segment": "coco8-seg.yaml",
"classify": "imagenet10",
"pose": "coco8-pose.yaml",
"obb": "dota8.yaml",
}
TASK2MODEL = {
"detect": "yolo11n.pt",
"segment": "yolo11n-seg.pt",
"classify": "yolo11n-cls.pt",
"pose": "yolo11n-pose.pt",
"obb": "yolo11n-obb.pt",
}
TASK2METRIC = {
"detect": "metrics/mAP50-95(B)",
"segment": "metrics/mAP50-95(M)",
"classify": "metrics/accuracy_top1",
"pose": "metrics/mAP50-95(P)",
"obb": "metrics/mAP50-95(B)",
}
MODELS = {TASK2MODEL[task] for task in TASKS}
ARGV = sys.argv or ["", ""] # sometimes sys.argv = []
SOLUTIONS_HELP_MSG = f"""
Arguments received: {str(["yolo"] + ARGV[1:])}. Ultralytics 'yolo solutions' usage overview:
yolo solutions SOLUTION ARGS
Where SOLUTION (optional) is one of {list(SOLUTION_MAP.keys())[:-1]}
ARGS (optional) are any number of custom 'arg=value' pairs like 'show_in=True' that override defaults
at https://docs.ultralytics.com/usage/cfg
1. Call object counting solution
yolo solutions count source="path/to/video/file.mp4" region=[(20, 400), (1080, 400), (1080, 360), (20, 360)]
2. Call heatmaps solution
yolo solutions heatmap colormap=cv2.COLORMAP_PARULA model=yolo11n.pt
3. Call queue management solution
yolo solutions queue region=[(20, 400), (1080, 400), (1080, 360), (20, 360)] model=yolo11n.pt
4. Call workouts monitoring solution for push-ups
yolo solutions workout model=yolo11n-pose.pt kpts=[6, 8, 10]
5. Generate analytical graphs
yolo solutions analytics analytics_type="pie"
6. Track objects within specific zones
yolo solutions trackzone source="path/to/video/file.mp4" region=[(150, 150), (1130, 150), (1130, 570), (150, 570)]
7. Streamlit real-time webcam inference GUI
yolo streamlit-predict
"""
CLI_HELP_MSG = f"""
Arguments received: {str(["yolo"] + ARGV[1:])}. Ultralytics 'yolo' commands use the following syntax:
yolo TASK MODE ARGS
Where TASK (optional) is one of {TASKS}
MODE (required) is one of {MODES}
ARGS (optional) are any number of custom 'arg=value' pairs like 'imgsz=320' that override defaults.
See all ARGS at https://docs.ultralytics.com/usage/cfg or with 'yolo cfg'
1. Train a detection model for 10 epochs with an initial learning_rate of 0.01
yolo train data=coco8.yaml model=yolo11n.pt epochs=10 lr0=0.01
2. Predict a YouTube video using a pretrained segmentation model at image size 320:
yolo predict model=yolo11n-seg.pt source='https://youtu.be/LNwODJXcvt4' imgsz=320
3. Val a pretrained detection model at batch-size 1 and image size 640:
yolo val model=yolo11n.pt data=coco8.yaml batch=1 imgsz=640
4. Export a YOLO11n classification model to ONNX format at image size 224 by 128 (no TASK required)
yolo export model=yolo11n-cls.pt format=onnx imgsz=224,128
5. Ultralytics solutions usage
yolo solutions count or in {list(SOLUTION_MAP.keys())[1:-1]} source="path/to/video/file.mp4"
6. Run special commands:
yolo help
yolo checks
yolo version
yolo settings
yolo copy-cfg
yolo cfg
yolo solutions help
Docs: https://docs.ultralytics.com
Solutions: https://docs.ultralytics.com/solutions/
Community: https://community.ultralytics.com
GitHub: https://github.com/ultralytics/ultralytics
"""
# Define keys for arg type checks
CFG_FLOAT_KEYS = { # integer or float arguments, i.e. x=2 and x=2.0
"warmup_epochs",
"box",
"cls",
"dfl",
"degrees",
"shear",
"time",
"workspace",
"batch",
}
CFG_FRACTION_KEYS = { # fractional float arguments with 0.0<=values<=1.0
"dropout",
"lr0",
"lrf",
"momentum",
"weight_decay",
"warmup_momentum",
"warmup_bias_lr",
"hsv_h",
"hsv_s",
"hsv_v",
"translate",
"scale",
"perspective",
"flipud",
"fliplr",
"bgr",
"mosaic",
"mixup",
"copy_paste",
"conf",
"iou",
"fraction",
}
CFG_INT_KEYS = { # integer-only arguments
"epochs",
"patience",
"workers",
"seed",
"close_mosaic",
"mask_ratio",
"max_det",
"vid_stride",
"line_width",
"nbs",
"save_period",
}
CFG_BOOL_KEYS = { # boolean-only arguments
"save",
"exist_ok",
"verbose",
"deterministic",
"single_cls",
"rect",
"cos_lr",
"overlap_mask",
"val",
"save_json",
"save_hybrid",
"half",
"dnn",
"plots",
"show",
"save_txt",
"save_conf",
"save_crop",
"save_frames",
"show_labels",
"show_conf",
"visualize",
"augment",
"agnostic_nms",
"retina_masks",
"show_boxes",
"keras",
"optimize",
"int8",
"dynamic",
"simplify",
"nms",
"profile",
"multi_scale",
}
def cfg2dict(cfg):
"""
Converts a configuration object to a dictionary.
Args:
cfg (str | Path | Dict | SimpleNamespace): Configuration object to be converted. Can be a file path,
a string, a dictionary, or a SimpleNamespace object.
Returns:
(Dict): Configuration object in dictionary format.
Examples:
Convert a YAML file path to a dictionary:
>>> config_dict = cfg2dict("config.yaml")
Convert a SimpleNamespace to a dictionary:
>>> from types import SimpleNamespace
>>> config_sn = SimpleNamespace(param1="value1", param2="value2")
>>> config_dict = cfg2dict(config_sn)
Pass through an already existing dictionary:
>>> config_dict = cfg2dict({"param1": "value1", "param2": "value2"})
Notes:
- If cfg is a path or string, it's loaded as YAML and converted to a dictionary.
- If cfg is a SimpleNamespace object, it's converted to a dictionary using vars().
- If cfg is already a dictionary, it's returned unchanged.
"""
if isinstance(cfg, (str, Path)):
cfg = yaml_load(cfg) # load dict
elif isinstance(cfg, SimpleNamespace):
cfg = vars(cfg) # convert to dict
return cfg
def get_cfg(cfg: Union[str, Path, Dict, SimpleNamespace] = DEFAULT_CFG_DICT, overrides: Dict = None):
"""
Load and merge configuration data from a file or dictionary, with optional overrides.
Args:
cfg (str | Path | Dict | SimpleNamespace): Configuration data source. Can be a file path, dictionary, or
SimpleNamespace object.
overrides (Dict | None): Dictionary containing key-value pairs to override the base configuration.
Returns:
(SimpleNamespace): Namespace containing the merged configuration arguments.
Examples:
>>> from ultralytics.cfg import get_cfg
>>> config = get_cfg() # Load default configuration
>>> config_with_overrides = get_cfg("path/to/config.yaml", overrides={"epochs": 50, "batch_size": 16})
Notes:
- If both `cfg` and `overrides` are provided, the values in `overrides` will take precedence.
- Special handling ensures alignment and correctness of the configuration, such as converting numeric
`project` and `name` to strings and validating configuration keys and values.
- The function performs type and value checks on the configuration data.
"""
cfg = cfg2dict(cfg)
# Merge overrides
if overrides:
overrides = cfg2dict(overrides)
if "save_dir" not in cfg:
overrides.pop("save_dir", None) # special override keys to ignore
check_dict_alignment(cfg, overrides)
cfg = {**cfg, **overrides} # merge cfg and overrides dicts (prefer overrides)
# Special handling for numeric project/name
for k in "project", "name":
if k in cfg and isinstance(cfg[k], (int, float)):
cfg[k] = str(cfg[k])
if cfg.get("name") == "model": # assign model to 'name' arg
cfg["name"] = str(cfg.get("model", "")).split(".")[0]
LOGGER.warning(f"WARNING ⚠️ 'name=model' automatically updated to 'name={cfg['name']}'.")
# Type and Value checks
check_cfg(cfg)
# Return instance
return IterableSimpleNamespace(**cfg)
def check_cfg(cfg, hard=True):
"""
Checks configuration argument types and values for the Ultralytics library.
This function validates the types and values of configuration arguments, ensuring correctness and converting
them if necessary. It checks for specific key types defined in global variables such as CFG_FLOAT_KEYS,
CFG_FRACTION_KEYS, CFG_INT_KEYS, and CFG_BOOL_KEYS.
Args:
cfg (Dict): Configuration dictionary to validate.
hard (bool): If True, raises exceptions for invalid types and values; if False, attempts to convert them.
Examples:
>>> config = {
... "epochs": 50, # valid integer
... "lr0": 0.01, # valid float
... "momentum": 1.2, # invalid float (out of 0.0-1.0 range)
... "save": "true", # invalid bool
... }
>>> check_cfg(config, hard=False)
>>> print(config)
{'epochs': 50, 'lr0': 0.01, 'momentum': 1.2, 'save': False} # corrected 'save' key
Notes:
- The function modifies the input dictionary in-place.
- None values are ignored as they may be from optional arguments.
- Fraction keys are checked to be within the range [0.0, 1.0].
"""
for k, v in cfg.items():
if v is not None: # None values may be from optional args
if k in CFG_FLOAT_KEYS and not isinstance(v, (int, float)):
if hard:
raise TypeError(
f"'{k}={v}' is of invalid type {type(v).__name__}. "
f"Valid '{k}' types are int (i.e. '{k}=0') or float (i.e. '{k}=0.5')"
)
cfg[k] = float(v)
elif k in CFG_FRACTION_KEYS:
if not isinstance(v, (int, float)):
if hard:
raise TypeError(
f"'{k}={v}' is of invalid type {type(v).__name__}. "
f"Valid '{k}' types are int (i.e. '{k}=0') or float (i.e. '{k}=0.5')"
)
cfg[k] = v = float(v)
if not (0.0 <= v <= 1.0):
raise ValueError(f"'{k}={v}' is an invalid value. Valid '{k}' values are between 0.0 and 1.0.")
elif k in CFG_INT_KEYS and not isinstance(v, int):
if hard:
raise TypeError(
f"'{k}={v}' is of invalid type {type(v).__name__}. '{k}' must be an int (i.e. '{k}=8')"
)
cfg[k] = int(v)
elif k in CFG_BOOL_KEYS and not isinstance(v, bool):
if hard:
raise TypeError(
f"'{k}={v}' is of invalid type {type(v).__name__}. "
f"'{k}' must be a bool (i.e. '{k}=True' or '{k}=False')"
)
cfg[k] = bool(v)
def get_save_dir(args, name=None):
"""
Returns the directory path for saving outputs, derived from arguments or default settings.
Args:
args (SimpleNamespace): Namespace object containing configurations such as 'project', 'name', 'task',
'mode', and 'save_dir'.
name (str | None): Optional name for the output directory. If not provided, it defaults to 'args.name'
or the 'args.mode'.
Returns:
(Path): Directory path where outputs should be saved.
Examples:
>>> from types import SimpleNamespace
>>> args = SimpleNamespace(project="my_project", task="detect", mode="train", exist_ok=True)
>>> save_dir = get_save_dir(args)
>>> print(save_dir)
my_project/detect/train
"""
if getattr(args, "save_dir", None):
save_dir = args.save_dir
else:
from ultralytics.utils.files import increment_path
project = args.project or (ROOT.parent / "tests/tmp/runs" if TESTS_RUNNING else RUNS_DIR) / args.task
name = name or args.name or f"{args.mode}"
save_dir = increment_path(Path(project) / name, exist_ok=args.exist_ok if RANK in {-1, 0} else True)
return Path(save_dir)
def _handle_deprecation(custom):
"""
Handles deprecated configuration keys by mapping them to current equivalents with deprecation warnings.
Args:
custom (Dict): Configuration dictionary potentially containing deprecated keys.
Examples:
>>> custom_config = {"boxes": True, "hide_labels": "False", "line_thickness": 2}
>>> _handle_deprecation(custom_config)
>>> print(custom_config)
{'show_boxes': True, 'show_labels': True, 'line_width': 2}
Notes:
This function modifies the input dictionary in-place, replacing deprecated keys with their current
equivalents. It also handles value conversions where necessary, such as inverting boolean values for
'hide_labels' and 'hide_conf'.
"""
for key in custom.copy().keys():
if key == "boxes":
deprecation_warn(key, "show_boxes")
custom["show_boxes"] = custom.pop("boxes")
if key == "hide_labels":
deprecation_warn(key, "show_labels")
custom["show_labels"] = custom.pop("hide_labels") == "False"
if key == "hide_conf":
deprecation_warn(key, "show_conf")
custom["show_conf"] = custom.pop("hide_conf") == "False"
if key == "line_thickness":
deprecation_warn(key, "line_width")
custom["line_width"] = custom.pop("line_thickness")
if key == "label_smoothing":
deprecation_warn(key)
custom.pop("label_smoothing")
return custom
def check_dict_alignment(base: Dict, custom: Dict, e=None):
"""
Checks alignment between custom and base configuration dictionaries, handling deprecated keys and providing error
messages for mismatched keys.
Args:
base (Dict): The base configuration dictionary containing valid keys.
custom (Dict): The custom configuration dictionary to be checked for alignment.
e (Exception | None): Optional error instance passed by the calling function.
Raises:
SystemExit: If mismatched keys are found between the custom and base dictionaries.
Examples:
>>> base_cfg = {"epochs": 50, "lr0": 0.01, "batch_size": 16}
>>> custom_cfg = {"epoch": 100, "lr": 0.02, "batch_size": 32}
>>> try:
... check_dict_alignment(base_cfg, custom_cfg)
... except SystemExit:
... print("Mismatched keys found")
Notes:
- Suggests corrections for mismatched keys based on similarity to valid keys.
- Automatically replaces deprecated keys in the custom configuration with updated equivalents.
- Prints detailed error messages for each mismatched key to help users correct their configurations.
"""
custom = _handle_deprecation(custom)
base_keys, custom_keys = (set(x.keys()) for x in (base, custom))
if mismatched := [k for k in custom_keys if k not in base_keys]:
from difflib import get_close_matches
string = ""
for x in mismatched:
matches = get_close_matches(x, base_keys) # key list
matches = [f"{k}={base[k]}" if base.get(k) is not None else k for k in matches]
match_str = f"Similar arguments are i.e. {matches}." if matches else ""
string += f"'{colorstr('red', 'bold', x)}' is not a valid YOLO argument. {match_str}\n"
raise SyntaxError(string + CLI_HELP_MSG) from e
def merge_equals_args(args: List[str]) -> List[str]:
"""
Merges arguments around isolated '=' in a list of strings and joins fragments with brackets.
This function handles the following cases:
1. ['arg', '=', 'val'] becomes ['arg=val']
2. ['arg=', 'val'] becomes ['arg=val']
3. ['arg', '=val'] becomes ['arg=val']
4. Joins fragments with brackets, e.g., ['imgsz=[3,', '640,', '640]'] becomes ['imgsz=[3,640,640]']
Args:
args (List[str]): A list of strings where each element represents an argument or fragment.
Returns:
List[str]: A list of strings where the arguments around isolated '=' are merged and fragments with brackets are joined.
Examples:
>>> args = ["arg1", "=", "value", "arg2=", "value2", "arg3", "=value3", "imgsz=[3,", "640,", "640]"]
>>> merge_and_join_args(args)
['arg1=value', 'arg2=value2', 'arg3=value3', 'imgsz=[3,640,640]']
"""
new_args = []
current = ""
depth = 0
i = 0
while i < len(args):
arg = args[i]
# Handle equals sign merging
if arg == "=" and 0 < i < len(args) - 1: # merge ['arg', '=', 'val']
new_args[-1] += f"={args[i + 1]}"
i += 2
continue
elif arg.endswith("=") and i < len(args) - 1 and "=" not in args[i + 1]: # merge ['arg=', 'val']
new_args.append(f"{arg}{args[i + 1]}")
i += 2
continue
elif arg.startswith("=") and i > 0: # merge ['arg', '=val']
new_args[-1] += arg
i += 1
continue
# Handle bracket joining
depth += arg.count("[") - arg.count("]")
current += arg
if depth == 0:
new_args.append(current)
current = ""
i += 1
# Append any remaining current string
if current:
new_args.append(current)
return new_args
def handle_yolo_hub(args: List[str]) -> None:
"""
Handles Ultralytics HUB command-line interface (CLI) commands for authentication.
This function processes Ultralytics HUB CLI commands such as login and logout. It should be called when executing a
script with arguments related to HUB authentication.
Args:
args (List[str]): A list of command line arguments. The first argument should be either 'login'
or 'logout'. For 'login', an optional second argument can be the API key.
Examples:
```bash
yolo login YOUR_API_KEY
```
Notes:
- The function imports the 'hub' module from ultralytics to perform login and logout operations.
- For the 'login' command, if no API key is provided, an empty string is passed to the login function.
- The 'logout' command does not require any additional arguments.
"""
from ultralytics import hub
if args[0] == "login":
key = args[1] if len(args) > 1 else ""
# Log in to Ultralytics HUB using the provided API key
hub.login(key)
elif args[0] == "logout":
# Log out from Ultralytics HUB
hub.logout()
def handle_yolo_settings(args: List[str]) -> None:
"""
Handles YOLO settings command-line interface (CLI) commands.
This function processes YOLO settings CLI commands such as reset and updating individual settings. It should be
called when executing a script with arguments related to YOLO settings management.
Args:
args (List[str]): A list of command line arguments for YOLO settings management.
Examples:
>>> handle_yolo_settings(["reset"]) # Reset YOLO settings
>>> handle_yolo_settings(["default_cfg_path=yolo11n.yaml"]) # Update a specific setting
Notes:
- If no arguments are provided, the function will display the current settings.
- The 'reset' command will delete the existing settings file and create new default settings.
- Other arguments are treated as key-value pairs to update specific settings.
- The function will check for alignment between the provided settings and the existing ones.
- After processing, the updated settings will be displayed.
- For more information on handling YOLO settings, visit:
https://docs.ultralytics.com/quickstart/#ultralytics-settings
"""
url = "https://docs.ultralytics.com/quickstart/#ultralytics-settings" # help URL
try:
if any(args):
if args[0] == "reset":
SETTINGS_FILE.unlink() # delete the settings file
SETTINGS.reset() # create new settings
LOGGER.info("Settings reset successfully") # inform the user that settings have been reset
else: # save a new setting
new = dict(parse_key_value_pair(a) for a in args)
check_dict_alignment(SETTINGS, new)
SETTINGS.update(new)
print(SETTINGS) # print the current settings
LOGGER.info(f"💡 Learn more about Ultralytics Settings at {url}")
except Exception as e:
LOGGER.warning(f"WARNING ⚠️ settings error: '{e}'. Please see {url} for help.")
def handle_yolo_solutions(args: List[str]) -> None:
"""
Processes YOLO solutions arguments and runs the specified computer vision solutions pipeline.
Args:
args (List[str]): Command-line arguments for configuring and running the Ultralytics YOLO
solutions: https://docs.ultralytics.com/solutions/, It can include solution name, source,
and other configuration parameters.
Returns:
None: The function processes video frames and saves the output but doesn't return any value.
Examples:
Run people counting solution with default settings:
>>> handle_yolo_solutions(["count"])
Run analytics with custom configuration:
>>> handle_yolo_solutions(["analytics", "conf=0.25", "source=path/to/video/file.mp4"])
Run inference with custom configuration, requires Streamlit version 1.29.0 or higher.
>>> handle_yolo_solutions(["inference", "model=yolo11n.pt"])
Notes:
- Default configurations are merged from DEFAULT_SOL_DICT and DEFAULT_CFG_DICT
- Arguments can be provided in the format 'key=value' or as boolean flags
- Available solutions are defined in SOLUTION_MAP with their respective classes and methods
- If an invalid solution is provided, defaults to 'count' solution
- Output videos are saved in 'runs/solution/{solution_name}' directory
- For 'analytics' solution, frame numbers are tracked for generating analytical graphs
- Video processing can be interrupted by pressing 'q'
- Processes video frames sequentially and saves output in .avi format
- If no source is specified, downloads and uses a default sample video\
- The inference solution will be launched using the 'streamlit run' command.
- The Streamlit app file is located in the Ultralytics package directory.
"""
full_args_dict = {**DEFAULT_SOL_DICT, **DEFAULT_CFG_DICT} # arguments dictionary
overrides = {}
# check dictionary alignment
for arg in merge_equals_args(args):
arg = arg.lstrip("-").rstrip(",")
if "=" in arg:
try:
k, v = parse_key_value_pair(arg)
overrides[k] = v
except (NameError, SyntaxError, ValueError, AssertionError) as e:
check_dict_alignment(full_args_dict, {arg: ""}, e)
elif arg in full_args_dict and isinstance(full_args_dict.get(arg), bool):
overrides[arg] = True
check_dict_alignment(full_args_dict, overrides) # dict alignment
# Get solution name
if args and args[0] in SOLUTION_MAP:
if args[0] != "help":
s_n = args.pop(0) # Extract the solution name directly
else:
LOGGER.info(SOLUTIONS_HELP_MSG)
else:
LOGGER.warning(
f"⚠️ No valid solution provided. Using default 'count'. Available: {', '.join(SOLUTION_MAP.keys())}"
)
s_n = "count" # Default solution if none provided
if args and args[0] == "help": # Add check for return if user call `yolo solutions help`
return
if s_n == "inference":
checks.check_requirements("streamlit>=1.29.0")
LOGGER.info("💡 Loading Ultralytics live inference app...")
subprocess.run(
[ # Run subprocess with Streamlit custom argument
"streamlit",
"run",
str(ROOT / "solutions/streamlit_inference.py"),
"--server.headless",
"true",
overrides.pop("model", "yolo11n.pt"),
]
)
else:
cls, method = SOLUTION_MAP[s_n] # solution class name, method name and default source
from ultralytics import solutions # import ultralytics solutions
solution = getattr(solutions, cls)(IS_CLI=True, **overrides) # get solution class i.e ObjectCounter
process = getattr(
solution, method
) # get specific function of class for processing i.e, count from ObjectCounter
cap = cv2.VideoCapture(solution.CFG["source"]) # read the video file
# extract width, height and fps of the video file, create save directory and initialize video writer
import os # for directory creation
from pathlib import Path
from ultralytics.utils.files import increment_path # for output directory path update
w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))
if s_n == "analytics": # analytical graphs follow fixed shape for output i.e w=1920, h=1080
w, h = 1920, 1080
save_dir = increment_path(Path("runs") / "solutions" / "exp", exist_ok=False)
save_dir.mkdir(parents=True, exist_ok=True) # create the output directory
vw = cv2.VideoWriter(os.path.join(save_dir, "solution.avi"), cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
try: # Process video frames
f_n = 0 # frame number, required for analytical graphs
while cap.isOpened():
success, frame = cap.read()
if not success:
break
frame = process(frame, f_n := f_n + 1) if s_n == "analytics" else process(frame)
vw.write(frame)
if cv2.waitKey(1) & 0xFF == ord("q"):
break
finally:
cap.release()
def parse_key_value_pair(pair: str = "key=value"):
"""
Parses a key-value pair string into separate key and value components.
Args:
pair (str): A string containing a key-value pair in the format "key=value".
Returns:
key (str): The parsed key.
value (str): The parsed value.
Raises:
AssertionError: If the value is missing or empty.
Examples:
>>> key, value = parse_key_value_pair("model=yolo11n.pt")
>>> print(f"Key: {key}, Value: {value}")
Key: model, Value: yolo11n.pt
>>> key, value = parse_key_value_pair("epochs=100")
>>> print(f"Key: {key}, Value: {value}")
Key: epochs, Value: 100
Notes:
- The function splits the input string on the first '=' character.
- Leading and trailing whitespace is removed from both key and value.
- An assertion error is raised if the value is empty after stripping.
"""
k, v = pair.split("=", 1) # split on first '=' sign
k, v = k.strip(), v.strip() # remove spaces
assert v, f"missing '{k}' value"
return k, smart_value(v)
def smart_value(v):
"""
Converts a string representation of a value to its appropriate Python type.
This function attempts to convert a given string into a Python object of the most appropriate type. It handles
conversions to None, bool, int, float, and other types that can be evaluated safely.
Args:
v (str): The string representation of the value to be converted.
Returns:
(Any): The converted value. The type can be None, bool, int, float, or the original string if no conversion
is applicable.
Examples:
>>> smart_value("42")
42
>>> smart_value("3.14")
3.14
>>> smart_value("True")
True
>>> smart_value("None")
None
>>> smart_value("some_string")
'some_string'
Notes:
- The function uses a case-insensitive comparison for boolean and None values.
- For other types, it attempts to use Python's eval() function, which can be unsafe if used on untrusted input.
- If no conversion is possible, the original string is returned.
"""
v_lower = v.lower()
if v_lower == "none":
return None
elif v_lower == "true":
return True
elif v_lower == "false":
return False
else:
try:
return eval(v)
except Exception:
return v
def entrypoint(debug=""):
"""
Ultralytics entrypoint function for parsing and executing command-line arguments.
This function serves as the main entry point for the Ultralytics CLI, parsing command-line arguments and
executing the corresponding tasks such as training, validation, prediction, exporting models, and more.
Args:
debug (str): Space-separated string of command-line arguments for debugging purposes.
Examples:
Train a detection model for 10 epochs with an initial learning_rate of 0.01:
>>> entrypoint("train data=coco8.yaml model=yolo11n.pt epochs=10 lr0=0.01")
Predict a YouTube video using a pretrained segmentation model at image size 320:
>>> entrypoint("predict model=yolo11n-seg.pt source='https://youtu.be/LNwODJXcvt4' imgsz=320")
Validate a pretrained detection model at batch-size 1 and image size 640:
>>> entrypoint("val model=yolo11n.pt data=coco8.yaml batch=1 imgsz=640")
Notes:
- If no arguments are passed, the function will display the usage help message.
- For a list of all available commands and their arguments, see the provided help messages and the
Ultralytics documentation at https://docs.ultralytics.com.
"""
args = (debug.split(" ") if debug else ARGV)[1:]
if not args: # no arguments passed
LOGGER.info(CLI_HELP_MSG)
return
special = {
"help": lambda: LOGGER.info(CLI_HELP_MSG),
"checks": checks.collect_system_info,
"version": lambda: LOGGER.info(__version__),
"settings": lambda: handle_yolo_settings(args[1:]),
"cfg": lambda: yaml_print(DEFAULT_CFG_PATH),
"hub": lambda: handle_yolo_hub(args[1:]),
"login": lambda: handle_yolo_hub(args),
"logout": lambda: handle_yolo_hub(args),
"copy-cfg": copy_default_cfg,
"solutions": lambda: handle_yolo_solutions(args[1:]),
}
full_args_dict = {**DEFAULT_CFG_DICT, **{k: None for k in TASKS}, **{k: None for k in MODES}, **special}
# Define common misuses of special commands, i.e. -h, -help, --help
special.update({k[0]: v for k, v in special.items()}) # singular
special.update({k[:-1]: v for k, v in special.items() if len(k) > 1 and k.endswith("s")}) # singular
special = {**special, **{f"-{k}": v for k, v in special.items()}, **{f"--{k}": v for k, v in special.items()}}
overrides = {} # basic overrides, i.e. imgsz=320
for a in merge_equals_args(args): # merge spaces around '=' sign
if a.startswith("--"):
LOGGER.warning(f"WARNING ⚠️ argument '{a}' does not require leading dashes '--', updating to '{a[2:]}'.")
a = a[2:]
if a.endswith(","):
LOGGER.warning(f"WARNING ⚠️ argument '{a}' does not require trailing comma ',', updating to '{a[:-1]}'.")
a = a[:-1]
if "=" in a:
try:
k, v = parse_key_value_pair(a)
if k == "cfg" and v is not None: # custom.yaml passed
LOGGER.info(f"Overriding {DEFAULT_CFG_PATH} with {v}")
overrides = {k: val for k, val in yaml_load(checks.check_yaml(v)).items() if k != "cfg"}
else:
overrides[k] = v
except (NameError, SyntaxError, ValueError, AssertionError) as e:
check_dict_alignment(full_args_dict, {a: ""}, e)
elif a in TASKS:
overrides["task"] = a
elif a in MODES:
overrides["mode"] = a
elif a.lower() in special:
special[a.lower()]()
return
elif a in DEFAULT_CFG_DICT and isinstance(DEFAULT_CFG_DICT[a], bool):
overrides[a] = True # auto-True for default bool args, i.e. 'yolo show' sets show=True
elif a in DEFAULT_CFG_DICT:
raise SyntaxError(
f"'{colorstr('red', 'bold', a)}' is a valid YOLO argument but is missing an '=' sign "
f"to set its value, i.e. try '{a}={DEFAULT_CFG_DICT[a]}'\n{CLI_HELP_MSG}"
)
else:
check_dict_alignment(full_args_dict, {a: ""})
# Check keys
check_dict_alignment(full_args_dict, overrides)
# Mode
mode = overrides.get("mode")
if mode is None:
mode = DEFAULT_CFG.mode or "predict"
LOGGER.warning(f"WARNING ⚠️ 'mode' argument is missing. Valid modes are {MODES}. Using default 'mode={mode}'.")
elif mode not in MODES:
raise ValueError(f"Invalid 'mode={mode}'. Valid modes are {MODES}.\n{CLI_HELP_MSG}")
# Task
task = overrides.pop("task", None)
if task:
if task == "classify" and mode == "track":
raise ValueError(
f"❌ Classification doesn't support 'mode=track'. Valid modes for classification are"
f" {MODES - {'track'}}.\n{CLI_HELP_MSG}"
)
elif task not in TASKS:
if task == "track":
LOGGER.warning(
"WARNING ⚠️ invalid 'task=track', setting 'task=detect' and 'mode=track'. Valid tasks are {TASKS}.\n{CLI_HELP_MSG}."
)
task, mode = "detect", "track"
else:
raise ValueError(f"Invalid 'task={task}'. Valid tasks are {TASKS}.\n{CLI_HELP_MSG}")
if "model" not in overrides:
overrides["model"] = TASK2MODEL[task]
# Model
model = overrides.pop("model", DEFAULT_CFG.model)
if model is None:
model = "yolo11n.pt"
LOGGER.warning(f"WARNING ⚠️ 'model' argument is missing. Using default 'model={model}'.")
overrides["model"] = model
stem = Path(model).stem.lower()
if "rtdetr" in stem: # guess architecture
from ultralytics import RTDETR
model = RTDETR(model) # no task argument
elif "fastsam" in stem:
from ultralytics import FastSAM
model = FastSAM(model)
elif "sam_" in stem or "sam2_" in stem or "sam2.1_" in stem:
from ultralytics import SAM
model = SAM(model)
else:
from ultralytics import YOLO
model = YOLO(model, task=task)
if isinstance(overrides.get("pretrained"), str):
model.load(overrides["pretrained"])
# Task Update
if task != model.task:
if task:
LOGGER.warning(
f"WARNING ⚠️ conflicting 'task={task}' passed with 'task={model.task}' model. "
f"Ignoring 'task={task}' and updating to 'task={model.task}' to match model."
)
task = model.task
# Mode
if mode in {"predict", "track"} and "source" not in overrides:
overrides["source"] = (
"https://ultralytics.com/images/boats.jpg" if task == "obb" else DEFAULT_CFG.source or ASSETS
)
LOGGER.warning(f"WARNING ⚠️ 'source' argument is missing. Using default 'source={overrides['source']}'.")
elif mode in {"train", "val"}:
if "data" not in overrides and "resume" not in overrides:
overrides["data"] = DEFAULT_CFG.data or TASK2DATA.get(task or DEFAULT_CFG.task, DEFAULT_CFG.data)
LOGGER.warning(f"WARNING ⚠️ 'data' argument is missing. Using default 'data={overrides['data']}'.")
elif mode == "export":
if "format" not in overrides:
overrides["format"] = DEFAULT_CFG.format or "torchscript"
LOGGER.warning(f"WARNING ⚠️ 'format' argument is missing. Using default 'format={overrides['format']}'.")
# Run command in python
getattr(model, mode)(**overrides) # default args from model
# Show help
LOGGER.info(f"💡 Learn more at https://docs.ultralytics.com/modes/{mode}")
# Recommend VS Code extension
if IS_VSCODE and SETTINGS.get("vscode_msg", True):
LOGGER.info(vscode_msg())
# Special modes --------------------------------------------------------------------------------------------------------
def copy_default_cfg():
"""
Copies the default configuration file and creates a new one with '_copy' appended to its name.
This function duplicates the existing default configuration file (DEFAULT_CFG_PATH) and saves it
with '_copy' appended to its name in the current working directory. It provides a convenient way
to create a custom configuration file based on the default settings.
Examples:
>>> copy_default_cfg()
# Output: default.yaml copied to /path/to/current/directory/default_copy.yaml
# Example YOLO command with this new custom cfg:
# yolo cfg='/path/to/current/directory/default_copy.yaml' imgsz=320 batch=8
Notes:
- The new configuration file is created in the current working directory.
- After copying, the function prints a message with the new file's location and an example
YOLO command demonstrating how to use the new configuration file.
- This function is useful for users who want to modify the default configuration without
altering the original file.
"""
new_file = Path.cwd() / DEFAULT_CFG_PATH.name.replace(".yaml", "_copy.yaml")
shutil.copy2(DEFAULT_CFG_PATH, new_file)
LOGGER.info(
f"{DEFAULT_CFG_PATH} copied to {new_file}\n"
f"Example YOLO command with this new custom cfg:\n yolo cfg='{new_file}' imgsz=320 batch=8"
)
if __name__ == "__main__":
# Example: entrypoint(debug='yolo predict model=yolo11n.pt')
entrypoint(debug="")
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment