Add docstrings to YOLOv5 functions

benchmarks.py

@@ -149,6 +149,7 @@ def test(
 
 
 def parse_opt():
+    """Parses command-line arguments for YOLOv5 model inference configuration."""
     parser = argparse.ArgumentParser()
     parser.add_argument("--weights", type=str, default=ROOT / "yolov5s.pt", help="weights path")
     parser.add_argument("--imgsz", "--img", "--img-size", type=int, default=640, help="inference size (pixels)")
@@ -166,6 +167,7 @@ def parse_opt():
 
 
 def main(opt):
+    """Executes a test run if `opt.test` is True, otherwise starts training or inference with provided options."""
     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 YOLOv5 inference settings including model, source, device, and image size."""
     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):
+    """Executes YOLOv5 model inference with options for ONNX DNN and video frame-rate stride adjustments."""
     check_requirements(ROOT / "requirements.txt", exclude=("tensorboard", "thop"))
     run(**vars(opt))
 

classify/train.py

@@ -76,6 +76,7 @@
 
 
 def train(opt, device):
+    """Trains a YOLOv5 model, managing datasets, model optimization, logging, and saving checkpoints."""
     init_seeds(opt.seed + 1 + RANK, deterministic=True)
     save_dir, data, bs, epochs, nw, imgsz, pretrained = (
         opt.save_dir,
@@ -306,6 +307,9 @@ def train(opt, device):
 
 
 def parse_opt(known=False):
+    """Parses command line arguments for YOLOv5 training including model path, dataset, epochs, and more, returning
+    parsed arguments.
+    """
     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
+    """Executes YOLOv5 training with given options, handling device setup and DDP mode; includes pre-training checks."""
     if RANK in {-1, 0}:
         print_args(vars(opt))
         check_git_status()
@@ -357,7 +361,11 @@ def main(opt):
 
 
 def run(**kwargs):
-    # Usage: from yolov5 import classify; classify.train.run(data=mnist, imgsz=320, model='yolov5m')
+    """
+    Executes YOLOv5 model training or inference with specified parameters, returning updated options.
+
+    Example: from yolov5 import classify; classify.train.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 and returns command line arguments for YOLOv5 model evaluation and inference settings."""
     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,7 @@ def parse_opt():
 
 
 def main(opt):
+    """Executes the YOLOv5 model prediction workflow, handling argument parsing and requirement checks."""
     check_requirements(ROOT / "requirements.txt", exclude=("tensorboard", "thop"))
     run(**vars(opt))
 

detect.py

@@ -166,6 +166,7 @@ def run(
 
         # Create or append to the CSV file
         def write_to_csv(image_name, prediction, confidence):
+            """Writes prediction data for an image to a CSV file, appending if the file exists."""
             data = {"Image Name": image_name, "Prediction": prediction, "Confidence": confidence}
             with open(csv_path, mode="a", newline="") as f:
                 writer = csv.DictWriter(f, fieldnames=data.keys())
@@ -264,6 +265,7 @@ def write_to_csv(image_name, prediction, confidence):
 
 
 def parse_opt():
+    """Parses command-line arguments for YOLOv5 detection, setting inference options and model configurations."""
     parser = argparse.ArgumentParser()
     parser.add_argument("--weights", nargs="+", type=str, default=ROOT / "yolov5s.pt", help="model path or triton URL")
     parser.add_argument("--source", type=str, default=ROOT / "data/images", help="file/dir/URL/glob/screen/0(webcam)")
@@ -300,6 +302,7 @@ def parse_opt():
 
 
 def main(opt):
+    """Executes YOLOv5 model inference with given options, checking requirements before running the model."""
     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 iOS compatible model with normalization based on image dimensions."""
         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):
+        """Runs forward pass on the input tensor, returning class confidences and normalized coordinates."""
         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():
-    # YOLOv5 export formats
+    """Returns a DataFrame of supported YOLOv5 model export formats and their properties."""
     x = [
         ["PyTorch", "-", ".pt", True, True],
         ["TorchScript", "torchscript", ".torchscript", True, True],
@@ -128,7 +130,7 @@ def export_formats():
 
 
 def try_export(inner_func):
-    # YOLOv5 export decorator, i..e @try_export
+    """Decorator @try_export for YOLOv5 model export functions that logs success/failure, time taken, and file size."""
     inner_args = get_default_args(inner_func)
 
     def outer_func(*args, **kwargs):
@@ -147,7 +149,9 @@ def outer_func(*args, **kwargs):
 
 @try_export
 def export_torchscript(model, im, file, optimize, prefix=colorstr("TorchScript:")):
-    # YOLOv5 TorchScript model export
+    """Exports YOLOv5 model to TorchScript format, optionally optimized for mobile, with image shape and stride
+    metadata.
+    """
     LOGGER.info(f"\n{prefix} starting export with torch {torch.__version__}...")
     f = file.with_suffix(".torchscript")
 
@@ -163,7 +167,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:")):
-    # YOLOv5 ONNX export
+    """Exports a YOLOv5 model to ONNX format with dynamic axes and optional simplification."""
     check_requirements("onnx>=1.12.0")
     import onnx
 
@@ -276,7 +280,9 @@ def transform_fn(data_item):
 
 @try_export
 def export_paddle(model, im, file, metadata, prefix=colorstr("PaddlePaddle:")):
-    # YOLOv5 Paddle export
+    """Exports a YOLOv5 model to PaddlePaddle format using X2Paddle, saving to `save_dir` and adding a metadata.yaml
+    file.
+    """
     check_requirements(("paddlepaddle", "x2paddle"))
     import x2paddle
     from x2paddle.convert import pytorch2paddle
@@ -291,7 +297,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:")):
-    # YOLOv5 CoreML export
+    """Exports YOLOv5 model to CoreML format with optional NMS, INT8, and FP16 support; requires coremltools."""
     check_requirements("coremltools")
     import coremltools as ct
 
@@ -316,7 +322,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:")):
-    # YOLOv5 TensorRT export https://developer.nvidia.com/tensorrt
+    """
+    Exports a YOLOv5 model to TensorRT engine format, requiring GPU and TensorRT>=7.0.0.
+
+    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
@@ -440,7 +450,7 @@ def export_saved_model(
 
 @try_export
 def export_pb(keras_model, file, prefix=colorstr("TensorFlow GraphDef:")):
-    # YOLOv5 TensorFlow GraphDef *.pb export https://github.com/leimao/Frozen_Graph_TensorFlow
+    """Exports YOLOv5 model to TensorFlow GraphDef *.pb format; 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
 
@@ -493,7 +503,11 @@ def export_tflite(
 
 @try_export
 def export_edgetpu(file, prefix=colorstr("Edge TPU:")):
-    # YOLOv5 Edge TPU export https://coral.ai/docs/edgetpu/models-intro/
+    """
+    Exports a YOLOv5 model to Edge TPU compatible TFLite format; requires Linux and Edge TPU compiler.
+
+    https://coral.ai/docs/edgetpu/models-intro/
+    """
     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 +545,7 @@ def export_edgetpu(file, prefix=colorstr("Edge TPU:")):
 
 @try_export
 def export_tfjs(file, int8, prefix=colorstr("TensorFlow.js:")):
-    # YOLOv5 TensorFlow.js export
+    """Exports a YOLOv5 model to TensorFlow.js format, optionally with uint8 quantization."""
     check_requirements("tensorflowjs")
     import tensorflowjs as tfjs
 
@@ -568,7 +582,11 @@ 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 metadata to a model file, supporting multiple outputs, as specified by TensorFlow guidelines.
+
+    https://www.tensorflow.org/lite/models/convert/metadata
+    """
     with contextlib.suppress(ImportError):
         # check_requirements('tflite_support')
         from tflite_support import flatbuffers
@@ -601,7 +619,9 @@ def add_tflite_metadata(file, metadata, num_outputs):
 
 
 def pipeline_coreml(model, im, file, names, y, prefix=colorstr("CoreML Pipeline:")):
-    # YOLOv5 CoreML pipeline
+    """Converts a PyTorch YOLOv5 model to CoreML format with NMS, handling different input/output shapes and saving the
+    model.
+    """
     import coremltools as ct
     from PIL import Image
 
@@ -869,6 +889,7 @@ def run(
 
 
 def parse_opt(known=False):
+    """Parses command-line arguments for YOLOv5 model export configurations, returning the parsed 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 / "yolov5s.pt", help="model.pt path(s)")
@@ -904,6 +925,7 @@ def parse_opt(known=False):
 
 
 def main(opt):
+    """Executes the YOLOv5 model inference or export with 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,77 @@ 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):
-    # YOLOv5 custom or local model
+    """Loads a custom or local YOLOv5 model from a given path with optional autoshaping and device specification."""
     return _create(path, autoshape=autoshape, verbose=_verbose, device=device)
 
 
 def yolov5n(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
-    # YOLOv5-nano model https://github.com/ultralytics/yolov5
+    """Instantiates the YOLOv5-nano model with options for pretraining, input channels, class count, autoshaping,
+    verbosity, and device.
+    """
     return _create("yolov5n", pretrained, channels, classes, autoshape, _verbose, device)
 
 
 def yolov5s(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
-    # YOLOv5-small model https://github.com/ultralytics/yolov5
+    """Creates YOLOv5-small model with options for pretraining, input channels, class count, autoshaping, verbosity, and
+    device.
+    """
     return _create("yolov5s", pretrained, channels, classes, autoshape, _verbose, device)
 
 
 def yolov5m(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
-    # YOLOv5-medium model https://github.com/ultralytics/yolov5
+    """Instantiates the YOLOv5-medium model with customizable pretraining, channel count, class count, autoshaping,
+    verbosity, and device.
+    """
     return _create("yolov5m", pretrained, channels, classes, autoshape, _verbose, device)
 
 
 def yolov5l(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
-    # YOLOv5-large model https://github.com/ultralytics/yolov5
+    """Creates YOLOv5-large model with options for pretraining, channels, classes, autoshaping, verbosity, 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):
-    # YOLOv5-xlarge model https://github.com/ultralytics/yolov5
+    """Instantiates the YOLOv5-xlarge model with customizable pretraining, channel count, class count, autoshaping,
+    verbosity, and device.
+    """
     return _create("yolov5x", pretrained, channels, classes, autoshape, _verbose, device)
 
 
 def yolov5n6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
-    # YOLOv5-nano-P6 model https://github.com/ultralytics/yolov5
+    """Creates YOLOv5-nano-P6 model with options for pretraining, channels, classes, autoshaping, verbosity, 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):
-    # YOLOv5-small-P6 model https://github.com/ultralytics/yolov5
+    """Instantiate YOLOv5-small-P6 model with options for pretraining, input channels, number of classes, autoshaping,
+    verbosity, and device selection.
+    """
     return _create("yolov5s6", pretrained, channels, classes, autoshape, _verbose, device)
 
 
 def yolov5m6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
-    # YOLOv5-medium-P6 model https://github.com/ultralytics/yolov5
+    """Creates YOLOv5-medium-P6 model with options for pretraining, channel count, class count, autoshaping, verbosity,
+    and device.
+    """
     return _create("yolov5m6", pretrained, channels, classes, autoshape, _verbose, device)
 
 
 def yolov5l6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
-    # YOLOv5-large-P6 model https://github.com/ultralytics/yolov5
+    """Instantiates the YOLOv5-large-P6 model with customizable pretraining, channel and class counts, autoshaping,
+    verbosity, and device selection.
+    """
     return _create("yolov5l6", pretrained, channels, classes, autoshape, _verbose, device)
 
 
 def yolov5x6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
-    # YOLOv5-xlarge-P6 model https://github.com/ultralytics/yolov5
+    """Creates YOLOv5-xlarge-P6 model with options for pretraining, channels, classes, autoshaping, verbosity, and
+    device.
+    """
     return _create("yolov5x6", pretrained, channels, classes, autoshape, _verbose, device)