dgenerate-ultralytics-headless 8.3.137__py3-none-any.whl → 8.3.224__py3-none-any.whl

ultralytics/utils/callbacks/base.py

@@ -175,13 +175,12 @@ default_callbacks = {
 
 
 def get_default_callbacks():
-    """
-    Get the default callbacks for Ultralytics training, validation, prediction, and export processes.
+    """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()
@@ -192,27 +191,27 @@ def get_default_callbacks():
 
 
 def add_integration_callbacks(instance):
-    """
-    Add integration callbacks to the instance's callbacks dictionary.
+    """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.
+    additional callbacks for various integrations like ClearML, Comet, DVC, MLflow, Neptune, Ray Tune, TensorBoard, and
+    Weights & Biases.
 
     Args:
-        instance (Trainer | Predictor | Validator | Exporter): The object instance to which callbacks will be added.
-            The type of instance determines which callbacks are loaded.
+        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
+    from .platform import callbacks as platform_cb
 
-    callbacks_list = [hub_cb]
+    # Load Ultralytics callbacks
+    callbacks_list = [hub_cb, platform_cb]
 
     # Load training callbacks
     if "Trainer" in instance.__class__.__name__:

ultralytics/utils/callbacks/clearml.py

@@ -15,11 +15,10 @@ except (ImportError, AssertionError):
 
 
 def _log_debug_samples(files, title: str = "Debug Samples") -> None:
-    """
-    Log files (images) as debug samples in the ClearML task.
+    """Log files (images) as debug samples in the ClearML task.
 
     Args:
-        files (List[Path]): A list of file paths in PosixPath format.
+        files (list[Path]): A list of file paths in PosixPath format.
         title (str): A title that groups together images with the same values.
     """
     import re
@@ -35,8 +34,7 @@ def _log_debug_samples(files, title: str = "Debug Samples") -> None:
 
 
 def _log_plot(title: str, plot_path: str) -> None:
-    """
-    Log an image as a plot in the plot section of ClearML.
+    """Log an image as a plot in the plot section of ClearML.
 
     Args:
         title (str): The title of the plot.
@@ -56,7 +54,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 +83,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,14 +96,15 @@ 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(
             title="Epoch Time", series="Epoch Time", value=trainer.epoch_time, iteration=trainer.epoch
         )
         for k, v in trainer.metrics.items():
-            task.get_logger().report_scalar("val", k, v, iteration=trainer.epoch)
+            title = k.split("/")[0]
+            task.get_logger().report_scalar(title, k, v, iteration=trainer.epoch)
         if trainer.epoch == 0:
             from ultralytics.utils.torch_utils import model_info_for_loggers
 
@@ -114,23 +113,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

@@ -1,8 +1,10 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
+from __future__ import annotations
+
 from collections.abc import Callable
 from types import SimpleNamespace
-from typing import Any, List, Optional
+from typing import Any
 
 import cv2
 import numpy as np
@@ -26,9 +28,12 @@ try:
     # Names of plots created by Ultralytics that are logged to Comet
     CONFUSION_MATRIX_PLOT_NAMES = "confusion_matrix", "confusion_matrix_normalized"
     EVALUATION_PLOT_NAMES = "F1_curve", "P_curve", "R_curve", "PR_curve"
-    LABEL_PLOT_NAMES = "labels", "labels_correlogram"
+    LABEL_PLOT_NAMES = ["labels"]
     SEGMENT_METRICS_PLOT_PREFIX = "Box", "Mask"
     POSE_METRICS_PLOT_PREFIX = "Box", "Pose"
+    DETECTION_METRICS_PLOT_PREFIX = ["Box"]
+    RESULTS_TABLE_NAME = "results.csv"
+    ARGS_YAML_NAME = "args.yaml"
 
     _comet_image_prediction_count = 0
 
@@ -37,7 +42,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 +57,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 +67,33 @@ 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 +123,14 @@ 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
@@ -133,11 +147,20 @@ def _fetch_trainer_metadata(trainer) -> dict:
 
 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.
+) -> list[float]:
+    """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.
 
-    This function rescales the bounding box labels to the original image shape.
+    Returns:
+        (list[float]): Scaled bounding box coordinates in xywh format with top-left corner adjustment.
     """
     resized_image_height, resized_image_width = resized_image_shape
 
@@ -154,9 +177,8 @@ def _scale_bounding_box_to_original_image_shape(
     return box
 
 
-def _format_ground_truth_annotations_for_detection(img_idx, image_path, batch, class_name_map=None) -> Optional[dict]:
-    """
-    Format ground truth annotations for object detection.
+def _format_ground_truth_annotations_for_detection(img_idx, image_path, batch, class_name_map=None) -> dict | None:
+    """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
@@ -172,7 +194,7 @@ 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:
@@ -209,8 +231,18 @@ def _format_ground_truth_annotations_for_detection(img_idx, image_path, batch, c
     return {"name": "ground_truth", "data": data}
 
 
-def _format_prediction_annotations(image_path, metadata, class_label_map=None, class_map=None) -> Optional[dict]:
-    """Format YOLO predictions for object detection visualization."""
+def _format_prediction_annotations(image_path, metadata, class_label_map=None, class_map=None) -> dict | None:
+    """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
 
@@ -224,7 +256,7 @@ def _format_prediction_annotations(image_path, metadata, class_label_map=None, c
         class_label_map = {class_map[k]: v for k, v in class_label_map.items()}
     try:
         # import pycotools utilities to decompress annotations for various tasks, e.g. segmentation
-        from pycocotools.mask import decode  # noqa
+        from faster_coco_eval.core.mask import decode
     except ImportError:
         decode = None
 
@@ -251,16 +283,15 @@ def _format_prediction_annotations(image_path, metadata, class_label_map=None, c
     return {"name": "prediction", "data": data}
 
 
-def _extract_segmentation_annotation(segmentation_raw: str, decode: Callable) -> Optional[List[List[Any]]]:
-    """
-    Extracts segmentation annotation from compressed segmentations as list of polygons.
+def _extract_segmentation_annotation(segmentation_raw: str, decode: Callable) -> list[list[Any]] | None:
+    """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)
@@ -272,10 +303,20 @@ def _extract_segmentation_annotation(segmentation_raw: str, decode: Callable) ->
     return None
 
 
-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."""
+def _fetch_annotations(img_idx, image_path, batch, prediction_metadata_map, class_label_map, class_map) -> list | None:
+    """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 +331,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"], [])
@@ -302,28 +343,24 @@ def _create_prediction_metadata_map(model_predictions) -> dict:
 def _log_confusion_matrix(experiment, trainer, curr_step, curr_epoch) -> None:
     """Log the confusion matrix to Comet experiment."""
     conf_mat = trainer.validator.confusion_matrix.matrix
-    names = list(trainer.data["names"].values()) + ["background"]
+    names = [*list(trainer.data["names"].values()), "background"]
     experiment.log_confusion_matrix(
         matrix=conf_mat, labels=names, max_categories=len(names), epoch=curr_epoch, step=curr_step
     )
 
 
-def _log_images(experiment, image_paths, curr_step, annotations=None) -> None:
-    """
-    Log images to the experiment with optional annotations.
+def _log_images(experiment, image_paths, curr_step: int | None, annotations=None) -> None:
+    """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.
+    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.
+        experiment (comet_ml.CometExperiment): 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
+        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):
@@ -335,15 +372,14 @@ def _log_images(experiment, image_paths, curr_step, annotations=None) -> None:
 
 
 def _log_image_predictions(experiment, validator, curr_step) -> None:
-    """
-    Log 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
+    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.
+        experiment (comet_ml.CometExperiment): 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.
 
@@ -398,15 +434,14 @@ def _log_image_predictions(experiment, validator, curr_step) -> None:
 
 
 def _log_plots(experiment, trainer) -> None:
-    """
-    Log 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.
+    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.
+        experiment (comet_ml.CometExperiment): The Comet ML experiment to log plots to.
         trainer (ultralytics.engine.trainer.BaseTrainer): The trainer object containing validation metrics and save
             directory information.
 
@@ -415,7 +450,7 @@ def _log_plots(experiment, trainer) -> None:
         >>> _log_plots(experiment, trainer)
     """
     plot_filenames = None
-    if isinstance(trainer.validator.metrics, SegmentMetrics) and trainer.validator.metrics.task == "segment":
+    if isinstance(trainer.validator.metrics, SegmentMetrics):
         plot_filenames = [
             trainer.save_dir / f"{prefix}{plots}.png"
             for plots in EVALUATION_PLOT_NAMES
@@ -428,7 +463,11 @@ def _log_plots(experiment, trainer) -> None:
             for prefix in POSE_METRICS_PLOT_PREFIX
         ]
     elif isinstance(trainer.validator.metrics, (DetMetrics, OBBMetrics)):
-        plot_filenames = [trainer.save_dir / f"{plots}.png" for plots in EVALUATION_PLOT_NAMES]
+        plot_filenames = [
+            trainer.save_dir / f"{prefix}{plots}.png"
+            for plots in EVALUATION_PLOT_NAMES
+            for prefix in DETECTION_METRICS_PLOT_PREFIX
+        ]
 
     if plot_filenames is not None:
         _log_images(experiment, plot_filenames, None)
@@ -448,13 +487,38 @@ 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 _log_asset(experiment, asset_path) -> None:
+    """Logs a specific asset file to the given experiment.
+
+    This function facilitates logging an asset, such as a file, to the provided
+    experiment. It enables integration with experiment tracking platforms.
+
+    Args:
+        experiment (comet_ml.CometExperiment): The experiment instance to which the asset will be logged.
+        asset_path (Path): The file path of the asset to log.
+    """
+    experiment.log_asset(asset_path)
+
+
+def _log_table(experiment, table_path) -> None:
+    """Logs a table to the provided experiment.
+
+    This function is used to log a table file to the given experiment. The table is identified by its file path.
+
+    Args:
+        experiment (comet_ml.CometExperiment): The experiment object where the table file will be logged.
+        table_path (Path): The file path of the table to be logged.
+    """
+    experiment.log_table(str(table_path))
+
+
 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)
 
 
@@ -472,16 +536,15 @@ def on_train_epoch_end(trainer) -> None:
 
 
 def on_fit_epoch_end(trainer) -> None:
-    """
-    Log model assets at the end of each epoch during training.
+    """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.
+    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).
+    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.
@@ -534,6 +597,16 @@ def on_train_end(trainer) -> None:
     _log_confusion_matrix(experiment, trainer, curr_step, curr_epoch)
     _log_image_predictions(experiment, trainer.validator, curr_step)
     _log_image_batches(experiment, trainer, curr_step)
+    # log results table
+    table_path = trainer.save_dir / RESULTS_TABLE_NAME
+    if table_path.exists():
+        _log_table(experiment, table_path)
+
+    # log arguments YAML
+    args_path = trainer.save_dir / ARGS_YAML_NAME
+    if args_path.exists():
+        _log_asset(experiment, args_path)
+
     experiment.end()
 
     global _comet_image_prediction_count

ultralytics/utils/callbacks/dvc.py

@@ -27,8 +27,7 @@ except (ImportError, AssertionError, TypeError):
 
 
 def _log_images(path: Path, prefix: str = "") -> None:
-    """
-    Log 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
@@ -36,7 +35,7 @@ def _log_images(path: Path, prefix: str = "") -> None:
 
     Args:
         path (Path): Path to the image file to be logged.
-        prefix (str): Optional prefix to add to the image name when logging.
+        prefix (str, optional): Optional prefix to add to the image name when logging.
 
     Examples:
         >>> from pathlib import Path
@@ -55,8 +54,7 @@ def _log_images(path: Path, prefix: str = "") -> None:
 
 
 def _log_plots(plots: dict, prefix: str = "") -> None:
-    """
-    Log 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.
@@ -70,18 +68,14 @@ def _log_plots(plots: dict, prefix: str = "") -> None:
 
 
 def _log_confusion_matrix(validator) -> None:
-    """
-    Log confusion matrix for a 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.
+    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
+        validator (BaseValidator): The validator object containing the confusion matrix and class names. Must have
+            attributes confusion_matrix.matrix, confusion_matrix.task, and names.
     """
     targets = []
     preds = []
@@ -99,7 +93,7 @@ def _log_confusion_matrix(validator) -> None:
 
 
 def on_pretrain_routine_start(trainer) -> None:
-    """Initializes DVCLive logger for training metadata during pre-training routine."""
+    """Initialize DVCLive logger for training metadata during pre-training routine."""
     try:
         global live
         live = dvclive.Live(save_dvc_exp=True, cache_images=True)
@@ -109,28 +103,27 @@ def on_pretrain_routine_start(trainer) -> None:
 
 
 def on_pretrain_routine_end(trainer) -> None:
-    """Logs plots related to the training process at the end of the pretraining routine."""
+    """Log plots related to the training process at the end of the pretraining routine."""
     _log_plots(trainer.plots, "train")
 
 
 def on_train_start(trainer) -> None:
-    """Logs the training parameters if DVCLive logging is active."""
+    """Log the training parameters if DVCLive logging is active."""
     if live:
         live.log_params(trainer.args)
 
 
 def on_train_epoch_start(trainer) -> None:
-    """Sets the global variable _training_epoch value to True at the start of training each epoch."""
+    """Set the global variable _training_epoch value to True at the start of training each epoch."""
     global _training_epoch
     _training_epoch = True
 
 
 def on_fit_epoch_end(trainer) -> None:
-    """
-    Log training metrics, model info, and advance to next step at 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
+    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:
@@ -160,12 +153,11 @@ def on_fit_epoch_end(trainer) -> None:
 
 
 def on_train_end(trainer) -> None:
-    """
-    Log best metrics, plots, and confusion matrix at the end of training.
+    """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.
+    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.