dgenerate-ultralytics-headless 8.3.190__py3-none-any.whl → 8.3.192__py3-none-any.whl

ultralytics/models/sam/modules/tiny_encoder.py

@@ -9,8 +9,9 @@
 # Build the TinyViT Model
 # --------------------------------------------------------
 
+from __future__ import annotations
+
 import itertools
-from typing import List, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -106,7 +107,7 @@ class PatchEmbed(nn.Module):
             activation (nn.Module): Activation function to use between convolutions.
         """
         super().__init__()
-        img_size: Tuple[int, int] = to_2tuple(resolution)
+        img_size: tuple[int, int] = to_2tuple(resolution)
         self.patches_resolution = (img_size[0] // 4, img_size[1] // 4)
         self.num_patches = self.patches_resolution[0] * self.patches_resolution[1]
         self.in_chans = in_chans
@@ -219,7 +220,7 @@ class PatchMerging(nn.Module):
         torch.Size([4, 3136, 128])
     """
 
-    def __init__(self, input_resolution: Tuple[int, int], dim: int, out_dim: int, activation):
+    def __init__(self, input_resolution: tuple[int, int], dim: int, out_dim: int, activation):
         """
         Initialize the PatchMerging module for merging and projecting neighboring patches in feature maps.
 
@@ -283,13 +284,13 @@ class ConvLayer(nn.Module):
     def __init__(
         self,
         dim: int,
-        input_resolution: Tuple[int, int],
+        input_resolution: tuple[int, int],
         depth: int,
         activation,
-        drop_path: Union[float, List[float]] = 0.0,
-        downsample: Optional[nn.Module] = None,
+        drop_path: float | list[float] = 0.0,
+        downsample: nn.Module | None = None,
         use_checkpoint: bool = False,
-        out_dim: Optional[int] = None,
+        out_dim: int | None = None,
         conv_expand_ratio: float = 4.0,
     ):
         """
@@ -370,8 +371,8 @@ class MLP(nn.Module):
     def __init__(
         self,
         in_features: int,
-        hidden_features: Optional[int] = None,
-        out_features: Optional[int] = None,
+        hidden_features: int | None = None,
+        out_features: int | None = None,
         activation=nn.GELU,
         drop: float = 0.0,
     ):
@@ -441,7 +442,7 @@ class Attention(torch.nn.Module):
         key_dim: int,
         num_heads: int = 8,
         attn_ratio: float = 4,
-        resolution: Tuple[int, int] = (14, 14),
+        resolution: tuple[int, int] = (14, 14),
     ):
         """
         Initialize the Attention module for multi-head attention with spatial awareness.
@@ -549,7 +550,7 @@ class TinyViTBlock(nn.Module):
     def __init__(
         self,
         dim: int,
-        input_resolution: Tuple[int, int],
+        input_resolution: tuple[int, int],
         num_heads: int,
         window_size: int = 7,
         mlp_ratio: float = 4.0,
@@ -690,18 +691,18 @@ class BasicLayer(nn.Module):
     def __init__(
         self,
         dim: int,
-        input_resolution: Tuple[int, int],
+        input_resolution: tuple[int, int],
         depth: int,
         num_heads: int,
         window_size: int,
         mlp_ratio: float = 4.0,
         drop: float = 0.0,
-        drop_path: Union[float, List[float]] = 0.0,
-        downsample: Optional[nn.Module] = None,
+        drop_path: float | list[float] = 0.0,
+        downsample: nn.Module | None = None,
         use_checkpoint: bool = False,
         local_conv_size: int = 3,
         activation=nn.GELU,
-        out_dim: Optional[int] = None,
+        out_dim: int | None = None,
     ):
         """
         Initialize a BasicLayer in the TinyViT architecture.
@@ -800,10 +801,10 @@ class TinyViT(nn.Module):
         img_size: int = 224,
         in_chans: int = 3,
         num_classes: int = 1000,
-        embed_dims: Tuple[int, int, int, int] = (96, 192, 384, 768),
-        depths: Tuple[int, int, int, int] = (2, 2, 6, 2),
-        num_heads: Tuple[int, int, int, int] = (3, 6, 12, 24),
-        window_sizes: Tuple[int, int, int, int] = (7, 7, 14, 7),
+        embed_dims: tuple[int, int, int, int] = (96, 192, 384, 768),
+        depths: tuple[int, int, int, int] = (2, 2, 6, 2),
+        num_heads: tuple[int, int, int, int] = (3, 6, 12, 24),
+        window_sizes: tuple[int, int, int, int] = (7, 7, 14, 7),
         mlp_ratio: float = 4.0,
         drop_rate: float = 0.0,
         drop_path_rate: float = 0.1,
@@ -980,7 +981,7 @@ class TinyViT(nn.Module):
         """Perform the forward pass through the TinyViT model, extracting features from the input image."""
         return self.forward_features(x)
 
-    def set_imgsz(self, imgsz: List[int] = [1024, 1024]):
+    def set_imgsz(self, imgsz: list[int] = [1024, 1024]):
         """Set image size to make model compatible with different image sizes."""
         imgsz = [s // 4 for s in imgsz]
         self.patches_resolution = imgsz

ultralytics/models/sam/modules/transformer.py

@@ -1,7 +1,8 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
+from __future__ import annotations
+
 import math
-from typing import Tuple, Type
 
 import torch
 from torch import Tensor, nn
@@ -44,7 +45,7 @@ class TwoWayTransformer(nn.Module):
         embedding_dim: int,
         num_heads: int,
         mlp_dim: int,
-        activation: Type[nn.Module] = nn.ReLU,
+        activation: type[nn.Module] = nn.ReLU,
         attention_downsample_rate: int = 2,
     ) -> None:
         """
@@ -85,7 +86,7 @@ class TwoWayTransformer(nn.Module):
         image_embedding: torch.Tensor,
         image_pe: torch.Tensor,
         point_embedding: torch.Tensor,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Process image and point embeddings through the Two-Way Transformer.
 
@@ -162,7 +163,7 @@ class TwoWayAttentionBlock(nn.Module):
         embedding_dim: int,
         num_heads: int,
         mlp_dim: int = 2048,
-        activation: Type[nn.Module] = nn.ReLU,
+        activation: type[nn.Module] = nn.ReLU,
         attention_downsample_rate: int = 2,
         skip_first_layer_pe: bool = False,
     ) -> None:
@@ -198,7 +199,7 @@ class TwoWayAttentionBlock(nn.Module):
 
     def forward(
         self, queries: torch.Tensor, keys: torch.Tensor, query_pe: torch.Tensor, key_pe: torch.Tensor
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Apply two-way attention to process query and key embeddings in a transformer block.
 

ultralytics/models/sam/modules/utils.py

@@ -1,12 +1,14 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
-from typing import Any, Dict, Tuple
+from __future__ import annotations
+
+from typing import Any
 
 import torch
 import torch.nn.functional as F
 
 
-def select_closest_cond_frames(frame_idx: int, cond_frame_outputs: Dict[int, Any], max_cond_frame_num: int):
+def select_closest_cond_frames(frame_idx: int, cond_frame_outputs: dict[int, Any], max_cond_frame_num: int):
     """
     Select the closest conditioning frames to a given frame index.
 
@@ -248,7 +250,7 @@ def window_partition(x: torch.Tensor, window_size: int):
     return windows, (Hp, Wp)
 
 
-def window_unpartition(windows: torch.Tensor, window_size: int, pad_hw: Tuple[int, int], hw: Tuple[int, int]):
+def window_unpartition(windows: torch.Tensor, window_size: int, pad_hw: tuple[int, int], hw: tuple[int, int]):
     """
     Unpartition windowed sequences into original sequences and remove padding.
 
@@ -333,8 +335,8 @@ def add_decomposed_rel_pos(
     q: torch.Tensor,
     rel_pos_h: torch.Tensor,
     rel_pos_w: torch.Tensor,
-    q_size: Tuple[int, int],
-    k_size: Tuple[int, int],
+    q_size: tuple[int, int],
+    k_size: tuple[int, int],
 ) -> torch.Tensor:
     """
     Add decomposed Relative Positional Embeddings to the attention map.

ultralytics/models/sam/predict.py

@@ -8,8 +8,10 @@ using SAM. It forms an integral part of the Ultralytics framework and is designe
 segmentation tasks.
 """
 
+from __future__ import annotations
+
 from collections import OrderedDict
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
 
 import cv2
 import numpy as np
@@ -1717,9 +1719,9 @@ class SAM2DynamicInteractivePredictor(SAM2Predictor):
     def __init__(
         self,
         cfg: Any = DEFAULT_CFG,
-        overrides: Optional[Dict[str, Any]] = None,
+        overrides: dict[str, Any] | None = None,
         max_obj_num: int = 3,
-        _callbacks: Optional[Dict[str, Any]] = None,
+        _callbacks: dict[str, Any] | None = None,
     ) -> None:
         """
         Initialize the predictor with configuration and optional overrides.
@@ -1759,14 +1761,14 @@ class SAM2DynamicInteractivePredictor(SAM2Predictor):
     @smart_inference_mode()
     def inference(
         self,
-        img: Union[torch.Tensor, np.ndarray],
-        bboxes: Optional[List[List[float]]] = None,
-        masks: Optional[Union[torch.Tensor, np.ndarray]] = None,
-        points: Optional[List[List[float]]] = None,
-        labels: Optional[List[int]] = None,
-        obj_ids: Optional[List[int]] = None,
+        img: torch.Tensor | np.ndarray,
+        bboxes: list[list[float]] | None = None,
+        masks: torch.Tensor | np.ndarray | None = None,
+        points: list[list[float]] | None = None,
+        labels: list[int] | None = None,
+        obj_ids: list[int] | None = None,
         update_memory: bool = False,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Perform inference on a single image with optional bounding boxes, masks, points and object IDs.
         It has two modes: one is to run inference on a single image without updating the memory,
@@ -1824,7 +1826,7 @@ class SAM2DynamicInteractivePredictor(SAM2Predictor):
         pred_scores = torch.clamp_(pred_scores / 32, min=0)
         return pred_masks.flatten(0, 1), pred_scores.flatten(0, 1)
 
-    def get_im_features(self, img: Union[torch.Tensor, np.ndarray]) -> None:
+    def get_im_features(self, img: torch.Tensor | np.ndarray) -> None:
         """
         Initialize the image state by processing the input image and extracting features.
 
@@ -1844,10 +1846,10 @@ class SAM2DynamicInteractivePredictor(SAM2Predictor):
     @smart_inference_mode()
     def update_memory(
         self,
-        obj_ids: List[int] = None,
-        points: Optional[torch.Tensor] = None,
-        labels: Optional[torch.Tensor] = None,
-        masks: Optional[torch.Tensor] = None,
+        obj_ids: list[int] = None,
+        points: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+        masks: torch.Tensor | None = None,
     ) -> None:
         """
         Append the imgState to the memory_bank and update the memory for the model.
@@ -1923,7 +1925,7 @@ class SAM2DynamicInteractivePredictor(SAM2Predictor):
         consolidated_out["maskmem_pos_enc"] = maskmem_pos_enc
         self.memory_bank.append(consolidated_out)
 
-    def _prepare_memory_conditioned_features(self, obj_idx: Optional[int]) -> torch.Tensor:
+    def _prepare_memory_conditioned_features(self, obj_idx: int | None) -> torch.Tensor:
         """
         Prepare the memory-conditioned features for the current image state. If obj_idx is provided, it supposes to
         prepare features for a specific prompted object in the image. If obj_idx is None, it prepares features for all
@@ -1958,7 +1960,7 @@ class SAM2DynamicInteractivePredictor(SAM2Predictor):
             *self.feat_sizes[-1],
         )
 
-    def get_maskmem_enc(self) -> Tuple[torch.Tensor, torch.Tensor]:
+    def get_maskmem_enc(self) -> tuple[torch.Tensor, torch.Tensor]:
         """Get the memory and positional encoding from the memory, which is used to condition the current image
         features.
         """
@@ -1973,7 +1975,7 @@ class SAM2DynamicInteractivePredictor(SAM2Predictor):
         memory_pos_embed = torch.cat(to_cat_memory_pos_embed, dim=0)
         return memory, memory_pos_embed
 
-    def _obj_id_to_idx(self, obj_id: int) -> Optional[int]:
+    def _obj_id_to_idx(self, obj_id: int) -> int | None:
         """
         Map client-side object id to model-side object index.
 
@@ -1987,11 +1989,11 @@ class SAM2DynamicInteractivePredictor(SAM2Predictor):
 
     def track_step(
         self,
-        obj_idx: Optional[int] = None,
-        point: Optional[torch.Tensor] = None,
-        label: Optional[torch.Tensor] = None,
-        mask: Optional[torch.Tensor] = None,
-    ) -> Dict[str, Any]:
+        obj_idx: int | None = None,
+        point: torch.Tensor | None = None,
+        label: torch.Tensor | None = None,
+        mask: torch.Tensor | None = None,
+    ) -> dict[str, Any]:
         """
         Tracking step for the current image state to predict masks.
 
@@ -2010,7 +2012,6 @@ class SAM2DynamicInteractivePredictor(SAM2Predictor):
             current_out (Dict[str, Any]): A dictionary containing the current output with mask predictions and object pointers.
                 Keys include 'point_inputs', 'mask_inputs', 'pred_masks', 'pred_masks_high_res', 'obj_ptr', 'object_score_logits'.
         """
-        current_out = {}
         if mask is not None and self.model.use_mask_input_as_output_without_sam:
             # When use_mask_input_as_output_without_sam=True, we directly output the mask input
             # (see it as a GT mask) without using a SAM prompt encoder + mask decoder.
@@ -2021,7 +2022,7 @@ class SAM2DynamicInteractivePredictor(SAM2Predictor):
             # fused the visual feature with previous memory features in the memory bank
             pix_feat_with_mem = self._prepare_memory_conditioned_features(obj_idx)
             # calculate the first feature if adding obj_idx exists(means adding prompts)
-            pix_feat_with_mem = pix_feat_with_mem[0:1] if obj_idx is not None else pix_feat_with_mem
+            pix_feat_with_mem = pix_feat_with_mem[:1] if obj_idx is not None else pix_feat_with_mem
             _, _, _, low_res_masks, high_res_masks, obj_ptr, object_score_logits = self.model._forward_sam_heads(
                 backbone_features=pix_feat_with_mem,
                 point_inputs={"point_coords": point, "point_labels": label} if obj_idx is not None else None,
@@ -2029,9 +2030,9 @@ class SAM2DynamicInteractivePredictor(SAM2Predictor):
                 multimask_output=False,
                 high_res_features=[feat[: pix_feat_with_mem.size(0)] for feat in self.high_res_features],
             )
-        current_out["pred_masks"] = low_res_masks
-        current_out["pred_masks_high_res"] = high_res_masks
-        current_out["obj_ptr"] = obj_ptr
-        current_out["object_score_logits"] = object_score_logits
-
-        return current_out
+        return {
+            "pred_masks": low_res_masks,
+            "pred_masks_high_res": high_res_masks,
+            "obj_ptr": obj_ptr,
+            "object_score_logits": object_score_logits,
+        }

ultralytics/models/utils/loss.py

@@ -1,6 +1,8 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
-from typing import Any, Dict, List, Optional, Tuple
+from __future__ import annotations
+
+from typing import Any
 
 import torch
 import torch.nn as nn
@@ -36,7 +38,7 @@ class DETRLoss(nn.Module):
     def __init__(
         self,
         nc: int = 80,
-        loss_gain: Optional[Dict[str, float]] = None,
+        loss_gain: dict[str, float] | None = None,
         aux_loss: bool = True,
         use_fl: bool = True,
         use_vfl: bool = False,
@@ -79,7 +81,7 @@ class DETRLoss(nn.Module):
 
     def _get_loss_class(
         self, pred_scores: torch.Tensor, targets: torch.Tensor, gt_scores: torch.Tensor, num_gts: int, postfix: str = ""
-    ) -> Dict[str, torch.Tensor]:
+    ) -> dict[str, torch.Tensor]:
         """
         Compute classification loss based on predictions, target values, and ground truth scores.
 
@@ -121,7 +123,7 @@ class DETRLoss(nn.Module):
 
     def _get_loss_bbox(
         self, pred_bboxes: torch.Tensor, gt_bboxes: torch.Tensor, postfix: str = ""
-    ) -> Dict[str, torch.Tensor]:
+    ) -> dict[str, torch.Tensor]:
         """
         Compute bounding box and GIoU losses for predicted and ground truth bounding boxes.
 
@@ -191,12 +193,12 @@ class DETRLoss(nn.Module):
         pred_scores: torch.Tensor,
         gt_bboxes: torch.Tensor,
         gt_cls: torch.Tensor,
-        gt_groups: List[int],
-        match_indices: Optional[List[Tuple]] = None,
+        gt_groups: list[int],
+        match_indices: list[tuple] | None = None,
         postfix: str = "",
-        masks: Optional[torch.Tensor] = None,
-        gt_mask: Optional[torch.Tensor] = None,
-    ) -> Dict[str, torch.Tensor]:
+        masks: torch.Tensor | None = None,
+        gt_mask: torch.Tensor | None = None,
+    ) -> dict[str, torch.Tensor]:
         """
         Get auxiliary losses for intermediate decoder layers.
 
@@ -258,7 +260,7 @@ class DETRLoss(nn.Module):
         return loss
 
     @staticmethod
-    def _get_index(match_indices: List[Tuple]) -> Tuple[Tuple[torch.Tensor, torch.Tensor], torch.Tensor]:
+    def _get_index(match_indices: list[tuple]) -> tuple[tuple[torch.Tensor, torch.Tensor], torch.Tensor]:
         """
         Extract batch indices, source indices, and destination indices from match indices.
 
@@ -275,8 +277,8 @@ class DETRLoss(nn.Module):
         return (batch_idx, src_idx), dst_idx
 
     def _get_assigned_bboxes(
-        self, pred_bboxes: torch.Tensor, gt_bboxes: torch.Tensor, match_indices: List[Tuple]
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        self, pred_bboxes: torch.Tensor, gt_bboxes: torch.Tensor, match_indices: list[tuple]
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Assign predicted bounding boxes to ground truth bounding boxes based on match indices.
 
@@ -309,12 +311,12 @@ class DETRLoss(nn.Module):
         pred_scores: torch.Tensor,
         gt_bboxes: torch.Tensor,
         gt_cls: torch.Tensor,
-        gt_groups: List[int],
-        masks: Optional[torch.Tensor] = None,
-        gt_mask: Optional[torch.Tensor] = None,
+        gt_groups: list[int],
+        masks: torch.Tensor | None = None,
+        gt_mask: torch.Tensor | None = None,
         postfix: str = "",
-        match_indices: Optional[List[Tuple]] = None,
-    ) -> Dict[str, torch.Tensor]:
+        match_indices: list[tuple] | None = None,
+    ) -> dict[str, torch.Tensor]:
         """
         Calculate losses for a single prediction layer.
 
@@ -358,10 +360,10 @@ class DETRLoss(nn.Module):
         self,
         pred_bboxes: torch.Tensor,
         pred_scores: torch.Tensor,
-        batch: Dict[str, Any],
+        batch: dict[str, Any],
         postfix: str = "",
         **kwargs: Any,
-    ) -> Dict[str, torch.Tensor]:
+    ) -> dict[str, torch.Tensor]:
         """
         Calculate loss for predicted bounding boxes and scores.
 
@@ -407,12 +409,12 @@ class RTDETRDetectionLoss(DETRLoss):
 
     def forward(
         self,
-        preds: Tuple[torch.Tensor, torch.Tensor],
-        batch: Dict[str, Any],
-        dn_bboxes: Optional[torch.Tensor] = None,
-        dn_scores: Optional[torch.Tensor] = None,
-        dn_meta: Optional[Dict[str, Any]] = None,
-    ) -> Dict[str, torch.Tensor]:
+        preds: tuple[torch.Tensor, torch.Tensor],
+        batch: dict[str, Any],
+        dn_bboxes: torch.Tensor | None = None,
+        dn_scores: torch.Tensor | None = None,
+        dn_meta: dict[str, Any] | None = None,
+    ) -> dict[str, torch.Tensor]:
         """
         Forward pass to compute detection loss with optional denoising loss.
 
@@ -448,8 +450,8 @@ class RTDETRDetectionLoss(DETRLoss):
 
     @staticmethod
     def get_dn_match_indices(
-        dn_pos_idx: List[torch.Tensor], dn_num_group: int, gt_groups: List[int]
-    ) -> List[Tuple[torch.Tensor, torch.Tensor]]:
+        dn_pos_idx: list[torch.Tensor], dn_num_group: int, gt_groups: list[int]
+    ) -> list[tuple[torch.Tensor, torch.Tensor]]:
         """
         Get match indices for denoising.
 

ultralytics/models/utils/ops.py

@@ -1,6 +1,8 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
-from typing import Any, Dict, List, Optional, Tuple
+from __future__ import annotations
+
+from typing import Any
 
 import torch
 import torch.nn as nn
@@ -47,7 +49,7 @@ class HungarianMatcher(nn.Module):
 
     def __init__(
         self,
-        cost_gain: Optional[Dict[str, float]] = None,
+        cost_gain: dict[str, float] | None = None,
         use_fl: bool = True,
         with_mask: bool = False,
         num_sample_points: int = 12544,
@@ -82,10 +84,10 @@ class HungarianMatcher(nn.Module):
         pred_scores: torch.Tensor,
         gt_bboxes: torch.Tensor,
         gt_cls: torch.Tensor,
-        gt_groups: List[int],
-        masks: Optional[torch.Tensor] = None,
-        gt_mask: Optional[List[torch.Tensor]] = None,
-    ) -> List[Tuple[torch.Tensor, torch.Tensor]]:
+        gt_groups: list[int],
+        masks: torch.Tensor | None = None,
+        gt_mask: list[torch.Tensor] | None = None,
+    ) -> list[tuple[torch.Tensor, torch.Tensor]]:
         """
         Compute optimal assignment between predictions and ground truth using Hungarian algorithm.
 
@@ -187,7 +189,7 @@ class HungarianMatcher(nn.Module):
 
 
 def get_cdn_group(
-    batch: Dict[str, Any],
+    batch: dict[str, Any],
     num_classes: int,
     num_queries: int,
     class_embed: torch.Tensor,
@@ -195,7 +197,7 @@ def get_cdn_group(
     cls_noise_ratio: float = 0.5,
     box_noise_scale: float = 1.0,
     training: bool = False,
-) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], Optional[Dict[str, Any]]]:
+) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, dict[str, Any] | None]:
     """
     Generate contrastive denoising training group with positive and negative samples from ground truths.
 

ultralytics/models/yolo/classify/train.py

@@ -1,7 +1,9 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
+from __future__ import annotations
+
 from copy import copy
-from typing import Any, Dict, Optional
+from typing import Any
 
 import torch
 
@@ -49,7 +51,7 @@ class ClassificationTrainer(BaseTrainer):
         >>> trainer.train()
     """
 
-    def __init__(self, cfg=DEFAULT_CFG, overrides: Optional[Dict[str, Any]] = None, _callbacks=None):
+    def __init__(self, cfg=DEFAULT_CFG, overrides: dict[str, Any] | None = None, _callbacks=None):
         """
         Initialize a ClassificationTrainer object.
 
@@ -162,10 +164,10 @@ class ClassificationTrainer(BaseTrainer):
                 self.model.transforms = loader.dataset.torch_transforms
         return loader
 
-    def preprocess_batch(self, batch: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    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)
+        batch["img"] = batch["img"].to(self.device, non_blocking=True)
+        batch["cls"] = batch["cls"].to(self.device, non_blocking=True)
         return batch
 
     def progress_string(self) -> str:
@@ -185,7 +187,7 @@ class ClassificationTrainer(BaseTrainer):
             self.test_loader, self.save_dir, args=copy(self.args), _callbacks=self.callbacks
         )
 
-    def label_loss_items(self, loss_items: Optional[torch.Tensor] = None, prefix: str = "train"):
+    def label_loss_items(self, loss_items: torch.Tensor | None = None, prefix: str = "train"):
         """
         Return a loss dict with labelled training loss items tensor.
 
@@ -220,7 +222,7 @@ class ClassificationTrainer(BaseTrainer):
                     self.metrics.pop("fitness", None)
                     self.run_callbacks("on_fit_epoch_end")
 
-    def plot_training_samples(self, batch: Dict[str, torch.Tensor], ni: int):
+    def plot_training_samples(self, batch: dict[str, torch.Tensor], ni: int):
         """
         Plot training samples with their annotations.
 

ultralytics/models/yolo/classify/val.py

@@ -1,7 +1,9 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
+from __future__ import annotations
+
 from pathlib import Path
-from typing import Any, Dict, List, Tuple, Union
+from typing import Any
 
 import torch
 
@@ -85,14 +87,14 @@ class ClassificationValidator(BaseValidator):
         self.targets = []
         self.confusion_matrix = ConfusionMatrix(names=model.names)
 
-    def preprocess(self, batch: Dict[str, Any]) -> Dict[str, Any]:
+    def preprocess(self, batch: dict[str, Any]) -> dict[str, Any]:
         """Preprocess input batch by moving data to device and converting to appropriate dtype."""
         batch["img"] = batch["img"].to(self.device, non_blocking=True)
         batch["img"] = batch["img"].half() if self.args.half else batch["img"].float()
-        batch["cls"] = batch["cls"].to(self.device)
+        batch["cls"] = batch["cls"].to(self.device, non_blocking=True)
         return batch
 
-    def update_metrics(self, preds: torch.Tensor, batch: Dict[str, Any]) -> None:
+    def update_metrics(self, preds: torch.Tensor, batch: dict[str, Any]) -> None:
         """
         Update running metrics with model predictions and batch targets.
 
@@ -131,11 +133,11 @@ class ClassificationValidator(BaseValidator):
         self.metrics.save_dir = self.save_dir
         self.metrics.confusion_matrix = self.confusion_matrix
 
-    def postprocess(self, preds: Union[torch.Tensor, List[torch.Tensor], Tuple[torch.Tensor]]) -> torch.Tensor:
+    def postprocess(self, preds: torch.Tensor | list[torch.Tensor] | tuple[torch.Tensor]) -> torch.Tensor:
         """Extract the primary prediction from model output if it's in a list or tuple format."""
         return preds[0] if isinstance(preds, (list, tuple)) else preds
 
-    def get_stats(self) -> Dict[str, float]:
+    def get_stats(self) -> dict[str, float]:
         """Calculate and return a dictionary of metrics by processing targets and predictions."""
         self.metrics.process(self.targets, self.pred)
         return self.metrics.results_dict
@@ -144,7 +146,7 @@ class ClassificationValidator(BaseValidator):
         """Create a ClassificationDataset instance for validation."""
         return ClassificationDataset(root=img_path, args=self.args, augment=False, prefix=self.args.split)
 
-    def get_dataloader(self, dataset_path: Union[Path, str], batch_size: int) -> torch.utils.data.DataLoader:
+    def get_dataloader(self, dataset_path: Path | str, batch_size: int) -> torch.utils.data.DataLoader:
         """
         Build and return a data loader for classification validation.
 
@@ -163,7 +165,7 @@ class ClassificationValidator(BaseValidator):
         pf = "%22s" + "%11.3g" * len(self.metrics.keys)  # print format
         LOGGER.info(pf % ("all", self.metrics.top1, self.metrics.top5))
 
-    def plot_val_samples(self, batch: Dict[str, Any], ni: int) -> None:
+    def plot_val_samples(self, batch: dict[str, Any], ni: int) -> None:
         """
         Plot validation image samples with their ground truth labels.
 
@@ -184,7 +186,7 @@ class ClassificationValidator(BaseValidator):
             on_plot=self.on_plot,
         )
 
-    def plot_predictions(self, batch: Dict[str, Any], preds: torch.Tensor, ni: int) -> None:
+    def plot_predictions(self, batch: dict[str, Any], preds: torch.Tensor, ni: int) -> None:
         """
         Plot images with their predicted class labels and save the visualization.
 

ultralytics/models/yolo/detect/predict.py

@@ -85,7 +85,7 @@ class DetectionPredictor(BasePredictor):
         """Extract object features from the feature maps."""
         import torch
 
-        s = min([x.shape[1] for x in feat_maps])  # find smallest vector length
+        s = min(x.shape[1] for x in feat_maps)  # find shortest vector length
         obj_feats = torch.cat(
             [x.permute(0, 2, 3, 1).reshape(x.shape[0], -1, s, x.shape[1] // s).mean(dim=-1) for x in feat_maps], dim=1
         )  # mean reduce all vectors to same length