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>

utils/segment/augmentations.py

@@ -12,7 +12,11 @@ from ..general import resample_segments, segment2box
 
 
 def mixup(im, labels, segments, im2, labels2, segments2):
-    # Applies MixUp augmentation https://arxiv.org/pdf/1710.09412.pdf
+    """
+    Applies MixUp augmentation blending two images, labels, and segments with a random ratio.
+
+    See https://arxiv.org/pdf/1710.09412.pdf
+    """
     r = np.random.beta(32.0, 32.0)  # mixup ratio, alpha=beta=32.0
     im = (im * r + im2 * (1 - r)).astype(np.uint8)
     labels = np.concatenate((labels, labels2), 0)

utils/segment/dataloaders.py

@@ -123,6 +123,7 @@ class LoadImagesAndLabelsAndMasks(LoadImagesAndLabels):  # for training/testing
         self.overlap = overlap
 
     def __getitem__(self, index):
+        """Returns a transformed item from the dataset at the specified index, handling indexing and image weighting."""
         index = self.indices[index]  # linear, shuffled, or image_weights
 
         hyp = self.hyp
@@ -230,7 +231,7 @@ class LoadImagesAndLabelsAndMasks(LoadImagesAndLabels):  # for training/testing
         return (torch.from_numpy(img), labels_out, self.im_files[index], shapes, masks)
 
     def load_mosaic(self, index):
-        # YOLOv5 4-mosaic loader. Loads 1 image + 3 random images into a 4-image mosaic
+        """Loads 1 image + 3 random images into a 4-image YOLOv5 mosaic, adjusting labels and segments accordingly."""
         labels4, segments4 = [], []
         s = self.img_size
         yc, xc = (int(random.uniform(-x, 2 * s + x)) for x in self.mosaic_border)  # mosaic center x, y
@@ -291,6 +292,7 @@ class LoadImagesAndLabelsAndMasks(LoadImagesAndLabels):  # for training/testing
 
     @staticmethod
     def collate_fn(batch):
+        """Custom collation function for DataLoader, batches images, labels, paths, shapes, and segmentation masks."""
         img, label, path, shapes, masks = zip(*batch)  # transposed
         batched_masks = torch.cat(masks, 0)
         for i, l in enumerate(label):

utils/segment/general.py

@@ -144,7 +144,9 @@ def masks_iou(mask1, mask2, eps=1e-7):
 
 
 def masks2segments(masks, strategy="largest"):
-    # Convert masks(n,160,160) into segments(n,xy)
+    """Converts binary (n,160,160) masks to polygon segments with options for concatenation or selecting the largest
+    segment.
+    """
     segments = []
     for x in masks.int().cpu().numpy().astype("uint8"):
         c = cv2.findContours(x, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]

utils/segment/loss.py

@@ -12,6 +12,9 @@ from .general import crop_mask
 class ComputeLoss:
     # Compute losses
     def __init__(self, model, autobalance=False, overlap=False):
+        """Initializes the compute loss function for YOLOv5 models with options for autobalancing and overlap
+        handling.
+        """
         self.sort_obj_iou = False
         self.overlap = overlap
         device = next(model.parameters()).device  # get model device
@@ -109,13 +112,15 @@ class ComputeLoss:
         return loss * bs, torch.cat((lbox, lseg, lobj, lcls)).detach()
 
     def single_mask_loss(self, gt_mask, pred, proto, xyxy, area):
-        # Mask loss for one image
+        """Calculates and normalizes single mask loss for YOLOv5 between predicted and ground truth masks."""
         pred_mask = (pred @ proto.view(self.nm, -1)).view(-1, *proto.shape[1:])  # (n,32) @ (32,80,80) -> (n,80,80)
         loss = F.binary_cross_entropy_with_logits(pred_mask, gt_mask, reduction="none")
         return (crop_mask(loss, xyxy).mean(dim=(1, 2)) / area).mean()
 
     def build_targets(self, p, targets):
-        # Build targets for compute_loss(), input targets(image,class,x,y,w,h)
+        """Prepares YOLOv5 targets for loss computation; inputs targets (image, class, x, y, w, h), output target
+        classes/boxes.
+        """
         na, nt = self.na, targets.shape[0]  # number of anchors, targets
         tcls, tbox, indices, anch, tidxs, xywhn = [], [], [], [], [], []
         gain = torch.ones(8, device=self.device)  # normalized to gridspace gain

utils/segment/metrics.py

@@ -7,7 +7,7 @@ from ..metrics import ap_per_class
 
 
 def fitness(x):
-    # Model fitness as a weighted combination of metrics
+    """Evaluates model fitness by a weighted sum of 8 metrics, `x`: [N,8] array, weights: [0.1, 0.9] for mAP and F1."""
     w = [0.0, 0.0, 0.1, 0.9, 0.0, 0.0, 0.1, 0.9]
     return (x[:, :8] * w).sum(1)
 
@@ -128,6 +128,7 @@ class Metric:
         return (self.p[i], self.r[i], self.ap50[i], self.ap[i])
 
     def get_maps(self, nc):
+        """Calculates and returns mean Average Precision (mAP) for each class given number of classes `nc`."""
         maps = np.zeros(nc) + self.map
         for i, c in enumerate(self.ap_class_index):
             maps[c] = self.ap[i]
@@ -162,17 +163,22 @@ class Metrics:
         self.metric_mask.update(list(results["masks"].values()))
 
     def mean_results(self):
+        """Computes and returns the mean results for both box and mask metrics by summing their individual means."""
         return self.metric_box.mean_results() + self.metric_mask.mean_results()
 
     def class_result(self, i):
+        """Returns the sum of box and mask metric results for a specified class index `i`."""
         return self.metric_box.class_result(i) + self.metric_mask.class_result(i)
 
     def get_maps(self, nc):
+        """Calculates and returns the sum of mean average precisions (mAPs) for both box and mask metrics for `nc`
+        classes.
+        """
         return self.metric_box.get_maps(nc) + self.metric_mask.get_maps(nc)
 
     @property
     def ap_class_index(self):
-        # boxes and masks have the same ap_class_index
+        """Returns the class index for average precision, shared by both box and mask metrics."""
         return self.metric_box.ap_class_index
 
 

utils/segment/plots.py

@@ -15,7 +15,7 @@ from ..plots import Annotator, colors
 
 @threaded
 def plot_images_and_masks(images, targets, masks, paths=None, fname="images.jpg", names=None):
-    # Plot image grid with labels
+    """Plots a grid of images, their labels, and masks with optional resizing and annotations, saving to fname."""
     if isinstance(images, torch.Tensor):
         images = images.cpu().float().numpy()
     if isinstance(targets, torch.Tensor):
@@ -111,7 +111,11 @@ def plot_images_and_masks(images, targets, masks, paths=None, fname="images.jpg"
 
 
 def plot_results_with_masks(file="path/to/results.csv", dir="", best=True):
-    # Plot training results.csv. Usage: from utils.plots import *; plot_results('path/to/results.csv')
+    """
+    Plots training results from CSV files, plotting best or last result highlights based on `best` parameter.
+
+    Example: from utils.plots import *; plot_results('path/to/results.csv')
+    """
     save_dir = Path(file).parent if file else Path(dir)
     fig, ax = plt.subplots(2, 8, figsize=(18, 6), tight_layout=True)
     ax = ax.ravel()