From b5b53bab0dbb04a30bf396aba2d15ecb6c5aa2c0 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 8 Jul 2024 23:56:23 +0200 Subject: [PATCH] Ultralytics Code Refactor https://ultralytics.com/actions (#2246) * Refactor code for speed and clarity * Update export.py * Update hubconf.py --- benchmarks.py | 138 +++++++++++- detect.py | 139 +++++++++++- export.py | 611 +++++++++++++++++++++++++++++++++++++++++++++++--- hubconf.py | 326 ++++++++++++++++++++++++--- train.py | 142 +++++++++++- val.py | 199 +++++++++++++++- 6 files changed, 1470 insertions(+), 85 deletions(-) diff --git a/benchmarks.py b/benchmarks.py index f01dc9b52d..06bde1e380 100644 --- a/benchmarks.py +++ b/benchmarks.py @@ -60,7 +60,40 @@ def run( pt_only=False, # test PyTorch only hard_fail=False, # throw error on benchmark failure ): - """Run YOLOv3 benchmarks on multiple export formats and validate performance metrics.""" + """ + Run YOLOv3 benchmarks on multiple export formats and validate performance metrics. + + Args: + weights (str | Path): Path to the weights file. Defaults to 'yolov5s.pt'. + imgsz (int): Inference image size in pixels. Defaults to 640. + batch_size (int): Batch size for inference. Defaults to 1. + data (str | Path): Path to the dataset configuration file (dataset.yaml). Defaults to 'data/coco128.yaml'. + device (str): Device to be used for inference, e.g., '0' or '0,1,2,3' for GPU or 'cpu' for CPU. Defaults to ''. + half (bool): Use FP16 half-precision for inference. Defaults to False. + test (bool): Test exports only without running benchmarks. Defaults to False. + pt_only (bool): Run benchmarks only for PyTorch format. Defaults to False. + hard_fail (bool): Raise an error if any benchmark test fails. Defaults to False. + + Returns: + None + + Notes: + This function iterates over multiple export formats, performs the export, and then validates the model's + performance using appropriate validation functions for detection and segmentation models. The results are logged, + and optionally, benchmarks can be configured to raise errors on failures using the `hard_fail` argument. + + Examples: + ```python + # Run benchmarks on the default 'yolov5s.pt' model with an image size of 640 pixels + run() + + # Run benchmarks on a specific model with GPU and half-precision enabled + run(weights='custom_model.pt', device='0', half=True) + + # Test only PyTorch export + run(pt_only=True) + ``` + """ y, t = [], time.time() device = select_device(device) model_type = type(attempt_load(weights, fuse=False)) # DetectionModel, SegmentationModel, etc. @@ -125,7 +158,45 @@ def test( pt_only=False, # test PyTorch only hard_fail=False, # throw error on benchmark failure ): - """Run YOLOv3 export tests for various formats and log the results, including export success status.""" + """ + Run YOLOv3 export tests for various formats and log the results, including export success status. + + Args: + weights (str | Path): Path to the weights file. Defaults to ROOT / "yolov5s.pt". + imgsz (int): Inference size in pixels. Defaults to 640. + batch_size (int): Number of images per batch. Defaults to 1. + data (str | Path): Path to the dataset yaml file. Defaults to ROOT / "data/coco128.yaml". + device (str): Device for inference. Accepts cuda device (e.g., "0" or "0,1,2,3") or "cpu". Defaults to "". + half (bool): Use FP16 half-precision inference. Defaults to False. + test (bool): Run export tests only, no inference. Defaults to False. + pt_only (bool): Run tests on PyTorch format only. Defaults to False. + hard_fail (bool): Raise an error on benchmark failure. Defaults to False. + + Returns: + pd.DataFrame: A DataFrame containing the export formats and their success status. (pd.DataFrame) + + Examples: + ```python + from ultralytics import test + + results = test( + weights="path/to/yolov5s.pt", + imgsz=640, + batch_size=1, + data="path/to/coco128.yaml", + device="0", + half=False, + test=True, + pt_only=False, + hard_fail=True, + ) + print(results) + ``` + + Notes: + Ensure all required packages are installed as specified in the Ultralytics YOLOv3 documentation: + https://github.com/ultralytics/ultralytics + """ 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,8 +222,38 @@ def test( def parse_opt(): - """Parses command line arguments for model inference configurations, including weights, image size, and device - options. + """ + Parses command line arguments for YOLOv3 inference and export configuration. + + Args: + --weights (str): Path to the weights file. Default is 'ROOT / "yolov3-tiny.pt"'. + --imgsz | --img | --img-size (int): Inference image size in pixels. Default is 640. + --batch-size (int): Batch size for inference. Default is 1. + --data (str): Path to the dataset configuration file (dataset.yaml). Default is 'ROOT / "data/coco128.yaml"'. + --device (str): CUDA device identifier, e.g., '0' for single GPU, '0,1,2,3' for multiple GPUs, or 'cpu' for CPU + inference. Default is "". + --half (bool): If set, use FP16 half-precision inference. Default is False. + --test (bool): If set, only test exports without running inference. Default is False. + --pt-only (bool): If set, test only the PyTorch model without exporting to other formats. Default is False. + --hard-fail (str | bool): If set, raise an exception on benchmark failure. Can also be a string representing + the minimum metric floor for success. Default is False. + + Returns: + argparse.Namespace: The parsed arguments as a namespace object. + + Example: + To run inference on the YOLOv3-tiny model with a different image size: + + ```python + $ python benchmarks.py --weights yolov3-tiny.pt --imgsz 512 --device 0 + ``` + + Notes: + The `--hard-fail` argument can be a boolean or a string. If a string is provided, it should be an expression that + represents the minimum acceptable metric value, such as '0.29' for mAP (mean Average Precision). + + Links: + https://github.com/ultralytics/ultralytics """ parser = argparse.ArgumentParser() parser.add_argument("--weights", type=str, default=ROOT / "yolov3-tiny.pt", help="weights path") @@ -171,7 +272,34 @@ def parse_opt(): def main(opt): - """Executes tests or main pipeline on provided options, determining behavior based on `opt.test`.""" + """ + Executes the export and benchmarking pipeline for YOLOv3 models, testing multiple export formats and validating + performance metrics. + + Args: + opt (argparse.Namespace): Parsed command line arguments, including options for weights, image size, batch size, + dataset path, device, half-precision inference, test mode, PyTorch-only testing, and hard fail conditions. + + Returns: + pd.DataFrame: A DataFrame containing benchmarking results with columns: + - Format: Name of the export format + - Size (MB): File size of the exported model + - mAP50-95: Mean Average Precision for the model + - Inference time (ms): Time taken for inference + + Notes: + The function runs the main pipeline by exporting the YOLOv3 model to various formats and running benchmarks to + evaluate performance. If `opt.test` is set to True, it only tests the export process and logs the results. + + Example: + Running the function from command line with required arguments: + + ```python + $ python benchmarks.py --weights yolov5s.pt --img 640 + ``` + + For more details, visit the Ultralytics YOLOv3 repository on [GitHub](https://github.com/ultralytics/ultralytics). + """ test(**vars(opt)) if opt.test else run(**vars(opt)) diff --git a/detect.py b/detect.py index ba11935c76..efe3f1d6e5 100644 --- a/detect.py +++ b/detect.py @@ -95,7 +95,60 @@ def run( dnn=False, # use OpenCV DNN for ONNX inference vid_stride=1, # video frame-rate stride ): - """Performs YOLOv3 detection on various input sources including images, videos, streams, and YouTube URLs.""" + """ + Performs YOLOv3 detection on various input sources including images, videos, streams, and YouTube URLs. + + Args: + weights (str | Path): Path to the model weights file or a Triton URL (default: 'yolov5s.pt'). + source (str | Path): Source of input data such as a file, directory, URL, glob pattern, or device identifier + (default: 'data/images'). + data (str | Path): Path to the dataset YAML file (default: 'data/coco128.yaml'). + imgsz (tuple[int, int]): Inference size as a tuple (height, width) (default: (640, 640)). + conf_thres (float): Confidence threshold for detection (default: 0.25). + iou_thres (float): Intersection Over Union (IOU) threshold for Non-Max Suppression (NMS) (default: 0.45). + max_det (int): Maximum number of detections per image (default: 1000). + device (str): CUDA device identifier, e.g., '0', '0,1,2,3', or 'cpu' (default: ''). + view_img (bool): Whether to display results during inference (default: False). + save_txt (bool): Whether to save detection results to text files (default: False). + save_conf (bool): Whether to save detection confidences in the text labels (default: False). + save_crop (bool): Whether to save cropped detection boxes (default: False). + nosave (bool): Whether to prevent saving images or videos with detections (default: False). + classes (list[int] | None): List of class indices to filter, e.g., [0, 2, 3] (default: None). + agnostic_nms (bool): Whether to perform class-agnostic NMS (default: False). + augment (bool): Whether to apply augmented inference (default: False). + visualize (bool): Whether to visualize feature maps (default: False). + update (bool): Whether to update all models (default: False). + project (str | Path): Path to the project directory where results will be saved (default: 'runs/detect'). + name (str): Name for the specific run within the project directory (default: 'exp'). + exist_ok (bool): Whether to allow existing project/name directory without incrementing run index (default: False). + line_thickness (int): Thickness of bounding box lines in pixels (default: 3). + hide_labels (bool): Whether to hide labels in the results (default: False). + hide_conf (bool): Whether to hide confidences in the results (default: False). + half (bool): Whether to use half-precision (FP16) for inference (default: False). + dnn (bool): Whether to use OpenCV DNN for ONNX inference (default: False). + vid_stride (int): Stride for video frame rate (default: 1). + + Returns: + None + + Notes: + This function supports a variety of input sources such as image files, video files, directories, URL patterns, + webcam streams, and YouTube links. It also supports multiple model formats including PyTorch, ONNX, OpenVINO, + TensorRT, CoreML, TensorFlow, PaddlePaddle, and others. The results can be visualized in real-time or saved to + specified directories. Use command-line arguments to modify the behavior of the function. + + Examples: + ```python + # Run YOLOv3 inference on an image + run(weights='yolov5s.pt', source='data/images/bus.jpg') + + # Run YOLOv3 inference on a video + run(weights='yolov5s.pt', source='data/videos/video.mp4', view_img=True) + + # Run YOLOv3 inference on a webcam + run(weights='yolov5s.pt', source='0', view_img=True) + ``` + """ 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) @@ -233,7 +286,41 @@ def run( def parse_opt(): - """Parses and returns command-line options for model inference configurations.""" + """ + Parses and returns command-line options for model inference configurations. + + Args: + --weights (list[str]): Model path or triton URL. Default: `ROOT / "yolov3-tiny.pt"`. + --source (str): File/directory/URL/glob/screen/0(webcam) for input data. Default: `ROOT / "data/images"`. + --data (str): (Optional) Path to dataset.yaml. Default: `ROOT / "data/coco128.yaml"`. + --imgsz (list[int]): Inference size as height, width. Accepts multiple values. Default: `[640]`. + --conf-thres (float): Confidence threshold for predictions. Default: `0.25`. + --iou-thres (float): IoU threshold for Non-Maximum Suppression (NMS). Default: `0.45`. + --max-det (int): Maximum number of detections per image. Default: `1000`. + --device (str): CUDA device to run the model on, e.g., `0` or `0,1,2,3` or `cpu`. Default: `""`. + --view-img (bool): Display results on the screen. Default: `False`. + --save-txt (bool): Save results to text files. Default: `False`. + --save-conf (bool): Save confidence scores in text labels. Default: `False`. + --save-crop (bool): Save cropped prediction boxes. Default: `False`. + --nosave (bool): Do not save images/videos. Default: `False`. + --classes (list[int]): Filter results by class, e.g., `--classes 0` or `--classes 0 2 3`. Default: `None`. + --agnostic-nms (bool): Perform class-agnostic NMS. Default: `False`. + --augment (bool): Perform augmented inference. Default: `False`. + --visualize (bool): Visualize features. Default: `False`. + --update (bool): Update all models. Default: `False`. + --project (str): Directory to save results, saved as "project/name". Default: `ROOT / "runs/detect"`. + --name (str): Directory name for saving results, e.g., "exp" in "project/name". Default: `"exp"`. + --exist-ok (bool): Allow results to be saved in an existing directory without incrementing. Default: `False`. + --line-thickness (int): Thickness of bounding box lines (in pixels). Default: `3`. + --hide-labels (bool): Hide labels on detected objects. Default: `False`. + --hide-conf (bool): Hide confidence scores on labels. Default: `False`. + --half (bool): Use FP16 half-precision inference. Default: `False`. + --dnn (bool): Use OpenCV DNN backend for ONNX inference. Default: `False`. + --vid-stride (int): Frame-rate stride for video input. Default: `1`. + + Returns: + argparse.Namespace: Parsed command-line arguments for inference configurations. + """ parser = argparse.ArgumentParser() parser.add_argument( "--weights", nargs="+", type=str, default=ROOT / "yolov3-tiny.pt", help="model path or triton URL" @@ -271,7 +358,53 @@ def parse_opt(): def main(opt): - """Entry point for running the model; checks requirements and calls `run` with parsed options.""" + """ + Entry point for running the model; checks requirements and calls `run` with parsed options. + + Args: + opt (argparse.Namespace): Parsed command-line options, which include: + - weights (str | list of str): Path to the model weights or Triton server URL. + - source (str): Input source, can be a file, directory, URL, glob, screen, or webcam index. + - data (str): Path to the dataset configuration file (.yaml). + - imgsz (tuple of int): Inference image size as (height, width). + - conf_thres (float): Confidence threshold for detections. + - iou_thres (float): Intersection over Union (IoU) threshold for Non-Maximum Suppression (NMS). + - max_det (int): Maximum number of detections per image. + - device (str): Device to run inference on; options are CUDA device id(s) or 'cpu' + - view_img (bool): Flag to display inference results. + - save_txt (bool): Save detection results in .txt format. + - save_conf (bool): Save detection confidences in .txt labels. + - save_crop (bool): Save cropped bounding box predictions. + - nosave (bool): Do not save images/videos with detections. + - classes (list of int): Filter results by class, e.g., --class 0 2 3. + - agnostic_nms (bool): Use class-agnostic NMS. + - augment (bool): Enable augmented inference. + - visualize (bool): Visualize feature maps. + - update (bool): Update the model during inference. + - project (str): Directory to save results. + - name (str): Name for the results directory. + - exist_ok (bool): Allow existing project/name directories without incrementing. + - line_thickness (int): Thickness of bounding box lines. + - hide_labels (bool): Hide class labels on bounding boxes. + - hide_conf (bool): Hide confidence scores on bounding boxes. + - half (bool): Use FP16 half-precision inference. + - dnn (bool): Use OpenCV DNN backend for ONNX inference. + - vid_stride (int): Video frame-rate stride. + + Returns: + None + + Example: + ```python + if __name__ == "__main__": + opt = parse_opt() + main(opt) + ``` + Notes: + Run this function as the entry point for using YOLOv3 for object detection on a variety of input sources such + as images, videos, directories, webcams, streams, etc. This function ensures all requirements are checked and + subsequently initiates the detection process by calling the `run` function with appropriate options. + """ check_requirements(ROOT / "requirements.txt", exclude=("tensorboard", "thop")) run(**vars(opt)) diff --git a/export.py b/export.py index 22ac1d5198..e30c4f0e7b 100644 --- a/export.py +++ b/export.py @@ -92,7 +92,25 @@ class iOSModel(torch.nn.Module): def __init__(self, model, im): - """Initializes an iOSModel with normalized input dimensions and number of classes from a torch model.""" + """ + Initializes an iOSModel with normalized input dimensions and number of classes from a PyTorch model. + + Args: + model (torch.nn.Module): The PyTorch model from which to initialize the iOS model. This should include + attributes like `nc` (number of classes) which will be used to configure the iOS model. + im (torch.Tensor): A Tensor representing a sample input image. The shape of this tensor should be (batch_size, + channels, height, width). This is used to extract dimensions for input normalization. + + Returns: + None + + Notes: + - This class is specifically designed for use in exporting a PyTorch model for deployment on iOS platforms, + optimizing input dimensions and class configurations to suit mobile requirements. + - Normalization factor is derived from the input image dimensions, which impacts the model's performance + during inference on iOS devices. + - Ensure the sample input image `im` provided has correct dimensions and shape for accurate model configuration. + """ super().__init__() b, c, h, w = im.shape # batch, channel, height, width self.model = model @@ -105,13 +123,61 @@ 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): - """Performs forward pass, returns scaled confidences and normalized coordinates given input tensor `x`.""" + """ + Performs a forward pass, returning scaled confidences and normalized coordinates given an input tensor. + + Args: + x (torch.Tensor): Input tensor representing a batch of images, with dimensions [batch_size, channels, height, width]. + + Returns: + tuple[torch.Tensor, torch.Tensor, torch.Tensor]: A tuple containing three elements: + - xywh (torch.Tensor): Tensor of shape [batch_size, num_detections, 4] containing normalized x, y, width, + and height coordinates. + - conf (torch.Tensor): Tensor of shape [batch_size, num_detections, 1] containing confidence scores for each detection. + - cls (torch.Tensor): Tensor of shape [batch_size, num_detections, num_classes] containing class probabilities. + + Notes: + The dimensions of `x` should match the input dimensions used during the model's initialization to ensure proper + scaling and normalization. + + Examples: + ```python + model = iOSModel(trained_model, input_image_tensor) + detection_results = model.forward(input_tensor) + xywh, conf, cls = detection_results + ``` + + Further reading on exporting models to different formats: + https://github.com/ultralytics/ultralytics + + See Also: + `export.py` for exporting a YOLOv3 PyTorch model to various formats. + https://github.com/zldrobit for TensorFlow export scripts. + """ 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(): - """Lists supported YOLOv3 model export formats including file suffixes and CPU/GPU compatibility.""" + """ + Lists supported YOLOv3 model export formats including file suffixes and CPU/GPU compatibility. + + Returns: + list: A list of lists where each sublist contains information about a specific export format. + Each sublist includes the following elements: + - str: The name of the format. + - str: The command-line argument for including this format. + - str: The file suffix used for this format. + - bool: Indicates if the format is compatible with CPU. + - bool: Indicates if the format is compatible with GPU. + + Examples: + ```python + formats = export_formats() + for format in formats: + print(f"Format: {format[0]}, Suffix: {format[2]}, CPU Compatible: {format[3]}, GPU Compatible: {format[4]}") + ``` + """ x = [ ["PyTorch", "-", ".pt", True, True], ["TorchScript", "torchscript", ".torchscript", True, True], @@ -130,7 +196,28 @@ def export_formats(): def try_export(inner_func): - """Decorator to profile and log the export process of YOLOv3 models, handling successes and failures.""" + """ + Profiles and logs the export process of YOLOv3 models, capturing success or failure details. + + Args: + inner_func (Callable): The function that performs the actual export process and returns the model file path + and the exported model. + + Returns: + Callable: A wrapped function that profiles and logs the export process, handling successes and failures. + + Usage example: + + ```python + @try_export + def export_onnx(py_model_path: str, output_path: str): + # Export logic here + return output_path, model + ``` + + After applying this decorator, `export_onnx` will log the export results, including export success or failure along + with associated time and file size details. + """ inner_args = get_default_args(inner_func) def outer_func(*args, **kwargs): @@ -151,9 +238,40 @@ def outer_func(*args, **kwargs): @try_export def export_torchscript(model, im, file, optimize, prefix=colorstr("TorchScript:")): """ - Exports a YOLOv3 model to TorchScript format, optionally optimizing for mobile. + Exports a YOLOv3 model to TorchScript format, with optional optimization for mobile deployment. + + Args: + model (torch.nn.Module): The YOLOv3 model to be exported. + im (torch.Tensor): A tensor representing the input image for the model. + file (pathlib.Path): The file path where the TorchScript model will be saved. + optimize (bool): A boolean flag indicating whether to optimize the model for mobile. + prefix (str): A prefix for logging messages. Defaults to `colorstr("TorchScript:")`. + + Returns: + (pathlib.Path | None, torch.nn.Module | None): Tuple containing the path to the saved TorchScript model and the + model itself. Returns `(None, None)` if the export fails. + + Raises: + Exception: If there is an error during export, it logs the error and returns `(None, None)`. + + Notes: + The function uses torch.jit.trace to trace the model with the input image tensor (`im`). Required metadata such as + input shape, stride, and class names are saved in an extra file included in the TorchScript model. + + Examples: + ```python + from pathlib import Path + import torch - Args: model, im (input image), file (path), optimize (bool). + model = ... # Assume model is loaded or created + im = torch.randn(1, 3, 640, 640) # A sample input tensor + file = Path("model.torchscript") + optimize = True + + export_torchscript(model, im, file, optimize) + ``` + + For more information, visit: https://ultralytics.com/. """ LOGGER.info(f"\n{prefix} starting export with torch {torch.__version__}...") f = file.with_suffix(".torchscript") @@ -170,7 +288,21 @@ 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 YOLOv3 model to ONNX format with dynamic shapes and simplification options.""" + """ + Exports a YOLOv3 model to ONNX format with dynamic shape and simplification options. + + Args: + model (torch.nn.Module): The YOLOv3 model to be exported. + im (torch.Tensor): A sample input tensor for tracing the model. + file (pathlib.Path): The file path where the ONNX model will be saved. + opset (int): The ONNX opset version to use for the export. + dynamic (bool): If `True`, enables dynamic shape support. + simplify (bool): If `True`, simplifies the ONNX model using onnx-simplifier. + prefix (str): A prefix for logging messages. + + Returns: + tuple[pathlib.Path, None]: The path to the saved ONNX model, None as the second tuple element (kept for consistency). + """ check_requirements("onnx>=1.12.0") import onnx @@ -227,7 +359,31 @@ 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 YOLOv3 model to OpenVINO format, with optional INT8 quantization and inference metadata.""" + """ + Exports a YOLOv3 model to OpenVINO format with optional INT8 quantization and inference metadata. + + Args: + file (Path): Path to the output file. + metadata (dict): Inference metadata to include in the exported model. + half (bool): Indicates if FP16 precision should be used. + int8 (bool): Indicates if INT8 quantization should be applied. + data (str): Path to the dataset file (.yaml) for post-training quantization. + + Returns: + tuple[Path | None, openvino.runtime.Model | None]: Tuple containing the path to the exported model and + the OpenVINO model object, or None if the export failed. + + Notes: + - Requires the `openvino-dev>=2023.0` and optional `nncf>=2.4.0` package for INT8 quantization. + - Refer to OpenVINO documentation for further details: https://docs.openvino.ai/latest/index.html. + + Examples: + ```python + model_file = Path('/path/to/model.onnx') + metadata = {'names': ['class1', 'class2'], 'stride': 32} + export_openvino(model_file, metadata, half=True, int8=False, data='/path/to/dataset.yaml') + ``` + """ 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 @@ -294,8 +450,42 @@ def transform_fn(data_item): @try_export def export_paddle(model, im, file, metadata, prefix=colorstr("PaddlePaddle:")): - """Exports a YOLOv3 model to PaddlePaddle format using X2Paddle, saving to `file` directory with metadata in - yaml. + """ + Exports a YOLOv3 model to PaddlePaddle format using X2Paddle, saving to a specified directory and including model + metadata. + + Args: + model (torch.nn.Module): The YOLOv3 model to be exported. + im (torch.Tensor): A sample input tensor used for tracing the model. + file (pathlib.Path): Destination file path for the exported model, with `.pt` suffix. + metadata (dict): Additional metadata to be saved in YAML format alongside the exported model. + prefix (str, optional): Log message prefix. Defaults to a colored "PaddlePaddle:" string. + + Returns: + tuple: A tuple containing the directory path (str) where the PaddlePaddle model is saved, and `None`. + + Requirements: + - paddlepaddle: Install via `pip install paddlepaddle`. + - x2paddle: Install via `pip install x2paddle`. + + Notes: + The function first checks for required packages `paddlepaddle` and `x2paddle`. It then uses X2Paddle to trace + the model and export it to a PaddlePaddle format, saving the resulting files in the specified directory + with included metadata in a YAML file. + + Example: + ```python + from pathlib import Path + import torch + from models.yolo import DetectionModel + + model = DetectionModel() # Example model initialization + im = torch.rand(1, 3, 640, 640) # Example input tensor + file = Path("path/to/save/model.pt") + metadata = {"nc": 80, "names": ["class1", "class2", ...]} # Example metadata + + export_paddle(model, im, file, metadata) + ``` """ check_requirements(("paddlepaddle", "x2paddle")) import x2paddle @@ -311,7 +501,21 @@ 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 YOLOv3 model to CoreML format with optional quantization and NMS; requires 'coremltools'.""" + """ + Exports a YOLOv3 model to CoreML format with optional quantization and Non-Maximum Suppression (NMS). + + Args: + model (torch.nn.Module): The YOLOv3 model to be exported. + im (torch.Tensor): Input tensor used for tracing the model. + file (pathlib.Path): Destination file path where the CoreML model will be saved. + int8 (bool): Whether to use INT8 quantization or not. + half (bool): Whether to use FP16 quantization or not. + nms (bool): Whether to include Non-Maximum Suppression in the CoreML model or not. + prefix (str): Prefix string to add context to log messages. + + Returns: + str: Path to the saved CoreML model file (.mlmodel). + """ check_requirements("coremltools") import coremltools as ct @@ -337,9 +541,41 @@ def export_coreml(model, im, file, int8, half, nms, prefix=colorstr("CoreML:")): @try_export def export_engine(model, im, file, half, dynamic, simplify, workspace=4, verbose=False, prefix=colorstr("TensorRT:")): """ - Exports a YOLOv3 model to TensorRT engine format. - - Nvidia TensorRT: https://developer.nvidia.com/tensorrt + Exports a YOLOv3 model to TensorRT engine format, optimizing it for GPU inference. + + Args: + model (torch.nn.Module): The YOLOv3 model to be exported. + im (torch.Tensor): Sample input tensor used for tracing the model. + file (Path): File path where the exported TensorRT engine will be saved. + half (bool): Whether to use FP16 precision. Requires a supported GPU. + dynamic (bool): Whether to use dynamic input shapes. + simplify (bool): Whether to simplify the model during the ONNX export. + workspace (int): The maximum workspace size in GB. Default is 4. + verbose (bool): Whether to print detailed export logs. + prefix (str): Prefix string for log messages. Default is "TensorRT:". + + Returns: + tuple: A tuple containing the output file path (Path) and None. + + Raises: + AssertionError: If the model is running on CPU instead of GPU. + RuntimeError: If the ONNX file failed to load. + + Notes: + Nvidia TensorRT: https://developer.nvidia.com/tensorrt + + Example: + ```python + from pathlib import Path + import torch + + # Initialize model and dummy input + model = YOLOv3(...) # or another correct initialization + im = torch.randn(1, 3, 640, 640) + + # Export the model + export_engine(model, im, Path("yolov3.engine"), half=True, dynamic=True, simplify=True) + ``` """ assert im.device.type != "cpu", "export running on CPU but must be on GPU, i.e. `python export.py --device 0`" try: @@ -415,7 +651,48 @@ def export_saved_model( keras=False, prefix=colorstr("TensorFlow SavedModel:"), ): - """Exports YOLOv3 model to TensorFlow SavedModel format; includes NMS and configuration options.""" + """ + Exports a YOLOv3 model to TensorFlow SavedModel format. The function includes options for post-processing like Non- + Max Suppression (NMS) and others. + + Args: + model (torch.nn.Module): The PyTorch model to be exported. + im (torch.Tensor): Input tensor with sample data to trace the model. + file (Path): Path to save the exported model. + dynamic (bool): Whether to export the model with dynamic shapes. + tf_nms (bool, optional): Whether to include TensorFlow NMS in the exported model. Defaults to False. + agnostic_nms (bool, optional): Whether to use class-agnostic NMS. Defaults to False. + topk_per_class (int, optional): Number of top-K predictions to keep per class after NMS. Defaults to 100. + topk_all (int, optional): Number of top-K predictions to keep overall after NMS. Defaults to 100. + iou_thres (float, optional): IoU threshold for NMS. Defaults to 0.45. + conf_thres (float, optional): Confidence threshold for NMS. Defaults to 0.25. + keras (bool, optional): Whether to save as a Keras model, retaining the original graph. Defaults to False. + prefix (str, optional): String prefix for log messages, formatted with color. Defaults to "TensorFlow SavedModel:". + + Returns: + Tuple[str, None]: Path to the saved model and None (to maintain consistency with other export functions). + + Raises: + ImportError: If the necessary TensorFlow libraries are not installed. + + Note: + - Ensure required libraries are installed: + `pip install tensorflow` or `tensorflow-cpu` or `tensorflow-macos` depending on your environment. + - Refer to https://github.com/ultralytics/yolov5 for more details and usage examples. + + Example: + ```python + from pathlib import Path + from models.common import DetectMultiBackend + import torch + + model = DetectMultiBackend(weights='yolov5s.pt') + im = torch.zeros(1, 3, 640, 640) + file = Path("output/saved_model") + + export_saved_model(model, im, file, dynamic=True) + ``` + """ # YOLOv3 TensorFlow SavedModel export try: import tensorflow as tf @@ -461,9 +738,30 @@ def export_saved_model( @try_export def export_pb(keras_model, file, prefix=colorstr("TensorFlow GraphDef:")): """ - Exports Keras model as TensorFlow GraphDef *.pb file, compatible with YOLOv3. - - See https://github.com/leimao/Frozen_Graph_TensorFlow for details. + Exports a Keras model to TensorFlow GraphDef (*.pb) format, which is compatible with YOLOv3. + + Args: + keras_model (tf.keras.Model): The trained Keras model to be exported. + file (pathlib.Path): The target file path for saving the exported model. + prefix (str, optional): Prefix string for logging. Defaults to colorstr("TensorFlow GraphDef:"). + + Returns: + tuple[pathlib.Path, None]: The file path where the model is saved and None. + + Example: + ```python + from tensorflow.keras.models import load_model + from pathlib import Path + export_pb(load_model('model.h5'), Path('model.pb')) + ``` + + See Also: + For more details on TensorFlow GraphDef, visit + https://github.com/leimao/Frozen_Graph_TensorFlow. + + Notes: + Ensure TensorFlow is properly installed in your environment as it is required for this function to execute. + TensorFlow's version should be compatible with the version used to train your model to avoid any compatibility issues. """ import tensorflow as tf from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2 @@ -481,7 +779,47 @@ def export_pb(keras_model, file, prefix=colorstr("TensorFlow GraphDef:")): @try_export def export_tflite(keras_model, im, file, int8, data, nms, agnostic_nms, prefix=colorstr("TensorFlow Lite:")): - """Exports a YOLOv3 model to TensorFlow Lite format, supports fp16/int8, NMS options, and custom data.""" + """ + Export a YOLOv3 model to TensorFlow Lite format, supporting fp16/int8, NMS options, and custom data. + + Args: + keras_model (tf.keras.Model): The Keras model to be exported. + im (torch.Tensor): The input tensor used for model tracing. + file (Path): The file path where the exported model will be saved. + int8 (bool): Flag to enable INT8 quantization for the model. + data (str): Path to the dataset YAML file for representative data generation. + nms (bool): Flag to include Non-Max Suppression (NMS) in the exported model. + agnostic_nms (bool): Flag to apply class-agnostic NMS. + prefix (str, optional): Prefix string for logging. Defaults to colorstr("TensorFlow Lite:"). + + Returns: + (str, NoneType): The file path of the saved TensorFlow Lite model and None. + + Notes: + - TensorFlow dependency is required for model export. + - INT8 quantization requires a representative dataset for calibration. + - Non-Max Suppression (NMS) can be optionally included in the exported model. + + Examples: + ```python + import torch + from pathlib import Path + from models.experimental import attempt_load + + # Load and prepare model + model = attempt_load('yolov5s.pt', map_location='cpu') + im = torch.zeros(1, 3, 640, 640) # Dummy input tensor + + # Export model + export_tflite(model, im, Path('yolov5s'), int8=False, data=None, nms=True, agnostic_nms=False) + ``` + + TensorFlow Lite Developer Guide: + https://www.tensorflow.org/lite/guide + + Model Conversion Reference: + https://github.com/leimao/Frozen_Graph_TensorFlow + """ import tensorflow as tf LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...") @@ -514,9 +852,35 @@ def export_tflite(keras_model, im, file, int8, data, nms, agnostic_nms, prefix=c @try_export def export_edgetpu(file, prefix=colorstr("Edge TPU:")): """ - Exports YOLOv3 model to Edge TPU compatible format; requires Linux and Edge TPU compiler. - - https://coral.ai/docs/edgetpu/compiler/ + Exports YOLOv3 model to Edge TPU compatible format. + + Args: + file (Path): The file path for the model to be exported, with a `.pt` suffix. + prefix (str): A prefix used for logging, default is 'Edge TPU:'. + + Returns: + Tuple[Path | None, None]: The tuple contains the file path of the exported model with `-int8_edgetpu.tflite` suffix + if successful, otherwise returns `None`. + + Raises: + AssertionError: If the export is not executed on a Linux system. + subprocess.CalledProcessError: If there are issues with subprocess execution, particularly around Edge TPU compiler + installation or execution. + + Notes: + This function requires Linux and the Edge TPU compiler. If the compiler is not installed, it attempts to install it. + For more details, refer to the Edge TPU compiler documentation: + https://coral.ai/docs/edgetpu/compiler/ + + Example: + ```python + from pathlib import Path + from ultralytics import export_edgetpu + + model_file = Path('yolov5s.pt') + exported_model, _ = export_edgetpu(model_file) + print(f"Model exported to {exported_model}") + ``` """ cmd = "edgetpu_compiler --version" help_url = "https://coral.ai/docs/edgetpu/compiler/" @@ -555,7 +919,48 @@ def export_edgetpu(file, prefix=colorstr("Edge TPU:")): @try_export def export_tfjs(file, int8, prefix=colorstr("TensorFlow.js:")): - """Exports a YOLOv3 model to TensorFlow.js format, optionally quantizing to uint8.""" + """ + Exports a YOLOv3 model to TensorFlow.js format, with an optional quantization to uint8. + + Args: + file (Path): The path to the model file to be exported. + int8 (bool): Boolean flag to determine if the model should be quantized to uint8. + prefix (str): String prefix for logging, by default "TensorFlow.js". + + Returns: + tuple: + str: Directory path where the TensorFlow.js files are saved. + None: Placeholder to match the expected return type from 'try_export' decorator. + + Raises: + ImportError: If the required 'tensorflowjs' package is not installed. + + Requirements: + - tensorflowjs: Install with `pip install tensorflowjs` + + Example: + ```python + from pathlib import Path + export_tfjs(file=Path("yolov5s.pt"), int8=False) + ``` + + Note: + Ensure that you have TensorFlow.js installed in your environment. For more details, visit the official + TensorFlow.js documentation at https://www.tensorflow.org/js. + + Usage: + The converted model can be used directly in JavaScript environments using the TensorFlow.js library. + + For usage in web applications: + - Clone the example repository: + $ cd .. && git clone https://github.com/zldrobit/tfjs-yolov5-example.git && cd tfjs-yolov5-example + - Install dependencies: + $ npm install + - Create a symbolic link to the exported web model: + $ ln -s ../../yolov5/yolov5s_web_model public/yolov5s_web_model + - Start the example application: + $ npm star + """ check_requirements("tensorflowjs") import tensorflowjs as tfjs @@ -592,7 +997,32 @@ def export_tfjs(file, int8, prefix=colorstr("TensorFlow.js:")): def add_tflite_metadata(file, metadata, num_outputs): - """Adds TFLite model metadata for enhanced model understanding and processing; requires `tflite_support`.""" + """ + Adds metadata to a TensorFlow Lite model to enhance model understanding and processing using `tflite_support`. + + Args: + file (str): The path to the TensorFlow Lite model file. + metadata (dict): A dictionary containing metadata information such as labels, input/output descriptions, etc. + num_outputs (int): The number of output tensors in the model. + + Returns: + None. + + Raises: + ImportError: If the `tflite_support` library is not available. + + Note: + This function uses the `tflite_support` library to add the metadata. The library must be installed separately. + + Example: + ```python + metadata = { + "input": {"description": "Input image tensor"}, + "output": [{"name": "scores", "description": "Detection scores"}], + } + add_tflite_metadata("/path/to/model.tflite", metadata, num_outputs=1) + ``` + """ with contextlib.suppress(ImportError): # check_requirements('tflite_support') from tflite_support import flatbuffers @@ -625,8 +1055,34 @@ def add_tflite_metadata(file, metadata, num_outputs): def pipeline_coreml(model, im, file, names, y, prefix=colorstr("CoreML Pipeline:")): - """Executes YOLOv3 CoreML pipeline, handling image preprocessing, prediction, and NMS, saving the model with - metadata. + """ + Executes YOLOv3 CoreML pipeline, handling image preprocessing, prediction, and NMS, saving the model with metadata. + + Args: + model (coremltools.models.MLModel): The CoreML model to be used for the prediction. + im (torch.Tensor): Input tensor in BCHW (Batch, Channel, Height, Width) format. + file (pathlib.Path): The file path where the exported CoreML model will be saved. + names (dict): Dictionary mapping class indices to class names. + y (torch.Tensor): YOLO model output tensor. + prefix (str): Optional prefix for logging messages; default is 'CoreML Pipeline:'. + + Returns: + (pathlib.Path | None): The file path of the saved CoreML model if successful; None otherwise. + + Example: + ```python + from pathlib import Path + import torch + from coremltools.models import MLModel + + model = MLModel('path/to/pretrained/model.mlmodel') + im = torch.randn(1, 3, 640, 640) # Example input tensor + file = Path('path/to/save/model.mlmodel') + names = {0: 'class0', 1: 'class1'} + y = torch.randn(1, 25200, 85) # Example output tensor + + pipeline_coreml(model, im, file, names, y) + ``` """ import coremltools as ct from PIL import Image @@ -782,7 +1238,55 @@ def run( iou_thres=0.45, # TF.js NMS: IoU threshold conf_thres=0.25, # TF.js NMS: confidence threshold ): - """Exports a PyTorch model to specified formats like ONNX, CoreML, TensorRT.""" + """ + Export a PyTorch model to various formats like ONNX, CoreML, and TensorRT. + + Args: + data (str | Path): Path to dataset configuration file. + weights (str | Path): Path to model weights file in PyTorch format. + imgsz (tuple[int, int]): Tuple specifying image height and width for input dimensions. + batch_size (int): Batch size for model inference. + device (str): Device to use for inference (e.g., '0', '0,1,2,3', 'cpu'). + include (tuple[str]): Formats to include for model export (e.g., 'torchscript', 'onnx', etc.). + half (bool): Whether to export model with FP16 precision. + inplace (bool): Set YOLOv3 Detect module inplace option to True. + keras (bool): Save Keras model when exporting TensorFlow SavedModel format. + optimize (bool): Optimize the TorchScript model for mobile inference. + int8 (bool): Apply INT8 quantization for CoreML/TF models. + dynamic (bool): Enable dynamic axes for ONNX/TF/TensorRT models. + simplify (bool): Simplify the ONNX model after export. + opset (int): ONNX opset version. + verbose (bool): Enable verbose logging for TensorRT engine export. + workspace (int): Workspace size in GB for TensorRT engine. + nms (bool): Enable Non-Maximum Suppression (NMS) in TensorFlow models. + agnostic_nms (bool): Enable class-agnostic NMS in TensorFlow models. + topk_per_class (int): Top-K per class to keep in TensorFlow JSON model. + topk_all (int): Top-K for all classes to keep in TensorFlow JSON model. + iou_thres (float): IOU threshold for TensorFlow JSON model. + conf_thres (float): Confidence threshold for TensorFlow JSON model. + + Returns: + None + + Notes: + - Requires various packages installed for different export formats, e.g., `onnx`, `coremltools`, etc. + - Some formats have additional dependencies (e.g., TensorFlow, TensorRT, etc.) + + Examples: + ```python + run( + data='data/coco128.yaml', + weights='yolov5s.pt', + imgsz=(640, 640), + batch_size=1, + device='cpu', + include=('torchscript', 'onnx'), + half=False, + dynamic=True, + opset=12 + ) + ``` + """ t = time.time() include = [x.lower() for x in include] # to lowercase fmts = tuple(export_formats()["Argument"][1:]) # --include arguments @@ -893,8 +1397,19 @@ def run( def parse_opt(known=False): - """Parses command line arguments for model export configuration, including data paths, weights, and export - options. + """ + Parses command line arguments for model export configuration, allowing users to specify various export options. + + Args: + known (bool): If True, parse only known arguments and ignore others. Default is False. + + Returns: + argparse.Namespace: An object containing all the configured arguments and their values. + + Notes: + This function uses `argparse` to define and process command line arguments. It supports multiple model export formats + including TorchScript, ONNX, OpenVINO, TensorRT, and others. Users can set various options like image size, batch size, + device, and thresholds for NMS. """ parser = argparse.ArgumentParser() parser.add_argument("--data", type=str, default=ROOT / "data/coco128.yaml", help="dataset.yaml path") @@ -930,7 +1445,43 @@ def parse_opt(known=False): def main(opt): - """Executes model inference, supporting various export formats based on specified weights and options.""" + """ + Main function for exporting a YOLOv3 PyTorch model to various formats. + + This function parses command line arguments to configure export options and then proceeds to export + the YOLOv3 model to specified formats such as TorchScript, ONNX, CoreML, TensorRT, and more. + + Args: + opt (argparse.Namespace): Parsed command line arguments containing export configurations such as: + data (str): Path to dataset YAML file. + weights (int | str): Path(s) to model weights file(s). + imgsz (tuple[int, int]): Image dimensions (height, width). + batch_size (int): Batch size for export. + device (str): Device to use for export, e.g., 'cpu' or '0' for GPU. + include (list[int | str]): List of formats to export, e.g., ['torchscript', 'onnx']. + half (bool): If True, export model in FP16 half-precision. + inplace (bool): If True, set YOLOv3 Detect() inplace=True. + keras (bool): If True, use Keras for TensorFlow SavedModel export. + optimize (bool): If True, optimize TorchScript model for mobile. + int8 (bool): If True, enable INT8 quantization for CoreML and TensorFlow. + dynamic (bool): If True, enable dynamic axes for ONNX, TensorRT, and TensorFlow. + simplify (bool): If True, simplify ONNX model using onnx-simplifier. + opset (int): ONNX opset version to use. + verbose (bool): If True, enable verbose logging for TensorRT export. + workspace (int): Workspace size in GB for TensorRT export. + nms (bool): If True, add Non-Max Suppression (NMS) to TensorFlow model. + agnostic_nms (bool): If True, add class-agnostic NMS to TensorFlow model. + topk_per_class (int): Top-k per class to keep for TensorFlow.js NMS. + topk_all (int): Top-k for all classes to keep for TensorFlow.js NMS. + iou_thres (float): IoU threshold for TensorFlow.js NMS. + conf_thres (float): Confidence threshold for TensorFlow.js NMS. + + Returns: + list[str]: List of paths to the exported model(s). + + Example: + $ python export.py --weights yolov5s.pt --include torchscript onnx openvino engine coreml tflite ... + """ 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 89c47a0721..d701c2d4c7 100644 --- a/hubconf.py +++ b/hubconf.py @@ -15,19 +15,31 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): """ - Creates or loads a YOLOv3 model. - - 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 YOLOv3 .autoshape() wrapper to model - verbose (bool): print all information to screen - device (str, torch.device, None): device to use for model parameters + Creates or loads a YOLOv3 model with specified configurations and optional pretrained weights. + + Args: + name (str): Model name such as 'yolov5s' or a path to a model checkpoint file, e.g., 'path/to/best.pt'. + pretrained (bool): Whether to 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): Whether to apply the YOLOv3 .autoshape() wrapper to the model for handling multiple input types. + Default is True. + verbose (bool): If True, print all information to the screen. Default is True. + device (str | torch.device | None): Device to use for model parameters ('cpu', 'cuda', etc.). If None, defaults to + the best available device. Returns: - YOLOv3 model + torch.nn.Module: YOLOv3 model loaded with or without pretrained weights. + + Example: + ```python + import torch + model = _create('yolov5s') + ``` + + Raises: + Exception: If an error occurs while loading the model, returns an error message with a helpful URL: + "https://docs.ultralytics.com/yolov5/tutorials/pytorch_hub_model_loading". """ from pathlib import Path @@ -84,71 +96,325 @@ 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 YOLOv3 model from a specified path, with options for autoshaping and device - assignment. + """ + Loads a custom or local YOLOv3 model from a specified path, with options for autoshaping and device assignment. + + Args: + path (str): Path to the model file. Supports both local and URL paths. + autoshape (bool): If True, applies the YOLOv3 `.autoshape()` wrapper to allow for various input formats. + _verbose (bool): If True, outputs detailed information. Otherwise, limits verbosity. + device (str | torch.device | None): Device to load the model on. Default is None, which uses the available GPU if + possible. + + Returns: + torch.nn.Module: The loaded YOLOv3 model, either with or without autoshaping applied. + + Raises: + Exception: If the model loading fails due to invalid path or incompatible model state, with helpful suggestions + including a reference to the troubleshooting page: + https://docs.ultralytics.com/yolov5/tutorials/pytorch_hub_model_loading + + Examples: + ```python + import torch + model = torch.hub.load('ultralytics/yolov5', 'custom', 'path/to/best.pt') + model = torch.hub.load('ultralytics/yolov5', 'custom', 'path/to/best.pt', 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 YOLOv5n model; optional pretrained weights, class/channel count, autoshaping, device selection.""" + """ + Instantiates a YOLOv5n model with optional pretrained weights, configurable input channels, classes, autoshaping, + and device selection. + + 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 detection classes. Defaults to 80. + autoshape (bool): If True, applies YOLOv5 .autoshape() wrapper to the model for various formats (file/URI/PIL/cv2/np) + and non-maximum suppression (NMS). Defaults to True. + _verbose (bool): If True, prints detailed information to the screen. Defaults to True. + device (str | torch.device | None): Device to use for model computations (e.g., 'cpu', 'cuda'). If None, the best + available device is automatically selected. Defaults to None. + + Returns: + torch.nn.Module: The instantiated YOLOv5n model. + + Example: + ```python + import torch + model = torch.hub.load('ultralytics/yolov5', 'yolov5n') # using official model + model = torch.hub.load('ultralytics/yolov5:master', 'yolov5n') # from specific branch + model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5n.pt') # using custom/local model + model = torch.hub.load('.', 'custom', 'yolov5n.pt', source='local') # from local repository + ``` + """ return _create("yolov5n", pretrained, channels, classes, autoshape, _verbose, device) def yolov5s(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None): - """Loads YOLOv5s model, with options for pretrained weights, channel/class customization, autoshaping, and device - choice. + """ + Loads the YOLOv5s model with customizable options for pretrained weights, input channels, number of classes, + autoshape functionality, and device selection. + + Args: + pretrained (bool, optional): If True, loads model with pretrained weights. Default is True. + channels (int, optional): Specifies the number of input channels. Default is 3. + classes (int, optional): Defines the number of model classes. Default is 80. + autoshape (bool, optional): Applies YOLOv5 .autoshape() wrapper to the model for enhanced usability. Default is + True. + _verbose (bool, optional): If True, prints detailed information during model loading. Default is True. + device (str | torch.device | None, optional): Specifies the device to load the model on. Accepts 'cpu', 'cuda', or + torch.device. Default is None, which automatically selects the best available option. + + Returns: + torch.nn.Module: The initialized YOLOv5s model loaded with the specified options. + + Example: + ```python + import torch + model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True) + ``` + + For more information, refer to [PyTorch Hub models](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): - """Loads YOLOv5m model with optional pretrained weights, channel/class customization, autoshaping, device + """ + Loads YOLOv5m model with options for pretrained weights, channel/class customization, autoshaping, and device selection. + + Args: + pretrained (bool, optional): If True, loads pretrained weights into the model. Default is True. + channels (int, optional): Number of input channels for the model. Default is 3. + classes (int, optional): Number of classes for the model. Default is 80. + autoshape (bool, optional): If True, applies the YOLOv3 .autoshape() wrapper to the model for convenient input + handling. Default is True. + _verbose (bool, optional): If True, prints all information to the screen. Default is True. + device (str | torch.device | None, optional): Device to use for model parameters. Default is None, which + automatically selects the available device. + + Returns: + torch.nn.Module: YOLOv5m model instance. + + Example: + ```python + import torch + model = yolov5m(pretrained=True, channels=3, classes=80, autoshape=True, device='cuda:0') + ``` """ return _create("yolov5m", pretrained, channels, classes, autoshape, _verbose, device) def yolov5l(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None): - """Loads YOLOv5l model with options for pretrained weights, channel/class customization, autoshaping, and device - selection. + """ + Loads the YOLOv5l model with options for pretrained weights, channel and class customization, autoshaping, and + device selection. + + Args: + pretrained (bool, optional): If True, loads the model with pretrained weights. 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): If True, applies YOLOv5's .autoshape() wrapper to the model. Defaults to True. + _verbose (bool, optional): If True, prints detailed information during model loading. Defaults to True. + device (str | torch.device | None, optional): Device to use for model parameters, e.g., 'cpu', 'cuda:0', + or torch.device object. If None, automatically selects the appropriate device. Defaults to None. + + Returns: + torch.nn.Module: The initialized YOLOv5l model. + + Example: + ```python + import torch + model = torch.hub.load('ultralytics/yolov5', 'yolov5l', pretrained=True) + ``` """ return _create("yolov5l", pretrained, channels, classes, autoshape, _verbose, device) def yolov5x(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None): - """Loads YOLOv5x model with customization options for weights, channels, classes, autoshaping, and device - selection. + """ + Loads YOLOv5x model with customization options for weights, channels, classes, autoshaping, and device selection. + + Args: + pretrained (bool, optional): Whether to load pretrained weights. 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 YOLOv3 .autoshape() wrapper to the model. Defaults to True. + _verbose (bool, optional): Whether to print all information to the screen. Defaults to True. + device (str | torch.device | None, optional): Device to use for model parameters, e.g., 'cpu', 'cuda', or + torch.device(). If None, default device will be selected. Defaults to None. + + Returns: + torch.nn.Module: The YOLOv5x model. + + Examples: + ```python + import torch + + # Load YOLOv5x model with default settings + model = torch.hub.load('ultralytics/yolov5', 'yolov5x') + + # Load YOLOv5x model with custom device + model = torch.hub.load('ultralytics/yolov5', 'yolov5x', device='cuda:0') + ``` """ return _create("yolov5x", pretrained, channels, classes, autoshape, _verbose, device) def yolov5n6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None): - """Loads the YOLOv5n6 model with options for weights, channels, classes, shaping, and device.""" + """ + Loads the YOLOv5n6 model with options for pretrained weights, input channels, classes, autoshaping, verbosity, and + device assignment. + + 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 YOLOv3 .autoshape() wrapper to the model. Default is True. + _verbose (bool, optional): If True, prints all information to the screen. Default is True. + device (str | torch.device | None, optional): Device to use for model parameters, e.g., 'cpu', '0', or torch.device. + Default is None. + + Returns: + torch.nn.Module: YOLOv5n6 model loaded on the specified device and configured as per the provided options. + + Notes: + For more information on PyTorch Hub models, refer to: https://pytorch.org/hub/ultralytics_yolov5 + + Example: + ```python + model = yolov5n6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device='cuda') + ``` + """ return _create("yolov5n6", pretrained, channels, classes, autoshape, _verbose, device) def yolov5s6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None): - """Loads the YOLOv5s6 model; see options for weights, channels, classes, - at https://github.com/ultralytics/yolov5.""" + """ + Loads the YOLOv5s6 model with options for weights, channels, classes, autoshaping, and device selection. + + Args: + pretrained (bool, optional): If True, loads 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): Apply YOLOv5 .autoshape() wrapper to model. Defaults to True. + _verbose (bool, optional): If True, prints detailed information to the screen. Defaults to True. + device (str | torch.device | None, optional): Device to use for model parameters, e.g., 'cpu', 'cuda:0'. + If None, it will select the appropriate device automatically. Defaults to None. + + Returns: + torch.nn.Module: The YOLOv5s6 model, ready for inference or further training. + + Example: + ```python + import torch + model = torch.hub.load('ultralytics/yolov5', 'yolov5s6', pretrained=True, channels=3, classes=80) + model.eval() # Set the model to evaluation mode + ``` + + For more details, see the official documentation at: + https://github.com/ultralytics/yolov5 + """ return _create("yolov5s6", pretrained, channels, classes, autoshape, _verbose, device) def yolov5m6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None): - """Loads YOLOv5m6 model; options for weights, channels, classes documented at - https://github.com/ultralytics/yolov5.""" + """ + Loads YOLOv5m6 model with options for pretrained weights, input channels, number of classes, autoshaping, and device + selection. + + Args: + pretrained (bool): Whether to 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): Whether to apply YOLOv3 .autoshape() wrapper to the model. Default is True. + _verbose (bool): Whether to print all information to the screen. Default is True. + device (str | torch.device | None): Device to use for model parameters, e.g., 'cpu', 'cuda', 'mps', or torch device. + Default is None. + + Returns: + YOLOv5m6 model (torch.nn.Module): The instantiated YOLOv5m6 model with specified options. + + Example: + ```python + import torch + model = torch.hub.load('ultralytics/yolov5', 'yolov5m6') + # Load custom YOLOv5m6 model from a local path with specific options + model = torch.hub.load('.', 'yolov5m6', pretrained=False, channels=1, classes=10, device='cuda') + ``` + + Notes: + For more detailed documentation, visit https://github.com/ultralytics/yolov5 + """ return _create("yolov5m6", pretrained, channels, classes, autoshape, _verbose, device) def yolov5l6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None): - """Loads YOLOv5l6 model with customization for pretraining, channels, classes; details at - https://github.com/ultralytics/yolov5.""" + """ + Loads the YOLOv5l6 model with options for pretrained weights, input channels, the number of classes, autoshaping, + and device selection. + + 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 for automatic shape + inference. Default is True. + _verbose (bool, optional): If True, prints all information to the screen. Default is True. + device (str | torch.device | None, optional): Device to use for the model parameters, e.g., 'cpu', 'cuda', or + a specific GPU like 'cuda:0'. Default is None, which means the best available device will be selected + automatically. + + Returns: + yolov5.models.yolo.DetectionModel: YOLOv5l6 model initialized with defined custom configurations. + + Examples: + ```python + import torch + model = torch.hub.load('ultralytics/yolov5', 'yolov5l6') # Load YOLOv5l6 model + ``` + + Note: + For more details, visit the [Ultralytics YOLOv5 GitHub repository](https://github.com/ultralytics/yolov5). + """ return _create("yolov5l6", pretrained, channels, classes, autoshape, _verbose, device) def yolov5x6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None): - """Loads the YOLOv5x6 model, allowing customization for channels and classes; more at - https://github.com/ultralytics/yolov5.""" + """ + Loads the YOLOv5x6 model, allowing customization for pretrained weights, input channels, and model classes. + + Args: + pretrained (bool): If True, loads the model with pretrained weights. Default is True. + channels (int): Number of input channels. Default is 3. + classes (int): Number of output classes for the model. Default is 80. + autoshape (bool): If True, applies the .autoshape() wrapper for inference on diverse input formats. Default is True. + _verbose (bool): If True, prints detailed information during model loading. Default is True. + device (str | torch.device | None): Specifies the device to load the model on ('cpu', 'cuda', etc.). Default is None, + which uses the best available device. + + Returns: + torch.nn.Module: The YOLOv5x6 model with the specified configurations. + + Example: + ```python + from ultralytics import yolov5x6 + + # Load the model with default settings + model = yolov5x6() + + # Load the model with custom configurations + model = yolov5x6(pretrained=False, channels=1, classes=10, autoshape=False, device='cuda') + ``` + + Notes: + For more information, refer to the YOLOv5 repository: https://github.com/ultralytics/yolov5 + """ return _create("yolov5x6", pretrained, channels, classes, autoshape, _verbose, device) diff --git a/train.py b/train.py index 708fc9670b..cbdce77f6f 100644 --- a/train.py +++ b/train.py @@ -101,8 +101,52 @@ def train(hyp, opt, device, callbacks): # hyp is path/to/hyp.yaml or hyp dictionary - """Trains a model with given hyperparameters, options, device, and callbacks, saving results and weights as - specified. + """ + Train a YOLOv3 model on a custom dataset and manage the training process. + + Args: + hyp (str | dict): Path to hyperparameters yaml file or hyperparameters dictionary. + opt (argparse.Namespace): Parsed command line arguments containing training options. + device (torch.device): Device to load and train the model on. + callbacks (Callbacks): Callbacks to handle various stages of the training lifecycle. + + Returns: + None + + Usage - Single-GPU training: + $ 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 + + Usage - Multi-GPU DDP training: + $ 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 + + 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 + + Examples: + ```python + from ultralytics import train + import argparse + import torch + from utils.callbacks import Callbacks + + # Example usage + args = argparse.Namespace( + data='coco128.yaml', + weights='yolov5s.pt', + cfg='yolov5s.yaml', + img_size=640, + epochs=50, + batch_size=16, + device='0' + ) + + device = torch.device(f'cuda:{args.device}' if torch.cuda.is_available() else 'cpu') + callbacks = Callbacks() + + train(hyp='hyp.scratch.yaml', opt=args, device=device, callbacks=callbacks) + ``` """ save_dir, epochs, batch_size, weights, single_cls, evolve, data, cfg, resume, noval, nosave, workers, freeze = ( Path(opt.save_dir), @@ -499,7 +543,30 @@ def lf(x): def parse_opt(known=False): - """Parses command line arguments for training configurations, with facilities for known arg parsing.""" + """ + Parses command line arguments for configuring the training of a YOLO model. + + Args: + known (bool): Flag to parse known arguments only, defaults to False. + + Returns: + argparse.Namespace: Parsed command line arguments. + + Examples: + ```python + options = parse_opt() + print(options.weights) + ``` + + Notes: + * The default weights path is 'yolov3-tiny.pt'. + * Set `known` to True for parsing only the known arguments, useful for partial arguments. + + References: + * Models: https://github.com/ultralytics/yolov5/tree/master/models + * Datasets: https://github.com/ultralytics/yolov5/tree/master/data + * Training Tutorial: https://docs.ultralytics.com/yolov5/tutorials/train_custom_data + """ parser = argparse.ArgumentParser() parser.add_argument("--weights", type=str, default=ROOT / "yolov3-tiny.pt", help="initial weights path") parser.add_argument("--cfg", type=str, default="", help="model.yaml path") @@ -546,7 +613,39 @@ def parse_opt(known=False): def main(opt, callbacks=Callbacks()): - """Main training/evolution script handling model checks, DDP setup, training, and hyperparameter evolution.""" + """ + Main training/evolution script handling model checks, DDP setup, training, and hyperparameter evolution. + + Args: + opt (argparse.Namespace): Parsed command-line options. + callbacks (Callbacks, optional): Callback object for handling training events. Defaults to Callbacks(). + + Returns: + None + + Raises: + AssertionError: If certain constraints are violated (e.g., when specific options are incompatible with DDP training). + + Notes: + - For tutorial on using Multi-GPU with DDP: https://docs.ultralytics.com/yolov5/tutorials/multi_gpu_training + + Example: + Single-GPU training: + ```python + $ 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: + ```python + $ 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 + ``` + + 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 + """ if RANK in {-1, 0}: print_args(vars(opt)) check_git_status() @@ -712,8 +811,39 @@ def main(opt, callbacks=Callbacks()): def run(**kwargs): - """Executes model training with specified configurations; see example: import train; train.run(data='coco128.yaml', - imgsz=320, weights='yolov5m.pt'). + """ + Executes model training with specified configurations. + + Args: + kwargs (dict): Configuration parameters for the training process. Supported parameters include but are not limited to: + - data (str): Path to the dataset YAML file. + - weights (str): Path to initial weights file. + - imgsz (int): Image size for training and validation. + - epochs (int): Number of training epochs. + - batch_size (int): Total batch size for all GPUs. + + Returns: + None + + Notes: + - Ensure that the provided dataset YAML file and the initial weights path are accessible. + - For multi-GPU training, use appropriate distributed training settings. + - Example usage: + ```python + run(data='coco128.yaml', imgsz=320, weights='yolov5m.pt') + ``` + + Example: + To execute model training with custom configurations: + ```python + from train import run + run(data='coco128.yaml', imgsz=320, weights='yolov5m.pt') + ``` + + References: + - 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 """ opt = parse_opt(True) for k, v in kwargs.items(): diff --git a/val.py b/val.py index 8d6c06aa68..775282d50d 100644 --- a/val.py +++ b/val.py @@ -62,8 +62,38 @@ def save_one_txt(predn, save_conf, shape, file): - """Saves detection results in txt format; includes labels and optionally confidence scores if `save_conf` is - True. + """ + Saves detection results in a text format, including labels and optionally confidence scores. + + Args: + predn (torch.Tensor): A tensor containing normalized prediction results in the format (x1, y1, x2, y2, conf, cls). + save_conf (bool): A flag indicating whether to save confidence scores. + shape (tuple[int, int]): Original image shape in the format (height, width). + file (str | Path): Path to the file where the results will be saved. + + Returns: + None + + Example: + ```python + from pathlib import Path + import torch + + predn = torch.tensor([ + [10, 20, 100, 200, 0.9, 1], + [30, 40, 150, 250, 0.8, 0], + ]) + save_conf = True + shape = (416, 416) + file = Path("results.txt") + + save_one_txt(predn, save_conf, shape, file) + ``` + + Notes: + - The function normalizes bounding box coordinates before saving. + - Each line in the output file will contain class, x-center, y-center, width, height and optionally confidence score. + - The format is compatible with YOLO training dataset format. """ gn = torch.tensor(shape)[[1, 0, 1, 0]] # normalization gain whwh for *xyxy, conf, cls in predn.tolist(): @@ -74,7 +104,28 @@ def save_one_txt(predn, save_conf, shape, file): def save_one_json(predn, jdict, path, class_map): - """Saves detection results in JSON format containing image_id, category_id, bbox, and score per detection.""" + """ + Saves detection results in JSON format containing image_id, category_id, bbox, and score per detection. + + Args: + predn (torch.Tensor): Normalized prediction tensor of shape (N, 6) where N is the number of detections. + Each detection should contain (x1, y1, x2, y2, confidence, class). + jdict (list): List to store the JSON serializable detections. + path (Path): Path object representing the image file path. + class_map (dict[int, int]): Dictionary mapping class indices to their respective category IDs. + + Returns: + None + + Example: + ```python + predn = torch.tensor([[50, 30, 200, 150, 0.9, 0], [30, 20, 180, 150, 0.8, 1]]) + jdict = [] + path = Path('images/000001.jpg') + class_map = {0: 1, 1: 2} + save_one_json(predn, jdict, path, class_map) + ``` + """ image_id = int(path.stem) if path.stem.isnumeric() else path.stem box = xyxy2xywh(predn[:, :4]) # xywh box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner @@ -91,12 +142,31 @@ def save_one_json(predn, jdict, path, class_map): def process_batch(detections, labels, iouv): """ - Return correct prediction matrix - Arguments: - detections (array[N, 6]), x1, y1, x2, y2, conf, class - labels (array[M, 5]), class, x1, y1, x2, y2 + Computes correct prediction matrix for detections against ground truth labels at various IoU thresholds. + + Args: + detections (np.ndarray): Array of detections with shape (N, 6), where each detection contains [x1, y1, x2, y2, + confidence, class]. + labels (np.ndarray): Array of ground truth labels with shape (M, 5), where each label contains [class, x1, y1, x2, y2]. + iouv (np.ndarray): Array of IoU thresholds to use for evaluation. + Returns: - correct (array[N, 10]), for 10 IoU levels + np.ndarray: Boolean array of shape (N, len(iouv)), indicating correct predictions at each IoU threshold. + + Notes: + - This function compares detections and ground truth labels to establish matches based on IoU and class. + - It supports multiple IoU thresholds to evaluate prediction accuracy flexibly. + + Example: + ```python + detections = np.array([[50, 50, 150, 150, 0.8, 0], + [30, 30, 120, 120, 0.7, 1]]) + labels = np.array([[0, 50, 50, 150, 150], + [1, 30, 30, 120, 120]]) + iouv = np.array([0.5, 0.6, 0.7]) + + correct = process_batch(detections, labels, iouv) + ``` """ correct = np.zeros((detections.shape[0], iouv.shape[0])).astype(bool) iou = box_iou(labels[:, 1:], detections[:, :4]) @@ -145,7 +215,45 @@ def run( callbacks=Callbacks(), compute_loss=None, ): - """Validates a trained YOLO model on a dataset and saves detection results in specified formats.""" + """ + Validates a trained YOLO model on a dataset and saves detection results in specified formats. + + Parameters: + data (str | dict): Path to the dataset configuration file (.yaml) or a dictionary containing the dataset paths. + weights (str | list, optional): Path to the trained model weights file(s). Default is None. + batch_size (int, optional): Batch size for inference. Default is 32. + imgsz (int, optional): Input image size for inference in 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, can be 'train', 'val', 'test', 'speed', or 'study'. Default is 'val'. + device (str, optional): Device for computation, e.g., '0' for GPU or 'cpu' for CPU. Default is "". + workers (int, optional): Number of dataloader workers. Default is 8. + single_cls (bool, optional): Whether to treat the dataset as a single-class dataset. Default is False. + augment (bool, optional): Whether to apply augmented inference. Default is False. + verbose (bool, optional): Whether to output verbose information. Default is False. + save_txt (bool, optional): Whether to save detection results in text format (*.txt). Default is False. + save_hybrid (bool, optional): Whether to save hybrid results (labels+predictions) in text format (*.txt). Default is False. + save_conf (bool, optional): Whether to save confidence scores in text format labels. Default is False. + save_json (bool, optional): Whether to save detection results in COCO JSON format. Default is False. + project (str | Path, optional): Directory path to save validation results. Default is ROOT / 'runs/val'. + name (str, optional): Directory name to save validation results. Default is 'exp'. + exist_ok (bool, optional): Whether to overwrite existing project/name directory. Default is False. + half (bool, optional): Whether to use half-precision (FP16) for inference. Default is True. + dnn (bool, optional): Whether to use OpenCV DNN for ONNX inference. Default is False. + model (nn.Module, optional): Existing model instance. Default is None. + dataloader (DataLoader, optional): Existing dataloader instance. Default is None. + save_dir (Path, optional): Path to directory to save results. Default is Path(""). + plots (bool, optional): Whether to generate plots for visual results. Default is True. + callbacks (Callbacks, optional): Callbacks instance for event handling. Default is Callbacks(). + compute_loss (Callable, optional): Loss function for computing training loss. Default is None. + + Returns: + Tuple[torch.Tensor, dict, dict, torch.Tensor]: + - metrics: Dictionary containing metrics such as precision, recall, mAP, F1 score, etc. + - time: Dictionary containing times for different parts of the pipeline (e.g., preprocessing, inference, NMS) + - samples: Torch tensor containing validation samples. + """ # Initialize/load model and set device training = model is not None if training: # called by train.py @@ -360,7 +468,47 @@ def run( def parse_opt(): - """Parses and returns command-line options for dataset paths, model parameters, and inference settings.""" + """ + Parses and returns command-line options for dataset paths, model parameters, and inference settings. + + Args: + --data (str): Path to the dataset YAML file. Default is 'data/coco128.yaml'. + --weights (list[str]): Paths to one or more model files. Default is 'yolov3-tiny.pt'. + --batch-size (int): Number of images per batch during inference. Default is 32. + --imgsz (int): Inference size (pixels). Default is 640. + --conf-thres (float): Confidence threshold for object detection. 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 to perform: 'train', 'val', 'test', 'speed', or 'study'. Default is 'val'. + --device (str): CUDA device identifier (e.g., '0' or '0,1,2,3') or 'cpu' for using CPU. Default is "". + --workers (int): Maximum number of dataloader workers (per RANK in DDP mode). Default is 8. + --single-cls (bool): Treat the dataset as a single-class dataset. Default is False. + --augment (bool): Apply test-time augmentation during inference. Default is False. + --verbose (bool): Print mAP by class. Default is False. + --save-txt (bool): Save detection results in '.txt' format. Default is False. + --save-hybrid (bool): Save hybrid results containing both label and prediction in '.txt' format. Default is False. + --save-conf (bool): Save confidence scores in the '--save-txt' labels. Default is False. + --save-json (bool): Save detection results in COCO JSON format. Default is False. + --project (str): Project directory to save results. Default is 'runs/val'. + --name (str): Name of the experiment to save results. Default is 'exp'. + --exist-ok (bool): Whether to overwrite existing project/name without incrementing. Default is False. + --half (bool): Use FP16 half-precision during inference. Default is False. + --dnn (bool): Use OpenCV DNN backend for ONNX inference. Default is False. + + Returns: + opt (argparse.Namespace): Parsed command-line options. + + Notes: + - The function uses `argparse` to handle command-line options. + - It also modifies some options based on specific conditions, such as appending additional flags for saving in + JSON format and checking for the `coco.yaml` dataset. + + Example: + Use the following command to run validation with custom settings: + ```python + $ python val.py --weights yolov5s.pt --data coco128.yaml --img 640 + ``` + """ 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 / "yolov3-tiny.pt", help="model path(s)") @@ -393,7 +541,36 @@ def parse_opt(): def main(opt): - """Executes model tasks including training, validation, and speed or study benchmarks based on specified options.""" + """ + Executes model tasks including training, validation, and speed or study benchmarks based on specified options. + + Args: + opt (argparse.Namespace): Parsed command-line options for dataset paths, model parameters, and inference settings. + + Returns: + None + + Note: + This function orchestrates different tasks based on the user input provided through command-line arguments. It + supports tasks like `train`, `val`, `test`, `speed`, and `study`. Depending on the task, it validates the model on a + dataset, performs speed benchmarks, or runs mAP benchmarks. + + Examples: + To validate a trained YOLOv3 model: + + ```bash + $ python val.py --weights yolov3.pt --data coco.yaml --img 640 --task val + ``` + + For running speed benchmarks: + + ```bash + $ python val.py --task speed --data coco.yaml --weights yolov3.pt --batch-size 1 + ``` + + Links: + For more information, visit the official repository: https://github.com/ultralytics/ultralytics + """ check_requirements(ROOT / "requirements.txt", exclude=("tensorboard", "thop")) if opt.task in ("train", "val", "test"): # run normally