Add docstrings to YOLOv5 functions (#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 <glenn.jocher@ultralytics.com>

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

---------

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>

models/common.py

@@ -22,7 +22,7 @@ import torch.nn as nn
 from PIL import Image
 from torch.cuda import amp
 
-# Import 'ultralytics' package or install if if missing
+# Import 'ultralytics' package or install if missing
 try:
     import ultralytics
 
@@ -71,15 +71,18 @@ class Conv(nn.Module):
     default_act = nn.SiLU()  # default activation
 
     def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True):
+        """Initializes a standard convolution layer with optional batch normalization and activation."""
         super().__init__()
         self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p, d), groups=g, dilation=d, bias=False)
         self.bn = nn.BatchNorm2d(c2)
         self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()
 
     def forward(self, x):
+        """Applies a convolution followed by batch normalization and an activation function to the input tensor `x`."""
         return self.act(self.bn(self.conv(x)))
 
     def forward_fuse(self, x):
+        """Applies a fused convolution and activation function to the input tensor `x`."""
         return self.act(self.conv(x))
 
 
@@ -98,6 +101,11 @@ class DWConvTranspose2d(nn.ConvTranspose2d):
 class TransformerLayer(nn.Module):
     # Transformer layer https://arxiv.org/abs/2010.11929 (LayerNorm layers removed for better performance)
     def __init__(self, c, num_heads):
+        """
+        Initializes a transformer layer, sans LayerNorm for performance, with multihead attention and linear layers.
+
+        See  as described in https://arxiv.org/abs/2010.11929.
+        """
         super().__init__()
         self.q = nn.Linear(c, c, bias=False)
         self.k = nn.Linear(c, c, bias=False)
@@ -107,6 +115,7 @@ class TransformerLayer(nn.Module):
         self.fc2 = nn.Linear(c, c, bias=False)
 
     def forward(self, x):
+        """Performs forward pass using MultiheadAttention and two linear transformations with residual connections."""
         x = self.ma(self.q(x), self.k(x), self.v(x))[0] + x
         x = self.fc2(self.fc1(x)) + x
         return x
@@ -115,6 +124,9 @@ class TransformerLayer(nn.Module):
 class TransformerBlock(nn.Module):
     # Vision Transformer https://arxiv.org/abs/2010.11929
     def __init__(self, c1, c2, num_heads, num_layers):
+        """Initializes a Transformer block for vision tasks, adapting dimensions if necessary and stacking specified
+        layers.
+        """
         super().__init__()
         self.conv = None
         if c1 != c2:
@@ -124,6 +136,9 @@ class TransformerBlock(nn.Module):
         self.c2 = c2
 
     def forward(self, x):
+        """Processes input through an optional convolution, followed by Transformer layers and position embeddings for
+        object detection.
+        """
         if self.conv is not None:
             x = self.conv(x)
         b, _, w, h = x.shape
@@ -141,6 +156,9 @@ class Bottleneck(nn.Module):
         self.add = shortcut and c1 == c2
 
     def forward(self, x):
+        """Processes input through two convolutions, optionally adds shortcut if channel dimensions match; input is a
+        tensor.
+        """
         return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))
 
 
@@ -158,6 +176,9 @@ class BottleneckCSP(nn.Module):
         self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))
 
     def forward(self, x):
+        """Performs forward pass by applying layers, activation, and concatenation on input x, returning feature-
+        enhanced output.
+        """
         y1 = self.cv3(self.m(self.cv1(x)))
         y2 = self.cv2(x)
         return self.cv4(self.act(self.bn(torch.cat((y1, y2), 1))))
@@ -166,7 +187,12 @@ class BottleneckCSP(nn.Module):
 class CrossConv(nn.Module):
     # Cross Convolution Downsample
     def __init__(self, c1, c2, k=3, s=1, g=1, e=1.0, shortcut=False):
-        # ch_in, ch_out, kernel, stride, groups, expansion, shortcut
+        """
+        Initializes CrossConv with downsampling, expanding, and optionally shortcutting; `c1` input, `c2` output
+        channels.
+
+        Inputs are ch_in, ch_out, kernel, stride, groups, expansion, shortcut.
+        """
         super().__init__()
         c_ = int(c2 * e)  # hidden channels
         self.cv1 = Conv(c1, c_, (1, k), (1, s))
@@ -174,6 +200,7 @@ class CrossConv(nn.Module):
         self.add = shortcut and c1 == c2
 
     def forward(self, x):
+        """Performs feature sampling, expanding, and applies shortcut if channels match; expects `x` input tensor."""
         return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))
 
 
@@ -188,12 +215,16 @@ class C3(nn.Module):
         self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))
 
     def forward(self, x):
+        """Performs forward propagation using concatenated outputs from two convolutions and a Bottleneck sequence."""
         return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), 1))
 
 
 class C3x(C3):
     # C3 module with cross-convolutions
     def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
+        """Initializes C3x module with cross-convolutions, extending C3 with customizable channel dimensions, groups,
+        and expansion.
+        """
         super().__init__(c1, c2, n, shortcut, g, e)
         c_ = int(c2 * e)
         self.m = nn.Sequential(*(CrossConv(c_, c_, 3, 1, g, 1.0, shortcut) for _ in range(n)))
@@ -202,6 +233,9 @@ class C3x(C3):
 class C3TR(C3):
     # C3 module with TransformerBlock()
     def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
+        """Initializes C3 module with TransformerBlock for enhanced feature extraction, accepts channel sizes, shortcut
+        config, group, and expansion.
+        """
         super().__init__(c1, c2, n, shortcut, g, e)
         c_ = int(c2 * e)
         self.m = TransformerBlock(c_, c_, 4, n)
@@ -210,6 +244,9 @@ class C3TR(C3):
 class C3SPP(C3):
     # C3 module with SPP()
     def __init__(self, c1, c2, k=(5, 9, 13), n=1, shortcut=True, g=1, e=0.5):
+        """Initializes a C3 module with SPP layer for advanced spatial feature extraction, given channel sizes, kernel
+        sizes, shortcut, group, and expansion ratio.
+        """
         super().__init__(c1, c2, n, shortcut, g, e)
         c_ = int(c2 * e)
         self.m = SPP(c_, c_, k)
@@ -218,6 +255,7 @@ class C3SPP(C3):
 class C3Ghost(C3):
     # C3 module with GhostBottleneck()
     def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
+        """Initializes YOLOv5's C3 module with Ghost Bottlenecks for efficient feature extraction."""
         super().__init__(c1, c2, n, shortcut, g, e)
         c_ = int(c2 * e)  # hidden channels
         self.m = nn.Sequential(*(GhostBottleneck(c_, c_) for _ in range(n)))
@@ -226,6 +264,7 @@ class C3Ghost(C3):
 class SPP(nn.Module):
     # Spatial Pyramid Pooling (SPP) layer https://arxiv.org/abs/1406.4729
     def __init__(self, c1, c2, k=(5, 9, 13)):
+        """Initializes SPP layer with Spatial Pyramid Pooling, ref: https://arxiv.org/abs/1406.4729, args: c1 (input channels), c2 (output channels), k (kernel sizes)."""
         super().__init__()
         c_ = c1 // 2  # hidden channels
         self.cv1 = Conv(c1, c_, 1, 1)
@@ -233,6 +272,9 @@ class SPP(nn.Module):
         self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k])
 
     def forward(self, x):
+        """Applies convolution and max pooling layers to the input tensor `x`, concatenates results, and returns output
+        tensor.
+        """
         x = self.cv1(x)
         with warnings.catch_warnings():
             warnings.simplefilter("ignore")  # suppress torch 1.9.0 max_pool2d() warning
@@ -249,6 +291,7 @@ class SPPF(nn.Module):
         self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)
 
     def forward(self, x):
+        """Processes input through a series of convolutions and max pooling operations for feature extraction."""
         x = self.cv1(x)
         with warnings.catch_warnings():
             warnings.simplefilter("ignore")  # suppress torch 1.9.0 max_pool2d() warning
@@ -278,6 +321,7 @@ class GhostConv(nn.Module):
         self.cv2 = Conv(c_, c_, 5, 1, None, c_, act=act)
 
     def forward(self, x):
+        """Performs forward pass, concatenating outputs of two convolutions on input `x`: shape (B,C,H,W)."""
         y = self.cv1(x)
         return torch.cat((y, self.cv2(y)), 1)
 
@@ -297,16 +341,23 @@ class GhostBottleneck(nn.Module):
         )
 
     def forward(self, x):
+        """Processes input through conv and shortcut layers, returning their summed output."""
         return self.conv(x) + self.shortcut(x)
 
 
 class Contract(nn.Module):
     # Contract width-height into channels, i.e. x(1,64,80,80) to x(1,256,40,40)
     def __init__(self, gain=2):
+        """Initializes a layer to contract spatial dimensions (width-height) into channels, e.g., input shape
+        (1,64,80,80) to (1,256,40,40).
+        """
         super().__init__()
         self.gain = gain
 
     def forward(self, x):
+        """Processes input tensor to expand channel dimensions by contracting spatial dimensions, yielding output shape
+        `(b, c*s*s, h//s, w//s)`.
+        """
         b, c, h, w = x.size()  # assert (h / s == 0) and (W / s == 0), 'Indivisible gain'
         s = self.gain
         x = x.view(b, c, h // s, s, w // s, s)  # x(1,64,40,2,40,2)
@@ -317,10 +368,19 @@ class Contract(nn.Module):
 class Expand(nn.Module):
     # Expand channels into width-height, i.e. x(1,64,80,80) to x(1,16,160,160)
     def __init__(self, gain=2):
+        """
+        Initializes the Expand module to increase spatial dimensions by redistributing channels, with an optional gain
+        factor.
+
+        Example: x(1,64,80,80) to x(1,16,160,160).
+        """
         super().__init__()
         self.gain = gain
 
     def forward(self, x):
+        """Processes input tensor x to expand spatial dimensions by redistributing channels, requiring C / gain^2 ==
+        0.
+        """
         b, c, h, w = x.size()  # assert C / s ** 2 == 0, 'Indivisible gain'
         s = self.gain
         x = x.view(b, s, s, c // s**2, h, w)  # x(1,2,2,16,80,80)
@@ -331,17 +391,21 @@ class Expand(nn.Module):
 class Concat(nn.Module):
     # Concatenate a list of tensors along dimension
     def __init__(self, dimension=1):
+        """Initializes a Concat module to concatenate tensors along a specified dimension."""
         super().__init__()
         self.d = dimension
 
     def forward(self, x):
+        """Concatenates a list of tensors along a specified dimension; `x` is a list of tensors, `dimension` is an
+        int.
+        """
         return torch.cat(x, self.d)
 
 
 class DetectMultiBackend(nn.Module):
     # YOLOv5 MultiBackend class for python inference on various backends
     def __init__(self, weights="yolov5s.pt", device=torch.device("cpu"), dnn=False, data=None, fp16=False, fuse=True):
-        # Usage:
+        """Initializes DetectMultiBackend with support for various inference backends, including PyTorch and ONNX."""
         #   PyTorch:              weights = *.pt
         #   TorchScript:                    *.torchscript
         #   ONNX Runtime:                   *.onnx
@@ -462,11 +526,13 @@ class DetectMultiBackend(nn.Module):
             import tensorflow as tf
 
             def wrap_frozen_graph(gd, inputs, outputs):
+                """Wraps a TensorFlow GraphDef for inference, returning a pruned function."""
                 x = tf.compat.v1.wrap_function(lambda: tf.compat.v1.import_graph_def(gd, name=""), [])  # wrapped
                 ge = x.graph.as_graph_element
                 return x.prune(tf.nest.map_structure(ge, inputs), tf.nest.map_structure(ge, outputs))
 
             def gd_outputs(gd):
+                """Generates a sorted list of graph outputs excluding NoOp nodes and inputs, formatted as '<name>:0'."""
                 name_list, input_list = [], []
                 for node in gd.node:  # tensorflow.core.framework.node_def_pb2.NodeDef
                     name_list.append(node.name)
@@ -540,7 +606,7 @@ class DetectMultiBackend(nn.Module):
         self.__dict__.update(locals())  # assign all variables to self
 
     def forward(self, im, augment=False, visualize=False):
-        # YOLOv5 MultiBackend inference
+        """Performs YOLOv5 inference on input images with options for augmentation and visualization."""
         b, ch, h, w = im.shape  # batch, channel, height, width
         if self.fp16 and im.dtype != torch.float16:
             im = im.half()  # to FP16
@@ -622,10 +688,11 @@ class DetectMultiBackend(nn.Module):
             return self.from_numpy(y)
 
     def from_numpy(self, x):
+        """Converts a NumPy array to a torch tensor, maintaining device compatibility."""
         return torch.from_numpy(x).to(self.device) if isinstance(x, np.ndarray) else x
 
     def warmup(self, imgsz=(1, 3, 640, 640)):
-        # Warmup model by running inference once
+        """Performs a single inference warmup to initialize model weights, accepting an `imgsz` tuple for image size."""
         warmup_types = self.pt, self.jit, self.onnx, self.engine, self.saved_model, self.pb, self.triton
         if any(warmup_types) and (self.device.type != "cpu" or self.triton):
             im = torch.empty(*imgsz, dtype=torch.half if self.fp16 else torch.float, device=self.device)  # input
@@ -634,7 +701,11 @@ class DetectMultiBackend(nn.Module):
 
     @staticmethod
     def _model_type(p="path/to/model.pt"):
-        # Return model type from model path, i.e. path='path/to/model.onnx' -> type=onnx
+        """
+        Determines model type from file path or URL, supporting various export formats.
+
+        Example: path='path/to/model.onnx' -> type=onnx
+        """
         # types = [pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle]
         from export import export_formats
         from utils.downloads import is_url
@@ -650,7 +721,7 @@ class DetectMultiBackend(nn.Module):
 
     @staticmethod
     def _load_metadata(f=Path("path/to/meta.yaml")):
-        # Load metadata from meta.yaml if it exists
+        """Loads metadata from a YAML file, returning strides and names if the file exists, otherwise `None`."""
         if f.exists():
             d = yaml_load(f)
             return d["stride"], d["names"]  # assign stride, names
@@ -668,6 +739,7 @@ class AutoShape(nn.Module):
     amp = False  # Automatic Mixed Precision (AMP) inference
 
     def __init__(self, model, verbose=True):
+        """Initializes YOLOv5 model for inference, setting up attributes and preparing model for evaluation."""
         super().__init__()
         if verbose:
             LOGGER.info("Adding AutoShape... ")
@@ -681,7 +753,11 @@ class AutoShape(nn.Module):
             m.export = True  # do not output loss values
 
     def _apply(self, fn):
-        # Apply to(), cpu(), cuda(), half() to model tensors that are not parameters or registered buffers
+        """
+        Applies to(), cpu(), cuda(), half() etc.
+
+        to model tensors excluding parameters or registered buffers.
+        """
         self = super()._apply(fn)
         if self.pt:
             m = self.model.model.model[-1] if self.dmb else self.model.model[-1]  # Detect()
@@ -693,7 +769,12 @@ class AutoShape(nn.Module):
 
     @smart_inference_mode()
     def forward(self, ims, size=640, augment=False, profile=False):
-        # Inference from various sources. For size(height=640, width=1280), RGB images example inputs are:
+        """
+        Performs inference on inputs with optional augment & profiling.
+
+        Supports various formats including file, URI, OpenCV, PIL, numpy, torch.
+        """
+        # For size(height=640, width=1280), RGB images example inputs are:
         #   file:        ims = 'data/images/zidane.jpg'  # str or PosixPath
         #   URI:             = 'https://ultralytics.com/images/zidane.jpg'
         #   OpenCV:          = cv2.imread('image.jpg')[:,:,::-1]  # HWC BGR to RGB x(640,1280,3)
@@ -761,6 +842,7 @@ class AutoShape(nn.Module):
 class Detections:
     # YOLOv5 detections class for inference results
     def __init__(self, ims, pred, files, times=(0, 0, 0), names=None, shape=None):
+        """Initializes the YOLOv5 Detections class with image info, predictions, filenames, timing and normalization."""
         super().__init__()
         d = pred[0].device  # device
         gn = [torch.tensor([*(im.shape[i] for i in [1, 0, 1, 0]), 1, 1], device=d) for im in ims]  # normalizations
@@ -778,6 +860,7 @@ class Detections:
         self.s = tuple(shape)  # inference BCHW shape
 
     def _run(self, pprint=False, show=False, save=False, crop=False, render=False, labels=True, save_dir=Path("")):
+        """Executes model predictions, displaying and/or saving outputs with optional crops and labels."""
         s, crops = "", []
         for i, (im, pred) in enumerate(zip(self.ims, self.pred)):
             s += f"\nimage {i + 1}/{len(self.pred)}: {im.shape[0]}x{im.shape[1]} "  # string
@@ -832,22 +915,42 @@ class Detections:
 
     @TryExcept("Showing images is not supported in this environment")
     def show(self, labels=True):
+        """
+        Displays detection results with optional labels.
+
+        Usage: show(labels=True)
+        """
         self._run(show=True, labels=labels)  # show results
 
     def save(self, labels=True, save_dir="runs/detect/exp", exist_ok=False):
+        """
+        Saves detection results with optional labels to a specified directory.
+
+        Usage: save(labels=True, save_dir='runs/detect/exp', exist_ok=False)
+        """
         save_dir = increment_path(save_dir, exist_ok, mkdir=True)  # increment save_dir
         self._run(save=True, labels=labels, save_dir=save_dir)  # save results
 
     def crop(self, save=True, save_dir="runs/detect/exp", exist_ok=False):
+        """
+        Crops detection results, optionally saves them to a directory.
+
+        Args: save (bool), save_dir (str), exist_ok (bool).
+        """
         save_dir = increment_path(save_dir, exist_ok, mkdir=True) if save else None
         return self._run(crop=True, save=save, save_dir=save_dir)  # crop results
 
     def render(self, labels=True):
+        """Renders detection results with optional labels on images; args: labels (bool) indicating label inclusion."""
         self._run(render=True, labels=labels)  # render results
         return self.ims
 
     def pandas(self):
-        # return detections as pandas DataFrames, i.e. print(results.pandas().xyxy[0])
+        """
+        Returns detections as pandas DataFrames for various box formats (xyxy, xyxyn, xywh, xywhn).
+
+        Example: print(results.pandas().xyxy[0]).
+        """
         new = copy(self)  # return copy
         ca = "xmin", "ymin", "xmax", "ymax", "confidence", "class", "name"  # xyxy columns
         cb = "xcenter", "ycenter", "width", "height", "confidence", "class", "name"  # xywh columns
@@ -857,7 +960,11 @@ class Detections:
         return new
 
     def tolist(self):
-        # return a list of Detections objects, i.e. 'for result in results.tolist():'
+        """
+        Converts a Detections object into a list of individual detection results for iteration.
+
+        Example: for result in results.tolist():
+        """
         r = range(self.n)  # iterable
         return [
             Detections(
@@ -872,6 +979,7 @@ class Detections:
         ]
 
     def print(self):
+        """Logs the string representation of the current object's state via the LOGGER."""
         LOGGER.info(self.__str__())
 
     def __len__(self):  # override len(results)
@@ -881,6 +989,7 @@ class Detections:
         return self._run(pprint=True)  # print results
 
     def __repr__(self):
+        """Returns a string representation of the YOLOv5 object, including its class and formatted results."""
         return f"YOLOv5 {self.__class__} instance\n" + self.__str__()
 
 
@@ -894,6 +1003,7 @@ class Proto(nn.Module):
         self.cv3 = Conv(c_, c2)
 
     def forward(self, x):
+        """Performs a forward pass using convolutional layers and upsampling on input tensor `x`."""
         return self.cv3(self.cv2(self.upsample(self.cv1(x))))
 
 
@@ -910,6 +1020,7 @@ class Classify(nn.Module):
         self.linear = nn.Linear(c_, c2)  # to x(b,c2)
 
     def forward(self, x):
+        """Processes input through conv, pool, drop, and linear layers; supports list concatenation input."""
         if isinstance(x, list):
             x = torch.cat(x, 1)
         return self.linear(self.drop(self.pool(self.conv(x)).flatten(1)))