Add Python docstrings (#2209)

* Add Python docstrings

* Auto-format by https://ultralytics.com/actions

* Reformat

* Auto-format by https://ultralytics.com/actions

* Reformat

* Reformat

---------

Co-authored-by: UltralyticsAssistant <[email protected]>

CITATION.cff

@@ -3,9 +3,9 @@ preferred-citation:
   type: software
   message: If you use YOLOv5, please cite it as below.
   authors:
-  - family-names: Jocher
-    given-names: Glenn
-    orcid: "https://orcid.org/0000-0001-5950-6979"
+    - family-names: Jocher
+      given-names: Glenn
+      orcid: "https://orcid.org/0000-0001-5950-6979"
   title: "YOLOv5 by Ultralytics"
   version: 7.0
   doi: 10.5281/zenodo.3908559

benchmarks.py

@@ -149,6 +149,9 @@ def test(
 
 
 def parse_opt():
+    """Parses command line arguments for model inference configurations, including weights, image size, and device
+    options.
+    """
     parser = argparse.ArgumentParser()
     parser.add_argument("--weights", type=str, default=ROOT / "yolov3-tiny.pt", help="weights path")
     parser.add_argument("--imgsz", "--img", "--img-size", type=int, default=640, help="inference size (pixels)")
@@ -166,6 +169,7 @@ def parse_opt():
 
 
 def main(opt):
+    """Executes tests or main pipeline on provided options, determining behavior based on `opt.test`."""
     test(**vars(opt)) if opt.test else run(**vars(opt))
 
 

classify/predict.py

@@ -204,6 +204,7 @@ def run(
 
 
 def parse_opt():
+    """Parses command line arguments for model inference settings, returns a Namespace of options."""
     parser = argparse.ArgumentParser()
     parser.add_argument("--weights", nargs="+", type=str, default=ROOT / "yolov5s-cls.pt", help="model path(s)")
     parser.add_argument("--source", type=str, default=ROOT / "data/images", help="file/dir/URL/glob/screen/0(webcam)")
@@ -229,6 +230,7 @@ def parse_opt():
 
 
 def main(opt):
+    """Entry point for running the model; checks requirements and calls `run` with options parsed from CLI."""
     check_requirements(ROOT / "requirements.txt", exclude=("tensorboard", "thop"))
     run(**vars(opt))
 

classify/train.py

@@ -76,6 +76,9 @@
 
 
 def train(opt, device):
+    """Trains a model on a given dataset using specified options and device, handling data loading, model optimization,
+    and logging.
+    """
     init_seeds(opt.seed + 1 + RANK, deterministic=True)
     save_dir, data, bs, epochs, nw, imgsz, pretrained = (
         opt.save_dir,
@@ -306,6 +309,7 @@ def train(opt, device):
 
 
 def parse_opt(known=False):
+    """Parses command line arguments for model configuration and training options."""
     parser = argparse.ArgumentParser()
     parser.add_argument("--model", type=str, default="yolov5s-cls.pt", help="initial weights path")
     parser.add_argument("--data", type=str, default="imagenette160", help="cifar10, cifar100, mnist, imagenet, ...")
@@ -333,7 +337,7 @@ def parse_opt(known=False):
 
 
 def main(opt):
-    # Checks
+    """Initializes training environment, checks, DDP mode setup, and starts training with given options."""
     if RANK in {-1, 0}:
         print_args(vars(opt))
         check_git_status()
@@ -357,7 +361,7 @@ def main(opt):
 
 
 def run(**kwargs):
-    # Usage: from yolov5 import classify; classify.train.run(data=mnist, imgsz=320, model='yolov5m')
+    """Executes YOLOv5 model training with dynamic options, e.g., `run(data='mnist', imgsz=320, model='yolov5m')`."""
     opt = parse_opt(True)
     for k, v in kwargs.items():
         setattr(opt, k, v)

classify/val.py

@@ -147,6 +147,7 @@ def run(
 
 
 def parse_opt():
+    """Parses command-line options for model configuration and returns an argparse.Namespace of options."""
     parser = argparse.ArgumentParser()
     parser.add_argument("--data", type=str, default=ROOT / "../datasets/mnist", help="dataset path")
     parser.add_argument("--weights", nargs="+", type=str, default=ROOT / "yolov5s-cls.pt", help="model.pt path(s)")
@@ -166,6 +167,9 @@ def parse_opt():
 
 
 def main(opt):
+    """Executes the main pipeline, checks and installs requirements, then runs inference or training based on provided
+    options.
+    """
     check_requirements(ROOT / "requirements.txt", exclude=("tensorboard", "thop"))
     run(**vars(opt))
 

detect.py

@@ -232,6 +232,7 @@ def run(
 
 
 def parse_opt():
+    """Parses and returns command-line options for model inference configurations."""
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "--weights", nargs="+", type=str, default=ROOT / "yolov3-tiny.pt", help="model path or triton URL"
@@ -269,6 +270,7 @@ def parse_opt():
 
 
 def main(opt):
+    """Entry point for running the model; checks requirements and calls `run` with parsed options."""
     check_requirements(ROOT / "requirements.txt", exclude=("tensorboard", "thop"))
     run(**vars(opt))
 

export.py

@@ -92,6 +92,7 @@
 
 class iOSModel(torch.nn.Module):
     def __init__(self, model, im):
+        """Initializes an iOSModel with normalized input dimensions and number of classes from a torch model."""
         super().__init__()
         b, c, h, w = im.shape  # batch, channel, height, width
         self.model = model
@@ -104,12 +105,13 @@ def __init__(self, model, im):
             # self.normalize = torch.tensor([1. / w, 1. / h, 1. / w, 1. / h]).expand(np, 4)  # explicit (faster, larger)
 
     def forward(self, x):
+        """Performs forward pass, returns scaled confidences and normalized coordinates given input tensor `x`."""
         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():
-    # YOLOv3 export formats
+    """Lists supported YOLOv3 model export formats including file suffixes and CPU/GPU compatibility."""
     x = [
         ["PyTorch", "-", ".pt", True, True],
         ["TorchScript", "torchscript", ".torchscript", True, True],
@@ -128,7 +130,7 @@ def export_formats():
 
 
 def try_export(inner_func):
-    # YOLOv3 export decorator, i..e @try_export
+    """Decorator to profile and log the export process of YOLOv3 models, handling successes and failures."""
     inner_args = get_default_args(inner_func)
 
     def outer_func(*args, **kwargs):
@@ -147,7 +149,11 @@ def outer_func(*args, **kwargs):
 
 @try_export
 def export_torchscript(model, im, file, optimize, prefix=colorstr("TorchScript:")):
-    # YOLOv3 TorchScript model export
+    """
+    Exports a YOLOv3 model to TorchScript format, optionally optimizing for mobile.
+
+    Args: model, im (input image), file (path), optimize (bool).
+    """
     LOGGER.info(f"\n{prefix} starting export with torch {torch.__version__}...")
     f = file.with_suffix(".torchscript")
 
@@ -163,7 +169,7 @@ def export_torchscript(model, im, file, optimize, prefix=colorstr("TorchScript:"
 
 @try_export
 def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr("ONNX:")):
-    # YOLOv3 ONNX export
+    """Exports a YOLOv3 model to ONNX format with dynamic shapes and simplification options."""
     check_requirements("onnx>=1.12.0")
     import onnx
 
@@ -285,7 +291,9 @@ def transform_fn(data_item):
 
 @try_export
 def export_paddle(model, im, file, metadata, prefix=colorstr("PaddlePaddle:")):
-    # YOLOv3 Paddle export
+    """Exports a YOLOv3 model to PaddlePaddle format using X2Paddle, saving to `file` directory with metadata in
+    yaml.
+    """
     check_requirements(("paddlepaddle", "x2paddle"))
     import x2paddle
     from x2paddle.convert import pytorch2paddle
@@ -300,7 +308,7 @@ def export_paddle(model, im, file, metadata, prefix=colorstr("PaddlePaddle:")):
 
 @try_export
 def export_coreml(model, im, file, int8, half, nms, prefix=colorstr("CoreML:")):
-    # YOLOv3 CoreML export
+    """Exports YOLOv3 model to CoreML format with optional quantization and NMS; requires 'coremltools'."""
     check_requirements("coremltools")
     import coremltools as ct
 
@@ -325,7 +333,11 @@ 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:")):
-    # YOLOv3 TensorRT export https://developer.nvidia.com/tensorrt
+    """
+    Exports a YOLOv3 model to TensorRT engine format.
+
+    Nvidia TensorRT: https://developer.nvidia.com/tensorrt
+    """
     assert im.device.type != "cpu", "export running on CPU but must be on GPU, i.e. `python export.py --device 0`"
     try:
         import tensorrt as trt
@@ -444,7 +456,11 @@ def export_saved_model(
 
 @try_export
 def export_pb(keras_model, file, prefix=colorstr("TensorFlow GraphDef:")):
-    # YOLOv3 TensorFlow GraphDef *.pb export https://github.com/leimao/Frozen_Graph_TensorFlow
+    """
+    Exports Keras model as TensorFlow GraphDef *.pb file, compatible with YOLOv3.
+
+    See https://github.com/leimao/Frozen_Graph_TensorFlow for details.
+    """
     import tensorflow as tf
     from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
 
@@ -461,7 +477,7 @@ 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:")):
-    # YOLOv3 TensorFlow Lite export
+    """Exports a YOLOv3 model to TensorFlow Lite format, supports fp16/int8, NMS options, and custom data."""
     import tensorflow as tf
 
     LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...")
@@ -493,7 +509,11 @@ def export_tflite(keras_model, im, file, int8, data, nms, agnostic_nms, prefix=c
 
 @try_export
 def export_edgetpu(file, prefix=colorstr("Edge TPU:")):
-    # YOLOv3 Edge TPU export https://coral.ai/docs/edgetpu/models-intro/
+    """
+    Exports YOLOv3 model to Edge TPU compatible format; requires Linux and Edge TPU compiler.
+
+    https://coral.ai/docs/edgetpu/compiler/
+    """
     cmd = "edgetpu_compiler --version"
     help_url = "https://coral.ai/docs/edgetpu/compiler/"
     assert platform.system() == "Linux", f"export only supported on Linux. See {help_url}"
@@ -531,7 +551,7 @@ def export_edgetpu(file, prefix=colorstr("Edge TPU:")):
 
 @try_export
 def export_tfjs(file, int8, prefix=colorstr("TensorFlow.js:")):
-    # YOLOv3 TensorFlow.js export
+    """Exports a YOLOv3 model to TensorFlow.js format, optionally quantizing to uint8."""
     check_requirements("tensorflowjs")
     import tensorflowjs as tfjs
 
@@ -568,7 +588,7 @@ def export_tfjs(file, int8, prefix=colorstr("TensorFlow.js:")):
 
 
 def add_tflite_metadata(file, metadata, num_outputs):
-    # Add metadata to *.tflite models per https://www.tensorflow.org/lite/models/convert/metadata
+    """Adds TFLite model metadata for enhanced model understanding and processing; requires `tflite_support`."""
     with contextlib.suppress(ImportError):
         # check_requirements('tflite_support')
         from tflite_support import flatbuffers
@@ -601,7 +621,9 @@ def add_tflite_metadata(file, metadata, num_outputs):
 
 
 def pipeline_coreml(model, im, file, names, y, prefix=colorstr("CoreML Pipeline:")):
-    # YOLOv3 CoreML pipeline
+    """Executes YOLOv3 CoreML pipeline, handling image preprocessing, prediction, and NMS, saving the model with
+    metadata.
+    """
     import coremltools as ct
     from PIL import Image
 
@@ -866,6 +888,9 @@ def run(
 
 
 def parse_opt(known=False):
+    """Parses command line arguments for model export configuration, including data paths, weights, and export
+    options.
+    """
     parser = argparse.ArgumentParser()
     parser.add_argument("--data", type=str, default=ROOT / "data/coco128.yaml", help="dataset.yaml path")
     parser.add_argument("--weights", nargs="+", type=str, default=ROOT / "yolov3-tiny.pt", help="model.pt path(s)")
@@ -900,6 +925,7 @@ def parse_opt(known=False):
 
 
 def main(opt):
+    """Executes model inference, supporting various export formats based on specified weights and options."""
     for opt.weights in opt.weights if isinstance(opt.weights, list) else [opt.weights]:
         run(**vars(opt))
 

hubconf.py

@@ -84,57 +84,71 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo
 
 
 def custom(path="path/to/model.pt", autoshape=True, _verbose=True, device=None):
-    # YOLOv3 custom or local model
+    """Loads a custom or local YOLOv3 model from a specified path, with options for autoshaping and device
+    assignment.
+    """
     return _create(path, autoshape=autoshape, verbose=_verbose, device=device)
 
 
 def yolov5n(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
-    # YOLOv3-nano model https://github.com/ultralytics/yolov5
+    """Instantiates YOLOv5n model; optional pretrained weights, class/channel count, autoshaping, device selection."""
     return _create("yolov5n", pretrained, channels, classes, autoshape, _verbose, device)
 
 
 def yolov5s(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
-    # YOLOv3-small model https://github.com/ultralytics/yolov5
+    """Loads YOLOv5s model, with options for pretrained weights, channel/class customization, autoshaping, and device
+    choice.
+    """
     return _create("yolov5s", pretrained, channels, classes, autoshape, _verbose, device)
 
 
 def yolov5m(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
-    # YOLOv3-medium model https://github.com/ultralytics/yolov5
+    """Loads YOLOv5m model with optional pretrained weights, channel/class customization, autoshaping, device
+    selection.
+    """
     return _create("yolov5m", pretrained, channels, classes, autoshape, _verbose, device)
 
 
 def yolov5l(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
-    # YOLOv3-large model https://github.com/ultralytics/yolov5
+    """Loads YOLOv5l model with options for pretrained weights, channel/class customization, autoshaping, and device
+    selection.
+    """
     return _create("yolov5l", pretrained, channels, classes, autoshape, _verbose, device)
 
 
 def yolov5x(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
-    # YOLOv3-xlarge model https://github.com/ultralytics/yolov5
+    """Loads YOLOv5x model with customization options for weights, channels, classes, autoshaping, and device
+    selection.
+    """
     return _create("yolov5x", pretrained, channels, classes, autoshape, _verbose, device)
 
 
 def yolov5n6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
-    # YOLOv3-nano-P6 model https://github.com/ultralytics/yolov5
+    """Loads the YOLOv5n6 model with options for weights, channels, classes, shaping, and device."""
     return _create("yolov5n6", pretrained, channels, classes, autoshape, _verbose, device)
 
 
 def yolov5s6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
-    # YOLOv3-small-P6 model https://github.com/ultralytics/yolov5
+    """Loads the YOLOv5s6 model; see options for weights, channels, classes,
+    at https://github.com/ultralytics/yolov5."""
     return _create("yolov5s6", pretrained, channels, classes, autoshape, _verbose, device)
 
 
 def yolov5m6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
-    # YOLOv3-medium-P6 model https://github.com/ultralytics/yolov5
+    """Loads YOLOv5m6 model; options for weights, channels, classes documented at
+    https://github.com/ultralytics/yolov5."""
     return _create("yolov5m6", pretrained, channels, classes, autoshape, _verbose, device)
 
 
 def yolov5l6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
-    # YOLOv3-large-P6 model https://github.com/ultralytics/yolov5
+    """Loads YOLOv5l6 model with customization for pretraining, channels, classes; details at
+    https://github.com/ultralytics/yolov5."""
     return _create("yolov5l6", pretrained, channels, classes, autoshape, _verbose, device)
 
 
 def yolov5x6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
-    # YOLOv3-xlarge-P6 model https://github.com/ultralytics/yolov5
+    """Loads the YOLOv5x6 model, allowing customization for channels and classes; more at
+    https://github.com/ultralytics/yolov5."""
     return _create("yolov5x6", pretrained, channels, classes, autoshape, _verbose, device)