diff --git a/benchmarks.py b/benchmarks.py
index c849eed6f..90cce4b3f 100644
--- a/benchmarks.py
+++ b/benchmarks.py
@@ -60,7 +60,41 @@ def run(
pt_only=False, # test PyTorch only
hard_fail=False, # throw error on benchmark failure
):
- """Run YOLOv5 benchmarks on multiple export formats and log results for model performance evaluation."""
+ """
+ Run YOLOv5 benchmarks on multiple export formats and log results for model performance evaluation.
+
+ Args:
+ weights (Path | str): Path to the model weights file (default: ROOT / "yolov5s.pt").
+ imgsz (int): Inference size in pixels (default: 640).
+ batch_size (int): Batch size for inference (default: 1).
+ data (Path | str): Path to the dataset.yaml file (default: ROOT / "data/coco128.yaml").
+ device (str): CUDA device, e.g., '0' or '0,1,2,3' or 'cpu' (default: None).
+ half (bool): Use FP16 half-precision inference (default: False).
+ test (bool): Test export formats only (default: False).
+ pt_only (bool): Test PyTorch format only (default: False).
+ hard_fail (bool): Throw an error on benchmark failure if True (default: False).
+
+ Returns:
+ None. Logs information about the benchmark results, including the format, size, mAP50-95, and inference time.
+
+ Notes:
+ Supported export formats and models include PyTorch, TorchScript, ONNX, OpenVINO, TensorRT, CoreML, TensorFlow
+ SavedModel, TensorFlow GraphDef, TensorFlow Lite, and TensorFlow Edge TPU. Edge TPU and TF.js are unsupported.
+
+ Examples:
+ ```python
+ $ python benchmarks.py --weights yolov5s.pt --img 640
+ ```
+
+ Usage:
+ Install required packages:
+ $ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime openvino-dev tensorflow-cpu # CPU support
+ $ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime-gpu openvino-dev tensorflow # GPU support
+ $ pip install -U nvidia-tensorrt --index-url https://pypi.ngc.nvidia.com # TensorRT
+
+ Run benchmarks:
+ $ python benchmarks.py --weights yolov5s.pt --img 640
+ """
y, t = [], time.time()
device = select_device(device)
model_type = type(attempt_load(weights, fuse=False)) # DetectionModel, SegmentationModel, etc.
@@ -125,7 +159,23 @@ def test(
pt_only=False, # test PyTorch only
hard_fail=False, # throw error on benchmark failure
):
- """Run YOLOv5 export tests for all supported formats and log the results, including inference speed and mAP."""
+ """
+ Run YOLOv5 export tests for all supported formats and log the results, including export statuses.
+
+ Args:
+ weights (Path | str): Path to the model weights file (.pt format). Default is 'ROOT / "yolov5s.pt"'.
+ imgsz (int): Inference image size (in pixels). Default is 640.
+ batch_size (int): Batch size for testing. Default is 1.
+ data (Path | str): Path to the dataset configuration file (.yaml format). Default is 'ROOT / "data/coco128.yaml"'.
+ device (str): Device for running the tests, can be 'cpu' or a specific CUDA device ('0', '0,1,2,3', etc.). Default is an empty string.
+ half (bool): Use FP16 half-precision for inference if True. Default is False.
+ test (bool): Test export formats only without running inference. Default is False.
+ pt_only (bool): Test only the PyTorch model if True. Default is False.
+ hard_fail (bool): Raise error on export or test failure if True. Default is False.
+
+ Returns:
+ pd.DataFrame: DataFrame containing the results of the export tests, including format names and export statuses.
+ """
y, t = [], time.time()
device = select_device(device)
for i, (name, f, suffix, gpu) in export.export_formats().iterrows(): # index, (name, file, suffix, gpu-capable)
@@ -151,7 +201,28 @@ def test(
def parse_opt():
- """Parses command-line arguments for YOLOv5 model inference configuration."""
+ """
+ Parses command-line arguments for YOLOv5 model inference configuration.
+
+ Args:
+ weights (str): The path to the weights file. Defaults to 'ROOT / "yolov5s.pt"'.
+ imgsz (int): Inference size in pixels. Defaults to 640.
+ batch_size (int): Batch size. Defaults to 1.
+ data (str): Path to the dataset YAML file. Defaults to 'ROOT / "data/coco128.yaml"'.
+ device (str): CUDA device, e.g., '0' or '0,1,2,3' or 'cpu'. Defaults to an empty string (auto-select).
+ half (bool): Use FP16 half-precision inference. This is a flag and defaults to False.
+ test (bool): Test exports only. This is a flag and defaults to False.
+ pt_only (bool): Test PyTorch only. This is a flag and defaults to False.
+ hard_fail (bool|str): Throw an error on benchmark failure. Can be a boolean or a string representing a minimum metric
+ floor, i.e., '0.29'. Defaults to False.
+
+ Returns:
+ argparse.Namespace: Parsed command-line arguments encapsulated in an argparse Namespace object.
+
+ Notes:
+ The function modifies the 'opt.data' by checking and validating the YAML path using 'check_yaml()'.
+ The parsed arguments are printed for reference using 'print_args()'.
+ """
parser = argparse.ArgumentParser()
parser.add_argument("--weights", type=str, default=ROOT / "yolov5s.pt", help="weights path")
parser.add_argument("--imgsz", "--img", "--img-size", type=int, default=640, help="inference size (pixels)")
@@ -169,7 +240,34 @@ def parse_opt():
def main(opt):
- """Executes a test run if `opt.test` is True, otherwise starts training or inference with provided options."""
+ """
+ Executes YOLOv5 benchmark tests or main training/inference routines based on the provided command-line arguments.
+
+ Args:
+ opt (argparse.Namespace): Parsed command-line arguments including options for weights, image size, batch size, data
+ configuration, device, and other flags for inference settings.
+
+ Returns:
+ None: This function does not return any value. It leverages side-effects such as logging and running benchmarks.
+
+ Example:
+ ```python
+ if __name__ == "__main__":
+ opt = parse_opt()
+ main(opt)
+ ```
+
+ Notes:
+ - For a complete list of supported export formats and their respective requirements, refer to the
+ [Ultralytics YOLOv5 Export Formats](https://github.com/ultralytics/yolov5#export-formats).
+ - Ensure that you have installed all necessary dependencies by following the installation instructions detailed in
+ the [main repository](https://github.com/ultralytics/yolov5#installation).
+
+ ```shell
+ # Running benchmarks on default weights and image size
+ $ python benchmarks.py --weights yolov5s.pt --img 640
+ ```
+ """
test(**vars(opt)) if opt.test else run(**vars(opt))
diff --git a/detect.py b/detect.py
index f791faa09..b774370f5 100644
--- a/detect.py
+++ b/detect.py
@@ -97,7 +97,56 @@ def run(
dnn=False, # use OpenCV DNN for ONNX inference
vid_stride=1, # video frame-rate stride
):
- """Runs YOLOv5 detection inference on various sources like images, videos, directories, streams, etc."""
+ """
+ Runs YOLOv5 detection inference on various sources like images, videos, directories, streams, etc.
+
+ Args:
+ weights (str | Path): Path to the model weights file or a Triton URL. Default is 'yolov5s.pt'.
+ source (str | Path): Input source, which can be a file, directory, URL, glob pattern, screen capture, or webcam index.
+ Default is 'data/images'.
+ data (str | Path): Path to the dataset YAML file. Default is 'data/coco128.yaml'.
+ imgsz (tuple[int, int]): Inference image size as a tuple (height, width). Default is (640, 640).
+ conf_thres (float): Confidence threshold for detections. Default is 0.25.
+ iou_thres (float): Intersection Over Union (IOU) threshold for non-max suppression. Default is 0.45.
+ max_det (int): Maximum number of detections per image. Default is 1000.
+ device (str): CUDA device identifier (e.g., '0' or '0,1,2,3') or 'cpu'. Default is an empty string, which
+ uses the best available device.
+ view_img (bool): If True, display inference results using OpenCV. Default is False.
+ save_txt (bool): If True, save results in a text file. Default is False.
+ save_csv (bool): If True, save results in a CSV file. Default is False.
+ save_conf (bool): If True, include confidence scores in the saved results. Default is False.
+ save_crop (bool): If True, save cropped prediction boxes. Default is False.
+ nosave (bool): If True, do not save inference images or videos. Default is False.
+ classes (list[int]): List of class indices to filter detections by. Default is None.
+ agnostic_nms (bool): If True, perform class-agnostic non-max suppression. Default is False.
+ augment (bool): If True, use augmented inference. Default is False.
+ visualize (bool): If True, visualize feature maps. Default is False.
+ update (bool): If True, update all models' weights. Default is False.
+ project (str | Path): Directory to save results. Default is 'runs/detect'.
+ name (str): Name of the current experiment; used to create a subdirectory within 'project'. Default is 'exp'.
+ exist_ok (bool): If True, existing directories with the same name are reused instead of being incremented. Default is
+ False.
+ line_thickness (int): Thickness of bounding box lines in pixels. Default is 3.
+ hide_labels (bool): If True, do not display labels on bounding boxes. Default is False.
+ hide_conf (bool): If True, do not display confidence scores on bounding boxes. Default is False.
+ half (bool): If True, use FP16 half-precision inference. Default is False.
+ dnn (bool): If True, use OpenCV DNN backend for ONNX inference. Default is False.
+ vid_stride (int): Stride for processing video frames, to skip frames between processing. Default is 1.
+
+ Returns:
+ None
+
+ Examples:
+ ```python
+ from ultralytics import run
+
+ # Run inference on an image
+ run(source='data/images/example.jpg', weights='yolov5s.pt', device='0')
+
+ # Run inference on a video with specific confidence threshold
+ run(source='data/videos/example.mp4', weights='yolov5s.pt', conf_thres=0.4, device='0')
+ ```
+ """
source = str(source)
save_img = not nosave and not source.endswith(".txt") # save inference images
is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)
@@ -266,7 +315,48 @@ def run(
def parse_opt():
- """Parses command-line arguments for YOLOv5 detection, setting inference options and model configurations."""
+ """
+ Parses command-line arguments for YOLOv5 detection, setting inference options and model configurations.
+
+ Args:
+ --weights (str | list[str], optional): Model path or Triton URL. Defaults to ROOT / 'yolov5s.pt'.
+ --source (str, optional): File/dir/URL/glob/screen/0(webcam). Defaults to ROOT / 'data/images'.
+ --data (str, optional): Dataset YAML path. Provides dataset configuration information.
+ --imgsz (list[int], optional): Inference size (height, width). Defaults to [640].
+ --conf-thres (float, optional): Confidence threshold. Defaults to 0.25.
+ --iou-thres (float, optional): NMS IoU threshold. Defaults to 0.45.
+ --max-det (int, optional): Maximum number of detections per image. Defaults to 1000.
+ --device (str, optional): CUDA device, i.e., '0' or '0,1,2,3' or 'cpu'. Defaults to "".
+ --view-img (bool, optional): Flag to display results. Defaults to False.
+ --save-txt (bool, optional): Flag to save results to *.txt files. Defaults to False.
+ --save-csv (bool, optional): Flag to save results in CSV format. Defaults to False.
+ --save-conf (bool, optional): Flag to save confidences in labels saved via --save-txt. Defaults to False.
+ --save-crop (bool, optional): Flag to save cropped prediction boxes. Defaults to False.
+ --nosave (bool, optional): Flag to prevent saving images/videos. Defaults to False.
+ --classes (list[int], optional): List of classes to filter results by, e.g., '--classes 0 2 3'. Defaults to None.
+ --agnostic-nms (bool, optional): Flag for class-agnostic NMS. Defaults to False.
+ --augment (bool, optional): Flag for augmented inference. Defaults to False.
+ --visualize (bool, optional): Flag for visualizing features. Defaults to False.
+ --update (bool, optional): Flag to update all models in the model directory. Defaults to False.
+ --project (str, optional): Directory to save results. Defaults to ROOT / 'runs/detect'.
+ --name (str, optional): Sub-directory name for saving results within --project. Defaults to 'exp'.
+ --exist-ok (bool, optional): Flag to allow overwriting if the project/name already exists. Defaults to False.
+ --line-thickness (int, optional): Thickness (in pixels) of bounding boxes. Defaults to 3.
+ --hide-labels (bool, optional): Flag to hide labels in the output. Defaults to False.
+ --hide-conf (bool, optional): Flag to hide confidences in the output. Defaults to False.
+ --half (bool, optional): Flag to use FP16 half-precision inference. Defaults to False.
+ --dnn (bool, optional): Flag to use OpenCV DNN for ONNX inference. Defaults to False.
+ --vid-stride (int, optional): Video frame-rate stride, determining the number of frames to skip in between consecutive frames. Defaults to 1.
+
+ Returns:
+ argparse.Namespace: Parsed command-line arguments as an argparse.Namespace object.
+
+ Example:
+ ```python
+ from ultralytics import YOLOv5
+ args = YOLOv5.parse_opt()
+ ```
+ """
parser = argparse.ArgumentParser()
parser.add_argument("--weights", nargs="+", type=str, default=ROOT / "yolov5s.pt", help="model path or triton URL")
parser.add_argument("--source", type=str, default=ROOT / "data/images", help="file/dir/URL/glob/screen/0(webcam)")
@@ -303,7 +393,28 @@ def parse_opt():
def main(opt):
- """Executes YOLOv5 model inference with given options, checking requirements before running the model."""
+ """
+ Executes YOLOv5 model inference based on provided command-line arguments, validating dependencies before running.
+
+ Args:
+ opt (argparse.Namespace): Command-line arguments for YOLOv5 detection. See function `parse_opt` for details.
+
+ Returns:
+ None
+
+ Note:
+ This function performs essential pre-execution checks and initiates the YOLOv5 detection process based on user-specified options.
+ Refer to the usage guide and examples for more information about different sources and formats at:
+ https://github.com/ultralytics/ultralytics
+
+ Example usage:
+
+ ```python
+ if __name__ == "__main__":
+ opt = parse_opt()
+ main(opt)
+ ```
+ """
check_requirements(ROOT / "requirements.txt", exclude=("tensorboard", "thop"))
run(**vars(opt))
diff --git a/export.py b/export.py
index 77c972ba2..bf68ae123 100644
--- a/export.py
+++ b/export.py
@@ -92,7 +92,22 @@ MACOS = platform.system() == "Darwin" # macOS environment
class iOSModel(torch.nn.Module):
def __init__(self, model, im):
- """Initializes an iOS compatible model with normalization based on image dimensions."""
+ """
+ Initializes an iOS compatible model with normalization based on image dimensions.
+
+ Args:
+ model (torch.nn.Module): The PyTorch model to be adapted for iOS compatibility.
+ im (torch.Tensor): An input tensor representing a batch of images with shape (batch, channel, height, width).
+
+ Returns:
+ None: This method does not return any value.
+
+ Notes:
+ This initializer configures normalization based on the input image dimensions, which is critical for
+ ensuring the model's compatibility and proper functionality on iOS devices. The normalization step
+ involves dividing by the image width if the image is square; otherwise, additional conditions might
+ apply (trimmed for brevity).
+ """
super().__init__()
b, c, h, w = im.shape # batch, channel, height, width
self.model = model
@@ -105,13 +120,48 @@ class iOSModel(torch.nn.Module):
# self.normalize = torch.tensor([1. / w, 1. / h, 1. / w, 1. / h]).expand(np, 4) # explicit (faster, larger)
def forward(self, x):
- """Runs forward pass on the input tensor, returning class confidences and normalized coordinates."""
+ """
+ Runs a forward pass on the input tensor, returning class confidences and normalized coordinates.
+
+ Args:
+ x (torch.Tensor): Input tensor containing the image data.
+
+ Returns:
+ torch.Tensor: Concatenated tensor with normalized coordinates (xywh), confidence scores (conf), and class probabilities (cls).
+
+ Examples:
+ ```python
+ model = iOSModel(pretrained_model, input_image)
+ output = model.forward(torch_input_tensor)
+ ```
+ """
xywh, conf, cls = self.model(x)[0].squeeze().split((4, 1, self.nc), 1)
return cls * conf, xywh * self.normalize # confidence (3780, 80), coordinates (3780, 4)
def export_formats():
- """Returns a DataFrame of supported YOLOv5 model export formats and their properties."""
+ """
+ Returns a DataFrame of supported YOLOv5 model export formats and their properties.
+
+ Returns:
+ pd.DataFrame: A pandas DataFrame containing supported export formats and their properties. The DataFrame includes
+ columns for format name, CLI argument suffix, file extension or directory name, and boolean flags indicating if the
+ export format supports training and detection.
+
+ Examples:
+ ```python
+ formats = export_formats()
+ print(formats)
+ ```
+
+ Notes:
+ The DataFrame contains the following columns:
+ - Format: The name of the model format (e.g., PyTorch, TorchScript, ONNX, etc.)
+ - Include Argument: The argument to use with the export script to include this format.
+ - File Suffix: File extension or directory name associated with the format.
+ - Supports Training: Whether the format supports training.
+ - Supports Detection: Whether the format supports detection.
+ """
x = [
["PyTorch", "-", ".pt", True, True],
["TorchScript", "torchscript", ".torchscript", True, True],
@@ -130,7 +180,28 @@ def export_formats():
def try_export(inner_func):
- """Decorator @try_export for YOLOv5 model export functions that logs success/failure, time taken, and file size."""
+ """
+ Logs success or failure, execution time, and file size for YOLOv5 model export functions wrapped with @try_export.
+
+ Args:
+ inner_func (Callable): The model export function to be wrapped by the decorator.
+
+ Returns:
+ Callable: The wrapped function that logs execution details. When executed, this wrapper function returns either:
+ - Tuple (str, torch.nn.Module): On success — the file path of the exported model and the model instance.
+ - Tuple (None, None): On failure — None values indicating export failed.
+
+ Examples:
+ @try_export
+ def export_onnx(model, filepath):
+ # implementation here
+ pass
+
+ exported_file, exported_model = export_onnx(yolo_model, 'path/to/save/model.onnx')
+
+ Notes:
+ For additional requirements and model export formats, refer to the [Ultralytics YOLOv5 GitHub repository](https://github.com/ultralytics/ultralytics).
+ """
inner_args = get_default_args(inner_func)
def outer_func(*args, **kwargs):
@@ -150,8 +221,46 @@ def try_export(inner_func):
@try_export
def export_torchscript(model, im, file, optimize, prefix=colorstr("TorchScript:")):
- """Exports YOLOv5 model to TorchScript format, optionally optimized for mobile, with image shape and stride
- metadata.
+ """
+ Exports a YOLOv5 model to the TorchScript format.
+
+ Args:
+ model (torch.nn.Module): The YOLOv5 model to be exported.
+ im (torch.Tensor): Example input tensor to be used for tracing the TorchScript model.
+ file (Path): File path where the exported TorchScript model will be saved.
+ optimize (bool): If True, applies optimizations for mobile deployment.
+ prefix (str): Optional prefix for log messages. Default is 'TorchScript:'.
+
+ Returns:
+ (str | None, torch.jit.ScriptModule | None): A tuple containing the file path of the exported model
+ (as a string) and the TorchScript model (as a torch.jit.ScriptModule). If the export fails, both elements
+ of the tuple will be None.
+
+ Notes:
+ - This function uses tracing to create the TorchScript model.
+ - Metadata, including the input image shape, model stride, and class names, is saved in an extra file (`config.txt`)
+ within the TorchScript model package.
+ - For mobile optimization, refer to the PyTorch tutorial: https://pytorch.org/tutorials/recipes/mobile_interpreter.html
+
+ Example:
+ ```python
+ from pathlib import Path
+ import torch
+ from models.experimental import attempt_load
+ from utils.torch_utils import select_device
+
+ # Load model
+ weights = 'yolov5s.pt'
+ device = select_device('')
+ model = attempt_load(weights, map_location=device)
+
+ # Example input tensor
+ im = torch.zeros(1, 3, 640, 640).to(device)
+
+ # Export model
+ file = Path('yolov5s.torchscript')
+ export_torchscript(model, im, file, optimize=False)
+ ```
"""
LOGGER.info(f"\n{prefix} starting export with torch {torch.__version__}...")
f = file.with_suffix(".torchscript")
@@ -168,7 +277,31 @@ def export_torchscript(model, im, file, optimize, prefix=colorstr("TorchScript:"
@try_export
def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr("ONNX:")):
- """Exports a YOLOv5 model to ONNX format with dynamic axes and optional simplification."""
+ """
+ Export a YOLOv5 model to ONNX format with dynamic axes support and optional model simplification.
+
+ Args:
+ model (torch.nn.Module): The YOLOv5 model to be exported.
+ im (torch.Tensor): A sample input tensor for model tracing, usually the shape is (1, 3, height, width).
+ file (pathlib.Path | str): The output file path where the ONNX model will be saved.
+ opset (int): The ONNX opset version to use for export.
+ dynamic (bool): If True, enables dynamic axes for batch, height, and width dimensions.
+ simplify (bool): If True, applies ONNX model simplification for optimization.
+ prefix (str): A prefix string for logging messages, defaults to 'ONNX:'.
+
+ Returns:
+ tuple[pathlib.Path | str, None]: The path to the saved ONNX model file and None (consistent with decorator).
+
+ Raises:
+ ImportError: If required libraries for export (e.g., 'onnx', 'onnx-simplifier') are not installed.
+ AssertionError: If the simplification check fails.
+
+ Notes:
+ The required packages for this function can be installed via:
+ ```
+ pip install onnx onnx-simplifier onnxruntime onnxruntime-gpu
+ ```
+ """
check_requirements("onnx>=1.12.0")
import onnx
@@ -225,7 +358,41 @@ def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr("ONNX
@try_export
def export_openvino(file, metadata, half, int8, data, prefix=colorstr("OpenVINO:")):
- """Exports a YOLOv5 model to OpenVINO format with optional FP16 and INT8 quantization; see https://pypi.org/project/openvino-dev/."""
+ """
+ Exports a YOLOv5 model to OpenVINO format with optional FP16 and INT8 quantization; see
+ https://pypi.org/project/openvino-dev/.
+
+ Args:
+ file (Path): The path to the output file where the OpenVINO model will be saved.
+ metadata (dict): Dictionary including model metadata such as names and strides.
+ half (bool): If True, export the model with FP16 precision.
+ int8 (bool): If True, export the model with INT8 quantization.
+ data (str): Path to the dataset YAML file required for INT8 quantization.
+ prefix (str): Prefix string for logging purposes (default is "OpenVINO:").
+
+ Returns:
+ (str, openvino.runtime.Model | None): Returns the OpenVINO model file path and openvino.runtime.Model object if
+ export is successful; otherwise, returns None.
+
+ Notes:
+ - Requires `openvino-dev` package version 2023.0 or higher. Install with:
+ `$ pip install openvino-dev>=2023.0`
+ - For INT8 quantization, also requires `nncf` library version 2.5.0 or higher. Install with:
+ `$ pip install nncf>=2.5.0`
+
+ Examples:
+ ```python
+ from pathlib import Path
+ from ultralytics import YOLOv5
+
+ model = YOLOv5('yolov5s.pt')
+ export_openvino(Path('yolov5s.onnx'), metadata={'names': model.names, 'stride': model.stride}, half=True,
+ int8=False, data='data.yaml')
+ ```
+
+ This will export the YOLOv5 model to OpenVINO with FP16 precision but without INT8 quantization, saving it to
+ the specified file path.
+ """
check_requirements("openvino-dev>=2023.0") # requires openvino-dev: https://pypi.org/project/openvino-dev/
import openvino.runtime as ov # noqa
from openvino.tools import mo # noqa
@@ -282,8 +449,39 @@ def export_openvino(file, metadata, half, int8, data, prefix=colorstr("OpenVINO:
@try_export
def export_paddle(model, im, file, metadata, prefix=colorstr("PaddlePaddle:")):
- """Exports a YOLOv5 model to PaddlePaddle format using X2Paddle, saving to `save_dir` and adding a metadata.yaml
- file.
+ """
+ Exports a YOLOv5 model to PaddlePaddle format using X2Paddle, saving the converted model and metadata.
+
+ Args:
+ model (torch.nn.Module): The YOLOv5 model to be exported.
+ im (torch.Tensor): Input tensor used for model tracing during export.
+ file (pathlib.Path): Path to the source file to be converted.
+ metadata (dict): Additional metadata to be saved alongside the model.
+ prefix (str): Prefix for logging information.
+
+ Returns:
+ tuple (str, None): A tuple where the first element is the path to the saved PaddlePaddle model, and the
+ second element is None.
+
+ Examples:
+ ```python
+ from pathlib import Path
+ import torch
+
+ # Assume 'model' is a pre-trained YOLOv5 model and 'im' is an example input tensor
+ model = ... # Load your model here
+ im = torch.randn((1, 3, 640, 640)) # Dummy input tensor for tracing
+ file = Path("yolov5s.pt")
+ metadata = {"stride": 32, "names": ["person", "bicycle", "car", "motorbike"]}
+
+ export_paddle(model=model, im=im, file=file, metadata=metadata)
+ ```
+ Notes:
+ Ensure that `paddlepaddle` and `x2paddle` are installed, as these are required for the export function. You can
+ install them via pip:
+ ```
+ $ pip install paddlepaddle x2paddle
+ ```
"""
check_requirements(("paddlepaddle", "x2paddle"))
import x2paddle
@@ -299,7 +497,36 @@ def export_paddle(model, im, file, metadata, prefix=colorstr("PaddlePaddle:")):
@try_export
def export_coreml(model, im, file, int8, half, nms, prefix=colorstr("CoreML:")):
- """Exports YOLOv5 model to CoreML format with optional NMS, INT8, and FP16 support; requires coremltools."""
+ """
+ Export a YOLOv5 model to CoreML format with optional NMS, INT8, and FP16 support.
+
+ Args:
+ model (torch.nn.Module): The YOLOv5 model to be exported.
+ im (torch.Tensor): Example input tensor to trace the model.
+ file (pathlib.Path): Path object where the CoreML model will be saved.
+ int8 (bool): Flag indicating whether to use INT8 quantization (default is False).
+ half (bool): Flag indicating whether to use FP16 quantization (default is False).
+ nms (bool): Flag indicating whether to include Non-Maximum Suppression (default is False).
+ prefix (str): Prefix string for logging purposes (default is 'CoreML:').
+
+ Returns:
+ tuple[pathlib.Path | None, None]: The path to the saved CoreML model file, or (None, None) if there is an error.
+
+ Notes:
+ The exported CoreML model will be saved with a .mlmodel extension.
+ Quantization is supported only on macOS.
+
+ Example:
+ ```python
+ from pathlib import Path
+ import torch
+ from models.yolo import Model
+ model = Model(cfg, ch=3, nc=80)
+ im = torch.randn(1, 3, 640, 640)
+ file = Path("yolov5s_coreml")
+ export_coreml(model, im, file, int8=False, half=False, nms=True)
+ ```
+ """
check_requirements("coremltools")
import coremltools as ct
@@ -327,7 +554,36 @@ def export_engine(model, im, file, half, dynamic, simplify, workspace=4, verbose
"""
Exports a YOLOv5 model to TensorRT engine format, requiring GPU and TensorRT>=7.0.0.
- https://developer.nvidia.com/tensorrt
+ Args:
+ model (torch.nn.Module): YOLOv5 model to be exported.
+ im (torch.Tensor): Input tensor of shape (B,C,H,W).
+ file (Path): Path to save the exported model.
+ half (bool): Set to True to export with FP16 precision.
+ dynamic (bool): Set to True to enable dynamic input shapes.
+ simplify (bool): Set to True to simplify the model during export.
+ workspace (int): Workspace size in GB (default is 4).
+ verbose (bool): Set to True for verbose logging output.
+ prefix (str): Log message prefix.
+
+ Returns:
+ (Path, None): Tuple containing the path to the exported model and None.
+
+ Raises:
+ AssertionError: If executed on CPU instead of GPU.
+ RuntimeError: If there is a failure in parsing the ONNX file.
+
+ Example:
+ ```python
+ from ultralytics import YOLOv5
+ import torch
+ from pathlib import Path
+
+ model = YOLOv5('yolov5s.pt') # Load a pre-trained YOLOv5 model
+ input_tensor = torch.randn(1, 3, 640, 640).cuda() # example input tensor on GPU
+ export_path = Path('yolov5s.engine') # export destination
+
+ export_engine(model.model, input_tensor, export_path, half=True, dynamic=True, simplify=True, workspace=8, verbose=True)
+ ```
"""
assert im.device.type != "cpu", "export running on CPU but must be on GPU, i.e. `python export.py --device 0`"
try:
@@ -407,8 +663,40 @@ def export_saved_model(
keras=False,
prefix=colorstr("TensorFlow SavedModel:"),
):
- """Exports a YOLOv5 model to TensorFlow SavedModel format, supporting dynamic axes and non-maximum suppression
- (NMS).
+ """
+ Exports a YOLOv5 model to TensorFlow SavedModel format, supporting dynamic axes and non-maximum suppression (NMS).
+
+ Args:
+ model (torch.nn.Module): The PyTorch model to convert.
+ im (torch.Tensor): Sample input tensor with shape (B, C, H, W) for tracing.
+ file (pathlib.Path): File path to save the exported model.
+ dynamic (bool): Flag to indicate whether dynamic axes should be used.
+ tf_nms (bool, optional): Enable TensorFlow non-maximum suppression (NMS). Default is False.
+ agnostic_nms (bool, optional): Enable class-agnostic NMS. Default is False.
+ topk_per_class (int, optional): Top K detections per class to keep before applying NMS. Default is 100.
+ topk_all (int, optional): Top K detections across all classes to keep before applying NMS. Default is 100.
+ iou_thres (float, optional): IoU threshold for NMS. Default is 0.45.
+ conf_thres (float, optional): Confidence threshold for detections. Default is 0.25.
+ keras (bool, optional): Save the model in Keras format if True. Default is False.
+ prefix (str, optional): Prefix for logging messages. Default is "TensorFlow SavedModel:".
+
+ Returns:
+ tuple: A tuple containing the path to the saved model folder (str) and the Keras model instance (tf.keras.Model | None).
+
+ Notes:
+ - The method supports TensorFlow versions up to 2.15.1.
+ - TensorFlow NMS may not be supported in older TensorFlow versions.
+ - If the TensorFlow version exceeds 2.13.1, it might cause issues when exporting to TFLite.
+ Refer to: https://github.com/ultralytics/yolov5/issues/12489
+
+ Raises:
+ Exception: If TensorFlow is not installed.
+
+ Example:
+ ```python
+ model, im = ... # Initialize your PyTorch model and input tensor
+ export_saved_model(model, im, Path("yolov5_saved_model"), dynamic=True)
+ ```
"""
# YOLOv5 TensorFlow SavedModel export
try:
@@ -460,7 +748,28 @@ def export_saved_model(
@try_export
def export_pb(keras_model, file, prefix=colorstr("TensorFlow GraphDef:")):
- """Exports YOLOv5 model to TensorFlow GraphDef *.pb format; see https://github.com/leimao/Frozen_Graph_TensorFlow for details."""
+ """
+ Exports YOLOv5 model to TensorFlow GraphDef (*.pb) format.
+
+ Args:
+ keras_model (tf.keras.Model): The Keras model to be converted.
+ file (Path): The output file path where the GraphDef will be saved.
+ prefix (str): Optional prefix string; defaults to a colored string indicating TensorFlow GraphDef export status.
+
+ Returns:
+ Tuple[Path, None]: The file path where the GraphDef model was saved and a None placeholder.
+
+ Notes:
+ For more details, refer to the guide on frozen graphs: https://github.com/leimao/Frozen_Graph_TensorFlow
+
+ Example:
+ ```python
+ from pathlib import Path
+ keras_model = ... # assume an existing Keras model
+ file = Path("model.pb")
+ export_pb(keras_model, file)
+ ```
+ """
import tensorflow as tf
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
@@ -480,7 +789,46 @@ def export_tflite(
keras_model, im, file, int8, per_tensor, data, nms, agnostic_nms, prefix=colorstr("TensorFlow Lite:")
):
# YOLOv5 TensorFlow Lite export
- """Exports YOLOv5 model to TensorFlow Lite format with optional FP16, INT8, and NMS support."""
+ """
+ Exports YOLOv5 model to TensorFlow Lite format with optional FP16, INT8, and NMS support.
+
+ Args:
+ keras_model (tf.keras.Model): The Keras model to be exported.
+ im (torch.Tensor): Image tensor for normalization and model tracing.
+ file (Path): The file path to save the exported TensorFlow Lite model.
+ int8 (bool): Enables INT8 quantization if True.
+ per_tensor (bool): If True, disable per-channel quantization (applicable when int8 is True).
+ data (str): Path to dataset for representative dataset generation in INT8 quantization.
+ nms (bool): Enables Non-Maximum Suppression (NMS) support if True.
+ agnostic_nms (bool): Enables class-agnostic NMS support if True.
+ prefix (str): Prefix for logging messages.
+
+ Returns:
+ (str | None, tf.lite.Model | None): The file path of the saved TFLite model, and the TFLite model instance if successful.
+
+ Example:
+ ```python
+ from pathlib import Path
+ import torch
+ import tensorflow as tf
+ from torchvision import models
+
+ # Load a pre-trained model from torchvision
+ model = models.yolov5() # Placeholder for actual YOLOv5 model loading
+ im = torch.zeros(1, 3, 640, 640) # Example image tensor
+
+ # Provide the Keras model wrapping the PyTorch YOLOv5 model
+ keras_model = tf.keras.models.load_model('path/to/keras_model.h5')
+
+ # Export the model to TensorFlow Lite format
+ file_path = export_tflite(keras_model, im, Path('model.tflite'), int8=False, per_tensor=False,
+ data='path/to/dataset.yaml', nms=False, agnostic_nms=False)
+ ```
+
+ Notes:
+ Ensure the TensorFlow and TensorFlow Lite dependencies are installed. The exported TFLite model can be used for
+ efficient inference on mobile and edge devices.
+ """
import tensorflow as tf
LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...")
@@ -515,9 +863,31 @@ def export_tflite(
@try_export
def export_edgetpu(file, prefix=colorstr("Edge TPU:")):
"""
+ Clear and concise summary line describing the function's purpose:
+
Exports a YOLOv5 model to Edge TPU compatible TFLite format; requires Linux and Edge TPU compiler.
- https://coral.ai/docs/edgetpu/models-intro/
+ Args:
+ file (Path): Path to the YOLOv5 model file to be exported (.pt format).
+ prefix (str, optional): Prefix for logging messages. Defaults to colorstr("Edge TPU:").
+
+ Returns:
+ tuple[Path, None]: Path to the exported Edge TPU compatible TFLite model, None.
+
+ Raises:
+ AssertionError: If the system is not Linux.
+ subprocess.CalledProcessError: If any subprocess call to install or run the Edge TPU compiler fails.
+
+ Notes:
+ To use this function, ensure you have the Edge TPU compiler installed on your Linux system. You can find
+ installation instructions here: https://coral.ai/docs/edgetpu/compiler/.
+
+ Example:
+ ```python
+ from pathlib import Path
+ file = Path('yolov5s.pt')
+ export_edgetpu(file)
+ ```
"""
cmd = "edgetpu_compiler --version"
help_url = "https://coral.ai/docs/edgetpu/compiler/"
@@ -556,7 +926,30 @@ def export_edgetpu(file, prefix=colorstr("Edge TPU:")):
@try_export
def export_tfjs(file, int8, prefix=colorstr("TensorFlow.js:")):
- """Exports a YOLOv5 model to TensorFlow.js format, optionally with uint8 quantization."""
+ """
+ Exports a YOLOv5 model to TensorFlow.js format, optionally with uint8 quantization.
+
+ Args:
+ file (Path): Path to the input model file.
+ int8 (bool): If True, applies uint8 quantization.
+ prefix (str): Prefix for logging information (default: colorstr("TensorFlow.js:")).
+
+ Returns:
+ tuple: Output directory path (str), None
+
+ Notes:
+ This function requires `tensorflowjs` to be installed. You can install it using:
+ ```shell
+ pip install tensorflowjs
+ ```
+
+ Example usage:
+ ```python
+ export_tfjs(Path('yolov5s.onnx'), int8=False)
+ ```
+
+ The TensorFlow.js converted model is saved in the directory specified by `file` with "_web_model" suffix.
+ """
check_requirements("tensorflowjs")
import tensorflowjs as tfjs
@@ -596,7 +989,30 @@ def add_tflite_metadata(file, metadata, num_outputs):
"""
Adds TFLite metadata to a model file, supporting multiple outputs, as specified by TensorFlow guidelines.
- https://www.tensorflow.org/lite/models/convert/metadata
+ Args:
+ file (str): The path to the TensorFlow Lite model file to which metadata will be added.
+ metadata (dict): Metadata information to be added to the model, structured as required by TFLite metadata schema.
+ num_outputs (int): Number of output tensors the model has, to properly configure the metadata.
+
+ Returns:
+ None
+
+ Example:
+ ```python
+ metadata = {
+ "name": "yolov5",
+ "description": "YOLOv5 object detection model",
+ "version": "1.0",
+ "author": "Ultralytics",
+ "license": "Apache License 2.0"
+ }
+ add_tflite_metadata("model.tflite", metadata, num_outputs=4)
+ ```
+
+ Note:
+ TFLite metadata can include information such as model name, version, author, and other relevant details.
+ For more details and structure of the metadata, refer to the TensorFlow Lite
+ [metadata guidelines](https://www.tensorflow.org/lite/models/convert/metadata).
"""
with contextlib.suppress(ImportError):
# check_requirements('tflite_support')
@@ -630,8 +1046,49 @@ def add_tflite_metadata(file, metadata, num_outputs):
def pipeline_coreml(model, im, file, names, y, prefix=colorstr("CoreML Pipeline:")):
- """Converts a PyTorch YOLOv5 model to CoreML format with NMS, handling different input/output shapes and saving the
- model.
+ """
+ Converts a PyTorch YOLOv5 model to CoreML format with Non-Maximum Suppression (NMS), handling different input/output
+ shapes and saving the model.
+
+ Args:
+ model (torch.nn.Module): The YOLOv5 PyTorch model.
+ im (torch.Tensor): Input tensor example with shape [N, C, H, W], where N is the batch size, C is the number of
+ channels, H is the height, and W is the width.
+ file (Path): Path to save the converted CoreML model.
+ names (dict[int, str]): Dictionary mapping class indices to class names.
+ y (torch.Tensor): Output tensor from the PyTorch model's forward pass.
+ prefix (str): Custom prefix for logging messages.
+
+ Returns:
+ Path: Path to the saved CoreML model (.mlmodel).
+
+ Raises:
+ AssertionError: If the number of class names does not match the number of classes in the model.
+
+ Notes:
+ - This function requires `coremltools` to be installed.
+ - Running this function on a non-macOS environment might not support some features.
+ - Flexible input shapes and additional NMS options can be customized within the function.
+
+ Examples:
+ ```python
+ from pathlib import Path
+ import torch
+
+ # Load YOLOv5 model and an example input tensor
+ model = torch.load("yolov5s.pt")
+ im = torch.zeros(1, 3, 640, 640) # Example input tensor
+
+ # Define class names
+ names = {0: "person", 1: "bicycle", 2: "car", ...}
+
+ # Perform forward pass to get model output
+ y = model(im)
+
+ # Convert to CoreML
+ output_file = Path("yolov5s.mlmodel")
+ pipeline_coreml(model, im, output_file, names, y)
+ ```
"""
import coremltools as ct
from PIL import Image
@@ -788,7 +1245,70 @@ def run(
iou_thres=0.45, # TF.js NMS: IoU threshold
conf_thres=0.25, # TF.js NMS: confidence threshold
):
- """Exports YOLOv5 model to specified formats including ONNX, TensorRT, CoreML, and TensorFlow; see https://github.com/ultralytics/yolov5."""
+ """
+ Exports a YOLOv5 model to specified formats including ONNX, TensorRT, CoreML, and TensorFlow.
+
+ Args:
+ data (str | Path): Path to the dataset YAML configuration file. Default is 'data/coco128.yaml'.
+ weights (str | Path): Path to the pretrained model weights file. Default is 'yolov5s.pt'.
+ imgsz (tuple): Image size as (height, width). Default is (640, 640).
+ batch_size (int): Batch size for exporting the model. Default is 1.
+ device (str): Device to run the export on, e.g., '0' for GPU, 'cpu' for CPU. Default is 'cpu'.
+ include (tuple): Formats to include in the export. Default is ('torchscript', 'onnx').
+ half (bool): Flag to export model with FP16 half-precision. Default is False.
+ inplace (bool): Set the YOLOv5 Detect() module inplace=True. Default is False.
+ keras (bool): Flag to use Keras for TensorFlow SavedModel export. Default is False.
+ optimize (bool): Optimize TorchScript model for mobile deployment. Default is False.
+ int8 (bool): Apply INT8 quantization for CoreML or TensorFlow models. Default is False.
+ per_tensor (bool): Apply per tensor quantization for TensorFlow models. Default is False.
+ dynamic (bool): Enable dynamic axes for ONNX, TensorFlow, or TensorRT exports. Default is False.
+ simplify (bool): Simplify the ONNX model during export. Default is False.
+ opset (int): ONNX opset version. Default is 12.
+ verbose (bool): Enable verbose logging for TensorRT export. Default is False.
+ workspace (int): TensorRT workspace size in GB. Default is 4.
+ nms (bool): Add non-maximum suppression (NMS) to the TensorFlow model. Default is False.
+ agnostic_nms (bool): Add class-agnostic NMS to the TensorFlow model. Default is False.
+ topk_per_class (int): Top-K boxes per class to keep for TensorFlow.js NMS. Default is 100.
+ topk_all (int): Top-K boxes for all classes to keep for TensorFlow.js NMS. Default is 100.
+ iou_thres (float): IoU threshold for NMS. Default is 0.45.
+ conf_thres (float): Confidence threshold for NMS. Default is 0.25.
+
+ Returns:
+ None
+
+ Notes:
+ - Model export is based on the specified formats in the 'include' argument.
+ - Be cautious of combinations where certain flags are mutually exclusive, such as `--half` and `--dynamic`.
+
+ Example:
+ ```python
+ run(
+ data="data/coco128.yaml",
+ weights="yolov5s.pt",
+ imgsz=(640, 640),
+ batch_size=1,
+ device="cpu",
+ include=("torchscript", "onnx"),
+ half=False,
+ inplace=False,
+ keras=False,
+ optimize=False,
+ int8=False,
+ per_tensor=False,
+ dynamic=False,
+ simplify=False,
+ opset=12,
+ verbose=False,
+ workspace=4,
+ nms=False,
+ agnostic_nms=False,
+ topk_per_class=100,
+ topk_all=100,
+ iou_thres=0.45,
+ conf_thres=0.25,
+ )
+ ```
+ """
t = time.time()
include = [x.lower() for x in include] # to lowercase
fmts = tuple(export_formats()["Argument"][1:]) # --include arguments
@@ -901,7 +1421,23 @@ def run(
def parse_opt(known=False):
- """Parses command-line arguments for YOLOv5 model export configurations, returning the parsed options."""
+ """
+ Parses command-line arguments for YOLOv5 model export configurations.
+
+ Args:
+ known (bool): If True, `argparse.ArgumentParser.parse_known_args` is used to parse command-line arguments; otherwise,
+ `argparse.ArgumentParser.parse_args` is used. Defaults to False.
+
+ Returns:
+ argparse.Namespace: An object containing parsed command-line arguments.
+
+ Example:
+ ```python
+ opts = parse_opt()
+ print(opts.data)
+ print(opts.weights)
+ ```
+ """
parser = argparse.ArgumentParser()
parser.add_argument("--data", type=str, default=ROOT / "data/coco128.yaml", help="dataset.yaml path")
parser.add_argument("--weights", nargs="+", type=str, default=ROOT / "yolov5s.pt", help="model.pt path(s)")
@@ -937,7 +1473,44 @@ def parse_opt(known=False):
def main(opt):
- """Executes the YOLOv5 model inference or export with specified weights and options."""
+ """
+ ```python Exports the YOLOv5 model to specified formats, including ONNX, TensorRT, CoreML, and TensorFlow.
+
+ Args:
+ opt (argparse.Namespace): Parsed command-line arguments containing the export configurations.
+ - data (str): Path to the dataset.yaml.
+ - weights (list[str]): Paths to model (.pt) file(s).
+ - imgsz (list[int]): Image size (height, width).
+ - batch_size (int): Batch size.
+ - device (str): CUDA device, e.g., '0' or '0,1,2,3' or 'cpu'.
+ - half (bool): FP16 half-precision export flag.
+ - inplace (bool): Set YOLOv5 Detect() inplace to True.
+ - keras (bool): Use Keras for TensorFlow models.
+ - optimize (bool): Optimize TorchScript model for mobile.
+ - int8 (bool): INT8 quantization flag.
+ - per_tensor (bool): Per tensor quantization for TensorFlow.
+ - dynamic (bool): Dynamic axes for ONNX/TF/TensorRT.
+ - simplify (bool): Simplify ONNX model.
+ - opset (int): ONNX opset version.
+ - verbose (bool): Verbose logging for TensorRT.
+ - workspace (int): Workspace size for TensorRT (in GB).
+ - nms (bool): Add NMS to TensorFlow model.
+ - agnostic_nms (bool): Add agnostic NMS to TensorFlow model.
+ - topk_per_class (int): Top-k per class for TensorFlow.js NMS.
+ - topk_all (int): Top-k for all classes for TensorFlow.js NMS.
+ - iou_thres (float): IoU threshold for TensorFlow.js NMS.
+ - conf_thres (float): Confidence threshold for TensorFlow.js NMS.
+ - include (list[str]): List of formats to include in export, e.g., ['torchscript', 'onnx'].
+
+ Returns:
+ list[str]: List of exported file paths.
+ ```python
+
+ # Example usage:
+ # opt = parse_opt()
+ # main(opt)
+ ```
+ """
for opt.weights in opt.weights if isinstance(opt.weights, list) else [opt.weights]:
run(**vars(opt))
diff --git a/hubconf.py b/hubconf.py
index 4b0c36b8d..de93c79bf 100644
--- a/hubconf.py
+++ b/hubconf.py
@@ -15,19 +15,38 @@ import torch
def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None):
"""
- Creates or loads a YOLOv5 model.
+ Creates or loads a YOLOv5 model, with options for pretrained weights and model customization.
- Arguments:
- name (str): model name 'yolov5s' or path 'path/to/best.pt'
- pretrained (bool): load pretrained weights into the model
- channels (int): number of input channels
- classes (int): number of model classes
- autoshape (bool): apply YOLOv5 .autoshape() wrapper to model
- verbose (bool): print all information to screen
- device (str, torch.device, None): device to use for model parameters
+ Args:
+ name (str): Model name (e.g., 'yolov5s') or path to the model checkpoint (e.g., 'path/to/best.pt').
+ pretrained (bool, optional): If True, loads pretrained weights into the model. Defaults to True.
+ channels (int, optional): Number of input channels the model expects. Defaults to 3.
+ classes (int, optional): Number of classes the model is expected to detect. Defaults to 80.
+ autoshape (bool, optional): If True, applies the YOLOv5 .autoshape() wrapper for various input formats. Defaults to True.
+ verbose (bool, optional): If True, prints detailed information during the model creation/loading process. Defaults to True.
+ device (str | torch.device | None, optional): Device to use for model parameters (e.g., 'cpu', 'cuda'). If None, selects the best available device. Defaults to None.
Returns:
- YOLOv5 model
+ DetectMultiBackend | AutoShape: The loaded YOLOv5 model, potentially wrapped with AutoShape if specified.
+
+ Examples:
+ ```python
+ import torch
+ from ultralytics import _create
+
+ # Load an official YOLOv5s model with pretrained weights
+ model = _create('yolov5s')
+
+ # Load a custom model from a local checkpoint
+ model = _create('path/to/custom_model.pt', pretrained=False)
+
+ # Load a model with specific input channels and classes
+ model = _create('yolov5s', channels=1, classes=10)
+ ```
+
+ Notes:
+ For more information on model loading and customization, visit the
+ [YOLOv5 PyTorch Hub Documentation](https://docs.ultralytics.com/yolov5/tutorials/pytorch_hub_model_loading).
"""
from pathlib import Path
@@ -84,76 +103,355 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo
def custom(path="path/to/model.pt", autoshape=True, _verbose=True, device=None):
- """Loads a custom or local YOLOv5 model from a given path with optional autoshaping and device specification."""
+ """
+ Loads a custom or local YOLOv5 model from a given path with optional autoshaping and device specification.
+
+ Args:
+ path (str): Path to the custom model file (e.g., 'path/to/model.pt').
+ autoshape (bool): Apply YOLOv5 .autoshape() wrapper to model if True, enabling compatibility with various input types
+ (default is True).
+ _verbose (bool): If True, prints all informational messages to the screen; otherwise, operates silently
+ (default is True).
+ device (str | torch.device | None): Device to load the model on, e.g., 'cpu', 'cuda', torch.device('cuda:0'), etc.
+ (default is None, which automatically selects the best available device).
+
+ Returns:
+ torch.nn.Module: A YOLOv5 model loaded with the specified parameters.
+
+ Notes:
+ For more details on loading models from PyTorch Hub:
+ https://docs.ultralytics.com/yolov5/tutorials/pytorch_hub_model_loading
+
+ Examples:
+ ```python
+ # Load model from a given path with autoshape enabled on the best available device
+ model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5s.pt')
+
+ # Load model from a local path without autoshape on the CPU device
+ model = torch.hub.load('.', 'custom', 'yolov5s.pt', source='local', autoshape=False, device='cpu')
+ ```
+ """
return _create(path, autoshape=autoshape, verbose=_verbose, device=device)
def yolov5n(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
- """Instantiates the YOLOv5-nano model with options for pretraining, input channels, class count, autoshaping,
+ """
+ Instantiates the YOLOv5-nano model with options for pretraining, input channels, class count, autoshaping,
verbosity, and device.
+
+ Args:
+ pretrained (bool): If True, loads pretrained weights into the model. Defaults to True.
+ channels (int): Number of input channels for the model. Defaults to 3.
+ classes (int): Number of classes for object detection. Defaults to 80.
+ autoshape (bool): If True, applies the YOLOv5 .autoshape() wrapper to the model for various formats (file/URI/PIL/
+ cv2/np) and non-maximum suppression (NMS) during inference. Defaults to True.
+ _verbose (bool): If True, prints detailed information to the screen. Defaults to True.
+ device (str | torch.device | None): Specifies the device to use for model computation. If None, uses the best device
+ available (i.e., GPU if available, otherwise CPU). Defaults to None.
+
+ Returns:
+ DetectionModel | ClassificationModel | SegmentationModel: The instantiated YOLOv5-nano model, potentially with
+ pretrained weights and autoshaping applied.
+
+ Notes:
+ For further details on loading models from PyTorch Hub, refer to [PyTorch Hub models](https://pytorch.org/hub/
+ ultralytics_yolov5).
+
+ Examples:
+ ```python
+ import torch
+ from ultralytics import yolov5n
+
+ # Load the YOLOv5-nano model with defaults
+ model = yolov5n()
+
+ # Load the YOLOv5-nano model with a specific device
+ model = yolov5n(device='cuda')
+ ```
"""
return _create("yolov5n", pretrained, channels, classes, autoshape, _verbose, device)
def yolov5s(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
- """Creates YOLOv5-small model with options for pretraining, input channels, class count, autoshaping, verbosity, and
+ """
+ Creates YOLOv5-small model with options for pretraining, input channels, class count, autoshaping, verbosity, and
device.
+
+ Args:
+ pretrained (bool, optional): Flag to load pretrained weights into the model. Defaults to True.
+ channels (int, optional): Number of input channels. Defaults to 3.
+ classes (int, optional): Number of model classes. Defaults to 80.
+ autoshape (bool, optional): Whether to apply YOLOv5 .autoshape() wrapper to the model for preprocessed inputs.
+ Defaults to True.
+ _verbose (bool, optional): Flag to print detailed information on model loading. Defaults to True.
+ device (str | torch.device | None, optional): Device to use for model parameters, e.g., 'cpu', 'cuda'. If None,
+ auto-select the best available device. Defaults to None.
+
+ Returns:
+ YOLOv5 model (torch.nn.Module): The YOLOv5-small model loaded with specified configurations and optionally
+ pretrained weights.
+
+ Usage:
+ ```python
+ import torch
+ model = torch.hub.load('ultralytics/yolov5', 'yolov5s') # official model
+ model = torch.hub.load('ultralytics/yolov5:master', 'yolov5s') # from branch
+ model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5s.pt') # custom/local model
+ model = torch.hub.load('.', 'custom', 'yolov5s.pt', source='local') # local repo
+ ```
+
+ For more information, visit https://pytorch.org/hub/ultralytics_yolov5.
"""
return _create("yolov5s", pretrained, channels, classes, autoshape, _verbose, device)
def yolov5m(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
- """Instantiates the YOLOv5-medium model with customizable pretraining, channel count, class count, autoshaping,
+ """
+ Instantiates the YOLOv5-medium model with customizable pretraining, channel count, class count, autoshaping,
verbosity, and device.
+
+ Args:
+ pretrained (bool, optional): Whether to load pretrained weights into the model. Default is True.
+ channels (int, optional): Number of input channels. Default is 3.
+ classes (int, optional): Number of model classes. Default is 80.
+ autoshape (bool, optional): Apply YOLOv5 .autoshape() wrapper to the model for handling various input formats. Default is True.
+ _verbose (bool, optional): Whether to print detailed information to the screen. Default is True.
+ device (str | torch.device | None, optional): Device specification to use for model parameters (e.g., 'cpu', 'cuda'). Default is None.
+
+ Returns:
+ torch.nn.Module: The instantiated YOLOv5-medium model.
+
+ Usage Example:
+ ```python
+ import torch
+
+ model = torch.hub.load('ultralytics/yolov5', 'yolov5m') # Load YOLOv5-medium from Ultralytics repository
+ model = torch.hub.load('ultralytics/yolov5:master', 'yolov5m') # Load from the master branch
+ model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5m.pt') # Load a custom/local YOLOv5-medium model
+ model = torch.hub.load('.', 'custom', 'yolov5m.pt', source='local') # Load from a local repository
+ ```
"""
return _create("yolov5m", pretrained, channels, classes, autoshape, _verbose, device)
def yolov5l(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
- """Creates YOLOv5-large model with options for pretraining, channels, classes, autoshaping, verbosity, and device
+ """
+ Creates YOLOv5-large model with options for pretraining, channels, classes, autoshaping, verbosity, and device
selection.
+
+ Args:
+ pretrained (bool): Load pretrained weights into the model. Default is True.
+ channels (int): Number of input channels. Default is 3.
+ classes (int): Number of model classes. Default is 80.
+ autoshape (bool): Apply YOLOv5 .autoshape() wrapper to model. Default is True.
+ _verbose (bool): Print all information to screen. Default is True.
+ device (str | torch.device | None): Device to use for model parameters, e.g., 'cpu', 'cuda', or a torch.device instance. Default is None.
+
+ Returns:
+ YOLOv5 model (torch.nn.Module).
+
+ Example:
+ ```python
+ import torch
+ model = torch.hub.load('ultralytics/yolov5', 'yolov5l')
+ ```
+
+ Notes:
+ For additional details, refer to the PyTorch Hub models documentation:
+ https://pytorch.org/hub/ultralytics_yolov5
"""
return _create("yolov5l", pretrained, channels, classes, autoshape, _verbose, device)
def yolov5x(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
- """Instantiates the YOLOv5-xlarge model with customizable pretraining, channel count, class count, autoshaping,
+ """
+ Instantiates the YOLOv5-xlarge model with customizable pretraining, channel count, class count, autoshaping,
verbosity, and device.
+
+ Args:
+ pretrained (bool): If True, loads pretrained weights into the model. Defaults to True.
+ channels (int): Number of input channels. Defaults to 3.
+ classes (int): Number of model classes. Defaults to 80.
+ autoshape (bool): If True, applies YOLOv5 .autoshape() wrapper to the model for easier image handling. Defaults to
+ True.
+ _verbose (bool): If True, prints detailed information to the screen. Defaults to True.
+ device (str | torch.device | None): Device for model parameters, e.g., 'cpu', 'cuda:0', or a torch.device object.
+ Defaults to None.
+
+ Returns:
+ torch.nn.Module: The instantiated YOLOv5-xlarge model.
+
+ Example:
+ ```python
+ import torch
+ model = torch.hub.load('ultralytics/yolov5', 'yolov5x')
+ ```
+
+ For more details and usage, refer to the official YOLOv5 PyTorch Hub models documentation:
+ https://pytorch.org/hub/ultralytics_yolov5
"""
return _create("yolov5x", pretrained, channels, classes, autoshape, _verbose, device)
def yolov5n6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
- """Creates YOLOv5-nano-P6 model with options for pretraining, channels, classes, autoshaping, verbosity, and
- device.
+ """
+ Creates YOLOv5-nano-P6 model with options for pretraining, channels, classes, autoshaping, verbosity, and device.
+
+ Args:
+ pretrained (bool, optional): If True, loads pretrained weights into the model. Default is True.
+ channels (int, optional): Number of input channels. Default is 3.
+ classes (int, optional): Number of model classes. Default is 80.
+ autoshape (bool, optional): If True, applies the YOLOv5 .autoshape() wrapper to the model. Default is True.
+ _verbose (bool, optional): If True, prints all information to screen. Default is True.
+ device (str | torch.device | None, optional): Device to use for model parameters. Can be 'cpu', 'cuda', or None.
+ Default is None.
+
+ Returns:
+ torch.nn.Module: YOLOv5 model loaded with the specified configurations.
+
+ Example:
+ ```python
+ import torch
+ model = yolov5n6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device='cuda')
+ ```
+
+ Notes:
+ For more information on PyTorch Hub models, visit: https://pytorch.org/hub/ultralytics_yolov5
"""
return _create("yolov5n6", pretrained, channels, classes, autoshape, _verbose, device)
def yolov5s6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
- """Instantiate YOLOv5-small-P6 model with options for pretraining, input channels, number of classes, autoshaping,
+ """
+ Instantiate the YOLOv5-small-P6 model with options for pretraining, input channels, number of classes, autoshaping,
verbosity, and device selection.
+
+ Args:
+ pretrained (bool): If True, loads pretrained weights. Default is True.
+ channels (int): Number of input channels. Default is 3.
+ classes (int): Number of object detection classes. Default is 80.
+ autoshape (bool): If True, applies YOLOv5 .autoshape() wrapper to the model, allowing for varied input formats.
+ Default is True.
+ _verbose (bool): If True, prints detailed information during model loading. Default is True.
+ device (str | torch.device | None): Device specification for model parameters (e.g., 'cpu', 'cuda', or torch.device).
+ Default is None, which selects an available device automatically.
+
+ Returns:
+ torch.nn.Module: The YOLOv5-small-P6 model instance.
+
+ Usage:
+ ```python
+ import torch
+
+ model = torch.hub.load('ultralytics/yolov5', 'yolov5s6')
+ model = torch.hub.load('ultralytics/yolov5:master', 'yolov5s6') # load from a specific branch
+ model = torch.hub.load('ultralytics/yolov5', 'custom', 'path/to/yolov5s6.pt') # custom/local model
+ model = torch.hub.load('.', 'custom', 'path/to/yolov5s6.pt', source='local') # local repo model
+ ```
+
+ Notes:
+ - For more information, refer to the PyTorch Hub models documentation at https://pytorch.org/hub/ultralytics_yolov5
+
+ Raises:
+ Exception: If there is an error during model creation or loading, with a suggestion to visit the YOLOv5
+ tutorials for help.
"""
return _create("yolov5s6", pretrained, channels, classes, autoshape, _verbose, device)
def yolov5m6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
- """Creates YOLOv5-medium-P6 model with options for pretraining, channel count, class count, autoshaping, verbosity,
- and device.
+ """
+ Creates YOLOv5-medium-P6 model with options for pretraining, channel count, class count, autoshaping, verbosity, and
+ device.
+
+ Args:
+ pretrained (bool): If True, loads pretrained weights. Default is True.
+ channels (int): Number of input channels. Default is 3.
+ classes (int): Number of model classes. Default is 80.
+ autoshape (bool): Apply YOLOv5 .autoshape() wrapper to the model for file/URI/PIL/cv2/np inputs and NMS. Default is True.
+ _verbose (bool): If True, prints detailed information to the screen. Default is True.
+ device (str | torch.device | None): Device to use for model parameters. Default is None, which uses the best available device.
+
+ Returns:
+ torch.nn.Module: The YOLOv5-medium-P6 model.
+
+ Refer to the PyTorch Hub models documentation: https://pytorch.org/hub/ultralytics_yolov5 for additional details.
+
+ Example:
+ ```python
+ import torch
+
+ # Load YOLOv5-medium-P6 model
+ model = torch.hub.load('ultralytics/yolov5', 'yolov5m6')
+ ```
+
+ Notes:
+ - The model can be loaded with pre-trained weights for better performance on specific tasks.
+ - The autoshape feature simplifies input handling by allowing various popular data formats.
"""
return _create("yolov5m6", pretrained, channels, classes, autoshape, _verbose, device)
def yolov5l6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
- """Instantiates the YOLOv5-large-P6 model with customizable pretraining, channel and class counts, autoshaping,
+ """
+ Instantiates the YOLOv5-large-P6 model with customizable pretraining, channel and class counts, autoshaping,
verbosity, and device selection.
+
+ Args:
+ pretrained (bool, optional): If True, load pretrained weights into the model. Default is True.
+ channels (int, optional): Number of input channels. Default is 3.
+ classes (int, optional): Number of model classes. Default is 80.
+ autoshape (bool, optional): If True, apply YOLOv5 .autoshape() wrapper to the model for input flexibility.
+ Default is True.
+ _verbose (bool, optional): If True, print all information to the screen. Default is True.
+ device (str | torch.device | None, optional): Device to use for model parameters, e.g., 'cpu', 'cuda', or
+ torch.device. If None, automatically selects the best available
+ device. Default is None.
+
+ Returns:
+ torch.nn.Module: The instantiated YOLOv5-large-P6 model.
+
+ Usage:
+ ```python
+ import torch
+ model = torch.hub.load('ultralytics/yolov5', 'yolov5l6') # official model
+ model = torch.hub.load('ultralytics/yolov5:master', 'yolov5l6') # from specific branch
+ model = torch.hub.load('ultralytics/yolov5', 'custom', 'path/to/yolov5l6.pt') # custom/local model
+ model = torch.hub.load('.', 'custom', 'path/to/yolov5l6.pt', source='local') # local repository
+ ```
+
+ Note:
+ Refer to [PyTorch Hub Documentation](https://pytorch.org/hub/ultralytics_yolov5) for additional usage instructions.
"""
return _create("yolov5l6", pretrained, channels, classes, autoshape, _verbose, device)
def yolov5x6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
- """Creates YOLOv5-xlarge-P6 model with options for pretraining, channels, classes, autoshaping, verbosity, and
- device.
+ """
+ Creates the YOLOv5-xlarge-P6 model with options for pretraining, number of input channels, class count, autoshaping,
+ verbosity, and device selection.
+
+ Args:
+ pretrained (bool): If True, loads pretrained weights into the model. Default is True.
+ channels (int): Number of input channels. Default is 3.
+ classes (int): Number of model classes. Default is 80.
+ autoshape (bool): If True, applies YOLOv5 .autoshape() wrapper to the model. Default is True.
+ _verbose (bool): If True, prints all information to the screen. Default is True.
+ device (str | torch.device | None): Device to use for model parameters, can be a string, torch.device object, or
+ None for default device selection. Default is None.
+
+ Returns:
+ torch.nn.Module: The instantiated YOLOv5-xlarge-P6 model.
+
+ Example:
+ ```python
+ import torch
+ model = torch.hub.load('ultralytics/yolov5', 'yolov5x6') # load the YOLOv5-xlarge-P6 model
+ ```
+
+ Note:
+ For more information on YOLOv5 models, visit the official documentation:
+ https://docs.ultralytics.com/yolov5
"""
return _create("yolov5x6", pretrained, channels, classes, autoshape, _verbose, device)
diff --git a/train.py b/train.py
index 472c1d395..65e6806b2 100644
--- a/train.py
+++ b/train.py
@@ -102,10 +102,37 @@ GIT_INFO = check_git_info()
def train(hyp, opt, device, callbacks):
"""
- Trains YOLOv5 model with given hyperparameters, options, and device, managing datasets, model architecture, loss
- computation, and optimizer steps.
+ Trains a YOLOv5 model on a custom dataset using specified hyperparameters, options, and device, managing datasets,
+ model architecture, loss computation, and optimizer steps.
- `hyp` argument is path/to/hyp.yaml or hyp dictionary.
+ Args:
+ hyp (str | dict): Path to the hyperparameters YAML file or a dictionary of hyperparameters.
+ opt (argparse.Namespace): Parsed command-line arguments containing training options.
+ device (torch.device): Device on which training occurs, e.g., 'cuda' or 'cpu'.
+ callbacks (Callbacks): Callback functions for various training events.
+
+ Returns:
+ None
+
+ Models and datasets download automatically from the latest YOLOv5 release.
+
+ Example:
+ Single-GPU training:
+ ```bash
+ $ python train.py --data coco128.yaml --weights yolov5s.pt --img 640 # from pretrained (recommended)
+ $ python train.py --data coco128.yaml --weights '' --cfg yolov5s.yaml --img 640 # from scratch
+ ```
+
+ Multi-GPU DDP training:
+ ```bash
+ $ python -m torch.distributed.run --nproc_per_node 4 --master_port 1 train.py --data coco128.yaml --weights
+ yolov5s.pt --img 640 --device 0,1,2,3
+ ```
+
+ For more usage details, refer to:
+ - Models: https://github.com/ultralytics/yolov5/tree/master/models
+ - Datasets: https://github.com/ultralytics/yolov5/tree/master/data
+ - Tutorial: https://docs.ultralytics.com/yolov5/tutorials/train_custom_data
"""
save_dir, epochs, batch_size, weights, single_cls, evolve, data, cfg, resume, noval, nosave, workers, freeze = (
Path(opt.save_dir),
@@ -515,7 +542,27 @@ def train(hyp, opt, device, callbacks):
def parse_opt(known=False):
- """Parses command-line arguments for YOLOv5 training, validation, and testing."""
+ """
+ Parses command-line arguments for YOLOv5 training, validation, and testing.
+
+ Args:
+ known (bool, optional): If True, parses known arguments, ignoring the unknown. Defaults to False.
+
+ Returns:
+ argparse.Namespace: Parsed command-line arguments.
+
+ Example:
+ ```python
+ from ultralytics.yolo import parse_opt
+ opt = parse_opt()
+ print(opt)
+ ```
+
+ Links:
+ Models: https://github.com/ultralytics/yolov5/tree/master/models
+ Datasets: https://github.com/ultralytics/yolov5/tree/master/data
+ Tutorial: https://docs.ultralytics.com/yolov5/tutorials/train_custom_data
+ """
parser = argparse.ArgumentParser()
parser.add_argument("--weights", type=str, default=ROOT / "yolov5s.pt", help="initial weights path")
parser.add_argument("--cfg", type=str, default="", help="model.yaml path")
@@ -570,7 +617,21 @@ def parse_opt(known=False):
def main(opt, callbacks=Callbacks()):
- """Runs training or hyperparameter evolution with specified options and optional callbacks."""
+ """
+ Runs training or hyperparameter evolution with specified options and optional callbacks.
+
+ Args:
+ opt (argparse.Namespace): The command-line arguments parsed for YOLOv5 training and evolution.
+ callbacks (ultralytics.utils.callbacks.Callbacks, optional): Callback functions for various training stages.
+ Defaults to Callbacks().
+
+ Returns:
+ None
+
+ Note:
+ For detailed usage, visit:
+ https://github.com/ultralytics/yolov5/tree/master/models
+ """
if RANK in {-1, 0}:
print_args(vars(opt))
check_git_status()
@@ -826,7 +887,25 @@ def main(opt, callbacks=Callbacks()):
def generate_individual(input_ranges, individual_length):
- """Generates a list of random values within specified input ranges for each gene in the individual."""
+ """
+ Generate a random individual with gene values within specified input ranges.
+
+ Args:
+ input_ranges (list[tuple[float, float]]): List of tuples where each tuple contains the lower and upper bounds
+ for the corresponding gene.
+ individual_length (int): The number of genes in the individual.
+
+ Returns:
+ list[float]: A list representing a generated individual with random gene values within the specified ranges.
+
+ Examples:
+ ```python
+ input_ranges = [(0.01, 0.1), (0.1, 1.0), (0.9, 2.0)]
+ individual_length = 3
+ individual = generate_individual(input_ranges, individual_length)
+ print(individual) # Output: [0.035, 0.678, 1.456] (example output)
+ ```
+ """
individual = []
for i in range(individual_length):
lower_bound, upper_bound = input_ranges[i]
@@ -836,9 +915,54 @@ def generate_individual(input_ranges, individual_length):
def run(**kwargs):
"""
- Executes YOLOv5 training with given options, overriding with any kwargs provided.
+ Executes YOLOv5 training with given options, allowing optional overrides through keyword arguments.
- Example: import train; train.run(data='coco128.yaml', imgsz=320, weights='yolov5m.pt')
+ Args:
+ weights (str): Path to initial weights. Defaults to ROOT / 'yolov5s.pt'.
+ cfg (str): Path to model YAML configuration. Defaults to an empty string.
+ data (str): Path to dataset YAML configuration. Defaults to ROOT / 'data/coco128.yaml'.
+ hyp (str): Path to hyperparameters YAML configuration. Defaults to ROOT / 'data/hyps/hyp.scratch-low.yaml'.
+ epochs (int): Total number of training epochs. Defaults to 100.
+ batch_size (int): Total batch size for all GPUs. Use -1 for automatic batch size determination. Defaults to 16.
+ imgsz (int): Image size (pixels) for training and validation. Defaults to 640.
+ rect (bool): Use rectangular training. Defaults to False.
+ resume (bool | str): Resume most recent training with an optional path. Defaults to False.
+ nosave (bool): Only save final checkpoint. Defaults to False.
+ noval (bool): Only validate at the final epoch. Defaults to False.
+ noautoanchor (bool): Disable AutoAnchor. Defaults to False.
+ noplots (bool): Do not save plot files. Defaults to False.
+ evolve (int): Evolve hyperparameters for a specified number of generations. Use 300 if provided without a value.
+ evolve_population (str): Directory for loading population during evolution. Defaults to ROOT / 'data/hyps'.
+ resume_evolve (str): Resume hyperparameter evolution from the last generation. Defaults to None.
+ bucket (str): gsutil bucket for saving checkpoints. Defaults to an empty string.
+ cache (str): Cache image data in 'ram' or 'disk'. Defaults to None.
+ image_weights (bool): Use weighted image selection for training. Defaults to False.
+ device (str): CUDA device identifier, e.g., '0', '0,1,2,3', or 'cpu'. Defaults to an empty string.
+ multi_scale (bool): Use multi-scale training, varying image size by ±50%. Defaults to False.
+ single_cls (bool): Train with multi-class data as single-class. Defaults to False.
+ optimizer (str): Optimizer type, choices are ['SGD', 'Adam', 'AdamW']. Defaults to 'SGD'.
+ sync_bn (bool): Use synchronized BatchNorm, only available in DDP mode. Defaults to False.
+ workers (int): Maximum dataloader workers per rank in DDP mode. Defaults to 8.
+ project (str): Directory for saving training runs. Defaults to ROOT / 'runs/train'.
+ name (str): Name for saving the training run. Defaults to 'exp'.
+ exist_ok (bool): Allow existing project/name without incrementing. Defaults to False.
+ quad (bool): Use quad dataloader. Defaults to False.
+ cos_lr (bool): Use cosine learning rate scheduler. Defaults to False.
+ label_smoothing (float): Label smoothing epsilon value. Defaults to 0.0.
+ patience (int): Patience for early stopping, measured in epochs without improvement. Defaults to 100.
+ freeze (list): Layers to freeze, e.g., backbone=10, first 3 layers = [0, 1, 2]. Defaults to [0].
+ save_period (int): Frequency in epochs to save checkpoints. Disabled if < 1. Defaults to -1.
+ seed (int): Global training random seed. Defaults to 0.
+ local_rank (int): Automatic DDP Multi-GPU argument. Do not modify. Defaults to -1.
+
+ Returns:
+ None: The function initiates YOLOv5 training or hyperparameter evolution based on the provided options.
+
+ Examples:
+ ```python
+ import train
+ train.run(data='coco128.yaml', imgsz=320, weights='yolov5m.pt')
+ ```
"""
opt = parse_opt(True)
for k, v in kwargs.items():
diff --git a/val.py b/val.py
index c1e8a6aa3..deb12624c 100644
--- a/val.py
+++ b/val.py
@@ -62,7 +62,30 @@ from utils.torch_utils import select_device, smart_inference_mode
def save_one_txt(predn, save_conf, shape, file):
- """Saves one detection result to a txt file in normalized xywh format, optionally including confidence."""
+ """
+ Saves one detection result to a txt file in normalized xywh format, optionally including confidence.
+
+ Args:
+ predn (torch.Tensor): Predicted bounding boxes and associated confidence scores and classes
+ in xyxy format, tensor of shape (N, 6) where N is the number of detections.
+ save_conf (bool): If True, saves the confidence scores along with the bounding box coordinates.
+ shape (tuple): Shape of the original image as (height, width).
+ file (str | Path): File path where the result will be saved.
+
+ Returns:
+ None
+
+ Notes:
+ The xyxy bounding box format represents the coordinates (xmin, ymin, xmax, ymax).
+ The xywh format represents the coordinates (center_x, center_y, width, height) and is
+ normalized by the width and height of the image.
+
+ Example:
+ ```python
+ predn = torch.tensor([[10, 20, 30, 40, 0.9, 1]]) # example prediction
+ save_one_txt(predn, save_conf=True, shape=(640, 480), file="output.txt")
+ ```
+ """
gn = torch.tensor(shape)[[1, 0, 1, 0]] # normalization gain whwh
for *xyxy, conf, cls in predn.tolist():
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
@@ -73,9 +96,36 @@ def save_one_txt(predn, save_conf, shape, file):
def save_one_json(predn, jdict, path, class_map):
"""
- Saves one JSON detection result with image ID, category ID, bounding box, and score.
+ Saves a single JSON detection result, including image ID, category ID, bounding box, and confidence score.
- Example: {"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}
+ Args:
+ predn (torch.Tensor): Predicted detections in xyxy format with shape (n, 6) where n is the number of detections.
+ The tensor should contain [x_min, y_min, x_max, y_max, confidence, class_id] for each detection.
+ jdict (list[dict]): List to collect JSON formatted detection results.
+ path (pathlib.Path): Path object of the image file, used to extract image_id.
+ class_map (dict[int, int]): Mapping from model class indices to dataset-specific category IDs.
+
+ Returns:
+ None: Appends detection results as dictionaries to `jdict` list in-place.
+
+ Example:
+ ```python
+ predn = torch.tensor([[100, 50, 200, 150, 0.9, 0], [50, 30, 100, 80, 0.8, 1]])
+ jdict = []
+ path = Path("42.jpg")
+ class_map = {0: 18, 1: 19}
+ save_one_json(predn, jdict, path, class_map)
+ ```
+ This will append to `jdict`:
+ ```
+ [
+ {'image_id': 42, 'category_id': 18, 'bbox': [125.0, 75.0, 100.0, 100.0], 'score': 0.9},
+ {'image_id': 42, 'category_id': 19, 'bbox': [75.0, 55.0, 50.0, 50.0], 'score': 0.8}
+ ]
+ ```
+
+ Notes:
+ The `bbox` values are formatted as [x, y, width, height], where x and y represent the top-left corner of the box.
"""
image_id = int(path.stem) if path.stem.isnumeric() else path.stem
box = xyxy2xywh(predn[:, :4]) # xywh
@@ -93,13 +143,30 @@ def save_one_json(predn, jdict, path, class_map):
def process_batch(detections, labels, iouv):
"""
- Return correct prediction matrix.
+ Return a correct prediction matrix given detections and labels at various IoU thresholds.
+
+ Args:
+ detections (np.ndarray): Array of shape (N, 6) where each row corresponds to a detection with
+ format [x1, y1, x2, y2, conf, class].
+ labels (np.ndarray): Array of shape (M, 5) where each row corresponds to a ground truth label with
+ format [class, x1, y1, x2, y2].
+ iouv (np.ndarray): Array of IoU thresholds to evaluate at.
- Arguments:
- detections (array[N, 6]), x1, y1, x2, y2, conf, class
- labels (array[M, 5]), class, x1, y1, x2, y2
Returns:
- correct (array[N, 10]), for 10 IoU levels
+ correct (np.ndarray): A binary array of shape (N, len(iouv)) indicating whether each detection
+ is a true positive for each IoU threshold. There are 10 IoU levels used in the evaluation.
+
+ Example:
+ ```python
+ detections = np.array([[50, 50, 200, 200, 0.9, 1], [30, 30, 150, 150, 0.7, 0]])
+ labels = np.array([[1, 50, 50, 200, 200]])
+ iouv = np.linspace(0.5, 0.95, 10)
+ correct = process_batch(detections, labels, iouv)
+ ```
+
+ Notes:
+ - This function is used as part of the evaluation pipeline for object detection models.
+ - IoU (Intersection over Union) is a common evaluation metric for object detection performance.
"""
correct = np.zeros((detections.shape[0], iouv.shape[0])).astype(bool)
iou = box_iou(labels[:, 1:], detections[:, :4])
@@ -148,7 +215,44 @@ def run(
callbacks=Callbacks(),
compute_loss=None,
):
- """Evaluates model on a dataset and logs performance metrics, results are saved to specific directories."""
+ """
+ Evaluates a YOLOv5 model on a dataset and logs performance metrics.
+
+ Args:
+ data (str | dict): Path to a dataset yaml file or a dataset dictionary.
+ weights (str | list[str], optional): Path to the model weights file(s). Supports various formats: PyTorch,
+ TorchScript, ONNX, OpenVINO, TensorRT, CoreML, TensorFlow SavedModel, TensorFlow GraphDef, TensorFlow Lite,
+ TensorFlow Edge TPU, and PaddlePaddle.
+ batch_size (int, optional): Batch size for inference. Default is 32.
+ imgsz (int, optional): Input image size (pixels). Default is 640.
+ conf_thres (float, optional): Confidence threshold for object detection. Default is 0.001.
+ iou_thres (float, optional): IoU threshold for Non-Maximum Suppression (NMS). Default is 0.6.
+ max_det (int, optional): Maximum number of detections per image. Default is 300.
+ task (str, optional): Task type - 'train', 'val', 'test', 'speed', or 'study'. Default is 'val'.
+ device (str, optional): Device to use for computation, e.g., '0' or '0,1,2,3' for CUDA or 'cpu' for CPU. Default is ''.
+ workers (int, optional): Number of dataloader workers. Default is 8.
+ single_cls (bool, optional): Treat dataset as a single class. Default is False.
+ augment (bool, optional): Enable augmented inference. Default is False.
+ verbose (bool, optional): Enable verbose output. Default is False.
+ save_txt (bool, optional): Save results to *.txt files. Default is False.
+ save_hybrid (bool, optional): Save label and prediction hybrid results to *.txt files. Default is False.
+ save_conf (bool, optional): Save confidences in --save-txt labels. Default is False.
+ save_json (bool, optional): Save a COCO-JSON results file. Default is False.
+ project (str | Path, optional): Directory to save results. Default is ROOT/'runs/val'.
+ name (str, optional): Name of the run. Default is 'exp'.
+ exist_ok (bool, optional): Overwrite existing project/name without incrementing. Default is False.
+ half (bool, optional): Use FP16 half-precision inference. Default is True.
+ dnn (bool, optional): Use OpenCV DNN for ONNX inference. Default is False.
+ model (torch.nn.Module, optional): Model object for training. Default is None.
+ dataloader (torch.utils.data.DataLoader, optional): Dataloader object. Default is None.
+ save_dir (Path, optional): Directory to save results. Default is Path('').
+ plots (bool, optional): Plot validation images and metrics. Default is True.
+ callbacks (utils.callbacks.Callbacks, optional): Callbacks for logging and monitoring. Default is Callbacks().
+ compute_loss (function, optional): Loss function for training. Default is None.
+
+ Returns:
+ dict: Contains performance metrics including precision, recall, mAP50, and mAP50-95.
+ """
# Initialize/load model and set device
training = model is not None
if training: # called by train.py
@@ -364,7 +468,53 @@ def run(
def parse_opt():
- """Parses command-line options for YOLOv5 model inference configuration."""
+ """
+ Parses command-line options for YOLOv5 model inference configuration.
+
+ Args:
+ data (str): Path to the dataset YAML file, default is 'data/coco128.yaml'.
+ weights (List[str]): List of paths to the model weight files, default is 'yolov5s.pt'.
+ batch_size (int): Batch size for inference, default is 32.
+ imgsz (int): Inference image size in pixels, default is 640.
+ conf_thres (float): Confidence threshold for predictions, default is 0.001.
+ iou_thres (float): IoU threshold for Non-Max Suppression (NMS), default is 0.6.
+ max_det (int): Maximum number of detections per image, default is 300.
+ task (str): Task type - options are 'train', 'val', 'test', 'speed', or 'study'. Default is 'val'.
+ device (str): Device to run the model on, e.g., '0' or '0,1,2,3' or 'cpu'. Default is empty to let
+ the system choose automatically.
+ workers (int): Maximum number of dataloader workers per rank in DDP mode, default is 8.
+ single_cls (bool): If set, treats the dataset as a single-class dataset. Default is False.
+ augment (bool): If set, performs augmented inference. Default is False.
+ verbose (bool): If set, reports mAP by class. Default is False.
+ save_txt (bool): If set, saves results to *.txt files. Default is False.
+ save_hybrid (bool): If set, saves label+prediction hybrid results to *.txt files. Default is False.
+ save_conf (bool): If set, saves confidences in --save-txt labels. Default is False.
+ save_json (bool): If set, saves results to a COCO-JSON file. Default is False.
+ project (str): Project directory to save results to. Default is 'runs/val'.
+ name (str): Name of the directory to save results to. Default is 'exp'.
+ exist_ok (bool): If set, existing directory will not be incremented. Default is False.
+ half (bool): If set, uses FP16 half-precision inference. Default is False.
+ dnn (bool): If set, uses OpenCV DNN for ONNX inference. Default is False.
+
+ Returns:
+ argparse.Namespace: Parsed command-line options
+
+ Notes:
+ - The '--data' parameter is checked to ensure it ends with 'coco.yaml' if '--save-json' is set.
+ - The '--save-txt' option is set to True if '--save-hybrid' is enabled.
+ - Args are printed using `print_args` to facilitate debugging.
+
+ Example:
+ To validate a trained YOLOv5 model on a COCO dataset:
+ ```python
+ $ python val.py --weights yolov5s.pt --data coco128.yaml --img 640
+ ```
+ Different model formats could be used instead of yolov5s.pt:
+ ```python
+ $ python val.py --weights yolov5s.pt yolov5s.torchscript yolov5s.onnx yolov5s_openvino_model yolov5s.engine
+ ```
+ Additional options include saving results in different formats, selecting devices, and more.
+ """
parser = argparse.ArgumentParser()
parser.add_argument("--data", type=str, default=ROOT / "data/coco128.yaml", help="dataset.yaml path")
parser.add_argument("--weights", nargs="+", type=str, default=ROOT / "yolov5s.pt", help="model path(s)")
@@ -397,8 +547,25 @@ def parse_opt():
def main(opt):
- """Executes YOLOv5 tasks like training, validation, testing, speed, and study benchmarks based on provided
- options.
+ """
+ Executes YOLOv5 tasks like training, validation, testing, speed, and study benchmarks based on provided options.
+
+ Args:
+ opt (argparse.Namespace): Parsed command-line options.
+ - This includes values for parameters like 'data', 'weights', 'batch_size', 'imgsz', 'conf_thres', 'iou_thres',
+ 'max_det', 'task', 'device', 'workers', 'single_cls', 'augment', 'verbose', 'save_txt', 'save_hybrid',
+ 'save_conf', 'save_json', 'project', 'name', 'exist_ok', 'half', and 'dnn', essential for configuring
+ the YOLOv5 tasks.
+
+ Returns:
+ None
+
+ Examples:
+ To validate a trained YOLOv5 model on the COCO dataset with a specific weights file, use:
+
+ ```python
+ $ python val.py --weights yolov5s.pt --data coco128.yaml --img 640
+ ```
"""
check_requirements(ROOT / "requirements.txt", exclude=("tensorboard", "thop"))