Add docstrings and improve comments (#11229)

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>

docs/en/guides/security-alarm-system.md

@@ -67,6 +67,7 @@ server.login(from_email, password)
 
 ```python
 def send_email(to_email, from_email, object_detected=1):
+    """Sends an email notification indicating the number of objects detected; defaults to 1 object."""
     message = MIMEMultipart()
     message['From'] = from_email
     message['To'] = to_email
@@ -83,7 +84,7 @@ def send_email(to_email, from_email, object_detected=1):
 ```python
 class ObjectDetection:
     def __init__(self, capture_index):
-        # default parameters
+        """Initializes an ObjectDetection instance with a given camera index."""
         self.capture_index = capture_index
         self.email_sent = False
 
@@ -99,10 +100,12 @@ class ObjectDetection:
         self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
     def predict(self, im0):
+        """Run prediction using a YOLO model for the input image `im0`."""
         results = self.model(im0)
         return results
 
     def display_fps(self, im0):
+        """Displays the FPS on an image `im0` by calculating and overlaying as white text on a black rectangle."""
         self.end_time = time()
         fps = 1 / np.round(self.end_time - self.start_time, 2)
         text = f'FPS: {int(fps)}'
@@ -112,6 +115,7 @@ class ObjectDetection:
         cv2.putText(im0, text, (20, 70), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 0), 2)
 
     def plot_bboxes(self, results, im0):
+        """Plots bounding boxes on an image given detection results; returns annotated image and class IDs."""
         class_ids = []
         self.annotator = Annotator(im0, 3, results[0].names)
         boxes = results[0].boxes.xyxy.cpu()
@@ -123,6 +127,7 @@ class ObjectDetection:
         return im0, class_ids
 
     def __call__(self):
+        """Executes object detection on video frames from a specified camera index, plotting bounding boxes and returning modified frames."""
         cap = cv2.VideoCapture(self.capture_index)
         assert cap.isOpened()
         cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)

docs/en/guides/yolo-thread-safe-inference.md

@@ -36,6 +36,7 @@ shared_model = YOLO("yolov8n.pt")
 
 
 def predict(image_path):
+    """Predicts objects in an image using a preloaded YOLO model, take path string to image as argument."""
     results = shared_model.predict(image_path)
     # Process results
 
@@ -62,6 +63,7 @@ shared_model_2 = YOLO("yolov8n_2.pt")
 
 
 def predict(model, image_path):
+    """Runs prediction on an image using a specified YOLO model, returning the results."""
     results = model.predict(image_path)
     # Process results
 
@@ -88,7 +90,7 @@ from threading import Thread
 
 
 def thread_safe_predict(image_path):
-    # Instantiate a new model inside the thread
+    """Predict on an image using a new YOLO model instance in a thread-safe manner; takes image path as input."""
     local_model = YOLO("yolov8n.pt")
     results = local_model.predict(image_path)
     # Process results

docs/en/integrations/amazon-sagemaker.md

@@ -118,6 +118,7 @@ After creating the AWS CloudFormation Stack, the next step is to deploy YOLOv8.
 import json
 
 def output_fn(prediction_output, content_type):
+    """Formats model outputs as JSON string according to content_type, extracting attributes like boxes, masks, keypoints."""
     print("Executing output_fn from inference.py ...")
     infer = {}
     for result in prediction_output:

docs/en/integrations/gradio.md

@@ -53,6 +53,7 @@ model = YOLO("yolov8n.pt")
 
 
 def predict_image(img, conf_threshold, iou_threshold):
+    """Predicts and plots labeled objects in an image using YOLOv8 model with adjustable confidence and IOU thresholds."""
     results = model.predict(
         source=img,
         conf=conf_threshold,

docs/en/modes/predict.md

@@ -731,7 +731,7 @@ When using YOLO models in a multi-threaded application, it's important to instan
     from threading import Thread
 
     def thread_safe_predict(image_path):
-        # Instantiate a new model inside the thread
+        """Performs thread-safe prediction on an image using a locally instantiated YOLO model."""
         local_model = YOLO("yolov8n.pt")
         results = local_model.predict(image_path)
         # Process results

docs/en/usage/callbacks.md

@@ -30,7 +30,7 @@ from ultralytics import YOLO
 
 
 def on_predict_batch_end(predictor):
-    # Retrieve the batch data
+    """Handle prediction batch end by combining results with corresponding frames; modifies predictor results."""
     _, image, _, _ = predictor.batch
 
     # Ensure that image is a list

docs/en/usage/engine.md

@@ -46,6 +46,7 @@ from ultralytics.models.yolo.detect import DetectionTrainer
 
 class CustomTrainer(DetectionTrainer):
     def get_model(self, cfg, weights):
+        """Loads a custom detection model given configuration and weight files."""
         ...
 
 
@@ -65,16 +66,19 @@ from ultralytics.nn.tasks import DetectionModel
 
 class MyCustomModel(DetectionModel):
     def init_criterion(self):
+        """Initializes the loss function and adds a callback for uploading the model to Google Drive every 10 epochs."""
         ...
 
 
 class CustomTrainer(DetectionTrainer):
     def get_model(self, cfg, weights):
+        """Returns a customized detection model instance configured with specified config and weights."""
         return MyCustomModel(...)
 
 
 # callback to upload model weights
 def log_model(trainer):
+    """Logs the path of the last model weight used by the trainer."""
     last_weight_path = trainer.last
     print(last_weight_path)
 

docs/en/yolov5/tutorials/architecture_description.md

@@ -38,12 +38,14 @@ import torch.nn as nn
 
 class SPP(nn.Module):
     def __init__(self):
+        """Initializes an SPP module with three different sizes of max pooling layers."""
         super().__init__()
         self.maxpool1 = nn.MaxPool2d(5, 1, padding=2)
         self.maxpool2 = nn.MaxPool2d(9, 1, padding=4)
         self.maxpool3 = nn.MaxPool2d(13, 1, padding=6)
 
     def forward(self, x):
+        """Applies three max pooling layers on input `x` and concatenates results along channel dimension."""
         o1 = self.maxpool1(x)
         o2 = self.maxpool2(x)
         o3 = self.maxpool3(x)
@@ -52,10 +54,12 @@ class SPP(nn.Module):
 
 class SPPF(nn.Module):
     def __init__(self):
+        """Initializes an SPPF module with a specific configuration of MaxPool2d layer."""
         super().__init__()
         self.maxpool = nn.MaxPool2d(5, 1, padding=2)
 
     def forward(self, x):
+        """Applies sequential max pooling and concatenates results with input tensor; expects input tensor x of any shape."""
         o1 = self.maxpool(x)
         o2 = self.maxpool(o1)
         o3 = self.maxpool(o2)
@@ -63,6 +67,7 @@ class SPPF(nn.Module):
 
 
 def main():
+    """Compares outputs and performance of SPP and SPPF on a random tensor (8, 32, 16, 16)."""
     input_tensor = torch.rand(8, 32, 16, 16)
     spp = SPP()
     sppf = SPPF()

docs/en/yolov5/tutorials/hyperparameter_evolution.md

@@ -65,7 +65,7 @@ Fitness is the value we seek to maximize. In YOLOv5 we define a default fitness
 
 ```python
 def fitness(x):
-    # Model fitness as a weighted combination of metrics
+    """Evaluates model fitness by summing weighted metrics [P, R, mAP@0.5, mAP@0.5:0.95], x is a numpy array of shape (n, 4)."""
     w = [0.0, 0.0, 0.1, 0.9]  # weights for [P, R, mAP@0.5, mAP@0.5:0.95]
     return (x[:, :4] * w).sum(1)
 ```

docs/en/yolov5/tutorials/pytorch_hub_model_loading.md

@@ -179,6 +179,7 @@ import threading
 
 
 def run(model, im):
+    """Performs inference on an image using a given model and saves the output; model must support `.save()` method."""
     results = model(im)
     results.save()
 

tests/test_cuda.py

@@ -19,6 +19,7 @@ def test_checks():
     assert torch.cuda.is_available() == CUDA_IS_AVAILABLE
     assert torch.cuda.device_count() == CUDA_DEVICE_COUNT
 
+
 @pytest.mark.slow
 @pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason="CUDA is not available")
 def test_export_engine():

tests/test_python.py

@@ -345,12 +345,12 @@ def test_labels_and_crops():
         labels = save_path / f"labels/{im_name}.txt"
         assert labels.exists()
         # Check detections match label count
-        assert len(r.boxes.data) == len([l for l in labels.read_text().splitlines() if l])
+        assert len(r.boxes.data) == len([line for line in labels.read_text().splitlines() if line])
         # Check crops path and files
-        crop_dirs = [p for p in (save_path / "crops").iterdir()]
+        crop_dirs = list((save_path / "crops").iterdir())
         crop_files = [f for p in crop_dirs for f in p.glob("*")]
         # Crop directories match detections
-        assert all([r.names.get(c) in {d.name for d in crop_dirs} for c in cls_idxs])
+        assert all(r.names.get(c) in {d.name for d in crop_dirs} for c in cls_idxs)
         # Same number of crops as detections
         assert len([f for f in crop_files if im_name in f.name]) == len(r.boxes.data)
 
@@ -643,8 +643,8 @@ def test_yolo_world():
     model(ASSETS / "bus.jpg", conf=0.01)
 
     model = YOLO("yolov8s-worldv2.pt")  # no YOLOv8n-world model yet
-    # Training from pretrain, evaluation process is included at the final stage of training.
+    # Training from a pretrained model. Eval is included at the final stage of training.
-    # Use dota8.yaml which has less categories to reduce the inference time of CLIP model
+    # Use dota8.yaml which has fewer categories to reduce the inference time of CLIP model
     model.train(
         data="dota8.yaml",
         epochs=1,

ultralytics/data/augment.py

@@ -917,7 +917,6 @@ class Albumentations:
         return labels
 
 
-# TODO: technically this is not an augmentation, maybe we should put this to another files
 class Format:
     """
     Formats image annotations for object detection, instance segmentation, and pose estimation tasks. The class

ultralytics/solutions/parking_management.py

@@ -14,7 +14,7 @@ import tkinter as tk
 
 class ParkingPtsSelection:
     def __init__(self, master):
-        # Initialize window and widgets.
+        """Initializes the UI for selecting parking zone points in a tkinter window."""
         self.master = master
         master.title("Ultralytics Parking Zones Points Selector")
         self.initialize_ui()
@@ -109,6 +109,7 @@ class ParkingPtsSelection:
             messagebox.showwarning("Warning", "No bounding boxes to remove.")
 
     def save_to_json(self):
+        """Saves rescaled bounding boxes to 'bounding_boxes.json' based on image-to-canvas size ratio."""
         canvas_width, canvas_height = self.canvas.winfo_width(), self.canvas.winfo_height()
         width_scaling_factor = self.img_width / canvas_width
         height_scaling_factor = self.img_height / canvas_height

ultralytics/trackers/README.md

@@ -268,6 +268,7 @@ from ultralytics import YOLO
 
 
 def run_tracker_in_thread(filename, model):
+    """Starts multi-thread tracking on video from `filename` using `model` and displays results frame by frame."""
     video = cv2.VideoCapture(filename)
     frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
     for _ in range(frames):

ultralytics/utils/instance.py

@@ -343,11 +343,7 @@ class Instances:
             self.keypoints[..., 1] = self.keypoints[..., 1].clip(0, h)
 
     def remove_zero_area_boxes(self):
-        """
+        """Remove zero-area boxes, i.e. after clipping some boxes may have zero width or height."""
-        Remove zero-area boxes, i.e. after clipping some boxes may have zero width or height.
-
-        This removes them.
-        """
         good = self.bbox_areas > 0
         if not all(good):
             self._bboxes = self._bboxes[good]