From 5e9aeb3e9ffe601613f0814712de7adf62f2a502 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Tue, 16 Jul 2024 00:18:12 +0200 Subject: [PATCH] Refactor code for speed and clarity --- benchmarks.py | 16 +- detect.py | 64 +++--- export.py | 569 +++++++++++++++++++++++++++++--------------------- hubconf.py | 106 +++++----- train.py | 33 ++- val.py | 47 +++-- 6 files changed, 477 insertions(+), 358 deletions(-) diff --git a/benchmarks.py b/benchmarks.py index 06bde1e380..127fecf4fb 100644 --- a/benchmarks.py +++ b/benchmarks.py @@ -78,9 +78,9 @@ def run( 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. + 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 @@ -173,7 +173,7 @@ def test( 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) + pd.DataFrame: A DataFrame containing the export formats and their success status. Examples: ```python @@ -223,7 +223,7 @@ def test( def parse_opt(): """ - Parses command line arguments for YOLOv3 inference and export configuration. + Parses command line arguments for YOLOv3 inference and export configurations. Args: --weights (str): Path to the weights file. Default is 'ROOT / "yolov3-tiny.pt"'. @@ -233,10 +233,10 @@ def parse_opt(): --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. + --test (bool): If set, test only 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. + --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. diff --git a/detect.py b/detect.py index efe3f1d6e5..0140fb6030 100644 --- a/detect.py +++ b/detect.py @@ -96,7 +96,7 @@ def run( vid_stride=1, # video frame-rate stride ): """ - Performs YOLOv3 detection on various input sources including images, videos, streams, and YouTube URLs. + Run YOLOv3 detection inference on various input sources such as images, videos, YouTube URLs, and streams. Args: weights (str | Path): Path to the model weights file or a Triton URL (default: 'yolov5s.pt'). @@ -290,33 +290,34 @@ def parse_opt(): 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`. + --weights (list[str]): Model path or Triton URL. Default: 'ROOT / "yolov3-tiny.pt"'. + --source (str | Path): Source of input data such as file, directory, URL, glob pattern, or device identifier + (default: 'ROOT / "data/images"'). + --data (str | Path): Optional path to dataset.yaml. Default: 'ROOT / "data/coco128.yaml"'. + --imgsz (list[int]): Inference size as height and 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', '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): Prevent saving images/videos. Default: False. + --classes (list[int] | None): Filter results by class, e.g., [0], or [0, 2, 3]. Default: None. + --agnostic-nms (bool): Perform class-agnostic Non-Maximum Suppression (NMS). Default: False. + --augment (bool): Apply augmented inference. Default: False. + --visualize (bool): Visualize features. Default: False. + --update (bool): Update all models. Default: False. + --project (str | Path): Directory to save results, in the format '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 existing project/name directory without incrementing. Default: False. + --line-thickness (int): Thickness of bounding box lines in pixels. Default: 3. + --hide-labels (bool): Hide labels in detected objects. Default: False. + --hide-conf (bool): Hide confidence scores in 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. @@ -359,7 +360,7 @@ 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 YOLO model; checks requirements and calls `run` with parsed options. Args: opt (argparse.Namespace): Parsed command-line options, which include: @@ -400,9 +401,10 @@ def main(opt): 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 + Run this function as the entry point for using YOLO 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")) diff --git a/export.py b/export.py index e30c4f0e7b..9ead8275a4 100644 --- a/export.py +++ b/export.py @@ -94,21 +94,21 @@ class iOSModel(torch.nn.Module): def __init__(self, model, im): """ 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. - + 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, + - 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. + - 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__() @@ -125,31 +125,34 @@ def __init__(self, model, im): def forward(self, 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]. - + 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. - + - 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. - + 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. @@ -161,16 +164,16 @@ def forward(self, x): def export_formats(): """ 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: + 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() @@ -198,23 +201,23 @@ def export_formats(): def try_export(inner_func): """ 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 - ``` - + + ```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. """ @@ -238,40 +241,42 @@ def outer_func(*args, **kwargs): @try_export def export_torchscript(model, im, file, optimize, prefix=colorstr("TorchScript:")): """ - Exports a YOLOv3 model to TorchScript format, with optional optimization for mobile deployment. - + Export 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. + im (torch.Tensor): A tensor representing the input image for the model, typically with shape (N, C, H, W). 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)`. - + Exception: Logs the error and returns `(None, None)` if any issue arises during export. + 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: + - The function leverages `torch.jit.trace` to trace the model using the provided input tensor (`im`). + - Metadata such as input shape, stride, and class names are embedded in an extra file stored with the + TorchScript model. + - When `optimize` is `True`, the exported model is optimized for mobile using `optimize_for_mobile`. + + Example: ```python from pathlib import Path import torch - + 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/. + + For more info, visit: https://ultralytics.com/. """ LOGGER.info(f"\n{prefix} starting export with torch {torch.__version__}...") f = file.with_suffix(".torchscript") @@ -289,8 +294,8 @@ 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 shape and simplification options. - + Export 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. @@ -299,9 +304,33 @@ def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr("ONNX 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). + + Example: + ```python + from pathlib import Path + import torch + + model = ... # Assume model is loaded or created + im = torch.randn(1, 3, 640, 640) # A sample input tensor + file = Path("model.onnx") + opset = 12 + dynamic = True + simplify = True + + export_onnx(model, im, file, opset, dynamic, simplify) + ``` + + Notes: + Ensure `onnx`, `onnx-simplifier`, and suitable runtime packages are installed: + ```sh + pip install onnx onnx-simplifier onnxruntime-gpu + ``` + This function also verifies the exported ONNX model and, if requested, simplifies it. + + For more information, refer to https://github.com/ultralytics/ultralytics. """ check_requirements("onnx>=1.12.0") import onnx @@ -360,23 +389,23 @@ 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. - + Export 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. - + 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') @@ -451,39 +480,39 @@ 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 a specified directory and including model + Export 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. + im (torch.Tensor): A sample input tensor used for tracing the model, with shape (B, C, H, W). 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`. - + tuple[str, None]: A tuple containing the directory path 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. - + 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) ``` """ @@ -503,7 +532,7 @@ def export_paddle(model, im, file, metadata, prefix=colorstr("PaddlePaddle:")): def export_coreml(model, im, file, int8, half, nms, prefix=colorstr("CoreML:")): """ 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. @@ -512,9 +541,25 @@ def export_coreml(model, im, file, int8, half, nms, prefix=colorstr("CoreML:")): 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). + + 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.mlmodel") + int8 = False + half = True + nms = True + + export_coreml(model, im, file, int8, half, nms) + ``` """ check_requirements("coremltools") import coremltools as ct @@ -541,8 +586,8 @@ 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, optimizing it for GPU inference. - + Export 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. @@ -553,26 +598,26 @@ def export_engine(model, im, file, half, dynamic, simplify, workspace=4, verbose 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. - + tuple (Path, None): The output file 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) ``` @@ -652,9 +697,9 @@ def export_saved_model( prefix=colorstr("TensorFlow SavedModel:"), ): """ - Exports a YOLOv3 model to TensorFlow SavedModel format. The function includes options for post-processing like Non- - Max Suppression (NMS) and others. - + 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. @@ -668,28 +713,28 @@ def export_saved_model( 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) ``` """ @@ -738,30 +783,31 @@ def export_saved_model( @try_export def export_pb(keras_model, file, prefix=colorstr("TensorFlow GraphDef:")): """ - Exports a Keras model to TensorFlow GraphDef (*.pb) format, which is compatible with YOLOv3. - + Export 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. + 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 @@ -780,45 +826,43 @@ 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:")): """ - Export a YOLOv3 model to TensorFlow Lite format, supporting fp16/int8, NMS options, and custom data. - + Export a YOLOv3 PyTorch model to TensorFlow Lite (TFLite) format. + 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. + keras_model (tf.keras.Model): The Keras model obtained after converting the PyTorch model. + im (torch.Tensor): Sample input tensor to determine model input size. + file (pathlib.Path): Desired file path for saving the exported TFLite model. + int8 (bool): Flag to enable INT8 quantization for the TFLite model. + data (str): Path to dataset YAML file for representative data generation used in quantization. + nms (bool): Flag to include Non-Max Suppression (NMS) in the TFLite model. agnostic_nms (bool): Flag to apply class-agnostic NMS. - prefix (str, optional): Prefix string for logging. Defaults to colorstr("TensorFlow Lite:"). - + prefix (str): Prefix string for log messages. 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. - + (str, None): The file path to the saved TensorFlow Lite model and None. + 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) + import torch + from tensorflow.keras.models import load_model + + # Load Keras model and create a dummy input tensor + keras_model = load_model('model.h5') + im = torch.randn(1, 3, 640, 640) + + # Export to TensorFlow Lite + export_tflite(keras_model, im, Path("model.tflite"), int8=True, data="data.yaml", nms=True, agnostic_nms=True) ``` - - TensorFlow Lite Developer Guide: - https://www.tensorflow.org/lite/guide - - Model Conversion Reference: - https://github.com/leimao/Frozen_Graph_TensorFlow + + Notes: + - TensorFlow and appropriate dependencies must be installed. + - INT8 quantization needs representative data for model calibration. + - For more details, refer to TensorFlow Lite Developer Guide: https://www.tensorflow.org/lite/guide + + See Also: + `export_onnx()`, `export_saved_model()` + ``` """ import tensorflow as tf @@ -852,35 +896,37 @@ 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. - + Export a YOLOv5 model to TensorFlow Edge TPU format with INT8 quantization. + 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`. - + 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. - + 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: + The function checks if the Edge TPU compiler is installed, and attempts installation if it isn't found. The + export process includes performing INT8 quantization on the model to make it Edge TPU-compatible. + + Examples: ```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}") ``` + + For more details, refer to the Edge TPU compiler documentation: + https://coral.ai/docs/edgetpu/compiler/ """ cmd = "edgetpu_compiler --version" help_url = "https://coral.ai/docs/edgetpu/compiler/" @@ -920,46 +966,55 @@ 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, with an optional quantization to uint8. - + Export 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. - + (tuple[str, None]): The directory path where the TensorFlow.js model files are saved and `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. - + Ensure that you have TensorFlow.js installed in your environment. Install the package via: + ```bash + pip install tensorflowjs + ``` + + For more details on using the converted model: + Refer to the official TensorFlow.js documentation: 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 + The converted model can be used directly in JavaScript environments using the TensorFlow.js library. + + For usage in web applications: + - Clone the example repository: + ```bash + cd .. && git clone https://github.com/zldrobit/tfjs-yolov5-example.git && cd tfjs-yolov5-example + ``` + - Install dependencies: + ```bash + npm install + ``` + - Create a symbolic link to the exported web model: + ```bash + ln -s ../../yolov5/yolov5s_web_model public/yolov5s_web_model + ``` + - Start the example application: + ```bash + npm start + ``` """ check_requirements("tensorflowjs") import tensorflowjs as tfjs @@ -999,21 +1054,21 @@ def export_tfjs(file, int8, prefix=colorstr("TensorFlow.js:")): def add_tflite_metadata(file, metadata, num_outputs): """ 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 = { @@ -1056,31 +1111,31 @@ 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. - + Pre-processes a YOLOv3 image input, runs CoreML model prediction, and applies Non-Maximum Suppression (NMS). + 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. + model (coremltools.models.MLModel): The pre-loaded CoreML model to use for making predictions. + im (torch.Tensor): Input image tensor in BCHW (batch, channel, height, width) format. + file (pathlib.Path): Path to save the processed CoreML model with .mlmodel extension. 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:'. - + y (torch.Tensor): YOLO model output tensor for inference shape purposes. + prefix (str): Optional log prefix, defaults to 'CoreML Pipeline:'. + Returns: - (pathlib.Path | None): The file path of the saved CoreML model if successful; None otherwise. - + (pathlib.Path | None): File path for the saved CoreML model if successful; None otherwise. + Example: ```python + from coremltools.models import MLModel from pathlib import Path import torch - from coremltools.models import MLModel - - model = MLModel('path/to/pretrained/model.mlmodel') + + model = MLModel('path/to/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 - + y = torch.randn(1, 25200, 85) # Example YOLO output tensor + pipeline_coreml(model, im, file, names, y) ``` """ @@ -1240,7 +1295,7 @@ def run( ): """ 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. @@ -1264,14 +1319,14 @@ def run( 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( @@ -1398,18 +1453,40 @@ def run( def parse_opt(known=False): """ - Parses command line arguments for model export configuration, allowing users to specify various export options. - + Parse command-line arguments for model export configuration. + 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. - + argparse.Namespace: Namespace object containing export configuration parameters. + - data (str): Path to dataset configuration file (default: 'data/coco128.yaml'). + - weights (list of str): List of model weight file paths (default: 'yolov3-tiny.pt'). + - imgsz (list of int): Image dimensions (height, width), default is [640, 640]. + - batch_size (int): Batch size for inference (default: 1). + - device (str): Device to use for export ('cpu', '0' for GPU). + - half (bool): Export the model with FP16 precision. + - inplace (bool): Set YOLOv3 Detect() inplace=True. + - keras (bool): Use Keras for TensorFlow SavedModel export. + - optimize (bool): Optimize TorchScript model for mobile. + - int8 (bool): Apply INT8 quantization for CoreML, TensorFlow, and OpenVINO. + - dynamic (bool): Enable dynamic axes for ONNX, TensorFlow, TensorRT models. + - simplify (bool): Simplify ONNX model after export. + - opset (int): ONNX opset version (default: 17). + - verbose (bool): Enable verbose logging for TensorRT export. + - workspace (int): Set workspace size for TensorRT in GB (default: 4). + - nms (bool): Add Non-Maximum Suppression (NMS) to TensorFlow models. + - agnostic_nms (bool): Enable class-agnostic NMS for TensorFlow models. + - topk_per_class (int): Keep top-K per class for TensorFlow JSON model (default: 100). + - topk_all (int): Keep top-K overall for TensorFlow JSON model (default: 100). + - iou_thres (float): IOU threshold for TensorFlow JSON model (default: 0.45). + - conf_thres (float): Confidence threshold for TensorFlow JSON model (default: 0.25). + - include (list of str): List of formats to include in export ('torchscript', 'onnx', etc.). + 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. + The function uses `argparse` to define and process command-line arguments for various model export formats, + including TorchScript, ONNX, OpenVINO, TensorRT, and more. Users can configure options such as 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") @@ -1446,41 +1523,53 @@ def parse_opt(known=False): def main(opt): """ - 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 ... + Main function to execute the export of a YOLOv3 PyTorch model to different 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 (Union[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[Union[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)) + + if __name__ == "__main__": + opt = parse_opt() + main(opt) + ``` + + Note: + - Ensure that any necessary packages are installed for specific export formats before running the function. + - Refer to the examples provided in the docstring for usage. """ 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 d701c2d4c7..cd9fb67eb0 100644 --- a/hubconf.py +++ b/hubconf.py @@ -22,11 +22,11 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo 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. + 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. + device (str | torch.device | None): Device to use for model parameters ('cpu', 'cuda', etc.). If None, defaults + to the best available device. Returns: torch.nn.Module: YOLOv3 model loaded with or without pretrained weights. @@ -101,13 +101,13 @@ def custom(path="path/to/model.pt", autoshape=True, _verbose=True, device=None): 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. + autoshape (bool): If True, applies the YOLOv3 `.autoshape()` wrapper to allow for various input formats. Default is True. + _verbose (bool): If True, outputs detailed information. Otherwise, limits verbosity. Default is True. 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. + (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 @@ -126,15 +126,15 @@ def custom(path="path/to/model.pt", autoshape=True, _verbose=True, device=None): def yolov5n(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None): """ - Instantiates a YOLOv5n model with optional pretrained weights, configurable input channels, classes, autoshaping, - and device selection. + Instantiates a YOLOv5n model with optional pretrained weights, configurable input channels, number of 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. + autoshape (bool): If True, applies YOLOv5 .autoshape() wrapper to the model for various input formats like file/URI/PIL/cv2/np + and adds 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. @@ -156,15 +156,14 @@ def yolov5n(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=Tr def yolov5s(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None): """ - Loads the YOLOv5s model with customizable options for pretrained weights, input channels, number of classes, + Load 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. + 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. @@ -212,45 +211,49 @@ def yolov5m(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=Tr def yolov5l(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None): """ - Loads the YOLOv5l model with options for pretrained weights, channel and class customization, autoshaping, and - device selection. + Load the YOLOv5l model with customizable options for pretrained weights, input channels, number of classes, + autoshape functionality, 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. + pretrained (bool, optional): If True, load 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): If True, applies YOLOv5's .autoshape() wrapper to the model. Default is True. + _verbose (bool, optional): If True, prints detailed information during model loading. Default is True. + device (str | torch.device | None, optional): Device to use for model parameters ('cpu', 'cuda', etc.). If None, + automatically selects the best available option. Default is None. Returns: - torch.nn.Module: The initialized YOLOv5l model. + (torch.nn.Module): The initialized YOLOv5l model with the specified options. Example: ```python import torch model = torch.hub.load('ultralytics/yolov5', 'yolov5l', pretrained=True) ``` + + For more information, refer to [PyTorch Hub models](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): """ - Loads YOLOv5x model with customization options for weights, channels, classes, autoshaping, and device selection. + Load the YOLOv5x model with options for pretrained weights, number of input channels, classes, autoshaping, and + device selection. Args: - pretrained (bool, optional): Whether to load pretrained weights. Defaults to True. + 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): 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. + classes (int, optional): Number of detection classes. Defaults to 80. + autoshape (bool, optional): If True, applies the YOLOv5 .autoshape() wrapper, enabling handling of various input formats + (file/URI/PIL/cv2/np) and non-maximum suppression (NMS). Defaults to True. + _verbose (bool, optional): Whether to print detailed information during model loading. Defaults to True. + device (str | torch.device | None, optional): 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 YOLOv5x model. + (torch.nn.Module): The instantiated YOLOv5x model. Examples: ```python @@ -262,23 +265,27 @@ def yolov5x(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=Tr # Load YOLOv5x model with custom device model = torch.hub.load('ultralytics/yolov5', 'yolov5x', device='cuda:0') ``` + + Refer to the [PyTorch Hub models](https://pytorch.org/hub/ultralytics_yolov5) for more details. + ``` """ 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 pretrained weights, input channels, classes, autoshaping, verbosity, and - device assignment. + Loads the YOLOv5n6 model with options for pretrained weights, input channels, number of 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. + autoshape (bool, optional): If True, applies the YOLOv5 .autoshape() wrapper to allow handling of multiple input types + (file/URI/PIL/cv2/np) and non-maximum suppression (NMS). Default is True. + _verbose (bool, optional): If True, prints detailed information to the screen. Default is True. + device (str | torch.device | None, optional): Device to use for model parameters (e.g., 'cpu', 'cuda'). If None, the best + available device is automatically selected. Default is None. Returns: torch.nn.Module: YOLOv5n6 model loaded on the specified device and configured as per the provided options. @@ -288,7 +295,11 @@ def yolov5n6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=T Example: ```python - model = yolov5n6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device='cuda') + import torch + model = torch.hub.load('ultralytics/yolov5', 'yolov5n6') + model = torch.hub.load('ultralytics/yolov5:master', 'yolov5n6') + model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5n6.pt') + model = torch.hub.load('.', 'custom', 'yolov5n6.pt', source='local') ``` """ return _create("yolov5n6", pretrained, channels, classes, autoshape, _verbose, device) @@ -302,7 +313,7 @@ def yolov5s6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=T 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. + autoshape (bool, optional): Apply YOLOv5 .autoshape() wrapper to the 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. @@ -325,20 +336,20 @@ def yolov5s6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=T def yolov5m6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None): """ - Loads YOLOv5m6 model with options for pretrained weights, input channels, number of classes, autoshaping, and device - selection. + Load the 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. + autoshape (bool): Whether to apply YOLOv5 .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. + 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. + torch.nn.Module: The instantiated YOLOv5m6 model with specified options. Example: ```python @@ -349,7 +360,7 @@ def yolov5m6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=T ``` Notes: - For more detailed documentation, visit https://github.com/ultralytics/yolov5 + For more detailed documentation, visit https://github.com/ultralytics/yolov5 """ return _create("yolov5m6", pretrained, channels, classes, autoshape, _verbose, device) @@ -387,7 +398,8 @@ def yolov5l6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=T def yolov5x6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None): """ - Loads the YOLOv5x6 model, allowing customization for pretrained weights, input channels, and model classes. + Loads the YOLOv5x6 model, allowing customization for pretrained weights, input channels, model classes, and device + selection. Args: pretrained (bool): If True, loads the model with pretrained weights. Default is True. diff --git a/train.py b/train.py index cbdce77f6f..5431e771c2 100644 --- a/train.py +++ b/train.py @@ -105,7 +105,7 @@ def train(hyp, opt, device, callbacks): # hyp is path/to/hyp.yaml or hyp dictio 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. + 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. @@ -118,7 +118,8 @@ def train(hyp, opt, device, callbacks): # hyp is path/to/hyp.yaml or hyp dictio $ 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 + $ 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 @@ -544,13 +545,13 @@ def lf(x): def parse_opt(known=False): """ - Parses command line arguments for configuring the training of a YOLO model. + Parse 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. + (argparse.Namespace): Parsed command line arguments. Examples: ```python @@ -627,7 +628,7 @@ def main(opt, callbacks=Callbacks()): 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 + - For a tutorial on using Multi-GPU with DDP: https://docs.ultralytics.com/yolov5/tutorials/multi_gpu_training Example: Single-GPU training: @@ -812,12 +813,13 @@ def main(opt, callbacks=Callbacks()): def run(**kwargs): """ - Executes model training with specified configurations. + Execute the training process for a YOLO model with the specified configuration. Args: - kwargs (dict): Configuration parameters for the training process. Supported parameters include but are not limited to: + 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. + - weights (str): Path to the 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. @@ -825,19 +827,12 @@ def run(**kwargs): 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: + Examples: ```python from train import run - run(data='coco128.yaml', imgsz=320, weights='yolov5m.pt') + + # Execute model training with custom configurations + run(data='coco128.yaml', imgsz=320, weights='yolov5m.pt', epochs=50) ``` References: diff --git a/val.py b/val.py index 775282d50d..a1646ef7fe 100644 --- a/val.py +++ b/val.py @@ -105,7 +105,7 @@ 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. + Save 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. @@ -125,6 +125,12 @@ def save_one_json(predn, jdict, path, class_map): class_map = {0: 1, 1: 2} save_one_json(predn, jdict, path, class_map) ``` + + Notes: + - The coordinates in `predn` are expected to be normalized. + - Each detection is converted to COCO-style format. + - `image_id` is derived from the filename stem and must be unique for each image. + - The function modifies `jdict` in place, appending detection results as dictionaries. """ image_id = int(path.stem) if path.stem.isnumeric() else path.stem box = xyxy2xywh(predn[:, :4]) # xywh @@ -218,7 +224,7 @@ def run( """ Validates a trained YOLO model on a dataset and saves detection results in specified formats. - Parameters: + Args: 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. @@ -241,18 +247,33 @@ def run( 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. + model (torch.nn.Module, optional): Existing model instance. Default is None. + dataloader (torch.utils.data.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. + (tuple): A tuple containing: + - metrics (torch.Tensor): Dictionary containing metrics such as precision, recall, mAP, F1 score, etc. + - times (dict): Dictionary containing times for different parts of the pipeline (e.g., preprocessing, inference, NMS). + - samples (torch.Tensor): Torch tensor containing validation samples. + + Example: + ```python + metrics, times, samples = run( + data='data/coco.yaml', + weights='yolov5s.pt', + batch_size=32, + imgsz=640, + conf_thres=0.001, + iou_thres=0.6, + max_det=300, + task='val', + device='cpu' + ) + ``` """ # Initialize/load model and set device training = model is not None @@ -500,8 +521,8 @@ def parse_opt(): 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. + - 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: @@ -551,9 +572,9 @@ def main(opt): 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. + 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: