diff --git a/benchmarks.py b/benchmarks.py index ea11ca1f8..e92a645fb 100644 --- a/benchmarks.py +++ b/benchmarks.py @@ -78,10 +78,11 @@ def run( None. Logs information about the benchmark results, including the format, size, mAP50-95, and inference time. Notes: - Supported export formats and models include PyTorch, TorchScript, ONNX, OpenVINO, TensorRT, CoreML, TensorFlow - SavedModel, TensorFlow GraphDef, TensorFlow Lite, and TensorFlow Edge TPU. Edge TPU and TF.js are unsupported. + Supported export formats and models include PyTorch, TorchScript, ONNX, OpenVINO, TensorRT, CoreML, + TensorFlow SavedModel, TensorFlow GraphDef, TensorFlow Lite, and TensorFlow Edge TPU. Edge TPU and TF.js + are unsupported. - Examples: + Example: ```python $ python benchmarks.py --weights yolov5s.pt --img 640 ``` @@ -190,7 +191,6 @@ def test( $ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime openvino-dev tensorflow-cpu # CPU support $ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime-gpu openvino-dev tensorflow # GPU support $ pip install -U nvidia-tensorrt --index-url https://pypi.ngc.nvidia.com # TensorRT - Run export tests: $ python benchmarks.py --weights yolov5s.pt --img 640 """ @@ -263,7 +263,7 @@ def main(opt): Args: opt (argparse.Namespace): Parsed command-line arguments including options for weights, image size, batch size, data - configuration, device, and other flags for inference settings. + configuration, device, and other flags for inference settings. Returns: None: This function does not return any value. It leverages side-effects such as logging and running benchmarks. diff --git a/detect.py b/detect.py index e35d220c2..61dd8d549 100644 --- a/detect.py +++ b/detect.py @@ -102,15 +102,15 @@ def run( Args: weights (str | Path): Path to the model weights file or a Triton URL. Default is 'yolov5s.pt'. - source (str | Path): Input source, which can be a file, directory, URL, glob pattern, screen capture, or webcam index. - Default is 'data/images'. + source (str | Path): Input source, which can be a file, directory, URL, glob pattern, screen capture, or webcam + index. Default is 'data/images'. data (str | Path): Path to the dataset YAML file. Default is 'data/coco128.yaml'. imgsz (tuple[int, int]): Inference image size as a tuple (height, width). Default is (640, 640). conf_thres (float): Confidence threshold for detections. Default is 0.25. iou_thres (float): Intersection Over Union (IOU) threshold for non-max suppression. Default is 0.45. max_det (int): Maximum number of detections per image. Default is 1000. - device (str): CUDA device identifier (e.g., '0' or '0,1,2,3') or 'cpu'. Default is an empty string, which - uses the best available device. + device (str): CUDA device identifier (e.g., '0' or '0,1,2,3') or 'cpu'. Default is an empty string, which uses the + best available device. view_img (bool): If True, display inference results using OpenCV. Default is False. save_txt (bool): If True, save results in a text file. Default is False. save_csv (bool): If True, save results in a CSV file. Default is False. @@ -137,15 +137,15 @@ def run( None Examples: - ```python - from ultralytics import run + ```python + from ultralytics import run - # Run inference on an image - run(source='data/images/example.jpg', weights='yolov5s.pt', device='0') + # Run inference on an image + run(source='data/images/example.jpg', weights='yolov5s.pt', device='0') - # Run inference on a video with specific confidence threshold - run(source='data/videos/example.mp4', weights='yolov5s.pt', conf_thres=0.4, device='0') - ``` + # Run inference on a video with specific confidence threshold + run(source='data/videos/example.mp4', weights='yolov5s.pt', conf_thres=0.4, device='0') + ``` """ source = str(source) save_img = not nosave and not source.endswith(".txt") # save inference images @@ -316,7 +316,7 @@ def run( def parse_opt(): """ - Parses command-line arguments for YOLOv5 detection, setting inference options and model configurations. + Parse command-line arguments for YOLOv5 detection, allowing custom inference options and model configurations. Args: --weights (str | list[str], optional): Model path or Triton URL. Defaults to ROOT / 'yolov5s.pt'. @@ -346,7 +346,8 @@ def parse_opt(): --hide-conf (bool, optional): Flag to hide confidences in the output. Defaults to False. --half (bool, optional): Flag to use FP16 half-precision inference. Defaults to False. --dnn (bool, optional): Flag to use OpenCV DNN for ONNX inference. Defaults to False. - --vid-stride (int, optional): Video frame-rate stride, determining the number of frames to skip in between consecutive frames. Defaults to 1. + --vid-stride (int, optional): Video frame-rate stride, determining the number of frames to skip in between + consecutive frames. Defaults to 1. Returns: argparse.Namespace: Parsed command-line arguments as an argparse.Namespace object. @@ -403,8 +404,8 @@ def main(opt): None Note: - This function performs essential pre-execution checks and initiates the YOLOv5 detection process based on user-specified options. - Refer to the usage guide and examples for more information about different sources and formats at: + This function performs essential pre-execution checks and initiates the YOLOv5 detection process based on user-specified + options. Refer to the usage guide and examples for more information about different sources and formats at: https://github.com/ultralytics/ultralytics Example usage: diff --git a/export.py b/export.py index 3688dfd84..30944a93a 100644 --- a/export.py +++ b/export.py @@ -97,15 +97,15 @@ class iOSModel(torch.nn.Module): Args: model (torch.nn.Module): The PyTorch model to be adapted for iOS compatibility. - im (torch.Tensor): An input tensor representing a batch of images with shape (batch, channel, height, width). + im (torch.Tensor): An input tensor representing a batch of images with shape (B, C, H, W). Returns: None: This method does not return any value. Notes: - This initializer configures normalization based on the input image dimensions, which is critical for ensuring the - model's compatibility and proper functionality on iOS devices. The normalization step involves dividing by the image - width if the image is square; otherwise, additional conditions might apply (trimmed for brevity). + This initializer configures normalization based on the input image dimensions, which is critical for + ensuring the model's compatibility and proper functionality on iOS devices. The normalization step + involves dividing by the image width if the image is square; otherwise, additional conditions might apply. """ super().__init__() b, c, h, w = im.shape # batch, channel, height, width @@ -120,14 +120,15 @@ class iOSModel(torch.nn.Module): def forward(self, x): """ - Runs a forward pass on the input tensor, returning class confidences and normalized coordinates. + Run a forward pass on the input tensor, returning class confidences and normalized coordinates. Args: - x (torch.Tensor): Input tensor containing the image data. + x (torch.Tensor): Input tensor containing the image data with shape (batch, channels, height, width). Returns: torch.Tensor: Concatenated tensor with normalized coordinates (xywh), confidence scores (conf), - and class probabilities (cls). + and class probabilities (cls), having shape (N, 4 + 1 + C), where N is the number of predictions, + and C is the number of classes. Examples: ```python @@ -144,9 +145,9 @@ def export_formats(): Returns a DataFrame of supported YOLOv5 model export formats and their properties. Returns: - pd.DataFrame: A pandas DataFrame containing supported export formats and their properties. The DataFrame includes - columns for format name, CLI argument suffix, file extension or directory name, and boolean flags indicating if the - export format supports training and detection. + pandas.DataFrame: A DataFrame containing supported export formats and their properties. The DataFrame + includes columns for format name, CLI argument suffix, file extension or directory name, and boolean flags + indicating if the export format supports training and detection. Examples: ```python @@ -156,7 +157,7 @@ def export_formats(): Notes: The DataFrame contains the following columns: - - Format: The name of the model format (e.g., PyTorch, TorchScript, ONNX, etc.) + - Format: The name of the model format (e.g., PyTorch, TorchScript, ONNX, etc.). - Include Argument: The argument to use with the export script to include this format. - File Suffix: File extension or directory name associated with the format. - Supports Training: Whether the format supports training. @@ -181,7 +182,7 @@ def export_formats(): def try_export(inner_func): """ - Logs success or failure, execution time, and file size for YOLOv5 model export functions wrapped with @try_export. + Log success or failure, execution time, and file size for YOLOv5 model export functions wrapped with @try_export. Args: inner_func (Callable): The model export function to be wrapped by the decorator. @@ -202,7 +203,8 @@ def try_export(inner_func): ``` Notes: - For additional requirements and model export formats, refer to the [Ultralytics YOLOv5 GitHub repository](https://github.com/ultralytics/ultralytics). + For additional requirements and model export formats, refer to the + [Ultralytics YOLOv5 GitHub repository](https://github.com/ultralytics/ultralytics). """ inner_args = get_default_args(inner_func) @@ -224,7 +226,7 @@ def try_export(inner_func): @try_export def export_torchscript(model, im, file, optimize, prefix=colorstr("TorchScript:")): """ - Exports a YOLOv5 model to the TorchScript format. + Export a YOLOv5 model to the TorchScript format. Args: model (torch.nn.Module): The YOLOv5 model to be exported. @@ -235,8 +237,8 @@ def export_torchscript(model, im, file, optimize, prefix=colorstr("TorchScript:" Returns: (str | None, torch.jit.ScriptModule | None): A tuple containing the file path of the exported model - (as a string) and the TorchScript model (as a torch.jit.ScriptModule). If the export fails, both elements - of the tuple will be None. + (as a string) and the TorchScript model (as a torch.jit.ScriptModule). If the export fails, both elements + of the tuple will be None. Notes: - This function uses tracing to create the TorchScript model. @@ -303,6 +305,26 @@ def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr("ONNX ``` pip install onnx onnx-simplifier onnxruntime onnxruntime-gpu ``` + + Example: + ```python + from pathlib import Path + import torch + from models.experimental import attempt_load + from utils.torch_utils import select_device + + # Load model + weights = 'yolov5s.pt' + device = select_device('') + model = attempt_load(weights, map_location=device) + + # Example input tensor + im = torch.zeros(1, 3, 640, 640).to(device) + + # Export model + file_path = Path('yolov5s.onnx') + export_onnx(model, im, file_path, opset=12, dynamic=True, simplify=True) + ``` """ check_requirements("onnx>=1.12.0") import onnx @@ -361,10 +383,10 @@ def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr("ONNX @try_export def export_openvino(file, metadata, half, int8, data, prefix=colorstr("OpenVINO:")): """ - Exports a YOLOv5 model to OpenVINO format with optional FP16 and INT8 quantization. + Export a YOLOv5 model to OpenVINO format with optional FP16 and INT8 quantization. Args: - file (Path): The path to the output file where the OpenVINO model will be saved. + file (Path): Path to the output file where the OpenVINO model will be saved. metadata (dict): Dictionary including model metadata such as names and strides. half (bool): If True, export the model with FP16 precision. int8 (bool): If True, export the model with INT8 quantization. @@ -372,8 +394,8 @@ def export_openvino(file, metadata, half, int8, data, prefix=colorstr("OpenVINO: prefix (str): Prefix string for logging purposes (default is "OpenVINO:"). Returns: - (str, openvino.runtime.Model | None): Returns the OpenVINO model file path and openvino.runtime.Model object if - export is successful; otherwise, returns None. + (str, openvino.runtime.Model | None): The OpenVINO model file path and openvino.runtime.Model object if export is + successful; otherwise, None. Notes: - Requires `openvino-dev` package version 2023.0 or higher. Install with: @@ -451,7 +473,7 @@ def export_openvino(file, metadata, half, int8, data, prefix=colorstr("OpenVINO: @try_export def export_paddle(model, im, file, metadata, prefix=colorstr("PaddlePaddle:")): """ - Export a YOLOv5 model to PaddlePaddle format using X2Paddle, saving the converted model and metadata. + Export a YOLOv5 PyTorch model to PaddlePaddle format using X2Paddle, saving the converted model and metadata. Args: model (torch.nn.Module): The YOLOv5 model to be exported. @@ -554,12 +576,12 @@ 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 YOLOv5 model to TensorRT engine format, requiring GPU and TensorRT>=7.0.0. + Export a YOLOv5 model to TensorRT engine format, requiring GPU and TensorRT>=7.0.0. Args: model (torch.nn.Module): YOLOv5 model to be exported. im (torch.Tensor): Input tensor of shape (B, C, H, W). - file (Path): Path to save the exported model. + file (pathlib.Path): Path to save the exported model. half (bool): Set to True to export with FP16 precision. dynamic (bool): Set to True to enable dynamic input shapes. simplify (bool): Set to True to simplify the model during export. @@ -568,7 +590,7 @@ def export_engine(model, im, file, half, dynamic, simplify, workspace=4, verbose prefix (str): Log message prefix. Returns: - (Path, None): Tuple containing the path to the exported model and None. + (pathlib.Path, None): Tuple containing the path to the exported model and None. Raises: AssertionError: If executed on CPU instead of GPU. @@ -666,7 +688,7 @@ def export_saved_model( prefix=colorstr("TensorFlow SavedModel:"), ): """ - Exports a YOLOv5 model to the TensorFlow SavedModel format, supporting dynamic axes and non-maximum suppression + Export a YOLOv5 model to the TensorFlow SavedModel format, supporting dynamic axes and non-maximum suppression (NMS). Args: @@ -791,7 +813,7 @@ def export_tflite( ): # YOLOv5 TensorFlow Lite export """ - Exports a YOLOv5 model to TensorFlow Lite format with optional INT8 quantization and NMS support. + Export a YOLOv5 model to TensorFlow Lite format with optional INT8 quantization and NMS support. Args: keras_model (tf.keras.Model): The Keras model to be exported. @@ -805,7 +827,8 @@ def export_tflite( prefix (str): Prefix for log messages. Returns: - (str | None, tflite.Model | None): The file path of the exported TFLite model and the TFLite model instance, or None if export failed. + (str | None, tflite.Model | None): The file path of the exported TFLite model and the TFLite model instance, or None + if the export failed. Example: ```python @@ -863,8 +886,6 @@ def export_tflite( @try_export def export_edgetpu(file, prefix=colorstr("Edge TPU:")): """ - Clear and concise summary line describing the function's purpose: - Exports a YOLOv5 model to Edge TPU compatible TFLite format; requires Linux and Edge TPU compiler. Args: @@ -927,28 +948,30 @@ def export_edgetpu(file, prefix=colorstr("Edge TPU:")): @try_export def export_tfjs(file, int8, prefix=colorstr("TensorFlow.js:")): """ - Exports a YOLOv5 model to TensorFlow.js format, optionally with uint8 quantization. + Convert a YOLOv5 model to TensorFlow.js format with optional uint8 quantization. Args: - file (Path): Path to the input model file. - int8 (bool): If True, applies uint8 quantization. - prefix (str): Prefix for logging information (default: colorstr("TensorFlow.js:")). + file (Path): Path to the YOLOv5 model file to be converted, typically having a ".pt" or ".onnx" extension. + int8 (bool): If True, applies uint8 quantization during the conversion process. + prefix (str): Optional prefix for logging messages, default is 'TensorFlow.js:' with color formatting. Returns: - (str, None): The output directory path as a string and None. + (str, None): Tuple containing the output directory path as a string and None. Notes: - This function requires `tensorflowjs` to be installed. You can install it using: - ```shell - pip install tensorflowjs - ``` + - This function requires the `tensorflowjs` package. Install it using: + ```shell + pip install tensorflowjs + ``` + - The converted TensorFlow.js model will be saved in a directory with the "_web_model" suffix appended to the original file name. + - The conversion involves running shell commands that invoke the TensorFlow.js converter tool. - Example usage: + Example: ```python - export_tfjs(Path('yolov5s.onnx'), int8=False) + from pathlib import Path + file = Path('yolov5.onnx') + export_tfjs(file, int8=False) ``` - - The TensorFlow.js converted model is saved in the directory specified by `file` with "_web_model" suffix. """ check_requirements("tensorflowjs") import tensorflowjs as tfjs @@ -987,11 +1010,13 @@ def export_tfjs(file, int8, prefix=colorstr("TensorFlow.js:")): def add_tflite_metadata(file, metadata, num_outputs): """ - Adds metadata to a TensorFlow Lite (TFLite) model file, supporting multiple outputs, based on TensorFlow guidelines. + Adds metadata to a TensorFlow Lite (TFLite) model file, supporting multiple outputs according to TensorFlow + guidelines. Args: file (str): Path to the TFLite model file to which metadata will be added. metadata (dict): Metadata information to be added to the model, structured as required by the TFLite metadata schema. + Common keys include "name", "description", "version", "author", and "license". num_outputs (int): Number of output tensors the model has, used to configure the metadata properly. Returns: @@ -1047,20 +1072,20 @@ def add_tflite_metadata(file, metadata, num_outputs): def pipeline_coreml(model, im, file, names, y, prefix=colorstr("CoreML Pipeline:")): """ - Converts a PyTorch YOLOv5 model to CoreML format with Non-Maximum Suppression (NMS), handling different input/output - shapes and saving the model. + Convert a PyTorch YOLOv5 model to CoreML format with Non-Maximum Suppression (NMS), handling different input/output + shapes, and saving the model. Args: - model (torch.nn.Module): The YOLOv5 PyTorch model. - im (torch.Tensor): Input tensor example with shape [N, C, H, W], where N is the batch size, C is the number of - channels, H is the height, and W is the width. + model (torch.nn.Module): The YOLOv5 PyTorch model to be converted. + im (torch.Tensor): Example input tensor with shape (N, C, H, W), where N is the batch size, C is the number of channels, + H is the height, and W is the width. file (Path): Path to save the converted CoreML model. names (dict[int, str]): Dictionary mapping class indices to class names. y (torch.Tensor): Output tensor from the PyTorch model's forward pass. prefix (str): Custom prefix for logging messages. Returns: - Path: Path to the saved CoreML model (.mlmodel). + (Path): Path to the saved CoreML model (.mlmodel). Raises: AssertionError: If the number of class names does not match the number of classes in the model. @@ -1075,18 +1100,14 @@ def pipeline_coreml(model, im, file, names, y, prefix=colorstr("CoreML Pipeline: from pathlib import Path import torch - # Load YOLOv5 model and an example input tensor - model = torch.load("yolov5s.pt") - im = torch.zeros(1, 3, 640, 640) # Example input tensor + model = torch.load('yolov5s.pt') # Load YOLOv5 model + im = torch.zeros((1, 3, 640, 640)) # Example input tensor - # Define class names - names = {0: "person", 1: "bicycle", 2: "car", ...} + names = {0: "person", 1: "bicycle", 2: "car", ...} # Define class names - # Perform forward pass to get model output - y = model(im) + y = model(im) # Perform forward pass to get model output - # Convert to CoreML - output_file = Path("yolov5s.mlmodel") + output_file = Path('yolov5s.mlmodel') # Convert to CoreML pipeline_coreml(model, im, output_file, names, y) ``` """ @@ -1246,8 +1267,6 @@ def run( conf_thres=0.25, # TF.js NMS: confidence threshold ): """ - Clear and concise summary line describing the function's purpose: - Exports a YOLOv5 model to specified formats including ONNX, TensorRT, CoreML, and TensorFlow. Args: @@ -1424,14 +1443,14 @@ def run( def parse_opt(known=False): """ - Parses command-line arguments for YOLOv5 model export configurations. + Parse command-line options for YOLOv5 model export configurations. Args: - known (bool): If True, `argparse.ArgumentParser.parse_known_args` is used to parse command-line arguments; otherwise, - `argparse.ArgumentParser.parse_args` is used. Defaults to False. + known (bool): If True, uses `argparse.ArgumentParser.parse_known_args`; otherwise, uses `argparse.ArgumentParser.parse_args`. + Default is False. Returns: - argparse.Namespace: An object containing parsed command-line arguments. + argparse.Namespace: Object containing parsed command-line arguments. Example: ```python @@ -1475,44 +1494,7 @@ def parse_opt(known=False): def main(opt): - """ - Exports the YOLOv5 model to specified formats, including ONNX, TensorRT, CoreML, and TensorFlow. - - Args: - opt (argparse.Namespace): Parsed command-line arguments containing the export configurations. - - data (str): Path to the dataset YAML configuration file (e.g., 'data/coco128.yaml'). - - weights (list[str] | str): Paths to the pretrained model weights file(s) (e.g., 'yolov5s.pt'). - - imgsz (list[int]): Image size as a list [height, width]. - - batch_size (int): Batch size for exporting the model. - - device (str): Device to run the export on, such as '0' for GPU, or 'cpu' for CPU. - - half (bool): Flag to export the model with FP16 half-precision. - - inplace (bool): Set the YOLOv5 Detect() module inplace mode to True. - - keras (bool): Flag to use Keras for TensorFlow SavedModel export. - - optimize (bool): Optimize TorchScript model for mobile deployment. - - int8 (bool): Apply INT8 quantization for CoreML or TensorFlow models. - - per_tensor (bool): Apply per-tensor quantization for TensorFlow models. - - dynamic (bool): Enable dynamic axes for ONNX, TensorFlow, or TensorRT exports. - - simplify (bool): Simplify ONNX model during export. - - opset (int): ONNX opset version. - - verbose (bool): Enable verbose logging for TensorRT export. - - workspace (int): TensorRT workspace size in GB. - - nms (bool): Add non-maximum suppression (NMS) to the TensorFlow model. - - agnostic_nms (bool): Add class-agnostic NMS to the TensorFlow model. - - topk_per_class (int): Top-K boxes per class to keep for TensorFlow.js NMS. - - topk_all (int): Top-K boxes for all classes to keep for TensorFlow.js NMS. - - iou_thres (float): IoU threshold for NMS. - - conf_thres (float): Confidence threshold for NMS. - - include (list[str]): List of formats to include in export (e.g., ['torchscript', 'onnx']). - - Returns: - list[str]: List of paths to the exported model files. - - Example: - ```python - opt = parse_opt() - main(opt) - ``` - """ + """Run(**vars(opt)) # Execute the run function with parsed options.""" 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 5815d9e67..98e399421 100644 --- a/hubconf.py +++ b/hubconf.py @@ -24,10 +24,11 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo classes (int, optional): Number of classes the model is expected to detect. Defaults to 80. autoshape (bool, optional): If True, applies the YOLOv5 .autoshape() wrapper for various input formats. Defaults to True. verbose (bool, optional): If True, prints detailed information during the model creation/loading process. Defaults to True. - device (str | torch.device | None, optional): Device to use for model parameters (e.g., 'cpu', 'cuda'). If None, selects the best available device. Defaults to None. + device (str | torch.device | None, optional): Device to use for model parameters (e.g., 'cpu', 'cuda'). If None, selects + the best available device. Defaults to None. Returns: - DetectMultiBackend | AutoShape: The loaded YOLOv5 model, potentially wrapped with AutoShape if specified. + (DetectMultiBackend | AutoShape): The loaded YOLOv5 model, potentially wrapped with AutoShape if specified. Examples: ```python @@ -108,8 +109,8 @@ def custom(path="path/to/model.pt", autoshape=True, _verbose=True, device=None): Args: path (str): Path to the custom model file (e.g., 'path/to/model.pt'). - autoshape (bool): Apply YOLOv5 .autoshape() wrapper to model if True, enabling compatibility with various input types - (default is True). + autoshape (bool): Apply YOLOv5 .autoshape() wrapper to model if True, enabling compatibility with various input + types (default is True). _verbose (bool): If True, prints all informational messages to the screen; otherwise, operates silently (default is True). device (str | torch.device | None): Device to load the model on, e.g., 'cpu', 'cuda', torch.device('cuda:0'), etc. @@ -174,32 +175,42 @@ 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): """ - Creates YOLOv5-small model with options for pretraining, input channels, class count, autoshaping, verbosity, and - device. + Create a YOLOv5-small (yolov5s) model with options for pretraining, input channels, class count, autoshaping, + verbosity, and device configuration. Args: pretrained (bool, optional): Flag to load pretrained weights into the model. Defaults to True. channels (int, optional): Number of input channels. Defaults to 3. classes (int, optional): Number of model classes. Defaults to 80. - autoshape (bool, optional): Whether to apply YOLOv5 .autoshape() wrapper to the model for preprocessed inputs. + autoshape (bool, optional): Whether to wrap the model with YOLOv5's .autoshape() for handling various input formats. Defaults to True. - _verbose (bool, optional): Flag to print detailed information on model loading. Defaults to True. - device (str | torch.device | None, optional): Device to use for model parameters, e.g., 'cpu', 'cuda'. If None, - auto-select the best available device. Defaults to None. + _verbose (bool, optional): Flag to print detailed information regarding model loading. Defaults to True. + device (str | torch.device | None, optional): Device to use for model computation, can be 'cpu', 'cuda', or + torch.device instances. If None, automatically selects the best available device. Defaults to None. Returns: - torch.nn.Module: The YOLOv5-small model loaded with specified configurations and optionally pretrained weights. + torch.nn.Module: The YOLOv5-small model configured and loaded according to the specified parameters. - Usage: + Example: ```python import torch - model = torch.hub.load('ultralytics/yolov5', 'yolov5s') # official model - model = torch.hub.load('ultralytics/yolov5:master', 'yolov5s') # from branch - model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5s.pt') # custom/local model - model = torch.hub.load('.', 'custom', 'yolov5s.pt', source='local') # local repo + + # Load the official YOLOv5-small model with pretrained weights + model = torch.hub.load('ultralytics/yolov5', 'yolov5s') + + # Load the YOLOv5-small model from a specific branch + model = torch.hub.load('ultralytics/yolov5:master', 'yolov5s') + + # Load a custom YOLOv5-small model from a local checkpoint + model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5s.pt') + + # Load a local YOLOv5-small model specifying source as local repository + model = torch.hub.load('.', 'custom', 'yolov5s.pt', source='local') ``` - For more information, visit https://pytorch.org/hub/ultralytics_yolov5. + Notes: + For more details on model loading and customization, visit + the [YOLOv5 PyTorch Hub Documentation](https://pytorch.org/hub/ultralytics_yolov5). """ return _create("yolov5s", pretrained, channels, classes, autoshape, _verbose, device) @@ -213,22 +224,26 @@ def yolov5m(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=Tr pretrained (bool, optional): Whether to load pretrained weights into the model. Default is True. channels (int, optional): Number of input channels. Default is 3. classes (int, optional): Number of model classes. Default is 80. - autoshape (bool, optional): Apply YOLOv5 .autoshape() wrapper to the model for handling various input formats. Default is True. + autoshape (bool, optional): Apply YOLOv5 .autoshape() wrapper to the model for handling various input formats. + Default is True. _verbose (bool, optional): Whether to print detailed information to the screen. Default is True. - device (str | torch.device | None, optional): Device specification to use for model parameters (e.g., 'cpu', 'cuda'). Default is None. + device (str | torch.device | None, optional): Device specification to use for model parameters (e.g., 'cpu', 'cuda'). + Default is None. Returns: torch.nn.Module: The instantiated YOLOv5-medium model. Usage Example: - ```python - import torch + ```python + import torch - model = torch.hub.load('ultralytics/yolov5', 'yolov5m') # Load YOLOv5-medium from Ultralytics repository - model = torch.hub.load('ultralytics/yolov5:master', 'yolov5m') # Load from the master branch - model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5m.pt') # Load a custom/local YOLOv5-medium model - model = torch.hub.load('.', 'custom', 'yolov5m.pt', source='local') # Load from a local repository - ``` + model = torch.hub.load('ultralytics/yolov5', 'yolov5m') # Load YOLOv5-medium from Ultralytics repository + model = torch.hub.load('ultralytics/yolov5:master', 'yolov5m') # Load from the master branch + model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5m.pt') # Load a custom/local YOLOv5-medium model + model = torch.hub.load('.', 'custom', 'yolov5m.pt', source='local') # Load from a local repository + ``` + + For more information, visit https://pytorch.org/hub/ultralytics_yolov5. """ return _create("yolov5m", pretrained, channels, classes, autoshape, _verbose, device) @@ -251,7 +266,7 @@ def yolov5l(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=Tr YOLOv5 model (torch.nn.Module): The YOLOv5-large model instantiated with specified configurations and possibly pretrained weights. - Example: + Examples: ```python import torch model = torch.hub.load('ultralytics/yolov5', 'yolov5l') @@ -266,21 +281,22 @@ def yolov5l(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=Tr def yolov5x(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None): """ - Instantiates the YOLOv5-xlarge model with customizable pretraining, channel count, class count, autoshaping, - verbosity, and device. + Perform object detection using the YOLOv5-xlarge model with options for pretraining, input channels, class count, + autoshaping, verbosity, and device specification. Args: pretrained (bool): If True, loads pretrained weights into the model. Defaults to True. - channels (int): Number of input channels. Defaults to 3. - classes (int): Number of model classes. Defaults to 80. - autoshape (bool): If True, applies YOLOv5 .autoshape() wrapper to the model for easier image handling. Defaults to + channels (int): Number of input channels for the model. Defaults to 3. + classes (int): Number of model classes for object detection. Defaults to 80. + autoshape (bool): If True, applies the YOLOv5 .autoshape() wrapper for handling different input formats. Defaults to True. - _verbose (bool): If True, prints detailed information to the screen. Defaults to True. - device (str | torch.device | None): Device for model parameters, e.g., 'cpu', 'cuda:0', or a torch.device object. + _verbose (bool): If True, prints detailed information during model loading. Defaults to True. + device (str | torch.device | None): Device specification for computing the model, e.g., 'cpu', 'cuda:0', torch.device('cuda'). Defaults to None. Returns: - torch.nn.Module: The instantiated YOLOv5-xlarge model. + torch.nn.Module: The YOLOv5-xlarge model loaded with the specified parameters, optionally with pretrained weights and + autoshaping applied. Example: ```python @@ -288,7 +304,7 @@ def yolov5x(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=Tr model = torch.hub.load('ultralytics/yolov5', 'yolov5x') ``` - For more details and usage, refer to the official YOLOv5 PyTorch Hub models documentation: + For additional details, refer to the official YOLOv5 PyTorch Hub models documentation: https://pytorch.org/hub/ultralytics_yolov5 """ return _create("yolov5x", pretrained, channels, classes, autoshape, _verbose, device) @@ -308,7 +324,7 @@ def yolov5n6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=T Default is None. Returns: - torch.nn.Module: YOLOv5 model loaded with the specified configurations. + torch.nn.Module: YOLOv5-nano-P6 model loaded with the specified configurations. Example: ```python @@ -362,7 +378,7 @@ 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): """ - Creates YOLOv5-medium-P6 model with options for pretraining, channel count, class count, autoshaping, verbosity, and + Create YOLOv5-medium-P6 model with options for pretraining, channel count, class count, autoshaping, verbosity, and device. Args: @@ -389,15 +405,15 @@ def yolov5m6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=T ``` Notes: - - The model can be loaded with pre-trained weights for better performance on specific tasks. - - The autoshape feature simplifies input handling by allowing various popular data formats. + - The model can be loaded with pre-trained weights for better performance on specific tasks. + - The autoshape feature simplifies input handling by allowing various popular data formats. """ return _create("yolov5m6", pretrained, channels, classes, autoshape, _verbose, device) def yolov5l6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None): """ - Instantiates the YOLOv5-large-P6 model with customizable pretraining, channel and class counts, autoshaping, + Instantiate the YOLOv5-large-P6 model with options for pretraining, channel and class counts, autoshaping, verbosity, and device selection. Args: @@ -407,12 +423,12 @@ def yolov5l6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=T autoshape (bool, optional): If True, apply YOLOv5 .autoshape() wrapper to the model for input flexibility. Default is True. _verbose (bool, optional): If True, print all information to the screen. Default is True. device (str | torch.device | None, optional): Device to use for model parameters, e.g., 'cpu', 'cuda', or torch.device. - If None, automatically selects the best available device. Default is None. + If None, automatically selects the best available device. Default is None. Returns: torch.nn.Module: The instantiated YOLOv5-large-P6 model. - Usage: + Example: ```python import torch model = torch.hub.load('ultralytics/yolov5', 'yolov5l6') # official model diff --git a/train.py b/train.py index 9e3382787..b4395d7e8 100644 --- a/train.py +++ b/train.py @@ -102,7 +102,7 @@ GIT_INFO = check_git_info() def train(hyp, opt, device, callbacks): """ - Trains a YOLOv5 model on a custom dataset using specified hyperparameters, options, and device, managing datasets, + Train a YOLOv5 model on a custom dataset using specified hyperparameters, options, and device, managing datasets, model architecture, loss computation, and optimizer steps. Args: @@ -543,13 +543,13 @@ def train(hyp, opt, device, callbacks): def parse_opt(known=False): """ - Parses command-line arguments for YOLOv5 training, validation, and testing. + Parse command-line arguments for YOLOv5 training, validation, and testing. Args: known (bool, optional): If True, parses known arguments, ignoring the unknown. Defaults to False. Returns: - argparse.Namespace: Parsed command-line arguments. + (argparse.Namespace): Parsed command-line arguments containing options for YOLOv5 execution. Example: ```python @@ -559,9 +559,9 @@ def parse_opt(known=False): ``` Links: - Models: https://github.com/ultralytics/yolov5/tree/master/models - Datasets: https://github.com/ultralytics/yolov5/tree/master/data - Tutorial: https://docs.ultralytics.com/yolov5/tutorials/train_custom_data + - Models: https://github.com/ultralytics/yolov5/tree/master/models + - Datasets: https://github.com/ultralytics/yolov5/tree/master/data + - Tutorial: https://docs.ultralytics.com/yolov5/tutorials/train_custom_data """ parser = argparse.ArgumentParser() parser.add_argument("--weights", type=str, default=ROOT / "yolov5s.pt", help="initial weights path") @@ -618,7 +618,7 @@ def parse_opt(known=False): def main(opt, callbacks=Callbacks()): """ - Runs training or hyperparameter evolution with specified options and optional callbacks. + Runs the main entry point for training or hyperparameter evolution with specified options and optional callbacks. Args: opt (argparse.Namespace): The command-line arguments parsed for YOLOv5 training and evolution. @@ -629,7 +629,7 @@ def main(opt, callbacks=Callbacks()): None Note: - For detailed usage, visit: + For detailed usage, refer to: https://github.com/ultralytics/yolov5/tree/master/models """ if RANK in {-1, 0}: @@ -888,23 +888,27 @@ def main(opt, callbacks=Callbacks()): def generate_individual(input_ranges, individual_length): """ - Generate a random individual with gene values within specified input ranges. + Generate an individual with random hyperparameters within specified ranges. Args: input_ranges (list[tuple[float, float]]): List of tuples where each tuple contains the lower and upper bounds - for the corresponding gene. - individual_length (int): The number of genes in the individual. + for the corresponding gene (hyperparameter). + individual_length (int): The number of genes (hyperparameters) in the individual. Returns: list[float]: A list representing a generated individual with random gene values within the specified ranges. - Examples: + Example: ```python input_ranges = [(0.01, 0.1), (0.1, 1.0), (0.9, 2.0)] individual_length = 3 individual = generate_individual(input_ranges, individual_length) print(individual) # Output: [0.035, 0.678, 1.456] (example output) ``` + + Note: + The individual returned will have a length equal to `individual_length`, with each gene value being a floating-point + number within its specified range in `input_ranges`. """ individual = [] for i in range(individual_length): @@ -915,7 +919,7 @@ def generate_individual(input_ranges, individual_length): def run(**kwargs): """ - Executes YOLOv5 training with given options, allowing optional overrides through keyword arguments. + Execute YOLOv5 training with specified options, allowing optional overrides through keyword arguments. Args: weights (str, optional): Path to initial weights. Defaults to ROOT / 'yolov5s.pt'. @@ -931,8 +935,9 @@ def run(**kwargs): noval (bool, optional): Only validate at the final epoch. Defaults to False. noautoanchor (bool, optional): Disable AutoAnchor. Defaults to False. noplots (bool, optional): Do not save plot files. Defaults to False. - evolve (int, optional): Evolve hyperparameters for a specified number of generations. Use 300 if provided without a value. - evolve_population (str, optional): Directory for loading population during evolution. Defaults to ROOT / 'data/hyps'. + evolve (int, optional): Evolve hyperparameters for a specified number of generations. Use 300 if provided without a + value. + evolve_population (str, optional): Directory for loading population during evolution. Defaults to ROOT / 'data/ hyps'. resume_evolve (str, optional): Resume hyperparameter evolution from the last generation. Defaults to None. bucket (str, optional): gsutil bucket for saving checkpoints. Defaults to an empty string. cache (str, optional): Cache image data in 'ram' or 'disk'. Defaults to None. diff --git a/val.py b/val.py index 218e0748d..c4e1e402e 100644 --- a/val.py +++ b/val.py @@ -66,8 +66,8 @@ def save_one_txt(predn, save_conf, shape, file): Saves one detection result to a txt file in normalized xywh format, optionally including confidence. Args: - predn (torch.Tensor): Predicted bounding boxes and associated confidence scores and classes in xyxy format, - tensor of shape (N, 6) where N is the number of detections. + predn (torch.Tensor): Predicted bounding boxes and associated confidence scores and classes in xyxy format, tensor + of shape (N, 6) where N is the number of detections. save_conf (bool): If True, saves the confidence scores along with the bounding box coordinates. shape (tuple): Shape of the original image as (height, width). file (str | Path): File path where the result will be saved. @@ -219,7 +219,7 @@ def run( Evaluates a YOLOv5 model on a dataset and logs performance metrics. Args: - data (str | dict): Path to a dataset yaml file or a dataset dictionary. + data (str | dict): Path to a dataset YAML file or a dataset dictionary. weights (str | list[str], optional): Path to the model weights file(s). Supports various formats including PyTorch, TorchScript, ONNX, OpenVINO, TensorRT, CoreML, TensorFlow SavedModel, TensorFlow GraphDef, TensorFlow Lite, TensorFlow Edge TPU, and PaddlePaddle. @@ -469,34 +469,34 @@ def run( def parse_opt(): """ - Parses command-line options for YOLOv5 model inference configuration. + Parse command-line options for configuring YOLOv5 model inference. Args: - data (str): Path to the dataset YAML file, default is 'data/coco128.yaml'. - weights (list[str]): List of paths to the model weight files, default is 'yolov5s.pt'. - batch_size (int): Batch size for inference, default is 32. - imgsz (int): Inference image size in pixels, default is 640. - conf_thres (float): Confidence threshold for predictions, default is 0.001. - iou_thres (float): IoU threshold for Non-Max Suppression (NMS), default is 0.6. - max_det (int): Maximum number of detections per image, default is 300. - task (str): Task type - options are 'train', 'val', 'test', 'speed', or 'study'. Default is 'val'. - device (str): Device to run the model on, e.g., '0' or '0,1,2,3' or 'cpu'. Default is empty to let the system choose automatically. - workers (int): Maximum number of dataloader workers per rank in DDP mode, default is 8. - single_cls (bool): If set, treats the dataset as a single-class dataset. Default is False. - augment (bool): If set, performs augmented inference. Default is False. - verbose (bool): If set, reports mAP by class. Default is False. - save_txt (bool): If set, saves results to *.txt files. Default is False. - save_hybrid (bool): If set, saves label+prediction hybrid results to *.txt files. Default is False. - save_conf (bool): If set, saves confidences in --save-txt labels. Default is False. - save_json (bool): If set, saves results to a COCO-JSON file. Default is False. - project (str): Project directory to save results to. Default is 'runs/val'. - name (str): Name of the directory to save results to. Default is 'exp'. - exist_ok (bool): If set, existing directory will not be incremented. Default is False. - half (bool): If set, uses FP16 half-precision inference. Default is False. - dnn (bool): If set, uses OpenCV DNN for ONNX inference. Default is False. + data (str, optional): Path to the dataset YAML file. Default is 'data/coco128.yaml'. + weights (list[str], optional): List of paths to model weight files. Default is 'yolov5s.pt'. + batch_size (int, optional): Batch size for inference. Default is 32. + imgsz (int, optional): Inference image size in pixels. Default is 640. + conf_thres (float, optional): Confidence threshold for predictions. Default is 0.001. + iou_thres (float, optional): IoU threshold for Non-Max Suppression (NMS). Default is 0.6. + max_det (int, optional): Maximum number of detections per image. Default is 300. + task (str, optional): Task type - options are 'train', 'val', 'test', 'speed', or 'study'. Default is 'val'. + device (str, optional): Device to run the model on. e.g., '0' or '0,1,2,3' or 'cpu'. Default is empty to let the system choose automatically. + workers (int, optional): Maximum number of dataloader workers per rank in DDP mode. Default is 8. + single_cls (bool, optional): If set, treats the dataset as a single-class dataset. Default is False. + augment (bool, optional): If set, performs augmented inference. Default is False. + verbose (bool, optional): If set, reports mAP by class. Default is False. + save_txt (bool, optional): If set, saves results to *.txt files. Default is False. + save_hybrid (bool, optional): If set, saves label+prediction hybrid results to *.txt files. Default is False. + save_conf (bool, optional): If set, saves confidences in --save-txt labels. Default is False. + save_json (bool, optional): If set, saves results to a COCO-JSON file. Default is False. + project (str, optional): Project directory to save results to. Default is 'runs/val'. + name (str, optional): Name of the directory to save results to. Default is 'exp'. + exist_ok (bool, optional): If set, existing directory will not be incremented. Default is False. + half (bool, optional): If set, uses FP16 half-precision inference. Default is False. + dnn (bool, optional): If set, uses OpenCV DNN for ONNX inference. Default is False. Returns: - argparse.Namespace: Parsed command-line options + argparse.Namespace: Parsed command-line options. Notes: - The '--data' parameter is checked to ensure it ends with 'coco.yaml' if '--save-json' is set. @@ -508,7 +508,7 @@ def parse_opt(): ```python $ python val.py --weights yolov5s.pt --data coco128.yaml --img 640 ``` - Different model formats could be used instead of yolov5s.pt: + Different model formats could be used instead of `yolov5s.pt`: ```python $ python val.py --weights yolov5s.pt yolov5s.torchscript yolov5s.onnx yolov5s_openvino_model yolov5s.engine ``` @@ -551,17 +551,16 @@ def main(opt): Args: opt (argparse.Namespace): Parsed command-line options. - This includes values for parameters like 'data', 'weights', 'batch_size', 'imgsz', 'conf_thres', 'iou_thres', - 'max_det', 'task', 'device', 'workers', 'single_cls', 'augment', 'verbose', 'save_txt', 'save_hybrid', - 'save_conf', 'save_json', 'project', 'name', 'exist_ok', 'half', and 'dnn', essential for configuring - the YOLOv5 tasks. + This includes values for parameters like 'data', 'weights', 'batch_size', 'imgsz', 'conf_thres', + 'iou_thres', 'max_det', 'task', 'device', 'workers', 'single_cls', 'augment', 'verbose', 'save_txt', + 'save_hybrid', 'save_conf', 'save_json', 'project', 'name', 'exist_ok', 'half', and 'dnn', essential + for configuring the YOLOv5 tasks. Returns: None Examples: To validate a trained YOLOv5 model on the COCO dataset with a specific weights file, use: - ```python $ python val.py --weights yolov5s.pt --data coco128.yaml --img 640 ```