ultralytics 8.3.89__py3-none-any.whl → 8.3.90__py3-none-any.whl

ultralytics/utils/benchmarks.py

@@ -201,7 +201,23 @@ def benchmark(
 
 
 class RF100Benchmark:
-    """Benchmark YOLO model performance across various formats for speed and accuracy."""
+    """
+    Benchmark YOLO model performance across various formats for speed and accuracy.
+
+    This class provides functionality to benchmark YOLO models on the RF100 dataset collection.
+
+    Attributes:
+        ds_names (List[str]): Names of datasets used for benchmarking.
+        ds_cfg_list (List[Path]): List of paths to dataset configuration files.
+        rf (Roboflow): Roboflow instance for accessing datasets.
+        val_metrics (List[str]): Metrics used for validation.
+
+    Methods:
+        set_key: Set Roboflow API key for accessing datasets.
+        parse_dataset: Parse dataset links and download datasets.
+        fix_yaml: Fix train and validation paths in YAML files.
+        evaluate: Evaluate model performance on validation results.
+    """
 
     def __init__(self):
         """Initialize the RF100Benchmark class for benchmarking YOLO model performance across various formats."""
@@ -234,6 +250,10 @@ class RF100Benchmark:
         Args:
             ds_link_txt (str): Path to the file containing dataset links.
 
+        Returns:
+            ds_names (List[str]): List of dataset names.
+            ds_cfg_list (List[Path]): List of paths to dataset configuration files.
+
         Examples:
             >>> benchmark = RF100Benchmark()
             >>> benchmark.set_key("api_key")
@@ -262,15 +282,7 @@ class RF100Benchmark:
 
     @staticmethod
     def fix_yaml(path):
-        """
-        Fixes the train and validation paths in a given YAML file.
-
-        Args:
-            path (str): Path to the YAML file to be fixed.
-
-        Examples:
-            >>> RF100Benchmark.fix_yaml("path/to/data.yaml")
-        """
+        """Fix the train and validation paths in a given YAML file."""
         with open(path, encoding="utf-8") as file:
             yaml_data = yaml.safe_load(file)
         yaml_data["train"] = "train/images"
@@ -353,6 +365,14 @@ class ProfileModels:
 
     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.
 
     Examples:
         Profile models and print results
@@ -404,7 +424,18 @@ class ProfileModels:
         self.device = device or torch.device(0 if torch.cuda.is_available() else "cpu")
 
     def profile(self):
-        """Profiles YOLO models for speed and accuracy across various formats including ONNX and TensorRT."""
+        """
+        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.
+
+        Examples:
+            Profile models and print results
+            >>> from ultralytics.utils.benchmarks import ProfileModels
+            >>> profiler = ProfileModels(["yolo11n.yaml", "yolov8s.yaml"])
+            >>> results = profiler.profile()
+        """
         files = self.get_files()
 
         if not files:
@@ -448,7 +479,12 @@ class ProfileModels:
         return output
 
     def get_files(self):
-        """Returns a list of paths for all relevant model files given by the user."""
+        """
+        Return a list of paths for all relevant model files given by the user.
+
+        Returns:
+            (List[Path]): List of Path objects for the model files.
+        """
         files = []
         for path in self.paths:
             path = Path(path)
@@ -470,7 +506,17 @@ class ProfileModels:
 
     @staticmethod
     def iterative_sigma_clipping(data, sigma=2, max_iters=3):
-        """Applies iterative sigma clipping to data to remove outliers based on specified sigma and iteration count."""
+        """
+        Apply iterative sigma clipping to data to remove outliers.
+
+        Args:
+            data (numpy.ndarray): Input data array.
+            sigma (float): Number of standard deviations to use for clipping.
+            max_iters (int): Maximum number of iterations for the clipping process.
+
+        Returns:
+            (numpy.ndarray): Clipped data array with outliers removed.
+        """
         data = np.array(data)
         for _ in range(max_iters):
             mean, std = np.mean(data), np.std(data)
@@ -481,7 +527,17 @@ class ProfileModels:
         return data
 
     def profile_tensorrt_model(self, engine_file: str, eps: float = 1e-3):
-        """Profiles YOLO model performance with TensorRT, measuring average run time and standard deviation."""
+        """
+        Profile YOLO model performance with TensorRT, measuring average run time and standard deviation.
+
+        Args:
+            engine_file (str): Path to the TensorRT engine file.
+            eps (float): Small epsilon value to prevent division by zero.
+
+        Returns:
+            mean_time (float): Mean inference time in milliseconds.
+            std_time (float): Standard deviation of inference time in milliseconds.
+        """
         if not self.trt or not Path(engine_file).is_file():
             return 0.0, 0.0
 
@@ -510,7 +566,17 @@ class ProfileModels:
         return np.mean(run_times), np.std(run_times)
 
     def profile_onnx_model(self, onnx_file: str, eps: float = 1e-3):
-        """Profiles an ONNX model, measuring average inference time and standard deviation across multiple runs."""
+        """
+        Profile an ONNX model, measuring average inference time and standard deviation across multiple runs.
+
+        Args:
+            onnx_file (str): Path to the ONNX model file.
+            eps (float): Small epsilon value to prevent division by zero.
+
+        Returns:
+            mean_time (float): Mean inference time in milliseconds.
+            std_time (float): Standard deviation of inference time in milliseconds.
+        """
         check_requirements("onnxruntime")
         import onnxruntime as ort
 
@@ -565,7 +631,18 @@ class ProfileModels:
         return np.mean(run_times), np.std(run_times)
 
     def generate_table_row(self, model_name, t_onnx, t_engine, model_info):
-        """Generates a table row string with model performance metrics including inference times and model details."""
+        """
+        Generate a table row string with model performance metrics.
+
+        Args:
+            model_name (str): Name of the model.
+            t_onnx (tuple): ONNX model inference time statistics (mean, std).
+            t_engine (tuple): TensorRT engine inference time statistics (mean, std).
+            model_info (tuple): Model information (layers, params, gradients, flops).
+
+        Returns:
+            (str): Formatted table row string with model metrics.
+        """
         layers, params, gradients, flops = model_info
         return (
             f"| {model_name:18s} | {self.imgsz} | - | {t_onnx[0]:.1f}±{t_onnx[1]:.1f} ms | {t_engine[0]:.1f}±"
@@ -574,7 +651,18 @@ class ProfileModels:
 
     @staticmethod
     def generate_results_dict(model_name, t_onnx, t_engine, model_info):
-        """Generates a dictionary of profiling results including model name, parameters, GFLOPs, and speed metrics."""
+        """
+        Generate a dictionary of profiling results.
+
+        Args:
+            model_name (str): Name of the model.
+            t_onnx (tuple): ONNX model inference time statistics (mean, std).
+            t_engine (tuple): TensorRT engine inference time statistics (mean, std).
+            model_info (tuple): Model information (layers, params, gradients, flops).
+
+        Returns:
+            (dict): Dictionary containing profiling results.
+        """
         layers, params, gradients, flops = model_info
         return {
             "model/name": model_name,
@@ -586,7 +674,12 @@ class ProfileModels:
 
     @staticmethod
     def print_table(table_rows):
-        """Prints a formatted table of model profiling results, including speed and accuracy metrics."""
+        """
+        Print a formatted table of model profiling results.
+
+        Args:
+            table_rows (List[str]): List of formatted table row strings.
+        """
         gpu = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "GPU"
         headers = [
             "Model",

ultralytics/utils/callbacks/base.py

@@ -1,5 +1,5 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
-"""Base callbacks."""
+"""Base callbacks for Ultralytics training, validation, prediction, and export processes."""
 
 from collections import defaultdict
 from copy import deepcopy
@@ -189,8 +189,8 @@ def add_integration_callbacks(instance):
     Add integration callbacks from various sources to the instance's callbacks.
 
     Args:
-        instance (Trainer, Predictor, Validator, Exporter): An object with a 'callbacks' attribute that is a dictionary
-            of callback lists.
+        instance (Trainer | Predictor | Validator | Exporter): An object with a 'callbacks' attribute that is a
+            dictionary of callback lists.
     """
     # Load HUB callbacks
     from .hub import callbacks as hub_cb

ultralytics/utils/callbacks/clearml.py

@@ -14,12 +14,12 @@ except (ImportError, AssertionError):
     clearml = None
 
 
-def _log_debug_samples(files, title="Debug Samples") -> None:
+def _log_debug_samples(files, title: str = "Debug Samples") -> None:
     """
     Log files (images) as debug samples in the ClearML task.
 
     Args:
-        files (list): 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
@@ -34,7 +34,7 @@ def _log_debug_samples(files, title="Debug Samples") -> None:
                 )
 
 
-def _log_plot(title, plot_path) -> None:
+def _log_plot(title: str, plot_path: str) -> None:
     """
     Log an image as a plot in the plot section of ClearML.
 
@@ -55,8 +55,8 @@ def _log_plot(title, plot_path) -> None:
     )
 
 
-def on_pretrain_routine_start(trainer):
-    """Runs at start of pretraining routine; initializes and connects/ logs task to ClearML."""
+def on_pretrain_routine_start(trainer) -> None:
+    """Runs at start of pretraining routine; initializes and connects/logs task to ClearML."""
     try:
         if task := Task.current_task():
             # WARNING: make sure the automatic pytorch and matplotlib bindings are disabled!
@@ -84,8 +84,8 @@ def on_pretrain_routine_start(trainer):
         LOGGER.warning(f"WARNING ⚠️ ClearML installed but not initialized correctly, not logging this run. {e}")
 
 
-def on_train_epoch_end(trainer):
-    """Logs debug samples for the first epoch of YOLO training and report current training progress."""
+def on_train_epoch_end(trainer) -> None:
+    """Logs debug samples for the first epoch of YOLO training and reports current training progress."""
     if task := Task.current_task():
         # Log debug samples
         if trainer.epoch == 1:
@@ -97,10 +97,10 @@ def on_train_epoch_end(trainer):
             task.get_logger().report_scalar("lr", k, v, iteration=trainer.epoch)
 
 
-def on_fit_epoch_end(trainer):
+def on_fit_epoch_end(trainer) -> None:
     """Reports model information to logger at the end of an epoch."""
     if task := Task.current_task():
-        # You should have access to the validation bboxes under jdict
+        # Report epoch time and validation metrics
         task.get_logger().report_scalar(
             title="Epoch Time", series="Epoch Time", value=trainer.epoch_time, iteration=trainer.epoch
         )
@@ -113,14 +113,14 @@ def on_fit_epoch_end(trainer):
                 task.get_logger().report_single_value(k, v)
 
 
-def on_val_end(validator):
+def on_val_end(validator) -> None:
     """Logs validation results including labels and predictions."""
     if Task.current_task():
         # Log val_labels and val_pred
         _log_debug_samples(sorted(validator.save_dir.glob("val*.jpg")), "Validation")
 
 
-def on_train_end(trainer):
+def on_train_end(trainer) -> None:
     """Logs final model and its name on training completion."""
     if task := Task.current_task():
         # Log final results, CM matrix + PR plots

ultralytics/utils/callbacks/comet.py

@@ -50,33 +50,33 @@ def _get_comet_mode() -> str:
     return "online"
 
 
-def _get_comet_model_name():
+def _get_comet_model_name() -> str:
     """Returns the model name for Comet from the environment variable COMET_MODEL_NAME or defaults to 'Ultralytics'."""
     return os.getenv("COMET_MODEL_NAME", "Ultralytics")
 
 
-def _get_eval_batch_logging_interval():
+def _get_eval_batch_logging_interval() -> int:
     """Get the evaluation batch logging interval from environment variable or use default value 1."""
     return int(os.getenv("COMET_EVAL_BATCH_LOGGING_INTERVAL", 1))
 
 
-def _get_max_image_predictions_to_log():
+def _get_max_image_predictions_to_log() -> int:
     """Get the maximum number of image predictions to log from the environment variables."""
     return int(os.getenv("COMET_MAX_IMAGE_PREDICTIONS", 100))
 
 
-def _scale_confidence_score(score):
+def _scale_confidence_score(score: float) -> float:
     """Scales the given confidence score by a factor specified in an environment variable."""
     scale = float(os.getenv("COMET_MAX_CONFIDENCE_SCORE", 100.0))
     return score * scale
 
 
-def _should_log_confusion_matrix():
+def _should_log_confusion_matrix() -> bool:
     """Determines if the confusion matrix should be logged based on the environment variable settings."""
     return os.getenv("COMET_EVAL_LOG_CONFUSION_MATRIX", "false").lower() == "true"
 
 
-def _should_log_image_predictions():
+def _should_log_image_predictions() -> bool:
     """Determines whether to log image predictions based on a specified environment variable."""
     return os.getenv("COMET_EVAL_LOG_IMAGE_PREDICTIONS", "true").lower() == "true"
 
@@ -114,7 +114,7 @@ def _resume_or_create_experiment(args: SimpleNamespace) -> None:
         LOGGER.warning(f"WARNING ⚠️ Comet installed but not initialized correctly, not logging this run. {e}")
 
 
-def _fetch_trainer_metadata(trainer):
+def _fetch_trainer_metadata(trainer) -> dict:
     """Returns metadata for YOLO training including epoch and asset saving status."""
     curr_epoch = trainer.epoch + 1
 
@@ -130,7 +130,9 @@ def _fetch_trainer_metadata(trainer):
     return dict(curr_epoch=curr_epoch, curr_step=curr_step, save_assets=save_assets, final_epoch=final_epoch)
 
 
-def _scale_bounding_box_to_original_image_shape(box, resized_image_shape, original_image_shape, ratio_pad):
+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.
 
@@ -151,7 +153,7 @@ def _scale_bounding_box_to_original_image_shape(box, resized_image_shape, origin
     return box
 
 
-def _format_ground_truth_annotations_for_detection(img_idx, image_path, batch, class_name_map=None):
+def _format_ground_truth_annotations_for_detection(img_idx, image_path, batch, class_name_map=None) -> Optional[dict]:
     """Format ground truth annotations for detection."""
     indices = batch["batch_idx"] == img_idx
     bboxes = batch["bboxes"][indices]
@@ -181,7 +183,7 @@ 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):
+def _format_prediction_annotations(image_path, metadata, class_label_map=None, class_map=None) -> Optional[dict]:
     """Format YOLO predictions for object detection visualization."""
     stem = image_path.stem
     image_id = int(stem) if stem.isnumeric() else stem
@@ -227,7 +229,16 @@ 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."""
+    """
+    Extracts 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.
+
+    Returns:
+        (Optional[List[List[Any]]]): List of polygon points or None if extraction fails.
+    """
     try:
         mask = decode(segmentation_raw)
         contours, _ = cv2.findContours(mask, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
@@ -238,7 +249,9 @@ 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):
+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."""
     ground_truth_annotations = _format_ground_truth_annotations_for_detection(
         img_idx, image_path, batch, class_label_map
@@ -253,7 +266,7 @@ def _fetch_annotations(img_idx, image_path, batch, prediction_metadata_map, clas
     return [annotations] if annotations else None
 
 
-def _create_prediction_metadata_map(model_predictions):
+def _create_prediction_metadata_map(model_predictions) -> dict:
     """Create metadata map for model predictions by groupings them based on image ID."""
     pred_metadata_map = {}
     for prediction in model_predictions:
@@ -263,7 +276,7 @@ def _create_prediction_metadata_map(model_predictions):
     return pred_metadata_map
 
 
-def _log_confusion_matrix(experiment, trainer, curr_step, curr_epoch):
+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"]
@@ -272,7 +285,7 @@ def _log_confusion_matrix(experiment, trainer, curr_step, curr_epoch):
     )
 
 
-def _log_images(experiment, image_paths, curr_step, annotations=None):
+def _log_images(experiment, image_paths, curr_step, annotations=None) -> None:
     """Logs images to the experiment with optional annotations."""
     if annotations:
         for image_path, annotation in zip(image_paths, annotations):
@@ -283,7 +296,7 @@ def _log_images(experiment, image_paths, curr_step, annotations=None):
             experiment.log_image(image_path, name=image_path.stem, step=curr_step)
 
 
-def _log_image_predictions(experiment, validator, curr_step):
+def _log_image_predictions(experiment, validator, curr_step) -> None:
     """Logs predicted boxes for a single image during training."""
     global _comet_image_prediction_count
 
@@ -330,7 +343,7 @@ def _log_image_predictions(experiment, validator, curr_step):
             _comet_image_prediction_count += 1
 
 
-def _log_plots(experiment, trainer):
+def _log_plots(experiment, trainer) -> None:
     """Logs evaluation plots and label plots for the experiment."""
     plot_filenames = None
     if isinstance(trainer.validator.metrics, SegmentMetrics) and trainer.validator.metrics.task == "segment":
@@ -359,18 +372,18 @@ def _log_plots(experiment, trainer):
         _log_images(experiment, label_plot_filenames, None)
 
 
-def _log_model(experiment, trainer):
+def _log_model(experiment, trainer) -> None:
     """Log the best-trained model to Comet.ml."""
     model_name = _get_comet_model_name()
     experiment.log_model(model_name, file_or_folder=str(trainer.best), file_name="best.pt", overwrite=True)
 
 
-def on_pretrain_routine_start(trainer):
+def on_pretrain_routine_start(trainer) -> None:
     """Creates or resumes a CometML experiment at the start of a YOLO pre-training routine."""
     _resume_or_create_experiment(trainer.args)
 
 
-def on_train_epoch_end(trainer):
+def on_train_epoch_end(trainer) -> None:
     """Log metrics and save batch images at the end of training epochs."""
     experiment = comet_ml.get_running_experiment()
     if not experiment:
@@ -383,7 +396,7 @@ def on_train_epoch_end(trainer):
     experiment.log_metrics(trainer.label_loss_items(trainer.tloss, prefix="train"), step=curr_step, epoch=curr_epoch)
 
 
-def on_fit_epoch_end(trainer):
+def on_fit_epoch_end(trainer) -> None:
     """Logs model assets at the end of each epoch."""
     experiment = comet_ml.get_running_experiment()
     if not experiment:
@@ -411,7 +424,7 @@ def on_fit_epoch_end(trainer):
         _log_image_predictions(experiment, trainer.validator, curr_step)
 
 
-def on_train_end(trainer):
+def on_train_end(trainer) -> None:
     """Perform operations at the end of training."""
     experiment = comet_ml.get_running_experiment()
     if not experiment:

ultralytics/utils/callbacks/dvc.py

@@ -1,5 +1,7 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
+from pathlib import Path
+
 from ultralytics.utils import LOGGER, SETTINGS, TESTS_RUNNING, checks
 
 try:
@@ -11,7 +13,6 @@ try:
 
     import os
     import re
-    from pathlib import Path
 
     # DVCLive logger instance
     live = None
@@ -25,7 +26,7 @@ except (ImportError, AssertionError, TypeError):
     dvclive = None
 
 
-def _log_images(path, prefix=""):
+def _log_images(path: Path, prefix: str = "") -> None:
     """Logs images at specified path with an optional prefix using DVCLive."""
     if live:
         name = path.name
@@ -39,7 +40,7 @@ def _log_images(path, prefix=""):
         live.log_image(os.path.join(prefix, name), path)
 
 
-def _log_plots(plots, prefix=""):
+def _log_plots(plots: dict, prefix: str = "") -> None:
     """Logs plot images for training progress if they have not been previously processed."""
     for name, params in plots.items():
         timestamp = params["timestamp"]
@@ -48,7 +49,7 @@ def _log_plots(plots, prefix=""):
             _processed_plots[name] = timestamp
 
 
-def _log_confusion_matrix(validator):
+def _log_confusion_matrix(validator) -> None:
     """Logs the confusion matrix for the given validator using DVCLive."""
     targets = []
     preds = []
@@ -65,7 +66,7 @@ def _log_confusion_matrix(validator):
     live.log_sklearn_plot("confusion_matrix", targets, preds, name="cf.json", normalized=True)
 
 
-def on_pretrain_routine_start(trainer):
+def on_pretrain_routine_start(trainer) -> None:
     """Initializes DVCLive logger for training metadata during pre-training routine."""
     try:
         global live
@@ -75,24 +76,24 @@ def on_pretrain_routine_start(trainer):
         LOGGER.warning(f"WARNING ⚠️ DVCLive installed but not initialized correctly, not logging this run. {e}")
 
 
-def on_pretrain_routine_end(trainer):
+def on_pretrain_routine_end(trainer) -> None:
     """Logs plots related to the training process at the end of the pretraining routine."""
     _log_plots(trainer.plots, "train")
 
 
-def on_train_start(trainer):
+def on_train_start(trainer) -> None:
     """Logs the training parameters if DVCLive logging is active."""
     if live:
         live.log_params(trainer.args)
 
 
-def on_train_epoch_start(trainer):
+def on_train_epoch_start(trainer) -> None:
     """Sets 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):
+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."""
     global _training_epoch
     if live and _training_epoch:
@@ -113,7 +114,7 @@ def on_fit_epoch_end(trainer):
         _training_epoch = False
 
 
-def on_train_end(trainer):
+def on_train_end(trainer) -> None:
     """Logs the best metrics, plots, and confusion matrix at the end of training if DVCLive is active."""
     if live:
         # At the end log the best metrics. It runs validator on the best model internally.

ultralytics/utils/callbacks/hub.py

@@ -14,16 +14,16 @@ def on_pretrain_routine_start(trainer):
 
 
 def on_pretrain_routine_end(trainer):
-    """Logs info before starting timer for upload rate limit."""
+    """Initialize timers for upload rate limiting before training begins."""
     if session := getattr(trainer, "hub_session", None):
         # Start timer for upload rate limit
-        session.timers = {"metrics": time(), "ckpt": time()}  # start timer on session.rate_limit
+        session.timers = {"metrics": time(), "ckpt": time()}  # start timer for session rate limiting
 
 
 def on_fit_epoch_end(trainer):
-    """Uploads training progress metrics at the end of each epoch."""
+    """Upload training progress metrics to Ultralytics HUB at the end of each epoch."""
     if session := getattr(trainer, "hub_session", None):
-        # Upload metrics after val end
+        # Upload metrics after validation ends
         all_plots = {
             **trainer.label_loss_items(trainer.tloss, prefix="train"),
             **trainer.metrics,
@@ -35,7 +35,7 @@ def on_fit_epoch_end(trainer):
 
         session.metrics_queue[trainer.epoch] = json.dumps(all_plots)
 
-        # If any metrics fail to upload, add them to the queue to attempt uploading again.
+        # If any metrics failed to upload previously, add them to the queue to attempt uploading again
         if session.metrics_upload_failed_queue:
             session.metrics_queue.update(session.metrics_upload_failed_queue)
 
@@ -46,7 +46,7 @@ def on_fit_epoch_end(trainer):
 
 
 def on_model_save(trainer):
-    """Saves checkpoints to Ultralytics HUB with rate limiting."""
+    """Upload model checkpoints to Ultralytics HUB with rate limiting."""
     if session := getattr(trainer, "hub_session", None):
         # Upload checkpoints with rate limiting
         is_best = trainer.best_fitness == trainer.fitness
@@ -77,7 +77,7 @@ def on_train_start(trainer):
 
 
 def on_val_start(validator):
-    """Runs events on validation start."""
+    """Run events on validation start."""
     events(validator.args)
 
 
@@ -105,4 +105,4 @@ callbacks = (
     }
     if SETTINGS["hub"] is True
     else {}
-)  # verify enabled
+)  # verify hub is enabled before registering callbacks

ultralytics/utils/callbacks/mlflow.py

@@ -39,7 +39,7 @@ except (ImportError, AssertionError):
     mlflow = None
 
 
-def sanitize_dict(x):
+def sanitize_dict(x: dict) -> dict:
     """Sanitize dictionary keys by removing parentheses and converting values to floats."""
     return {k.replace("(", "").replace(")", ""): float(v) for k, v in x.items()}