Add docstrings to YOLOv5 functions (ultralytics#12760)

* Add docstrings to top level files * Auto-format by https://ultralytics.com/actions * Add docstrings * Auto-format by https://ultralytics.com/actions * Add docstrings * Auto-format by https://ultralytics.com/actions * Add docstrings * Add docstrings * Auto-format by https://ultralytics.com/actions * Add docstrings * Add docstrings * Auto-format by https://ultralytics.com/actions * Add docstrings * Auto-format by https://ultralytics.com/actions * Add docstrings * Auto-format by https://ultralytics.com/actions * Add docstrings * Auto-format by https://ultralytics.com/actions * Add docstrings * Auto-format by https://ultralytics.com/actions * Add docstrings * Auto-format by https://ultralytics.com/actions * Add docstrings * Add docstrings * Auto-format by https://ultralytics.com/actions * Add docstrings * Auto-format by https://ultralytics.com/actions * Add docstrings * Auto-format by https://ultralytics.com/actions * Add docstrings * Auto-format by https://ultralytics.com/actions * Add docstrings * Auto-format by https://ultralytics.com/actions * Add docstrings * Auto-format by https://ultralytics.com/actions * Add docstrings * Auto-format by https://ultralytics.com/actions * Update activations.py Signed-off-by: Glenn Jocher <[email protected]> * Auto-format by https://ultralytics.com/actions --------- Signed-off-by: Glenn Jocher <[email protected]> Co-authored-by: UltralyticsAssistant <[email protected]>
SamSamhuns · Feb 25, 2024 · 41603da · 41603da
1 parent ca00c34
commit 41603da
Show file tree

Hide file tree

Showing 42 changed files with 983 additions and 322 deletions.
diff --git a/benchmarks.py b/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))
 
 

diff --git a/classify/predict.py b/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))
 

diff --git a/classify/train.py b/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)

diff --git a/classify/val.py b/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))
 

diff --git a/detect.py b/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))
 

diff --git a/export.py b/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))
 

diff --git a/hubconf.py b/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)