ultralytics 8.3.142__py3-none-any.whl → 8.3.144__py3-none-any.whl

ultralytics/utils/autodevice.py

@@ -1,5 +1,7 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
+from typing import Any, Dict, List, Optional
+
 from ultralytics.utils import LOGGER
 from ultralytics.utils.checks import check_requirements
 
@@ -19,17 +21,32 @@ class GPUInfo:
         pynvml (module | None): The `pynvml` module if successfully imported and initialized, otherwise `None`.
         nvml_available (bool): Indicates if `pynvml` is ready for use. True if import and `nvmlInit()` succeeded,
             False otherwise.
-        gpu_stats (list[dict]): A list of dictionaries, each holding stats for one GPU. Populated on initialization
-            and by `refresh_stats()`. Keys include: 'index', 'name', 'utilization' (%), 'memory_used' (MiB),
-            'memory_total' (MiB), 'memory_free' (MiB), 'temperature' (C), 'power_draw' (W),
+        gpu_stats (List[Dict[str, Any]]): A list of dictionaries, each holding stats for one GPU. Populated on
+            initialization and by `refresh_stats()`. Keys include: 'index', 'name', 'utilization' (%),
+            'memory_used' (MiB), 'memory_total' (MiB), 'memory_free' (MiB), 'temperature' (C), 'power_draw' (W),
             'power_limit' (W or 'N/A'). Empty if NVML is unavailable or queries fail.
+
+    Methods:
+        refresh_stats: Refresh the internal gpu_stats list by querying NVML.
+        print_status: Print GPU status in a compact table format using current stats.
+        select_idle_gpu: Select the most idle GPUs based on utilization and free memory.
+        shutdown: Shut down NVML if it was initialized.
+
+    Examples:
+        Initialize GPUInfo and print status
+        >>> gpu_info = GPUInfo()
+        >>> gpu_info.print_status()
+
+        Select idle GPUs with minimum memory requirements
+        >>> selected = gpu_info.select_idle_gpu(count=2, min_memory_fraction=0.2)
+        >>> print(f"Selected GPU indices: {selected}")
     """
 
     def __init__(self):
-        """Initializes GPUInfo, attempting to import and initialize pynvml."""
-        self.pynvml = None
-        self.nvml_available = False
-        self.gpu_stats = []
+        """Initialize GPUInfo, attempting to import and initialize pynvml."""
+        self.pynvml: Optional[Any] = None
+        self.nvml_available: bool = False
+        self.gpu_stats: List[Dict[str, Any]] = []
 
         try:
             check_requirements("pynvml>=12.0.0")
@@ -41,11 +58,11 @@ class GPUInfo:
             LOGGER.warning(f"Failed to initialize pynvml, GPU stats disabled: {e}")
 
     def __del__(self):
-        """Ensures NVML is shut down when the object is garbage collected."""
+        """Ensure NVML is shut down when the object is garbage collected."""
         self.shutdown()
 
     def shutdown(self):
-        """Shuts down NVML if it was initialized."""
+        """Shut down NVML if it was initialized."""
         if self.nvml_available and self.pynvml:
             try:
                 self.pynvml.nvmlShutdown()
@@ -54,7 +71,7 @@ class GPUInfo:
             self.nvml_available = False
 
     def refresh_stats(self):
-        """Refreshes the internal gpu_stats list by querying NVML."""
+        """Refresh the internal gpu_stats list by querying NVML."""
         self.gpu_stats = []
         if not self.nvml_available or not self.pynvml:
             return
@@ -67,8 +84,8 @@ class GPUInfo:
             LOGGER.warning(f"Error during device query: {e}")
             self.gpu_stats = []
 
-    def _get_device_stats(self, index):
-        """Gets stats for a single GPU device."""
+    def _get_device_stats(self, index: int) -> Dict[str, Any]:
+        """Get stats for a single GPU device."""
         handle = self.pynvml.nvmlDeviceGetHandleByIndex(index)
         memory = self.pynvml.nvmlDeviceGetMemoryInfo(handle)
         util = self.pynvml.nvmlDeviceGetUtilizationRates(handle)
@@ -86,16 +103,16 @@ class GPUInfo:
             "index": index,
             "name": self.pynvml.nvmlDeviceGetName(handle),
             "utilization": util.gpu if util else -1,
-            "memory_used": memory.used >> 20 if memory else -1,
+            "memory_used": memory.used >> 20 if memory else -1,  # Convert bytes to MiB
             "memory_total": memory.total >> 20 if memory else -1,
             "memory_free": memory.free >> 20 if memory else -1,
             "temperature": safe_get(self.pynvml.nvmlDeviceGetTemperature, handle, temp_type),
-            "power_draw": safe_get(self.pynvml.nvmlDeviceGetPowerUsage, handle, divisor=1000),
+            "power_draw": safe_get(self.pynvml.nvmlDeviceGetPowerUsage, handle, divisor=1000),  # Convert mW to W
             "power_limit": safe_get(self.pynvml.nvmlDeviceGetEnforcedPowerLimit, handle, divisor=1000),
         }
 
     def print_status(self):
-        """Prints GPU status in a compact table format using current stats."""
+        """Print GPU status in a compact table format using current stats."""
         self.refresh_stats()
         if not self.gpu_stats:
             LOGGER.warning("No GPU stats available.")
@@ -116,23 +133,29 @@ class GPUInfo:
 
         LOGGER.info(f"{'-' * len(hdr)}\n")
 
-    def select_idle_gpu(self, count=1, min_memory_fraction=0):
+    def select_idle_gpu(
+        self, count: int = 1, min_memory_fraction: float = 0, min_util_fraction: float = 0
+    ) -> List[int]:
         """
-        Selects the 'count' most idle GPUs based on utilization and free memory.
+        Select the most idle GPUs based on utilization and free memory.
 
         Args:
-            count (int): The number of idle GPUs to select. Defaults to 1.
-            min_memory_fraction (float): Minimum free memory required (fraction). Defaults to 0.
+            count (int): The number of idle GPUs to select.
+            min_memory_fraction (float): Minimum free memory required as a fraction of total memory.
+            min_util_fraction (float): Minimum free utilization rate required from 0.0 - 1.0.
 
         Returns:
-            (list[int]): Indices of the selected GPUs, sorted by idleness.
+            (List[int]): Indices of the selected GPUs, sorted by idleness (lowest utilization first).
 
         Notes:
              Returns fewer than 'count' if not enough qualify or exist.
              Returns basic CUDA indices if NVML fails. Empty list if no GPUs found.
         """
         assert min_memory_fraction <= 1.0, f"min_memory_fraction must be <= 1.0, got {min_memory_fraction}"
-        LOGGER.info(f"Searching for {count} idle GPUs with >= {min_memory_fraction * 100:.1f}% free memory...")
+        assert min_util_fraction <= 1.0, f"min_util_fraction must be <= 1.0, got {min_util_fraction}"
+        LOGGER.info(
+            f"Searching for {count} idle GPUs with free memory >= {min_memory_fraction * 100:.1f}% and free utilization >= {min_util_fraction * 100:.1f}%..."
+        )
 
         if count <= 0:
             return []
@@ -147,7 +170,7 @@ class GPUInfo:
             gpu
             for gpu in self.gpu_stats
             if gpu.get("memory_free", 0) / gpu.get("memory_total", 1) >= min_memory_fraction
-            and gpu.get("utilization", -1) != -1
+            and (100 - gpu.get("utilization", 100)) >= min_util_fraction * 100
         ]
         eligible_gpus.sort(key=lambda x: (x.get("utilization", 101), -x.get("memory_free", 0)))
 
@@ -157,19 +180,26 @@ class GPUInfo:
         if selected:
             LOGGER.info(f"Selected idle CUDA devices {selected}")
         else:
-            LOGGER.warning(f"No GPUs met criteria (Util != -1, Free Mem >= {min_memory_fraction * 100:.1f}%).")
+            LOGGER.warning(
+                f"No GPUs met criteria (Free Mem >= {min_memory_fraction * 100:.1f}% and Free Util >= {min_util_fraction * 100:.1f}%)."
+            )
 
         return selected
 
 
 if __name__ == "__main__":
     required_free_mem_fraction = 0.2  # Require 20% free VRAM
+    required_free_util_fraction = 0.2  # Require 20% free utilization
     num_gpus_to_select = 1
 
     gpu_info = GPUInfo()
     gpu_info.print_status()
 
-    selected = gpu_info.select_idle_gpu(count=num_gpus_to_select, min_memory_fraction=required_free_mem_fraction)
+    selected = gpu_info.select_idle_gpu(
+        count=num_gpus_to_select,
+        min_memory_fraction=required_free_mem_fraction,
+        min_util_fraction=required_free_util_fraction,
+    )
     if selected:
         print(f"\n==> Using selected GPU indices: {selected}")
         devices = [f"cuda:{idx}" for idx in selected]

ultralytics/utils/benchmarks.py

@@ -34,6 +34,7 @@ import re
 import shutil
 import time
 from pathlib import Path
+from typing import List, Optional, Tuple, Union
 
 import numpy as np
 import torch.cuda
@@ -226,7 +227,7 @@ class RF100Benchmark:
         self.rf = None
         self.val_metrics = ["class", "images", "targets", "precision", "recall", "map50", "map95"]
 
-    def set_key(self, api_key):
+    def set_key(self, api_key: str):
         """
         Set Roboflow API key for processing.
 
@@ -243,7 +244,7 @@ class RF100Benchmark:
 
         self.rf = Roboflow(api_key=api_key)
 
-    def parse_dataset(self, ds_link_txt="datasets_links.txt"):
+    def parse_dataset(self, ds_link_txt: str = "datasets_links.txt"):
         """
         Parse dataset links and download datasets.
 
@@ -281,14 +282,14 @@ class RF100Benchmark:
         return self.ds_names, self.ds_cfg_list
 
     @staticmethod
-    def fix_yaml(path):
+    def fix_yaml(path: Path):
         """Fix the train and validation paths in a given YAML file."""
         yaml_data = YAML.load(path)
         yaml_data["train"] = "train/images"
         yaml_data["val"] = "valid/images"
         YAML.dump(yaml_data, path)
 
-    def evaluate(self, yaml_path, val_log_file, eval_log_file, list_ind):
+    def evaluate(self, yaml_path: str, val_log_file: str, eval_log_file: str, list_ind: int):
         """
         Evaluate model performance on validation results.
 
@@ -343,6 +344,8 @@ class RF100Benchmark:
         with open(eval_log_file, "a", encoding="utf-8") as f:
             f.write(f"{self.ds_names[list_ind]}: {map_val}\n")
 
+        return float(map_val)
+
 
 class ProfileModels:
     """
@@ -361,15 +364,15 @@ class ProfileModels:
         device (torch.device): Device used for profiling.
 
     Methods:
-        profile: Profiles the models and prints the result.
-        get_files: Gets all relevant model files.
-        get_onnx_model_info: Extracts metadata from an ONNX model.
-        iterative_sigma_clipping: Applies sigma clipping to remove outliers.
-        profile_tensorrt_model: Profiles a TensorRT model.
-        profile_onnx_model: Profiles an ONNX model.
-        generate_table_row: Generates a table row with model metrics.
-        generate_results_dict: Generates a dictionary of profiling results.
-        print_table: Prints a formatted table of results.
+        run: Profile YOLO models for speed and accuracy across various formats.
+        get_files: Get all relevant model files.
+        get_onnx_model_info: Extract metadata from an ONNX model.
+        iterative_sigma_clipping: Apply sigma clipping to remove outliers.
+        profile_tensorrt_model: Profile a TensorRT model.
+        profile_onnx_model: Profile an ONNX model.
+        generate_table_row: Generate a table row with model metrics.
+        generate_results_dict: Generate a dictionary of profiling results.
+        print_table: Print a formatted table of results.
 
     Examples:
         Profile models and print results
@@ -380,14 +383,14 @@ class ProfileModels:
 
     def __init__(
         self,
-        paths: list,
-        num_timed_runs=100,
-        num_warmup_runs=10,
-        min_time=60,
-        imgsz=640,
-        half=True,
-        trt=True,
-        device=None,
+        paths: List[str],
+        num_timed_runs: int = 100,
+        num_warmup_runs: int = 10,
+        min_time: float = 60,
+        imgsz: int = 640,
+        half: bool = True,
+        trt: bool = True,
+        device: Optional[Union[torch.device, str]] = None,
     ):
         """
         Initialize the ProfileModels class for profiling models.
@@ -425,7 +428,7 @@ class ProfileModels:
         Profile YOLO models for speed and accuracy across various formats including ONNX and TensorRT.
 
         Returns:
-            (List[Dict]): List of dictionaries containing profiling results for each model.
+            (List[dict]): List of dictionaries containing profiling results for each model.
 
         Examples:
             Profile models and print results
@@ -437,7 +440,7 @@ class ProfileModels:
 
         if not files:
             LOGGER.warning("No matching *.pt or *.onnx files found.")
-            return
+            return []
 
         table_rows = []
         output = []
@@ -498,11 +501,11 @@ class ProfileModels:
 
     @staticmethod
     def get_onnx_model_info(onnx_file: str):
-        """Extracts metadata from an ONNX model file including parameters, GFLOPs, and input shape."""
+        """Extract metadata from an ONNX model file including parameters, GFLOPs, and input shape."""
         return 0.0, 0.0, 0.0, 0.0  # return (num_layers, num_params, num_gradients, num_flops)
 
     @staticmethod
-    def iterative_sigma_clipping(data, sigma=2, max_iters=3):
+    def iterative_sigma_clipping(data: np.ndarray, sigma: float = 2, max_iters: int = 3):
         """
         Apply iterative sigma clipping to data to remove outliers.
 
@@ -627,7 +630,13 @@ class ProfileModels:
         run_times = self.iterative_sigma_clipping(np.array(run_times), sigma=2, max_iters=5)  # sigma clipping
         return np.mean(run_times), np.std(run_times)
 
-    def generate_table_row(self, model_name, t_onnx, t_engine, model_info):
+    def generate_table_row(
+        self,
+        model_name: str,
+        t_onnx: Tuple[float, float],
+        t_engine: Tuple[float, float],
+        model_info: Tuple[float, float, float, float],
+    ):
         """
         Generate a table row string with model performance metrics.
 
@@ -647,7 +656,12 @@ class ProfileModels:
         )
 
     @staticmethod
-    def generate_results_dict(model_name, t_onnx, t_engine, model_info):
+    def generate_results_dict(
+        model_name: str,
+        t_onnx: Tuple[float, float],
+        t_engine: Tuple[float, float],
+        model_info: Tuple[float, float, float, float],
+    ):
         """
         Generate a dictionary of profiling results.
 
@@ -670,7 +684,7 @@ class ProfileModels:
         }
 
     @staticmethod
-    def print_table(table_rows):
+    def print_table(table_rows: List[str]):
         """
         Print a formatted table of model profiling results.
 

ultralytics/utils/callbacks/base.py

@@ -179,9 +179,9 @@ def get_default_callbacks():
     Get the default callbacks for Ultralytics training, validation, prediction, and export processes.
 
     Returns:
-        (dict): Dictionary of default callbacks for various training events. Each key in the dictionary represents an
-            event during the training process, and the corresponding value is a list of callback functions that are
-            executed when that event occurs.
+        (dict): Dictionary of default callbacks for various training events. Each key represents an event during the
+            training process, and the corresponding value is a list of callback functions executed when that event
+            occurs.
 
     Examples:
         >>> callbacks = get_default_callbacks()

ultralytics/utils/callbacks/clearml.py

@@ -56,7 +56,7 @@ def _log_plot(title: str, plot_path: str) -> None:
 
 
 def on_pretrain_routine_start(trainer) -> None:
-    """Runs at start of pretraining routine; initializes and connects/logs task to ClearML."""
+    """Initialize and connect ClearML task at the start of pretraining routine."""
     try:
         if task := Task.current_task():
             # WARNING: make sure the automatic pytorch and matplotlib bindings are disabled!
@@ -85,9 +85,9 @@ def on_pretrain_routine_start(trainer) -> None:
 
 
 def on_train_epoch_end(trainer) -> None:
-    """Logs debug samples for the first epoch of YOLO training and reports current training progress."""
+    """Log debug samples for the first epoch and report current training progress."""
     if task := Task.current_task():
-        # Log debug samples
+        # Log debug samples for first epoch only
         if trainer.epoch == 1:
             _log_debug_samples(sorted(trainer.save_dir.glob("train_batch*.jpg")), "Mosaic")
         # Report the current training progress
@@ -98,7 +98,7 @@ def on_train_epoch_end(trainer) -> None:
 
 
 def on_fit_epoch_end(trainer) -> None:
-    """Reports model information to logger at the end of an epoch."""
+    """Report model information and metrics to logger at the end of an epoch."""
     if task := Task.current_task():
         # Report epoch time and validation metrics
         task.get_logger().report_scalar(
@@ -114,23 +114,23 @@ def on_fit_epoch_end(trainer) -> None:
 
 
 def on_val_end(validator) -> None:
-    """Logs validation results including labels and predictions."""
+    """Log validation results including labels and predictions."""
     if Task.current_task():
-        # Log val_labels and val_pred
+        # Log validation labels and predictions
         _log_debug_samples(sorted(validator.save_dir.glob("val*.jpg")), "Validation")
 
 
 def on_train_end(trainer) -> None:
-    """Logs final model and its name on training completion."""
+    """Log final model and training results on training completion."""
     if task := Task.current_task():
-        # Log final results, CM matrix + PR plots
+        # Log final results, confusion matrix and PR plots
         files = [
             "results.png",
             "confusion_matrix.png",
             "confusion_matrix_normalized.png",
             *(f"{x}_curve.png" for x in ("F1", "PR", "P", "R")),
         ]
-        files = [(trainer.save_dir / f) for f in files if (trainer.save_dir / f).exists()]  # filter
+        files = [(trainer.save_dir / f) for f in files if (trainer.save_dir / f).exists()]  # filter existing files
         for f in files:
             _log_plot(title=f.stem, plot_path=f)
         # Report final metrics

ultralytics/utils/callbacks/comet.py

@@ -37,7 +37,7 @@ except (ImportError, AssertionError):
 
 
 def _get_comet_mode() -> str:
-    """Returns the mode of comet set in the environment variables, defaults to 'online' if not set."""
+    """Return the Comet mode from environment variables, defaulting to 'online'."""
     comet_mode = os.getenv("COMET_MODE")
     if comet_mode is not None:
         LOGGER.warning(
@@ -52,7 +52,7 @@ def _get_comet_mode() -> str:
 
 
 def _get_comet_model_name() -> str:
-    """Returns the model name for Comet from the environment variable COMET_MODEL_NAME or defaults to 'Ultralytics'."""
+    """Return the Comet model name from environment variable or default to 'Ultralytics'."""
     return os.getenv("COMET_MODEL_NAME", "Ultralytics")
 
 
@@ -62,31 +62,34 @@ def _get_eval_batch_logging_interval() -> int:
 
 
 def _get_max_image_predictions_to_log() -> int:
-    """Get the maximum number of image predictions to log from the environment variables."""
+    """Get the maximum number of image predictions to log from environment variables."""
     return int(os.getenv("COMET_MAX_IMAGE_PREDICTIONS", 100))
 
 
 def _scale_confidence_score(score: float) -> float:
-    """Scales the given confidence score by a factor specified in an environment variable."""
+    """Scale the confidence score by a factor specified in environment variable."""
     scale = float(os.getenv("COMET_MAX_CONFIDENCE_SCORE", 100.0))
     return score * scale
 
 
 def _should_log_confusion_matrix() -> bool:
-    """Determines if the confusion matrix should be logged based on the environment variable settings."""
+    """Determine if the confusion matrix should be logged based on environment variable settings."""
     return os.getenv("COMET_EVAL_LOG_CONFUSION_MATRIX", "false").lower() == "true"
 
 
 def _should_log_image_predictions() -> bool:
-    """Determines whether to log image predictions based on a specified environment variable."""
+    """Determine whether to log image predictions based on environment variable."""
     return os.getenv("COMET_EVAL_LOG_IMAGE_PREDICTIONS", "true").lower() == "true"
 
 
 def _resume_or_create_experiment(args: SimpleNamespace) -> None:
     """
-    Resumes CometML experiment or creates a new experiment based on args.
+    Resume CometML experiment or create a new experiment based on args.
 
     Ensures that the experiment object is only created in a single process during distributed training.
+
+    Args:
+        args (SimpleNamespace): Training arguments containing project configuration and other parameters.
     """
     if RANK not in {-1, 0}:
         return
@@ -116,7 +119,15 @@ def _resume_or_create_experiment(args: SimpleNamespace) -> None:
 
 
 def _fetch_trainer_metadata(trainer) -> dict:
-    """Returns metadata for YOLO training including epoch and asset saving status."""
+    """
+    Return metadata for YOLO training including epoch and asset saving status.
+
+    Args:
+        trainer (ultralytics.engine.trainer.BaseTrainer): The YOLO trainer object containing training state and config.
+
+    Returns:
+        (dict): Dictionary containing current epoch, step, save assets flag, and final epoch flag.
+    """
     curr_epoch = trainer.epoch + 1
 
     train_num_steps_per_epoch = len(trainer.train_loader.dataset) // trainer.batch_size
@@ -135,9 +146,19 @@ def _scale_bounding_box_to_original_image_shape(
     box, resized_image_shape, original_image_shape, ratio_pad
 ) -> List[float]:
     """
-    YOLO resizes images during training and the label values are normalized based on this resized shape.
+    Scale bounding box from resized image coordinates to original image coordinates.
 
+    YOLO resizes images during training and the label values are normalized based on this resized shape.
     This function rescales the bounding box labels to the original image shape.
+
+    Args:
+        box (torch.Tensor): Bounding box in normalized xywh format.
+        resized_image_shape (tuple): Shape of the resized image (height, width).
+        original_image_shape (tuple): Shape of the original image (height, width).
+        ratio_pad (tuple): Ratio and padding information for scaling.
+
+    Returns:
+        (List[float]): Scaled bounding box coordinates in xywh format with top-left corner adjustment.
     """
     resized_image_height, resized_image_width = resized_image_shape
 
@@ -172,14 +193,14 @@ def _format_ground_truth_annotations_for_detection(img_idx, image_path, batch, c
             - 'ori_shape': Original image shapes
             - 'resized_shape': Resized image shapes
             - 'ratio_pad': Ratio and padding information
-        class_name_map (dict | None, optional): Mapping from class indices to class names.
+        class_name_map (dict, optional): Mapping from class indices to class names.
 
     Returns:
         (dict | None): Formatted ground truth annotations with the following structure:
             - 'boxes': List of box coordinates [x, y, width, height]
             - 'label': Label string with format "gt_{class_name}"
             - 'score': Confidence score (always 1.0, scaled by _scale_confidence_score)
-            Returns None if no bounding boxes are found for the image.
+        Returns None if no bounding boxes are found for the image.
     """
     indices = batch["batch_idx"] == img_idx
     bboxes = batch["bboxes"][indices]
@@ -210,7 +231,18 @@ def _format_ground_truth_annotations_for_detection(img_idx, image_path, batch, c
 
 
 def _format_prediction_annotations(image_path, metadata, class_label_map=None, class_map=None) -> Optional[dict]:
-    """Format YOLO predictions for object detection visualization."""
+    """
+    Format YOLO predictions for object detection visualization.
+
+    Args:
+        image_path (Path): Path to the image file.
+        metadata (dict): Prediction metadata containing bounding boxes and class information.
+        class_label_map (dict, optional): Mapping from class indices to class names.
+        class_map (dict, optional): Additional class mapping for label conversion.
+
+    Returns:
+        (dict | None): Formatted prediction annotations or None if no predictions exist.
+    """
     stem = image_path.stem
     image_id = int(stem) if stem.isnumeric() else stem
 
@@ -253,14 +285,14 @@ def _format_prediction_annotations(image_path, metadata, class_label_map=None, c
 
 def _extract_segmentation_annotation(segmentation_raw: str, decode: Callable) -> Optional[List[List[Any]]]:
     """
-    Extracts segmentation annotation from compressed segmentations as list of polygons.
+    Extract segmentation annotation from compressed segmentations as list of polygons.
 
     Args:
-        segmentation_raw: Raw segmentation data in compressed format.
-        decode: Function to decode the compressed segmentation data.
+        segmentation_raw (str): Raw segmentation data in compressed format.
+        decode (Callable): Function to decode the compressed segmentation data.
 
     Returns:
-        (Optional[List[List[Any]]]): List of polygon points or None if extraction fails.
+        (List[List[Any]] | None): List of polygon points or None if extraction fails.
     """
     try:
         mask = decode(segmentation_raw)
@@ -275,7 +307,20 @@ def _extract_segmentation_annotation(segmentation_raw: str, decode: Callable) ->
 def _fetch_annotations(
     img_idx, image_path, batch, prediction_metadata_map, class_label_map, class_map
 ) -> Optional[List]:
-    """Join the ground truth and prediction annotations if they exist."""
+    """
+    Join the ground truth and prediction annotations if they exist.
+
+    Args:
+        img_idx (int): Index of the image in the batch.
+        image_path (Path): Path to the image file.
+        batch (dict): Batch data containing ground truth annotations.
+        prediction_metadata_map (dict): Map of prediction metadata by image ID.
+        class_label_map (dict): Mapping from class indices to class names.
+        class_map (dict): Additional class mapping for label conversion.
+
+    Returns:
+        (List | None): List of annotation dictionaries or None if no annotations exist.
+    """
     ground_truth_annotations = _format_ground_truth_annotations_for_detection(
         img_idx, image_path, batch, class_label_map
     )
@@ -290,7 +335,7 @@ def _fetch_annotations(
 
 
 def _create_prediction_metadata_map(model_predictions) -> dict:
-    """Create metadata map for model predictions by groupings them based on image ID."""
+    """Create metadata map for model predictions by grouping them based on image ID."""
     pred_metadata_map = {}
     for prediction in model_predictions:
         pred_metadata_map.setdefault(prediction["image_id"], [])
@@ -319,11 +364,8 @@ def _log_images(experiment, image_paths, curr_step, annotations=None) -> None:
         experiment (comet_ml.Experiment): The Comet ML experiment to log images to.
         image_paths (List[Path]): List of paths to images that will be logged.
         curr_step (int): Current training step/iteration for tracking in the experiment timeline.
-        annotations (List[List[dict]], optional): Nested list of annotation dictionaries for each image. Each annotation
-            contains visualization data like bounding boxes, labels, and confidence scores.
-
-    Returns:
-        None
+        annotations (List[List[dict]], optional): Nested list of annotation dictionaries for each image. Each
+            annotation contains visualization data like bounding boxes, labels, and confidence scores.
     """
     if annotations:
         for image_path, annotation in zip(image_paths, annotations):
@@ -448,13 +490,13 @@ def _log_model(experiment, trainer) -> None:
 
 
 def _log_image_batches(experiment, trainer, curr_step: int) -> None:
-    """Log samples of images batches for train, validation, and test."""
+    """Log samples of image batches for train, validation, and test."""
     _log_images(experiment, trainer.save_dir.glob("train_batch*.jpg"), curr_step)
     _log_images(experiment, trainer.save_dir.glob("val_batch*.jpg"), curr_step)
 
 
 def on_pretrain_routine_start(trainer) -> None:
-    """Creates or resumes a CometML experiment at the start of a YOLO pre-training routine."""
+    """Create or resume a CometML experiment at the start of a YOLO pre-training routine."""
     _resume_or_create_experiment(trainer.args)