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

ultralytics/models/yolo/classify/predict.py

@@ -44,7 +44,7 @@ class ClassificationPredictor(BasePredictor):
         tasks. It ensures the task is set to 'classify' regardless of input configuration.
 
         Args:
-            cfg (dict): Default configuration dictionary containing prediction settings. Defaults to DEFAULT_CFG.
+            cfg (dict): Default configuration dictionary containing prediction settings.
             overrides (dict, optional): Configuration overrides that take precedence over cfg.
             _callbacks (list, optional): List of callback functions to be executed during prediction.
         """
@@ -53,7 +53,7 @@ class ClassificationPredictor(BasePredictor):
         self._legacy_transform_name = "ultralytics.yolo.data.augment.ToTensor"
 
     def setup_source(self, source):
-        """Sets up source and inference mode and classify transforms."""
+        """Set up source and inference mode and classify transforms."""
         super().setup_source(source)
         updated = (
             self.model.model.transforms.transforms[0].size != max(self.imgsz)
@@ -68,14 +68,14 @@ class ClassificationPredictor(BasePredictor):
             is_legacy_transform = any(
                 self._legacy_transform_name in str(transform) for transform in self.transforms.transforms
             )
-            if is_legacy_transform:  # to handle legacy transforms
+            if is_legacy_transform:  # Handle legacy transforms
                 img = torch.stack([self.transforms(im) for im in img], dim=0)
             else:
                 img = torch.stack(
                     [self.transforms(Image.fromarray(cv2.cvtColor(im, cv2.COLOR_BGR2RGB))) for im in img], dim=0
                 )
         img = (img if isinstance(img, torch.Tensor) else torch.from_numpy(img)).to(self.model.device)
-        return img.half() if self.model.fp16 else img.float()  # uint8 to fp16/32
+        return img.half() if self.model.fp16 else img.float()  # Convert uint8 to fp16/32
 
     def postprocess(self, preds, img, orig_imgs):
         """
@@ -89,7 +89,7 @@ class ClassificationPredictor(BasePredictor):
         Returns:
             (List[Results]): List of Results objects containing classification results for each image.
         """
-        if not isinstance(orig_imgs, list):  # input images are a torch.Tensor, not a list
+        if not isinstance(orig_imgs, list):  # Input images are a torch.Tensor, not a list
             orig_imgs = ops.convert_torch2numpy_batch(orig_imgs)
 
         preds = preds[0] if isinstance(preds, (list, tuple)) else preds

ultralytics/models/yolo/classify/train.py

@@ -1,6 +1,7 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
 from copy import copy
+from typing import Any, Dict, Optional
 
 import torch
 
@@ -15,14 +16,14 @@ from ultralytics.utils.torch_utils import is_parallel, strip_optimizer, torch_di
 
 class ClassificationTrainer(BaseTrainer):
     """
-    A class extending the BaseTrainer class for training based on a classification model.
+    A trainer class extending BaseTrainer for training image classification models.
 
     This trainer handles the training process for image classification tasks, supporting both YOLO classification models
-    and torchvision models.
+    and torchvision models with comprehensive dataset handling and validation.
 
     Attributes:
         model (ClassificationModel): The classification model to be trained.
-        data (dict): Dictionary containing dataset information including class names and number of classes.
+        data (Dict[str, Any]): Dictionary containing dataset information including class names and number of classes.
         loss_names (List[str]): Names of the loss functions used during training.
         validator (ClassificationValidator): Validator instance for model evaluation.
 
@@ -41,13 +42,14 @@ class ClassificationTrainer(BaseTrainer):
         plot_training_samples: Plot training samples with their annotations.
 
     Examples:
+        Initialize and train a classification model
         >>> from ultralytics.models.yolo.classify import ClassificationTrainer
         >>> args = dict(model="yolo11n-cls.pt", data="imagenet10", epochs=3)
         >>> trainer = ClassificationTrainer(overrides=args)
         >>> trainer.train()
     """
 
-    def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
+    def __init__(self, cfg=DEFAULT_CFG, overrides: Optional[Dict[str, Any]] = None, _callbacks=None):
         """
         Initialize a ClassificationTrainer object.
 
@@ -55,11 +57,12 @@ class ClassificationTrainer(BaseTrainer):
         image size if not specified.
 
         Args:
-            cfg (dict, optional): Default configuration dictionary containing training parameters.
-            overrides (dict, optional): Dictionary of parameter overrides for the default configuration.
-            _callbacks (list, optional): List of callback functions to be executed during training.
+            cfg (Dict[str, Any], optional): Default configuration dictionary containing training parameters.
+            overrides (Dict[str, Any], optional): Dictionary of parameter overrides for the default configuration.
+            _callbacks (List[Any], optional): List of callback functions to be executed during training.
 
         Examples:
+            Create a trainer with custom configuration
             >>> from ultralytics.models.yolo.classify import ClassificationTrainer
             >>> args = dict(model="yolo11n-cls.pt", data="imagenet10", epochs=3)
             >>> trainer = ClassificationTrainer(overrides=args)
@@ -76,14 +79,14 @@ class ClassificationTrainer(BaseTrainer):
         """Set the YOLO model's class names from the loaded dataset."""
         self.model.names = self.data["names"]
 
-    def get_model(self, cfg=None, weights=None, verbose=True):
+    def get_model(self, cfg=None, weights=None, verbose: bool = True):
         """
-        Return a modified PyTorch model configured for training YOLO.
+        Return a modified PyTorch model configured for training YOLO classification.
 
         Args:
-            cfg (Any): Model configuration.
-            weights (Any): Pre-trained model weights.
-            verbose (bool): Whether to display model information.
+            cfg (Any, optional): Model configuration.
+            weights (Any, optional): Pre-trained model weights.
+            verbose (bool, optional): Whether to display model information.
 
         Returns:
             (ClassificationModel): Configured PyTorch model for classification.
@@ -120,29 +123,29 @@ class ClassificationTrainer(BaseTrainer):
         ClassificationModel.reshape_outputs(self.model, self.data["nc"])
         return ckpt
 
-    def build_dataset(self, img_path, mode="train", batch=None):
+    def build_dataset(self, img_path: str, mode: str = "train", batch=None):
         """
         Create a ClassificationDataset instance given an image path and mode.
 
         Args:
             img_path (str): Path to the dataset images.
-            mode (str): Dataset mode ('train', 'val', or 'test').
-            batch (Any): Batch information (unused in this implementation).
+            mode (str, optional): Dataset mode ('train', 'val', or 'test').
+            batch (Any, optional): Batch information (unused in this implementation).
 
         Returns:
             (ClassificationDataset): Dataset for the specified mode.
         """
         return ClassificationDataset(root=img_path, args=self.args, augment=mode == "train", prefix=mode)
 
-    def get_dataloader(self, dataset_path, batch_size=16, rank=0, mode="train"):
+    def get_dataloader(self, dataset_path: str, batch_size: int = 16, rank: int = 0, mode: str = "train"):
         """
         Return PyTorch DataLoader with transforms to preprocess images.
 
         Args:
             dataset_path (str): Path to the dataset.
-            batch_size (int): Number of images per batch.
-            rank (int): Process rank for distributed training.
-            mode (str): 'train', 'val', or 'test' mode.
+            batch_size (int, optional): Number of images per batch.
+            rank (int, optional): Process rank for distributed training.
+            mode (str, optional): 'train', 'val', or 'test' mode.
 
         Returns:
             (torch.utils.data.DataLoader): DataLoader for the specified dataset and mode.
@@ -159,14 +162,14 @@ class ClassificationTrainer(BaseTrainer):
                 self.model.transforms = loader.dataset.torch_transforms
         return loader
 
-    def preprocess_batch(self, batch):
-        """Preprocesses a batch of images and classes."""
+    def preprocess_batch(self, batch: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        """Preprocess a batch of images and classes."""
         batch["img"] = batch["img"].to(self.device)
         batch["cls"] = batch["cls"].to(self.device)
         return batch
 
-    def progress_string(self):
-        """Returns a formatted string showing training progress."""
+    def progress_string(self) -> str:
+        """Return a formatted string showing training progress."""
         return ("\n" + "%11s" * (4 + len(self.loss_names))) % (
             "Epoch",
             "GPU_mem",
@@ -176,22 +179,23 @@ class ClassificationTrainer(BaseTrainer):
         )
 
     def get_validator(self):
-        """Returns an instance of ClassificationValidator for validation."""
+        """Return an instance of ClassificationValidator for validation."""
         self.loss_names = ["loss"]
         return yolo.classify.ClassificationValidator(
             self.test_loader, self.save_dir, args=copy(self.args), _callbacks=self.callbacks
         )
 
-    def label_loss_items(self, loss_items=None, prefix="train"):
+    def label_loss_items(self, loss_items: Optional[torch.Tensor] = None, prefix: str = "train"):
         """
         Return a loss dict with labelled training loss items tensor.
 
         Args:
             loss_items (torch.Tensor, optional): Loss tensor items.
-            prefix (str): Prefix to prepend to loss names.
+            prefix (str, optional): Prefix to prepend to loss names.
 
         Returns:
-            (Dict[str, float] | List[str]): Dictionary of loss items or list of loss keys if loss_items is None.
+            keys (List[str]): List of loss keys if loss_items is None.
+            loss_dict (Dict[str, float]): Dictionary of loss items if loss_items is provided.
         """
         keys = [f"{prefix}/{x}" for x in self.loss_names]
         if loss_items is None:
@@ -216,7 +220,7 @@ class ClassificationTrainer(BaseTrainer):
                     self.metrics.pop("fitness", None)
                     self.run_callbacks("on_fit_epoch_end")
 
-    def plot_training_samples(self, batch, ni):
+    def plot_training_samples(self, batch: Dict[str, torch.Tensor], ni: int):
         """
         Plot training samples with their annotations.
 

ultralytics/models/yolo/classify/val.py

@@ -52,9 +52,6 @@ class ClassificationValidator(BaseValidator):
         """
         Initialize ClassificationValidator with dataloader, save directory, and other parameters.
 
-        This validator handles the validation process for classification models, including metrics calculation,
-        confusion matrix generation, and visualization of results.
-
         Args:
             dataloader (torch.utils.data.DataLoader, optional): Dataloader to use for validation.
             save_dir (str | Path, optional): Directory to save results.
@@ -101,8 +98,9 @@ class ClassificationValidator(BaseValidator):
             preds (torch.Tensor): Model predictions, typically logits or probabilities for each class.
             batch (dict): Batch data containing images and class labels.
 
-        This method appends the top-N predictions (sorted by confidence in descending order) to the
-        prediction list for later evaluation. N is limited to the minimum of 5 and the number of classes.
+        Notes:
+            This method appends the top-N predictions (sorted by confidence in descending order) to the
+            prediction list for later evaluation. N is limited to the minimum of 5 and the number of classes.
         """
         n5 = min(len(self.names), 5)
         self.pred.append(preds.argsort(1, descending=True)[:, :n5].type(torch.int32).cpu())
@@ -112,13 +110,14 @@ class ClassificationValidator(BaseValidator):
         """
         Finalize metrics including confusion matrix and processing speed.
 
-        This method processes the accumulated predictions and targets to generate the confusion matrix,
-        optionally plots it, and updates the metrics object with speed information.
-
         Args:
             *args (Any): Variable length argument list.
             **kwargs (Any): Arbitrary keyword arguments.
 
+        Notes:
+            This method processes the accumulated predictions and targets to generate the confusion matrix,
+            optionally plots it, and updates the metrics object with speed information.
+
         Examples:
             >>> validator = ClassificationValidator()
             >>> validator.pred = [torch.tensor([[0, 1, 2]])]  # Top-3 predictions for one sample

ultralytics/models/yolo/detect/predict.py

@@ -21,6 +21,7 @@ class DetectionPredictor(BasePredictor):
         postprocess: Process raw model predictions into detection results.
         construct_results: Build Results objects from processed predictions.
         construct_result: Create a single Result object from a prediction.
+        get_obj_feats: Extract object features from the feature maps.
 
     Examples:
         >>> from ultralytics.utils import ASSETS

ultralytics/models/yolo/detect/train.py

@@ -3,6 +3,7 @@
 import math
 import random
 from copy import copy
+from typing import Dict, List, Optional
 
 import numpy as np
 import torch.nn as nn
@@ -21,12 +22,12 @@ class DetectionTrainer(BaseTrainer):
     A class extending the BaseTrainer class for training based on a detection model.
 
     This trainer specializes in object detection tasks, handling the specific requirements for training YOLO models
-    for object detection.
+    for object detection including dataset building, data loading, preprocessing, and model configuration.
 
     Attributes:
         model (DetectionModel): The YOLO detection model being trained.
-        data (dict): Dictionary containing dataset information including class names and number of classes.
-        loss_names (Tuple[str]): Names of the loss components used in training (box_loss, cls_loss, dfl_loss).
+        data (Dict): Dictionary containing dataset information including class names and number of classes.
+        loss_names (tuple): Names of the loss components used in training (box_loss, cls_loss, dfl_loss).
 
     Methods:
         build_dataset: Build YOLO dataset for training or validation.
@@ -49,14 +50,14 @@ class DetectionTrainer(BaseTrainer):
         >>> trainer.train()
     """
 
-    def build_dataset(self, img_path, mode="train", batch=None):
+    def build_dataset(self, img_path: str, mode: str = "train", batch: Optional[int] = None):
         """
         Build YOLO Dataset for training or validation.
 
         Args:
             img_path (str): Path to the folder containing images.
-            mode (str): `train` mode or `val` mode, users are able to customize different augmentations for each mode.
-            batch (int, optional): Size of batches, this is for `rect`.
+            mode (str): 'train' mode or 'val' mode, users are able to customize different augmentations for each mode.
+            batch (int, optional): Size of batches, this is for 'rect' mode.
 
         Returns:
             (Dataset): YOLO dataset object configured for the specified mode.
@@ -64,7 +65,7 @@ class DetectionTrainer(BaseTrainer):
         gs = max(int(de_parallel(self.model).stride.max() if self.model else 0), 32)
         return build_yolo_dataset(self.args, img_path, batch, self.data, mode=mode, rect=mode == "val", stride=gs)
 
-    def get_dataloader(self, dataset_path, batch_size=16, rank=0, mode="train"):
+    def get_dataloader(self, dataset_path: str, batch_size: int = 16, rank: int = 0, mode: str = "train"):
         """
         Construct and return dataloader for the specified mode.
 
@@ -87,15 +88,15 @@ class DetectionTrainer(BaseTrainer):
         workers = self.args.workers if mode == "train" else self.args.workers * 2
         return build_dataloader(dataset, batch_size, workers, shuffle, rank)  # return dataloader
 
-    def preprocess_batch(self, batch):
+    def preprocess_batch(self, batch: Dict) -> Dict:
         """
         Preprocess a batch of images by scaling and converting to float.
 
         Args:
-            batch (dict): Dictionary containing batch data with 'img' tensor.
+            batch (Dict): Dictionary containing batch data with 'img' tensor.
 
         Returns:
-            (dict): Preprocessed batch with normalized images.
+            (Dict): Preprocessed batch with normalized images.
         """
         batch["img"] = batch["img"].to(self.device, non_blocking=True).float() / 255
         if self.args.multi_scale:
@@ -125,7 +126,7 @@ class DetectionTrainer(BaseTrainer):
         self.model.args = self.args  # attach hyperparameters to model
         # TODO: self.model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc
 
-    def get_model(self, cfg=None, weights=None, verbose=True):
+    def get_model(self, cfg: Optional[str] = None, weights: Optional[str] = None, verbose: bool = True):
         """
         Return a YOLO detection model.
 
@@ -149,7 +150,7 @@ class DetectionTrainer(BaseTrainer):
             self.test_loader, save_dir=self.save_dir, args=copy(self.args), _callbacks=self.callbacks
         )
 
-    def label_loss_items(self, loss_items=None, prefix="train"):
+    def label_loss_items(self, loss_items: Optional[List[float]] = None, prefix: str = "train"):
         """
         Return a loss dict with labeled training loss items tensor.
 
@@ -177,12 +178,12 @@ class DetectionTrainer(BaseTrainer):
             "Size",
         )
 
-    def plot_training_samples(self, batch, ni):
+    def plot_training_samples(self, batch: Dict, ni: int):
         """
         Plot training samples with their annotations.
 
         Args:
-            batch (dict): Dictionary containing batch data.
+            batch (Dict): Dictionary containing batch data.
             ni (int): Number of iterations.
         """
         plot_images(

ultralytics/models/yolo/detect/val.py

@@ -2,6 +2,7 @@
 
 import os
 from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
 
 import numpy as np
 import torch
@@ -26,13 +27,13 @@ class DetectionValidator(BaseValidator):
         nt_per_image (np.ndarray): Number of targets per image.
         is_coco (bool): Whether the dataset is COCO.
         is_lvis (bool): Whether the dataset is LVIS.
-        class_map (list): Mapping from model class indices to dataset class indices.
+        class_map (List[int]): Mapping from model class indices to dataset class indices.
         metrics (DetMetrics): Object detection metrics calculator.
         iouv (torch.Tensor): IoU thresholds for mAP calculation.
         niou (int): Number of IoU thresholds.
-        lb (list): List for storing ground truth labels for hybrid saving.
-        jdict (list): List for storing JSON detection results.
-        stats (dict): Dictionary for storing statistics during validation.
+        lb (List[Any]): List for storing ground truth labels for hybrid saving.
+        jdict (List[Dict[str, Any]]): List for storing JSON detection results.
+        stats (Dict[str, List[torch.Tensor]]): Dictionary for storing statistics during validation.
 
     Examples:
         >>> from ultralytics.models.yolo.detect import DetectionValidator
@@ -49,8 +50,8 @@ class DetectionValidator(BaseValidator):
             dataloader (torch.utils.data.DataLoader, optional): Dataloader to use for validation.
             save_dir (Path, optional): Directory to save results.
             pbar (Any, optional): Progress bar for displaying progress.
-            args (dict, optional): Arguments for the validator.
-            _callbacks (list, optional): List of callback functions.
+            args (Dict[str, Any], optional): Arguments for the validator.
+            _callbacks (List[Any], optional): List of callback functions.
         """
         super().__init__(dataloader, save_dir, pbar, args, _callbacks)
         self.nt_per_class = None
@@ -63,15 +64,15 @@ class DetectionValidator(BaseValidator):
         self.iouv = torch.linspace(0.5, 0.95, 10)  # IoU vector for mAP@0.5:0.95
         self.niou = self.iouv.numel()
 
-    def preprocess(self, batch):
+    def preprocess(self, batch: Dict[str, Any]) -> Dict[str, Any]:
         """
         Preprocess batch of images for YOLO validation.
 
         Args:
-            batch (dict): Batch containing images and annotations.
+            batch (Dict[str, Any]): Batch containing images and annotations.
 
         Returns:
-            (dict): Preprocessed batch.
+            (Dict[str, Any]): Preprocessed batch.
         """
         batch["img"] = batch["img"].to(self.device, non_blocking=True)
         batch["img"] = (batch["img"].half() if self.args.half else batch["img"].float()) / 255
@@ -80,7 +81,7 @@ class DetectionValidator(BaseValidator):
 
         return batch
 
-    def init_metrics(self, model):
+    def init_metrics(self, model: torch.nn.Module) -> None:
         """
         Initialize evaluation metrics for YOLO detection validation.
 
@@ -106,11 +107,11 @@ class DetectionValidator(BaseValidator):
         self.jdict = []
         self.stats = dict(tp=[], conf=[], pred_cls=[], target_cls=[], target_img=[])
 
-    def get_desc(self):
+    def get_desc(self) -> str:
         """Return a formatted string summarizing class metrics of YOLO model."""
         return ("%22s" + "%11s" * 6) % ("Class", "Images", "Instances", "Box(P", "R", "mAP50", "mAP50-95)")
 
-    def postprocess(self, preds):
+    def postprocess(self, preds: torch.Tensor) -> List[torch.Tensor]:
         """
         Apply Non-maximum suppression to prediction outputs.
 
@@ -132,16 +133,16 @@ class DetectionValidator(BaseValidator):
             rotated=self.args.task == "obb",
         )
 
-    def _prepare_batch(self, si, batch):
+    def _prepare_batch(self, si: int, batch: Dict[str, Any]) -> Dict[str, Any]:
         """
         Prepare a batch of images and annotations for validation.
 
         Args:
             si (int): Batch index.
-            batch (dict): Batch data containing images and annotations.
+            batch (Dict[str, Any]): Batch data containing images and annotations.
 
         Returns:
-            (dict): Prepared batch with processed annotations.
+            (Dict[str, Any]): Prepared batch with processed annotations.
         """
         idx = batch["batch_idx"] == si
         cls = batch["cls"][idx].squeeze(-1)
@@ -154,13 +155,13 @@ class DetectionValidator(BaseValidator):
             ops.scale_boxes(imgsz, bbox, ori_shape, ratio_pad=ratio_pad)  # native-space labels
         return {"cls": cls, "bbox": bbox, "ori_shape": ori_shape, "imgsz": imgsz, "ratio_pad": ratio_pad}
 
-    def _prepare_pred(self, pred, pbatch):
+    def _prepare_pred(self, pred: torch.Tensor, pbatch: Dict[str, Any]) -> torch.Tensor:
         """
         Prepare predictions for evaluation against ground truth.
 
         Args:
             pred (torch.Tensor): Model predictions.
-            pbatch (dict): Prepared batch information.
+            pbatch (Dict[str, Any]): Prepared batch information.
 
         Returns:
             (torch.Tensor): Prepared predictions in native space.
@@ -171,13 +172,13 @@ class DetectionValidator(BaseValidator):
         )  # native-space pred
         return predn
 
-    def update_metrics(self, preds, batch):
+    def update_metrics(self, preds: List[torch.Tensor], batch: Dict[str, Any]) -> None:
         """
         Update metrics with new predictions and ground truth.
 
         Args:
             preds (List[torch.Tensor]): List of predictions from the model.
-            batch (dict): Batch data containing ground truth.
+            batch (Dict[str, Any]): Batch data containing ground truth.
         """
         for si, pred in enumerate(preds):
             self.seen += 1
@@ -226,7 +227,7 @@ class DetectionValidator(BaseValidator):
                     self.save_dir / "labels" / f"{Path(batch['im_file'][si]).stem}.txt",
                 )
 
-    def finalize_metrics(self, *args, **kwargs):
+    def finalize_metrics(self, *args: Any, **kwargs: Any) -> None:
         """
         Set final values for metrics speed and confusion matrix.
 
@@ -237,12 +238,12 @@ class DetectionValidator(BaseValidator):
         self.metrics.speed = self.speed
         self.metrics.confusion_matrix = self.confusion_matrix
 
-    def get_stats(self):
+    def get_stats(self) -> Dict[str, Any]:
         """
         Calculate and return metrics statistics.
 
         Returns:
-            (dict): Dictionary containing metrics results.
+            (Dict[str, Any]): Dictionary containing metrics results.
         """
         stats = {k: torch.cat(v, 0).cpu().numpy() for k, v in self.stats.items()}  # to numpy
         self.nt_per_class = np.bincount(stats["target_cls"].astype(int), minlength=self.nc)
@@ -252,7 +253,7 @@ class DetectionValidator(BaseValidator):
             self.metrics.process(**stats, on_plot=self.on_plot)
         return self.metrics.results_dict
 
-    def print_results(self):
+    def print_results(self) -> None:
         """Print training/validation set metrics per class."""
         pf = "%22s" + "%11i" * 2 + "%11.3g" * len(self.metrics.keys)  # print format
         LOGGER.info(pf % ("all", self.seen, self.nt_per_class.sum(), *self.metrics.mean_results()))
@@ -272,7 +273,7 @@ class DetectionValidator(BaseValidator):
                     save_dir=self.save_dir, names=self.names.values(), normalize=normalize, on_plot=self.on_plot
                 )
 
-    def _process_batch(self, detections, gt_bboxes, gt_cls):
+    def _process_batch(self, detections: torch.Tensor, gt_bboxes: torch.Tensor, gt_cls: torch.Tensor) -> torch.Tensor:
         """
         Return correct prediction matrix.
 
@@ -289,7 +290,7 @@ class DetectionValidator(BaseValidator):
         iou = box_iou(gt_bboxes, detections[:, :4])
         return self.match_predictions(detections[:, 5], gt_cls, iou)
 
-    def build_dataset(self, img_path, mode="val", batch=None):
+    def build_dataset(self, img_path: str, mode: str = "val", batch: Optional[int] = None):
         """
         Build YOLO Dataset.
 
@@ -303,7 +304,7 @@ class DetectionValidator(BaseValidator):
         """
         return build_yolo_dataset(self.args, img_path, batch, self.data, mode=mode, stride=self.stride)
 
-    def get_dataloader(self, dataset_path, batch_size):
+    def get_dataloader(self, dataset_path: str, batch_size: int) -> torch.utils.data.DataLoader:
         """
         Construct and return dataloader.
 
@@ -317,12 +318,12 @@ class DetectionValidator(BaseValidator):
         dataset = self.build_dataset(dataset_path, batch=batch_size, mode="val")
         return build_dataloader(dataset, batch_size, self.args.workers, shuffle=False, rank=-1)  # return dataloader
 
-    def plot_val_samples(self, batch, ni):
+    def plot_val_samples(self, batch: Dict[str, Any], ni: int) -> None:
         """
         Plot validation image samples.
 
         Args:
-            batch (dict): Batch containing images and annotations.
+            batch (Dict[str, Any]): Batch containing images and annotations.
             ni (int): Batch index.
         """
         plot_images(
@@ -336,12 +337,12 @@ class DetectionValidator(BaseValidator):
             on_plot=self.on_plot,
         )
 
-    def plot_predictions(self, batch, preds, ni):
+    def plot_predictions(self, batch: Dict[str, Any], preds: List[torch.Tensor], ni: int) -> None:
         """
         Plot predicted bounding boxes on input images and save the result.
 
         Args:
-            batch (dict): Batch containing images and annotations.
+            batch (Dict[str, Any]): Batch containing images and annotations.
             preds (List[torch.Tensor]): List of predictions from the model.
             ni (int): Batch index.
         """
@@ -354,14 +355,14 @@ class DetectionValidator(BaseValidator):
             on_plot=self.on_plot,
         )  # pred
 
-    def save_one_txt(self, predn, save_conf, shape, file):
+    def save_one_txt(self, predn: torch.Tensor, save_conf: bool, shape: Tuple[int, int], file: Path) -> None:
         """
         Save YOLO detections to a txt file in normalized coordinates in a specific format.
 
         Args:
             predn (torch.Tensor): Predictions in the format (x1, y1, x2, y2, conf, class).
             save_conf (bool): Whether to save confidence scores.
-            shape (tuple): Shape of the original image.
+            shape (Tuple[int, int]): Shape of the original image.
             file (Path): File path to save the detections.
         """
         from ultralytics.engine.results import Results
@@ -373,7 +374,7 @@ class DetectionValidator(BaseValidator):
             boxes=predn[:, :6],
         ).save_txt(file, save_conf=save_conf)
 
-    def pred_to_json(self, predn, filename):
+    def pred_to_json(self, predn: torch.Tensor, filename: str) -> None:
         """
         Serialize YOLO predictions to COCO json format.
 
@@ -395,15 +396,15 @@ class DetectionValidator(BaseValidator):
                 }
             )
 
-    def eval_json(self, stats):
+    def eval_json(self, stats: Dict[str, Any]) -> Dict[str, Any]:
         """
         Evaluate YOLO output in JSON format and return performance statistics.
 
         Args:
-            stats (dict): Current statistics dictionary.
+            stats (Dict[str, Any]): Current statistics dictionary.
 
         Returns:
-            (dict): Updated statistics dictionary with COCO/LVIS evaluation results.
+            (Dict[str, Any]): Updated statistics dictionary with COCO/LVIS evaluation results.
         """
         if self.args.save_json and (self.is_coco or self.is_lvis) and len(self.jdict):
             pred_json = self.save_dir / "predictions.json"  # predictions