ultralytics 8.3.101__py3-none-any.whl → 8.3.103__py3-none-any.whl

ultralytics/nn/text_model.py

@@ -6,7 +6,7 @@ from pathlib import Path
 import torch
 import torch.nn as nn
 
-from ultralytics.utils import LOGGER, checks
+from ultralytics.utils import checks
 from ultralytics.utils.torch_utils import smart_inference_mode
 
 try:
@@ -59,9 +59,10 @@ class TextModel(nn.Module):
 
 class CLIP(TextModel):
     """
-    OpenAI CLIP text encoder implementation.
+    Implements OpenAI's CLIP (Contrastive Language-Image Pre-training) text encoder.
 
-    This class implements the TextModel interface using OpenAI's CLIP model for text encoding.
+    This class provides a text encoder based on OpenAI's CLIP model, which can convert text into feature vectors
+    that are aligned with corresponding image features in a shared embedding space.
 
     Attributes:
         model (clip.model.CLIP): The loaded CLIP model.
@@ -70,15 +71,33 @@ class CLIP(TextModel):
     Methods:
         tokenize: Convert input texts to CLIP tokens.
         encode_text: Encode tokenized texts into normalized feature vectors.
+
+    Examples:
+        >>> from ultralytics.models.sam import CLIP
+        >>> import torch
+        >>> device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        >>> clip_model = CLIP(size="ViT-B/32", device=device)
+        >>> tokens = clip_model.tokenize(["a photo of a cat", "a photo of a dog"])
+        >>> text_features = clip_model.encode_text(tokens)
+        >>> print(text_features.shape)
     """
 
     def __init__(self, size, device):
         """
         Initialize the CLIP text encoder.
 
+        This class implements the TextModel interface using OpenAI's CLIP model for text encoding. It loads
+        a pre-trained CLIP model of the specified size and prepares it for text encoding tasks.
+
         Args:
             size (str): Model size identifier (e.g., 'ViT-B/32').
             device (torch.device): Device to load the model on.
+
+        Examples:
+            >>> import torch
+            >>> from ultralytics.models.sam.modules.clip import CLIP
+            >>> clip_model = CLIP("ViT-B/32", device=torch.device("cuda:0"))
+            >>> text_features = clip_model.encode_text(["a photo of a cat", "a photo of a dog"])
         """
         super().__init__()
         self.model = clip.load(size, device=device)[0]
@@ -87,7 +106,20 @@ class CLIP(TextModel):
         self.eval()
 
     def tokenize(self, texts):
-        """Convert input texts to CLIP tokens."""
+        """
+        Convert input texts to CLIP tokens.
+
+        Args:
+            texts (str | List[str]): Input text or list of texts to tokenize.
+
+        Returns:
+            (torch.Tensor): Tokenized text tensor with shape (batch_size, context_length) ready for model processing.
+
+        Examples:
+            >>> model = CLIP("ViT-B/32", device="cpu")
+            >>> tokens = model.tokenize("a photo of a cat")
+            >>> print(tokens.shape)  # torch.Size([1, 77])
+        """
         return clip.tokenize(texts).to(self.device)
 
     @smart_inference_mode()
@@ -95,12 +127,22 @@ class CLIP(TextModel):
         """
         Encode tokenized texts into normalized feature vectors.
 
+        This method processes tokenized text inputs through the CLIP model to generate feature vectors, which are then
+        normalized to unit length. These normalized vectors can be used for text-image similarity comparisons.
+
         Args:
-            texts (torch.Tensor): Tokenized text inputs.
-            dtype (torch.dtype): Data type for output features.
+            texts (torch.Tensor): Tokenized text inputs, typically created using the tokenize() method.
+            dtype (torch.dtype, optional): Data type for output features. Default is torch.float32.
 
         Returns:
-            (torch.Tensor): Normalized text feature vectors.
+            (torch.Tensor): Normalized text feature vectors with unit length (L2 norm = 1).
+
+        Examples:
+            >>> clip_model = CLIP("ViT-B/32", device="cuda")
+            >>> tokens = clip_model.tokenize(["a photo of a cat", "a photo of a dog"])
+            >>> features = clip_model.encode_text(tokens)
+            >>> features.shape
+            torch.Size([2, 512])
         """
         txt_feats = self.model.encode_text(texts).to(dtype)
         txt_feats = txt_feats / txt_feats.norm(p=2, dim=-1, keepdim=True)
@@ -109,9 +151,10 @@ class CLIP(TextModel):
 
 class MobileCLIP(TextModel):
     """
-    Apple MobileCLIP text encoder implementation.
+    Implement Apple's MobileCLIP text encoder for efficient text encoding.
 
-    This class implements the TextModel interface using Apple's MobileCLIP model for efficient text encoding.
+    This class implements the TextModel interface using Apple's MobileCLIP model, providing efficient text encoding
+    capabilities for vision-language tasks.
 
     Attributes:
         model (mobileclip.model.MobileCLIP): The loaded MobileCLIP model.
@@ -122,6 +165,12 @@ class MobileCLIP(TextModel):
     Methods:
         tokenize: Convert input texts to MobileCLIP tokens.
         encode_text: Encode tokenized texts into normalized feature vectors.
+
+    Examples:
+        >>> device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        >>> text_encoder = MobileCLIP(size="s0", device=device)
+        >>> tokens = text_encoder.tokenize(["a photo of a cat", "a photo of a dog"])
+        >>> features = text_encoder.encode_text(tokens)
     """
 
     config_size_map = {"s0": "s0", "s1": "s1", "s2": "s2", "b": "b", "blt": "b"}
@@ -130,9 +179,18 @@ class MobileCLIP(TextModel):
         """
         Initialize the MobileCLIP text encoder.
 
+        This class implements the TextModel interface using Apple's MobileCLIP model for efficient text encoding.
+
         Args:
             size (str): Model size identifier (e.g., 's0', 's1', 's2', 'b', 'blt').
             device (torch.device): Device to load the model on.
+
+        Examples:
+            >>> from ultralytics.nn.modules import MobileCLIP
+            >>> import torch
+            >>> model = MobileCLIP("s0", device=torch.device("cpu"))
+            >>> tokens = model.tokenize(["a photo of a cat", "a photo of a dog"])
+            >>> features = model.encode_text(tokens)
         """
         super().__init__()
         config = self.config_size_map[size]
@@ -148,7 +206,19 @@ class MobileCLIP(TextModel):
         self.eval()
 
     def tokenize(self, texts):
-        """Convert input texts to MobileCLIP tokens."""
+        """
+        Convert input texts to MobileCLIP tokens.
+
+        Args:
+            texts (list[str]): List of text strings to tokenize.
+
+        Returns:
+            (torch.Tensor): Tokenized text inputs with shape (batch_size, sequence_length).
+
+        Examples:
+            >>> model = MobileCLIP("s0", "cpu")
+            >>> tokens = model.tokenize(["a photo of a cat", "a photo of a dog"])
+        """
         return self.tokenizer(texts).to(self.device)
 
     @smart_inference_mode()
@@ -158,10 +228,17 @@ class MobileCLIP(TextModel):
 
         Args:
             texts (torch.Tensor): Tokenized text inputs.
-            dtype (torch.dtype): Data type for output features.
+            dtype (torch.dtype, optional): Data type for output features.
 
         Returns:
-            (torch.Tensor): Normalized text feature vectors.
+            (torch.Tensor): Normalized text feature vectors with L2 normalization applied.
+
+        Examples:
+            >>> model = MobileCLIP("s0", device="cpu")
+            >>> tokens = model.tokenize(["a photo of a cat", "a photo of a dog"])
+            >>> features = model.encode_text(tokens)
+            >>> features.shape
+            torch.Size([2, 512])  # Actual dimension depends on model size
         """
         text_features = self.model.encode_text(texts).to(dtype)
         text_features /= text_features.norm(p=2, dim=-1, keepdim=True)
@@ -179,15 +256,14 @@ def build_text_model(variant, device=None):
     Returns:
         (TextModel): Instantiated text encoding model.
 
-    Raises:
-        AssertionError: If the specified variant is not supported.
+    Examples:
+        >>> model = build_text_model("clip:ViT-B/32", device=torch.device("cuda"))
+        >>> model = build_text_model("mobileclip:s0", device=torch.device("cpu"))
     """
-    LOGGER.info(f"Build text model {variant}")
     base, size = variant.split(":")
     if base == "clip":
         return CLIP(size, device)
     elif base == "mobileclip":
         return MobileCLIP(size, device)
     else:
-        print("Variant not found")
-        assert False
+        raise ValueError(f"Unrecognized base model: '{base}'. Supported base models: 'clip', 'mobileclip'.")

ultralytics/solutions/instance_segmentation.py

@@ -1,7 +1,7 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
-from ultralytics.solutions.solutions import BaseSolution, SolutionAnnotator, SolutionResults
-from ultralytics.utils.plotting import colors
+from ultralytics.engine.results import Results
+from ultralytics.solutions.solutions import BaseSolution, SolutionResults
 
 
 class InstanceSegmentation(BaseSolution):
@@ -41,6 +41,10 @@ class InstanceSegmentation(BaseSolution):
         kwargs["model"] = kwargs.get("model", "yolo11n-seg.pt")
         super().__init__(**kwargs)
 
+        self.show_conf = self.CFG.get("show_conf", True)
+        self.show_labels = self.CFG.get("show_labels", True)
+        self.show_boxes = self.CFG.get("show_boxes", True)
+
     def process(self, im0):
         """
         Perform instance segmentation on the input image and annotate the results.
@@ -58,17 +62,21 @@ class InstanceSegmentation(BaseSolution):
             >>> print(summary)
         """
         self.extract_tracks(im0)  # Extract tracks (bounding boxes, classes, and masks)
-        annotator = SolutionAnnotator(im0, self.line_width)
 
         # Iterate over detected classes, track IDs, and segmentation masks
         if self.masks is None:
             self.LOGGER.warning("⚠️ No masks detected! Ensure you're using a supported Ultralytics segmentation model.")
+            plot_im = im0
         else:
-            for cls, t_id, mask in zip(self.clss, self.track_ids, self.masks):
-                # Annotate the image with segmentation mask, mask color, and label
-                annotator.segmentation_mask(mask=mask, mask_color=colors(t_id, True), label=self.names[cls])
+            results = Results(im0, path=None, names=self.names, boxes=self.track_data.data, masks=self.masks.data)
+            plot_im = results.plot(
+                line_width=self.line_width,
+                boxes=self.show_boxes,
+                conf=self.show_conf,
+                labels=self.show_labels,
+                color_mode="instance",
+            )
 
-        plot_im = annotator.result()
         self.display_output(plot_im)  # Display the annotated output using the base class function
 
         # Return SolutionResults

ultralytics/solutions/solutions.py

@@ -52,7 +52,7 @@ class BaseSolution:
             is_cli (bool): Enables CLI mode if set to True.
             **kwargs (Any): Additional configuration parameters that override defaults.
         """
-        check_requirements("shapely>=2.0.0")
+        check_requirements("shapely>=2.0.0,<2.1.0")
         from shapely.geometry import LineString, Point, Polygon
         from shapely.prepared import prep
 
@@ -122,7 +122,7 @@ class BaseSolution:
         self.track_data = self.tracks[0].obb or self.tracks[0].boxes  # Extract tracks for OBB or object detection
 
         self.masks = (
-            self.tracks[0].masks.xy if hasattr(self.tracks[0], "masks") and self.tracks[0].masks is not None else None
+            self.tracks[0].masks if hasattr(self.tracks[0], "masks") and self.tracks[0].masks is not None else None
         )
 
         if self.track_data and self.track_data.id is not None:
@@ -225,7 +225,6 @@ class SolutionAnnotator(Annotator):
         plot_angle_and_count_and_stage: Visualizes angle, step count, and stage for workout monitoring.
         plot_distance_and_line: Displays the distance between centroids and connects them with a line.
         display_objects_labels: Annotates bounding boxes with object class labels.
-        segmentation_mask: Draws mask for segmented objects and optionally labels them.
         sweep_annotator: Visualizes a vertical sweep line and optional label.
         visioneye: Maps and connects object centroids to a visual "eye" point.
         circle_label: Draws a circular label within a bounding box.
@@ -519,50 +518,6 @@ class SolutionAnnotator(Annotator):
             lineType=cv2.LINE_AA,
         )
 
-    def segmentation_mask(self, mask, mask_color=(255, 0, 255), label=None, alpha=0.5):
-        """
-        Draw an optimized segmentation mask with smooth corners, highlighted edge, and dynamic text box size.
-
-        Args:
-            mask (np.ndarray): A 2D array of shape (N, 2) containing the object mask.
-            mask_color (Tuple[int, int, int]): RGB color for the mask.
-            label (str, optional): Text label for the object.
-            alpha (float): Transparency level (0 = fully transparent, 1 = fully opaque).
-        """
-        if mask.size == 0:
-            return
-
-        overlay = self.im.copy()
-        mask = np.int32([mask])
-
-        # Approximate polygon for smooth corners with epsilon
-        refined_mask = cv2.approxPolyDP(mask, 0.002 * cv2.arcLength(mask, True), True)
-
-        # Apply a highlighter effect by drawing a thick outer shadow
-        cv2.polylines(overlay, [refined_mask], isClosed=True, color=mask_color, thickness=self.lw * 3)
-        cv2.fillPoly(overlay, [refined_mask], mask_color)  # draw mask with primary color
-
-        # Apply an inner glow effect for extra clarity
-        cv2.polylines(overlay, [refined_mask], isClosed=True, color=mask_color, thickness=self.lw)
-
-        self.im = cv2.addWeighted(overlay, alpha, self.im, 1 - alpha, 0)  # blend overlay with the original image
-
-        # Draw label if provided
-        if label:
-            text_size, _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, self.sf, self.tf)
-            text_x, text_y = refined_mask[0][0][0], refined_mask[0][0][1]
-            rect_start, rect_end = (text_x - 5, text_y - text_size[1] - 5), (text_x + text_size[0] + 5, text_y + 5)
-            cv2.rectangle(self.im, rect_start, rect_end, mask_color, -1)
-            cv2.putText(
-                self.im,
-                label,
-                (text_x, text_y),
-                cv2.FONT_HERSHEY_SIMPLEX,
-                self.sf,
-                self.get_txt_color(mask_color),
-                self.tf,
-            )
-
     def sweep_annotator(self, line_x=0, line_y=0, label=None, color=(221, 0, 186), txt_color=(255, 255, 255)):
         """
         Draw a sweep annotation line and an optional label.

ultralytics/utils/benchmarks.py

@@ -126,7 +126,7 @@ def benchmark(
                 assert not isinstance(model, YOLOWorld), "YOLOWorldv2 TensorFlow exports not supported by onnx2tf yet"
             if i == 11:  # Paddle
                 assert not isinstance(model, YOLOWorld), "YOLOWorldv2 Paddle exports not supported yet"
-                assert not model.task == "obb", "Paddle OBB bug https://github.com/PaddlePaddle/Paddle/issues/72024"
+                assert model.task != "obb", "Paddle OBB bug https://github.com/PaddlePaddle/Paddle/issues/72024"
                 assert not is_end2end, "End-to-end models not supported by PaddlePaddle yet"
                 assert LINUX or MACOS, "Windows Paddle exports not supported yet"
             if i == 12:  # MNN

ultralytics/utils/callbacks/base.py

@@ -176,21 +176,38 @@ default_callbacks = {
 
 def get_default_callbacks():
     """
-    Return a copy of the default_callbacks dictionary with lists as default values.
+    Get the default callbacks for Ultralytics training, validation, prediction, and export processes.
 
     Returns:
-        (defaultdict): A defaultdict with keys from default_callbacks and empty lists as default values.
+        (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.
+
+    Examples:
+        >>> callbacks = get_default_callbacks()
+        >>> print(list(callbacks.keys()))  # show all available callback events
+        ['on_pretrain_routine_start', 'on_pretrain_routine_end', ...]
     """
     return defaultdict(list, deepcopy(default_callbacks))
 
 
 def add_integration_callbacks(instance):
     """
-    Add integration callbacks from various sources to the instance's callbacks.
+    Add integration callbacks to the instance's callbacks dictionary.
+
+    This function loads and adds various integration callbacks to the provided instance. The specific callbacks added
+    depend on the type of instance provided. All instances receive HUB callbacks, while Trainer instances also receive
+    additional callbacks for various integrations like ClearML, Comet, DVC, MLflow, Neptune, Ray Tune, TensorBoard,
+    and Weights & Biases.
 
     Args:
-        instance (Trainer | Predictor | Validator | Exporter): An object with a 'callbacks' attribute that is a
-            dictionary of callback lists.
+        instance (Trainer | Predictor | Validator | Exporter): The object instance to which callbacks will be added.
+            The type of instance determines which callbacks are loaded.
+
+    Examples:
+        >>> from ultralytics.engine.trainer import BaseTrainer
+        >>> trainer = BaseTrainer()
+        >>> add_integration_callbacks(trainer)
     """
     # Load HUB callbacks
     from .hub import callbacks as hub_cb

ultralytics/utils/callbacks/comet.py

@@ -155,7 +155,32 @@ def _scale_bounding_box_to_original_image_shape(
 
 
 def _format_ground_truth_annotations_for_detection(img_idx, image_path, batch, class_name_map=None) -> Optional[dict]:
-    """Format ground truth annotations for detection."""
+    """
+    Format ground truth annotations for object detection.
+
+    This function processes ground truth annotations from a batch of images for object detection tasks. It extracts
+    bounding boxes, class labels, and other metadata for a specific image in the batch, and formats them for
+    visualization or evaluation.
+
+    Args:
+        img_idx (int): Index of the image in the batch to process.
+        image_path (str | Path): Path to the image file.
+        batch (dict): Batch dictionary containing detection data with keys:
+            - 'batch_idx': Tensor of batch indices
+            - 'bboxes': Tensor of bounding boxes in normalized xywh format
+            - 'cls': Tensor of class labels
+            - '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.
+
+    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.
+    """
     indices = batch["batch_idx"] == img_idx
     bboxes = batch["bboxes"][indices]
     if len(bboxes) == 0:
@@ -284,7 +309,22 @@ def _log_confusion_matrix(experiment, trainer, curr_step, curr_epoch) -> None:
 
 
 def _log_images(experiment, image_paths, curr_step, annotations=None) -> None:
-    """Logs images to the experiment with optional annotations."""
+    """
+    Log images to the experiment with optional annotations.
+
+    This function logs images to a Comet ML experiment, optionally including annotation data for visualization
+    such as bounding boxes or segmentation masks.
+
+    Args:
+        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
+    """
     if annotations:
         for image_path, annotation in zip(image_paths, annotations):
             experiment.log_image(image_path, name=image_path.stem, step=curr_step, annotations=annotation)
@@ -295,7 +335,23 @@ def _log_images(experiment, image_paths, curr_step, annotations=None) -> None:
 
 
 def _log_image_predictions(experiment, validator, curr_step) -> None:
-    """Logs predicted boxes for a single image during training."""
+    """
+    Log predicted boxes for a single image during training.
+
+    This function logs image predictions to a Comet ML experiment during model validation. It processes
+    validation data and formats both ground truth and prediction annotations for visualization in the Comet
+    dashboard. The function respects configured limits on the number of images to log.
+
+    Args:
+        experiment (comet_ml.Experiment): The Comet ML experiment to log to.
+        validator (BaseValidator): The validator instance containing validation data and predictions.
+        curr_step (int): The current training step for logging timeline.
+
+    Notes:
+        This function uses global state to track the number of logged predictions across calls.
+        It only logs predictions for supported tasks defined in COMET_SUPPORTED_TASKS.
+        The number of logged images is limited by the COMET_MAX_IMAGE_PREDICTIONS environment variable.
+    """
     global _comet_image_prediction_count
 
     task = validator.args.task
@@ -342,7 +398,22 @@ def _log_image_predictions(experiment, validator, curr_step) -> None:
 
 
 def _log_plots(experiment, trainer) -> None:
-    """Logs evaluation plots and label plots for the experiment."""
+    """
+    Log evaluation plots and label plots for the experiment.
+
+    This function logs various evaluation plots and confusion matrices to the experiment tracking system. It handles
+    different types of metrics (SegmentMetrics, PoseMetrics, DetMetrics, OBBMetrics) and logs the appropriate plots
+    for each type.
+
+    Args:
+        experiment (comet_ml.Experiment): The Comet ML experiment to log plots to.
+        trainer (ultralytics.engine.trainer.BaseTrainer): The trainer object containing validation metrics and save
+            directory information.
+
+    Examples:
+        >>> from ultralytics.utils.callbacks.comet import _log_plots
+        >>> _log_plots(experiment, trainer)
+    """
     plot_filenames = None
     if isinstance(trainer.validator.metrics, SegmentMetrics) and trainer.validator.metrics.task == "segment":
         plot_filenames = [
@@ -401,7 +472,24 @@ def on_train_epoch_end(trainer) -> None:
 
 
 def on_fit_epoch_end(trainer) -> None:
-    """Logs model assets at the end of each epoch."""
+    """
+    Log model assets at the end of each epoch during training.
+
+    This function is called at the end of each training epoch to log metrics, learning rates, and model information
+    to a Comet ML experiment. It also logs model assets, confusion matrices, and image predictions based on
+    configuration settings.
+
+    The function retrieves the current Comet ML experiment and logs various training metrics. If it's the first epoch,
+    it also logs model information. On specified save intervals, it logs the model, confusion matrix (if enabled),
+    and image predictions (if enabled).
+
+    Args:
+        trainer (BaseTrainer): The YOLO trainer object containing training state, metrics, and configuration.
+
+    Examples:
+        >>> # Inside a training loop
+        >>> on_fit_epoch_end(trainer)  # Log metrics and assets to Comet ML
+    """
     experiment = comet_ml.get_running_experiment()
     if not experiment:
         return

ultralytics/utils/callbacks/dvc.py

@@ -27,7 +27,21 @@ except (ImportError, AssertionError, TypeError):
 
 
 def _log_images(path: Path, prefix: str = "") -> None:
-    """Logs images at specified path with an optional prefix using DVCLive."""
+    """
+    Log images at specified path with an optional prefix using DVCLive.
+
+    This function logs images found at the given path to DVCLive, organizing them by batch to enable slider
+    functionality in the UI. It processes image filenames to extract batch information and restructures the path
+    accordingly.
+
+    Args:
+        path (Path): Path to the image file to be logged.
+        prefix (str): Optional prefix to add to the image name when logging.
+
+    Examples:
+        >>> from pathlib import Path
+        >>> _log_images(Path("runs/train/exp/val_batch0_pred.jpg"), prefix="validation")
+    """
     if live:
         name = path.name
 
@@ -41,7 +55,13 @@ def _log_images(path: Path, prefix: str = "") -> None:
 
 
 def _log_plots(plots: dict, prefix: str = "") -> None:
-    """Logs plot images for training progress if they have not been previously processed."""
+    """
+    Log plot images for training progress if they have not been previously processed.
+
+    Args:
+        plots (dict): Dictionary containing plot information with timestamps.
+        prefix (str, optional): Optional prefix to add to the logged image paths.
+    """
     for name, params in plots.items():
         timestamp = params["timestamp"]
         if _processed_plots.get(name) != timestamp:
@@ -50,7 +70,19 @@ def _log_plots(plots: dict, prefix: str = "") -> None:
 
 
 def _log_confusion_matrix(validator) -> None:
-    """Logs the confusion matrix for the given validator using DVCLive."""
+    """
+    Log confusion matrix for a validator using DVCLive.
+
+    This function processes the confusion matrix from a validator object and logs it to DVCLive by converting
+    the matrix into lists of target and prediction labels.
+
+    Args:
+        validator (BaseValidator): The validator object containing the confusion matrix and class names.
+            Must have attributes: confusion_matrix.matrix, confusion_matrix.task, and names.
+
+    Returns:
+        None
+    """
     targets = []
     preds = []
     matrix = validator.confusion_matrix.matrix
@@ -94,7 +126,20 @@ def on_train_epoch_start(trainer) -> None:
 
 
 def on_fit_epoch_end(trainer) -> None:
-    """Logs training metrics and model info, and advances to next step on the end of each fit epoch."""
+    """
+    Log training metrics, model info, and advance to next step at the end of each fit epoch.
+
+    This function is called at the end of each fit epoch during training. It logs various metrics including
+    training loss items, validation metrics, and learning rates. On the first epoch, it also logs model
+    information. Additionally, it logs training and validation plots and advances the DVCLive step counter.
+
+    Args:
+        trainer (BaseTrainer): The trainer object containing training state, metrics, and plots.
+
+    Notes:
+        This function only performs logging operations when DVCLive logging is active and during a training epoch.
+        The global variable _training_epoch is used to track whether the current epoch is a training epoch.
+    """
     global _training_epoch
     if live and _training_epoch:
         all_metrics = {**trainer.label_loss_items(trainer.tloss, prefix="train"), **trainer.metrics, **trainer.lr}
@@ -115,7 +160,21 @@ def on_fit_epoch_end(trainer) -> None:
 
 
 def on_train_end(trainer) -> None:
-    """Logs the best metrics, plots, and confusion matrix at the end of training if DVCLive is active."""
+    """
+    Log best metrics, plots, and confusion matrix at the end of training.
+
+    This function is called at the conclusion of the training process to log final metrics, visualizations, and
+    model artifacts if DVCLive logging is active. It captures the best model performance metrics, training plots,
+    validation plots, and confusion matrix for later analysis.
+
+    Args:
+        trainer (BaseTrainer): The trainer object containing training state, metrics, and validation results.
+
+    Examples:
+        >>> # Inside a custom training loop
+        >>> from ultralytics.utils.callbacks.dvc import on_train_end
+        >>> on_train_end(trainer)  # Log final metrics and artifacts
+    """
     if live:
         # At the end log the best metrics. It runs validator on the best model internally.
         all_metrics = {**trainer.label_loss_items(trainer.tloss, prefix="train"), **trainer.metrics, **trainer.lr}

ultralytics/utils/callbacks/neptune.py

@@ -19,14 +19,37 @@ except (ImportError, AssertionError):
 
 
 def _log_scalars(scalars: dict, step: int = 0) -> None:
-    """Log scalars to the NeptuneAI experiment logger."""
+    """
+    Log scalars to the NeptuneAI experiment logger.
+
+    Args:
+        scalars (dict): Dictionary of scalar values to log to NeptuneAI.
+        step (int): The current step or iteration number for logging.
+
+    Examples:
+        >>> metrics = {"mAP": 0.85, "loss": 0.32}
+        >>> _log_scalars(metrics, step=100)
+    """
     if run:
         for k, v in scalars.items():
             run[k].append(value=v, step=step)
 
 
 def _log_images(imgs_dict: dict, group: str = "") -> None:
-    """Log images to the NeptuneAI experiment logger."""
+    """
+    Log images to the NeptuneAI experiment logger.
+
+    This function logs image data to Neptune.ai when a valid Neptune run is active. Images are organized
+    under the specified group name.
+
+    Args:
+        imgs_dict (dict): Dictionary of images to log, with keys as image names and values as image data.
+        group (str, optional): Group name to organize images under in the Neptune UI.
+
+    Examples:
+        >>> # Log validation images
+        >>> _log_images({"val_batch": img_tensor}, group="validation")
+    """
     if run:
         for k, v in imgs_dict.items():
             run[f"{group}/{k}"].upload(File(v))