From 8257e0a68cc1d192c6ae1392bce1ed3917119795 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 8 Jul 2024 21:19:04 +0200 Subject: [PATCH] Ultralytics Code Refactor https://ultralytics.com/actions (#13174) Refactor code for speed and clarity --- benchmarks.py | 106 ++++++++- detect.py | 117 +++++++++- export.py | 621 ++++++++++++++++++++++++++++++++++++++++++++++++-- hubconf.py | 346 ++++++++++++++++++++++++++-- train.py | 144 +++++++++++- val.py | 191 +++++++++++++++- 6 files changed, 1448 insertions(+), 77 deletions(-) diff --git a/benchmarks.py b/benchmarks.py index c849eed6f078..90cce4b3ff3c 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 f791faa09087..b774370f54ed 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 write_to_csv(image_name, prediction, confidence): 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 77c972ba27f6..bf68ae123ab3 100644 --- a/export.py +++ b/export.py @@ -92,7 +92,22 @@ 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 @@ def __init__(self, model, im): # 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 outer_func(*args, **kwargs): @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 transform_fn(data_item): @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 4b0c36b8daed..de93c79bfae5 100644 --- a/hubconf.py +++ b/hubconf.py @@ -15,19 +15,38 @@ 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 472c1d395981..65e6806b22e0 100644 --- a/train.py +++ b/train.py @@ -102,10 +102,37 @@ def train(hyp, opt, device, callbacks): """ - Trains YOLOv5 model with given hyperparameters, options, and device, managing datasets, model architecture, loss - computation, and optimizer steps. - - `hyp` argument is path/to/hyp.yaml or hyp dictionary. + Trains a YOLOv5 model on a custom dataset using specified hyperparameters, options, and device, managing datasets, + model architecture, loss computation, and optimizer steps. + + 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 lf(x): 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. - - Example: import train; train.run(data='coco128.yaml', imgsz=320, weights='yolov5m.pt') + Executes YOLOv5 training with given options, allowing optional overrides through keyword arguments. + + 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 c1e8a6aa3094..deb12624c32f 100644 --- a/val.py +++ b/val.py @@ -62,7 +62,30 @@ 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. + + 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. - Example: {"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236} + 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"))