Ultralytics Code Refactor https://ultralytics.com/actions (#13174)

Refactor code for speed and clarity

export.py

@@ -92,7 +92,22 @@ MACOS = platform.system() == "Darwin"  # macOS environment
 
 class iOSModel(torch.nn.Module):
     def __init__(self, model, im):
-        """Initializes an iOS compatible model with normalization based on image dimensions."""
+        """
+        Initializes an iOS compatible model with normalization based on image dimensions.
+
+        Args:
+            model (torch.nn.Module): The PyTorch model to be adapted for iOS compatibility.
+            im (torch.Tensor): An input tensor representing a batch of images with shape (batch, channel, height, width).
+
+        Returns:
+            None: This method does not return any value.
+
+        Notes:
+            This initializer configures normalization based on the input image dimensions, which is critical for
+            ensuring the model's compatibility and proper functionality on iOS devices. The normalization step
+            involves dividing by the image width if the image is square; otherwise, additional conditions might
+            apply (trimmed for brevity).
+        """
         super().__init__()
         b, c, h, w = im.shape  # batch, channel, height, width
         self.model = model
@@ -105,13 +120,48 @@ class iOSModel(torch.nn.Module):
             # self.normalize = torch.tensor([1. / w, 1. / h, 1. / w, 1. / h]).expand(np, 4)  # explicit (faster, larger)
 
     def forward(self, x):
-        """Runs forward pass on the input tensor, returning class confidences and normalized coordinates."""
+        """
+        Runs a forward pass on the input tensor, returning class confidences and normalized coordinates.
+
+        Args:
+          x (torch.Tensor): Input tensor containing the image data.
+
+        Returns:
+          torch.Tensor: Concatenated tensor with normalized coordinates (xywh), confidence scores (conf), and class probabilities (cls).
+
+        Examples:
+          ```python
+          model = iOSModel(pretrained_model, input_image)
+          output = model.forward(torch_input_tensor)
+          ```
+        """
         xywh, conf, cls = self.model(x)[0].squeeze().split((4, 1, self.nc), 1)
         return cls * conf, xywh * self.normalize  # confidence (3780, 80), coordinates (3780, 4)
 
 
 def export_formats():
-    """Returns a DataFrame of supported YOLOv5 model export formats and their properties."""
+    """
+    Returns a DataFrame of supported YOLOv5 model export formats and their properties.
+
+    Returns:
+        pd.DataFrame: A pandas DataFrame containing supported export formats and their properties. The DataFrame includes
+        columns for format name, CLI argument suffix, file extension or directory name, and boolean flags indicating if the
+        export format supports training and detection.
+
+    Examples:
+        ```python
+        formats = export_formats()
+        print(formats)
+        ```
+
+    Notes:
+        The DataFrame contains the following columns:
+        - Format: The name of the model format (e.g., PyTorch, TorchScript, ONNX, etc.)
+        - Include Argument: The argument to use with the export script to include this format.
+        - File Suffix: File extension or directory name associated with the format.
+        - Supports Training: Whether the format supports training.
+        - Supports Detection: Whether the format supports detection.
+    """
     x = [
         ["PyTorch", "-", ".pt", True, True],
         ["TorchScript", "torchscript", ".torchscript", True, True],
@@ -130,7 +180,28 @@ def export_formats():
 
 
 def try_export(inner_func):
-    """Decorator @try_export for YOLOv5 model export functions that logs success/failure, time taken, and file size."""
+    """
+    Logs success or failure, execution time, and file size for YOLOv5 model export functions wrapped with @try_export.
+
+    Args:
+        inner_func (Callable): The model export function to be wrapped by the decorator.
+
+    Returns:
+        Callable: The wrapped function that logs execution details. When executed, this wrapper function returns either:
+            - Tuple (str, torch.nn.Module): On success — the file path of the exported model and the model instance.
+            - Tuple (None, None): On failure — None values indicating export failed.
+
+    Examples:
+        @try_export
+        def export_onnx(model, filepath):
+            # implementation here
+            pass
+
+        exported_file, exported_model = export_onnx(yolo_model, 'path/to/save/model.onnx')
+
+    Notes:
+        For additional requirements and model export formats, refer to the [Ultralytics YOLOv5 GitHub repository](https://github.com/ultralytics/ultralytics).
+    """
     inner_args = get_default_args(inner_func)
 
     def outer_func(*args, **kwargs):
@@ -150,8 +221,46 @@ def try_export(inner_func):
 
 @try_export
 def export_torchscript(model, im, file, optimize, prefix=colorstr("TorchScript:")):
-    """Exports YOLOv5 model to TorchScript format, optionally optimized for mobile, with image shape and stride
-    metadata.
+    """
+    Exports a YOLOv5 model to the TorchScript format.
+
+    Args:
+        model (torch.nn.Module): The YOLOv5 model to be exported.
+        im (torch.Tensor): Example input tensor to be used for tracing the TorchScript model.
+        file (Path): File path where the exported TorchScript model will be saved.
+        optimize (bool): If True, applies optimizations for mobile deployment.
+        prefix (str): Optional prefix for log messages. Default is 'TorchScript:'.
+
+    Returns:
+        (str | None, torch.jit.ScriptModule | None): A tuple containing the file path of the exported model
+        (as a string) and the TorchScript model (as a torch.jit.ScriptModule). If the export fails, both elements
+        of the tuple will be None.
+
+    Notes:
+        - This function uses tracing to create the TorchScript model.
+        - Metadata, including the input image shape, model stride, and class names, is saved in an extra file (`config.txt`)
+        within the TorchScript model package.
+        - For mobile optimization, refer to the PyTorch tutorial: https://pytorch.org/tutorials/recipes/mobile_interpreter.html
+
+    Example:
+        ```python
+        from pathlib import Path
+        import torch
+        from models.experimental import attempt_load
+        from utils.torch_utils import select_device
+
+        # Load model
+        weights = 'yolov5s.pt'
+        device = select_device('')
+        model = attempt_load(weights, map_location=device)
+
+        # Example input tensor
+        im = torch.zeros(1, 3, 640, 640).to(device)
+
+        # Export model
+        file = Path('yolov5s.torchscript')
+        export_torchscript(model, im, file, optimize=False)
+        ```
     """
     LOGGER.info(f"\n{prefix} starting export with torch {torch.__version__}...")
     f = file.with_suffix(".torchscript")
@@ -168,7 +277,31 @@ def export_torchscript(model, im, file, optimize, prefix=colorstr("TorchScript:"
 
 @try_export
 def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr("ONNX:")):
-    """Exports a YOLOv5 model to ONNX format with dynamic axes and optional simplification."""
+    """
+    Export a YOLOv5 model to ONNX format with dynamic axes support and optional model simplification.
+
+    Args:
+        model (torch.nn.Module): The YOLOv5 model to be exported.
+        im (torch.Tensor): A sample input tensor for model tracing, usually the shape is (1, 3, height, width).
+        file (pathlib.Path | str): The output file path where the ONNX model will be saved.
+        opset (int): The ONNX opset version to use for export.
+        dynamic (bool): If True, enables dynamic axes for batch, height, and width dimensions.
+        simplify (bool): If True, applies ONNX model simplification for optimization.
+        prefix (str): A prefix string for logging messages, defaults to 'ONNX:'.
+
+    Returns:
+        tuple[pathlib.Path | str, None]: The path to the saved ONNX model file and None (consistent with decorator).
+
+    Raises:
+        ImportError: If required libraries for export (e.g., 'onnx', 'onnx-simplifier') are not installed.
+        AssertionError: If the simplification check fails.
+
+    Notes:
+        The required packages for this function can be installed via:
+        ```
+        pip install onnx onnx-simplifier onnxruntime onnxruntime-gpu
+        ```
+    """
     check_requirements("onnx>=1.12.0")
     import onnx
 
@@ -225,7 +358,41 @@ def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr("ONNX
 
 @try_export
 def export_openvino(file, metadata, half, int8, data, prefix=colorstr("OpenVINO:")):
-    """Exports a YOLOv5 model to OpenVINO format with optional FP16 and INT8 quantization; see https://pypi.org/project/openvino-dev/."""
+    """
+    Exports a YOLOv5 model to OpenVINO format with optional FP16 and INT8 quantization; see
+    https://pypi.org/project/openvino-dev/.
+
+    Args:
+        file (Path): The path to the output file where the OpenVINO model will be saved.
+        metadata (dict): Dictionary including model metadata such as names and strides.
+        half (bool): If True, export the model with FP16 precision.
+        int8 (bool): If True, export the model with INT8 quantization.
+        data (str): Path to the dataset YAML file required for INT8 quantization.
+        prefix (str): Prefix string for logging purposes (default is "OpenVINO:").
+
+    Returns:
+        (str, openvino.runtime.Model | None): Returns the OpenVINO model file path and openvino.runtime.Model object if
+        export is successful; otherwise, returns None.
+
+    Notes:
+        - Requires `openvino-dev` package version 2023.0 or higher. Install with:
+          `$ pip install openvino-dev>=2023.0`
+        - For INT8 quantization, also requires `nncf` library version 2.5.0 or higher. Install with:
+          `$ pip install nncf>=2.5.0`
+
+    Examples:
+        ```python
+        from pathlib import Path
+        from ultralytics import YOLOv5
+
+        model = YOLOv5('yolov5s.pt')
+        export_openvino(Path('yolov5s.onnx'), metadata={'names': model.names, 'stride': model.stride}, half=True,
+                        int8=False, data='data.yaml')
+        ```
+
+        This will export the YOLOv5 model to OpenVINO with FP16 precision but without INT8 quantization, saving it to
+        the specified file path.
+    """
     check_requirements("openvino-dev>=2023.0")  # requires openvino-dev: https://pypi.org/project/openvino-dev/
     import openvino.runtime as ov  # noqa
     from openvino.tools import mo  # noqa
@@ -282,8 +449,39 @@ def export_openvino(file, metadata, half, int8, data, prefix=colorstr("OpenVINO:
 
 @try_export
 def export_paddle(model, im, file, metadata, prefix=colorstr("PaddlePaddle:")):
-    """Exports a YOLOv5 model to PaddlePaddle format using X2Paddle, saving to `save_dir` and adding a metadata.yaml
-    file.
+    """
+    Exports a YOLOv5 model to PaddlePaddle format using X2Paddle, saving the converted model and metadata.
+
+    Args:
+        model (torch.nn.Module): The YOLOv5 model to be exported.
+        im (torch.Tensor): Input tensor used for model tracing during export.
+        file (pathlib.Path): Path to the source file to be converted.
+        metadata (dict): Additional metadata to be saved alongside the model.
+        prefix (str): Prefix for logging information.
+
+    Returns:
+        tuple (str, None): A tuple where the first element is the path to the saved PaddlePaddle model, and the
+        second element is None.
+
+    Examples:
+        ```python
+        from pathlib import Path
+        import torch
+
+        # Assume 'model' is a pre-trained YOLOv5 model and 'im' is an example input tensor
+        model = ...  # Load your model here
+        im = torch.randn((1, 3, 640, 640))  # Dummy input tensor for tracing
+        file = Path("yolov5s.pt")
+        metadata = {"stride": 32, "names": ["person", "bicycle", "car", "motorbike"]}
+
+        export_paddle(model=model, im=im, file=file, metadata=metadata)
+        ```
+    Notes:
+        Ensure that `paddlepaddle` and `x2paddle` are installed, as these are required for the export function. You can
+        install them via pip:
+        ```
+        $ pip install paddlepaddle x2paddle
+        ```
     """
     check_requirements(("paddlepaddle", "x2paddle"))
     import x2paddle
@@ -299,7 +497,36 @@ def export_paddle(model, im, file, metadata, prefix=colorstr("PaddlePaddle:")):
 
 @try_export
 def export_coreml(model, im, file, int8, half, nms, prefix=colorstr("CoreML:")):
-    """Exports YOLOv5 model to CoreML format with optional NMS, INT8, and FP16 support; requires coremltools."""
+    """
+    Export a YOLOv5 model to CoreML format with optional NMS, INT8, and FP16 support.
+
+    Args:
+        model (torch.nn.Module): The YOLOv5 model to be exported.
+        im (torch.Tensor): Example input tensor to trace the model.
+        file (pathlib.Path): Path object where the CoreML model will be saved.
+        int8 (bool): Flag indicating whether to use INT8 quantization (default is False).
+        half (bool): Flag indicating whether to use FP16 quantization (default is False).
+        nms (bool): Flag indicating whether to include Non-Maximum Suppression (default is False).
+        prefix (str): Prefix string for logging purposes (default is 'CoreML:').
+
+    Returns:
+        tuple[pathlib.Path | None, None]: The path to the saved CoreML model file, or (None, None) if there is an error.
+
+    Notes:
+        The exported CoreML model will be saved with a .mlmodel extension.
+        Quantization is supported only on macOS.
+
+    Example:
+        ```python
+        from pathlib import Path
+        import torch
+        from models.yolo import Model
+        model = Model(cfg, ch=3, nc=80)
+        im = torch.randn(1, 3, 640, 640)
+        file = Path("yolov5s_coreml")
+        export_coreml(model, im, file, int8=False, half=False, nms=True)
+        ```
+    """
     check_requirements("coremltools")
     import coremltools as ct
 
@@ -327,7 +554,36 @@ def export_engine(model, im, file, half, dynamic, simplify, workspace=4, verbose
     """
     Exports a YOLOv5 model to TensorRT engine format, requiring GPU and TensorRT>=7.0.0.
 
-    https://developer.nvidia.com/tensorrt
+    Args:
+        model (torch.nn.Module): YOLOv5 model to be exported.
+        im (torch.Tensor): Input tensor of shape (B,C,H,W).
+        file (Path): Path to save the exported model.
+        half (bool): Set to True to export with FP16 precision.
+        dynamic (bool): Set to True to enable dynamic input shapes.
+        simplify (bool): Set to True to simplify the model during export.
+        workspace (int): Workspace size in GB (default is 4).
+        verbose (bool): Set to True for verbose logging output.
+        prefix (str): Log message prefix.
+
+    Returns:
+        (Path, None): Tuple containing the path to the exported model and None.
+
+    Raises:
+        AssertionError: If executed on CPU instead of GPU.
+        RuntimeError: If there is a failure in parsing the ONNX file.
+
+    Example:
+        ```python
+        from ultralytics import YOLOv5
+        import torch
+        from pathlib import Path
+
+        model = YOLOv5('yolov5s.pt')  # Load a pre-trained YOLOv5 model
+        input_tensor = torch.randn(1, 3, 640, 640).cuda()  # example input tensor on GPU
+        export_path = Path('yolov5s.engine')  # export destination
+
+        export_engine(model.model, input_tensor, export_path, half=True, dynamic=True, simplify=True, workspace=8, verbose=True)
+        ```
     """
     assert im.device.type != "cpu", "export running on CPU but must be on GPU, i.e. `python export.py --device 0`"
     try:
@@ -407,8 +663,40 @@ def export_saved_model(
     keras=False,
     prefix=colorstr("TensorFlow SavedModel:"),
 ):
-    """Exports a YOLOv5 model to TensorFlow SavedModel format, supporting dynamic axes and non-maximum suppression
-    (NMS).
+    """
+    Exports a YOLOv5 model to TensorFlow SavedModel format, supporting dynamic axes and non-maximum suppression (NMS).
+
+    Args:
+        model (torch.nn.Module): The PyTorch model to convert.
+        im (torch.Tensor): Sample input tensor with shape (B, C, H, W) for tracing.
+        file (pathlib.Path): File path to save the exported model.
+        dynamic (bool): Flag to indicate whether dynamic axes should be used.
+        tf_nms (bool, optional): Enable TensorFlow non-maximum suppression (NMS). Default is False.
+        agnostic_nms (bool, optional): Enable class-agnostic NMS. Default is False.
+        topk_per_class (int, optional): Top K detections per class to keep before applying NMS. Default is 100.
+        topk_all (int, optional): Top K detections across all classes to keep before applying NMS. Default is 100.
+        iou_thres (float, optional): IoU threshold for NMS. Default is 0.45.
+        conf_thres (float, optional): Confidence threshold for detections. Default is 0.25.
+        keras (bool, optional): Save the model in Keras format if True. Default is False.
+        prefix (str, optional): Prefix for logging messages. Default is "TensorFlow SavedModel:".
+
+    Returns:
+        tuple: A tuple containing the path to the saved model folder (str) and the Keras model instance (tf.keras.Model | None).
+
+    Notes:
+        - The method supports TensorFlow versions up to 2.15.1.
+        - TensorFlow NMS may not be supported in older TensorFlow versions.
+        - If the TensorFlow version exceeds 2.13.1, it might cause issues when exporting to TFLite.
+          Refer to: https://github.com/ultralytics/yolov5/issues/12489
+
+    Raises:
+        Exception: If TensorFlow is not installed.
+
+    Example:
+        ```python
+        model, im = ...  # Initialize your PyTorch model and input tensor
+        export_saved_model(model, im, Path("yolov5_saved_model"), dynamic=True)
+        ```
     """
     # YOLOv5 TensorFlow SavedModel export
     try:
@@ -460,7 +748,28 @@ def export_saved_model(
 
 @try_export
 def export_pb(keras_model, file, prefix=colorstr("TensorFlow GraphDef:")):
-    """Exports YOLOv5 model to TensorFlow GraphDef *.pb format; see https://github.com/leimao/Frozen_Graph_TensorFlow for details."""
+    """
+    Exports YOLOv5 model to TensorFlow GraphDef (*.pb) format.
+
+    Args:
+        keras_model (tf.keras.Model): The Keras model to be converted.
+        file (Path): The output file path where the GraphDef will be saved.
+        prefix (str): Optional prefix string; defaults to a colored string indicating TensorFlow GraphDef export status.
+
+    Returns:
+        Tuple[Path, None]: The file path where the GraphDef model was saved and a None placeholder.
+
+    Notes:
+        For more details, refer to the guide on frozen graphs: https://github.com/leimao/Frozen_Graph_TensorFlow
+
+    Example:
+        ```python
+        from pathlib import Path
+        keras_model = ...  # assume an existing Keras model
+        file = Path("model.pb")
+        export_pb(keras_model, file)
+        ```
+    """
     import tensorflow as tf
     from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
 
@@ -480,7 +789,46 @@ def export_tflite(
     keras_model, im, file, int8, per_tensor, data, nms, agnostic_nms, prefix=colorstr("TensorFlow Lite:")
 ):
     # YOLOv5 TensorFlow Lite export
-    """Exports YOLOv5 model to TensorFlow Lite format with optional FP16, INT8, and NMS support."""
+    """
+    Exports YOLOv5 model to TensorFlow Lite format with optional FP16, INT8, and NMS support.
+
+    Args:
+        keras_model (tf.keras.Model): The Keras model to be exported.
+        im (torch.Tensor): Image tensor for normalization and model tracing.
+        file (Path): The file path to save the exported TensorFlow Lite model.
+        int8 (bool): Enables INT8 quantization if True.
+        per_tensor (bool): If True, disable per-channel quantization (applicable when int8 is True).
+        data (str): Path to dataset for representative dataset generation in INT8 quantization.
+        nms (bool): Enables Non-Maximum Suppression (NMS) support if True.
+        agnostic_nms (bool): Enables class-agnostic NMS support if True.
+        prefix (str): Prefix for logging messages.
+
+    Returns:
+        (str | None, tf.lite.Model | None): The file path of the saved TFLite model, and the TFLite model instance if successful.
+
+    Example:
+        ```python
+        from pathlib import Path
+        import torch
+        import tensorflow as tf
+        from torchvision import models
+
+        # Load a pre-trained model from torchvision
+        model = models.yolov5()  # Placeholder for actual YOLOv5 model loading
+        im = torch.zeros(1, 3, 640, 640)  # Example image tensor
+
+        # Provide the Keras model wrapping the PyTorch YOLOv5 model
+        keras_model = tf.keras.models.load_model('path/to/keras_model.h5')
+
+        # Export the model to TensorFlow Lite format
+        file_path = export_tflite(keras_model, im, Path('model.tflite'), int8=False, per_tensor=False,
+                                  data='path/to/dataset.yaml', nms=False, agnostic_nms=False)
+        ```
+
+    Notes:
+        Ensure the TensorFlow and TensorFlow Lite dependencies are installed. The exported TFLite model can be used for
+        efficient inference on mobile and edge devices.
+    """
     import tensorflow as tf
 
     LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...")
@@ -515,9 +863,31 @@ def export_tflite(
 @try_export
 def export_edgetpu(file, prefix=colorstr("Edge TPU:")):
     """
+    Clear and concise summary line describing the function's purpose:
+
     Exports a YOLOv5 model to Edge TPU compatible TFLite format; requires Linux and Edge TPU compiler.
 
-    https://coral.ai/docs/edgetpu/models-intro/
+    Args:
+        file (Path): Path to the YOLOv5 model file to be exported (.pt format).
+        prefix (str, optional): Prefix for logging messages. Defaults to colorstr("Edge TPU:").
+
+    Returns:
+        tuple[Path, None]: Path to the exported Edge TPU compatible TFLite model, None.
+
+    Raises:
+        AssertionError: If the system is not Linux.
+        subprocess.CalledProcessError: If any subprocess call to install or run the Edge TPU compiler fails.
+
+    Notes:
+        To use this function, ensure you have the Edge TPU compiler installed on your Linux system. You can find
+        installation instructions here: https://coral.ai/docs/edgetpu/compiler/.
+
+    Example:
+        ```python
+        from pathlib import Path
+        file = Path('yolov5s.pt')
+        export_edgetpu(file)
+        ```
     """
     cmd = "edgetpu_compiler --version"
     help_url = "https://coral.ai/docs/edgetpu/compiler/"
@@ -556,7 +926,30 @@ def export_edgetpu(file, prefix=colorstr("Edge TPU:")):
 
 @try_export
 def export_tfjs(file, int8, prefix=colorstr("TensorFlow.js:")):
-    """Exports a YOLOv5 model to TensorFlow.js format, optionally with uint8 quantization."""
+    """
+    Exports a YOLOv5 model to TensorFlow.js format, optionally with uint8 quantization.
+
+    Args:
+        file (Path): Path to the input model file.
+        int8 (bool): If True, applies uint8 quantization.
+        prefix (str): Prefix for logging information (default: colorstr("TensorFlow.js:")).
+
+    Returns:
+        tuple: Output directory path (str), None
+
+    Notes:
+        This function requires `tensorflowjs` to be installed. You can install it using:
+            ```shell
+            pip install tensorflowjs
+            ```
+
+        Example usage:
+            ```python
+            export_tfjs(Path('yolov5s.onnx'), int8=False)
+            ```
+
+        The TensorFlow.js converted model is saved in the directory specified by `file` with "_web_model" suffix.
+    """
     check_requirements("tensorflowjs")
     import tensorflowjs as tfjs
 
@@ -596,7 +989,30 @@ def add_tflite_metadata(file, metadata, num_outputs):
     """
     Adds TFLite metadata to a model file, supporting multiple outputs, as specified by TensorFlow guidelines.
 
-    https://www.tensorflow.org/lite/models/convert/metadata
+    Args:
+        file (str): The path to the TensorFlow Lite model file to which metadata will be added.
+        metadata (dict): Metadata information to be added to the model, structured as required by TFLite metadata schema.
+        num_outputs (int): Number of output tensors the model has, to properly configure the metadata.
+
+    Returns:
+        None
+
+    Example:
+        ```python
+        metadata = {
+            "name": "yolov5",
+            "description": "YOLOv5 object detection model",
+            "version": "1.0",
+            "author": "Ultralytics",
+            "license": "Apache License 2.0"
+        }
+        add_tflite_metadata("model.tflite", metadata, num_outputs=4)
+        ```
+
+    Note:
+        TFLite metadata can include information such as model name, version, author, and other relevant details.
+        For more details and structure of the metadata, refer to the TensorFlow Lite
+        [metadata guidelines](https://www.tensorflow.org/lite/models/convert/metadata).
     """
     with contextlib.suppress(ImportError):
         # check_requirements('tflite_support')
@@ -630,8 +1046,49 @@ def add_tflite_metadata(file, metadata, num_outputs):
 
 
 def pipeline_coreml(model, im, file, names, y, prefix=colorstr("CoreML Pipeline:")):
-    """Converts a PyTorch YOLOv5 model to CoreML format with NMS, handling different input/output shapes and saving the
-    model.
+    """
+    Converts a PyTorch YOLOv5 model to CoreML format with Non-Maximum Suppression (NMS), handling different input/output
+    shapes and saving the model.
+
+    Args:
+        model (torch.nn.Module): The YOLOv5 PyTorch model.
+        im (torch.Tensor): Input tensor example with shape [N, C, H, W], where N is the batch size, C is the number of
+            channels, H is the height, and W is the width.
+        file (Path): Path to save the converted CoreML model.
+        names (dict[int, str]): Dictionary mapping class indices to class names.
+        y (torch.Tensor): Output tensor from the PyTorch model's forward pass.
+        prefix (str): Custom prefix for logging messages.
+
+    Returns:
+        Path: Path to the saved CoreML model (.mlmodel).
+
+    Raises:
+        AssertionError: If the number of class names does not match the number of classes in the model.
+
+    Notes:
+        - This function requires `coremltools` to be installed.
+        - Running this function on a non-macOS environment might not support some features.
+        - Flexible input shapes and additional NMS options can be customized within the function.
+
+    Examples:
+    ```python
+    from pathlib import Path
+    import torch
+
+    # Load YOLOv5 model and an example input tensor
+    model = torch.load("yolov5s.pt")
+    im = torch.zeros(1, 3, 640, 640)  # Example input tensor
+
+    # Define class names
+    names = {0: "person", 1: "bicycle", 2: "car", ...}
+
+    # Perform forward pass to get model output
+    y = model(im)
+
+    # Convert to CoreML
+    output_file = Path("yolov5s.mlmodel")
+    pipeline_coreml(model, im, output_file, names, y)
+    ```
     """
     import coremltools as ct
     from PIL import Image
@@ -788,7 +1245,70 @@ def run(
     iou_thres=0.45,  # TF.js NMS: IoU threshold
     conf_thres=0.25,  # TF.js NMS: confidence threshold
 ):
-    """Exports YOLOv5 model to specified formats including ONNX, TensorRT, CoreML, and TensorFlow; see https://github.com/ultralytics/yolov5."""
+    """
+    Exports a YOLOv5 model to specified formats including ONNX, TensorRT, CoreML, and TensorFlow.
+
+    Args:
+        data (str | Path): Path to the dataset YAML configuration file. Default is 'data/coco128.yaml'.
+        weights (str | Path): Path to the pretrained model weights file. Default is 'yolov5s.pt'.
+        imgsz (tuple): Image size as (height, width). Default is (640, 640).
+        batch_size (int): Batch size for exporting the model. Default is 1.
+        device (str): Device to run the export on, e.g., '0' for GPU, 'cpu' for CPU. Default is 'cpu'.
+        include (tuple): Formats to include in the export. Default is ('torchscript', 'onnx').
+        half (bool): Flag to export model with FP16 half-precision. Default is False.
+        inplace (bool): Set the YOLOv5 Detect() module inplace=True. Default is False.
+        keras (bool): Flag to use Keras for TensorFlow SavedModel export. Default is False.
+        optimize (bool): Optimize TorchScript model for mobile deployment. Default is False.
+        int8 (bool): Apply INT8 quantization for CoreML or TensorFlow models. Default is False.
+        per_tensor (bool): Apply per tensor quantization for TensorFlow models. Default is False.
+        dynamic (bool): Enable dynamic axes for ONNX, TensorFlow, or TensorRT exports. Default is False.
+        simplify (bool): Simplify the ONNX model during export. Default is False.
+        opset (int): ONNX opset version. Default is 12.
+        verbose (bool): Enable verbose logging for TensorRT export. Default is False.
+        workspace (int): TensorRT workspace size in GB. Default is 4.
+        nms (bool): Add non-maximum suppression (NMS) to the TensorFlow model. Default is False.
+        agnostic_nms (bool): Add class-agnostic NMS to the TensorFlow model. Default is False.
+        topk_per_class (int): Top-K boxes per class to keep for TensorFlow.js NMS. Default is 100.
+        topk_all (int): Top-K boxes for all classes to keep for TensorFlow.js NMS. Default is 100.
+        iou_thres (float): IoU threshold for NMS. Default is 0.45.
+        conf_thres (float): Confidence threshold for NMS. Default is 0.25.
+
+    Returns:
+        None
+
+    Notes:
+        - Model export is based on the specified formats in the 'include' argument.
+        - Be cautious of combinations where certain flags are mutually exclusive, such as `--half` and `--dynamic`.
+
+    Example:
+        ```python
+        run(
+            data="data/coco128.yaml",
+            weights="yolov5s.pt",
+            imgsz=(640, 640),
+            batch_size=1,
+            device="cpu",
+            include=("torchscript", "onnx"),
+            half=False,
+            inplace=False,
+            keras=False,
+            optimize=False,
+            int8=False,
+            per_tensor=False,
+            dynamic=False,
+            simplify=False,
+            opset=12,
+            verbose=False,
+            workspace=4,
+            nms=False,
+            agnostic_nms=False,
+            topk_per_class=100,
+            topk_all=100,
+            iou_thres=0.45,
+            conf_thres=0.25,
+        )
+        ```
+    """
     t = time.time()
     include = [x.lower() for x in include]  # to lowercase
     fmts = tuple(export_formats()["Argument"][1:])  # --include arguments
@@ -901,7 +1421,23 @@ def run(
 
 
 def parse_opt(known=False):
-    """Parses command-line arguments for YOLOv5 model export configurations, returning the parsed options."""
+    """
+    Parses command-line arguments for YOLOv5 model export configurations.
+
+    Args:
+        known (bool): If True, `argparse.ArgumentParser.parse_known_args` is used to parse command-line arguments; otherwise,
+                     `argparse.ArgumentParser.parse_args` is used. Defaults to False.
+
+    Returns:
+        argparse.Namespace: An object containing parsed command-line arguments.
+
+    Example:
+        ```python
+        opts = parse_opt()
+        print(opts.data)
+        print(opts.weights)
+        ```
+    """
     parser = argparse.ArgumentParser()
     parser.add_argument("--data", type=str, default=ROOT / "data/coco128.yaml", help="dataset.yaml path")
     parser.add_argument("--weights", nargs="+", type=str, default=ROOT / "yolov5s.pt", help="model.pt path(s)")
@@ -937,7 +1473,44 @@ def parse_opt(known=False):
 
 
 def main(opt):
-    """Executes the YOLOv5 model inference or export with specified weights and options."""
+    """
+    ```python Exports the YOLOv5 model to specified formats, including ONNX, TensorRT, CoreML, and TensorFlow.
+
+    Args:
+        opt (argparse.Namespace): Parsed command-line arguments containing the export configurations.
+              - data (str): Path to the dataset.yaml.
+              - weights (list[str]): Paths to model (.pt) file(s).
+              - imgsz (list[int]): Image size (height, width).
+              - batch_size (int): Batch size.
+              - device (str): CUDA device, e.g., '0' or '0,1,2,3' or 'cpu'.
+              - half (bool): FP16 half-precision export flag.
+              - inplace (bool): Set YOLOv5 Detect() inplace to True.
+              - keras (bool): Use Keras for TensorFlow models.
+              - optimize (bool): Optimize TorchScript model for mobile.
+              - int8 (bool): INT8 quantization flag.
+              - per_tensor (bool): Per tensor quantization for TensorFlow.
+              - dynamic (bool): Dynamic axes for ONNX/TF/TensorRT.
+              - simplify (bool): Simplify ONNX model.
+              - opset (int): ONNX opset version.
+              - verbose (bool): Verbose logging for TensorRT.
+              - workspace (int): Workspace size for TensorRT (in GB).
+              - nms (bool): Add NMS to TensorFlow model.
+              - agnostic_nms (bool): Add agnostic NMS to TensorFlow model.
+              - topk_per_class (int): Top-k per class for TensorFlow.js NMS.
+              - topk_all (int): Top-k for all classes for TensorFlow.js NMS.
+              - iou_thres (float): IoU threshold for TensorFlow.js NMS.
+              - conf_thres (float): Confidence threshold for TensorFlow.js NMS.
+              - include (list[str]): List of formats to include in export, e.g., ['torchscript', 'onnx'].
+
+    Returns:
+        list[str]: List of exported file paths.
+    ```python
+
+    # Example usage:
+    # opt = parse_opt()
+    # main(opt)
+    ```
+    """
     for opt.weights in opt.weights if isinstance(opt.weights, list) else [opt.weights]:
         run(**vars(opt))