ultralytics 8.3.189__py3-none-any.whl → 8.3.191__py3-none-any.whl

ultralytics/models/sam/modules/memory_attention.py

@@ -1,7 +1,8 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
+from __future__ import annotations
+
 import copy
-from typing import Optional
 
 import torch
 from torch import nn
@@ -103,7 +104,7 @@ class MemoryAttentionLayer(nn.Module):
         self.pos_enc_at_cross_attn_queries = pos_enc_at_cross_attn_queries
         self.pos_enc_at_cross_attn_keys = pos_enc_at_cross_attn_keys
 
-    def _forward_sa(self, tgt: torch.Tensor, query_pos: Optional[torch.Tensor]) -> torch.Tensor:
+    def _forward_sa(self, tgt: torch.Tensor, query_pos: torch.Tensor | None) -> torch.Tensor:
         """Perform self-attention on input tensor using positional encoding and RoPE attention mechanism."""
         tgt2 = self.norm1(tgt)
         q = k = tgt2 + query_pos if self.pos_enc_at_attn else tgt2
@@ -115,8 +116,8 @@ class MemoryAttentionLayer(nn.Module):
         self,
         tgt: torch.Tensor,
         memory: torch.Tensor,
-        query_pos: Optional[torch.Tensor],
-        pos: Optional[torch.Tensor],
+        query_pos: torch.Tensor | None,
+        pos: torch.Tensor | None,
         num_k_exclude_rope: int = 0,
     ) -> torch.Tensor:
         """Perform cross-attention between target and memory tensors using RoPEAttention mechanism."""
@@ -140,8 +141,8 @@ class MemoryAttentionLayer(nn.Module):
         self,
         tgt: torch.Tensor,
         memory: torch.Tensor,
-        pos: Optional[torch.Tensor] = None,
-        query_pos: Optional[torch.Tensor] = None,
+        pos: torch.Tensor | None = None,
+        query_pos: torch.Tensor | None = None,
         num_k_exclude_rope: int = 0,
     ) -> torch.Tensor:
         """
@@ -242,8 +243,8 @@ class MemoryAttention(nn.Module):
         self,
         curr: torch.Tensor,  # self-attention inputs
         memory: torch.Tensor,  # cross-attention inputs
-        curr_pos: Optional[torch.Tensor] = None,  # pos_enc for self-attention inputs
-        memory_pos: Optional[torch.Tensor] = None,  # pos_enc for cross-attention inputs
+        curr_pos: torch.Tensor | None = None,  # pos_enc for self-attention inputs
+        memory_pos: torch.Tensor | None = None,  # pos_enc for cross-attention inputs
         num_obj_ptr_tokens: int = 0,  # number of object pointer *tokens*
     ) -> torch.Tensor:
         """

ultralytics/models/sam/modules/sam.py

@@ -3,10 +3,7 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 
-# This source code is licensed under the license found in the
-# LICENSE file in the root directory of this source tree.
-
-from typing import List
+from __future__ import annotations
 
 import torch
 import torch.nn.functional as F
@@ -61,8 +58,8 @@ class SAMModel(nn.Module):
         image_encoder: ImageEncoderViT,
         prompt_encoder: PromptEncoder,
         mask_decoder: MaskDecoder,
-        pixel_mean: List[float] = (123.675, 116.28, 103.53),
-        pixel_std: List[float] = (58.395, 57.12, 57.375),
+        pixel_mean: list[float] = (123.675, 116.28, 103.53),
+        pixel_std: list[float] = (58.395, 57.12, 57.375),
     ) -> None:
         """
         Initialize the SAMModel class to predict object masks from an image and input prompts.
@@ -959,7 +956,6 @@ class SAM2Model(torch.nn.Module):
         prev_sam_mask_logits=None,
     ):
         """Perform a single tracking step, updating object masks and memory features based on current frame inputs."""
-        current_out = {}
         sam_outputs, _, _ = self._track_step(
             frame_idx,
             is_init_cond_frame,
@@ -975,9 +971,11 @@ class SAM2Model(torch.nn.Module):
         )
         _, _, _, low_res_masks, high_res_masks, obj_ptr, object_score_logits = sam_outputs
 
-        current_out["pred_masks"] = low_res_masks
-        current_out["pred_masks_high_res"] = high_res_masks
-        current_out["obj_ptr"] = obj_ptr
+        current_out = {
+            "pred_masks": low_res_masks,
+            "pred_masks_high_res": high_res_masks,
+            "obj_ptr": obj_ptr,
+        }
         if not self.training:
             # Only add this in inference (to avoid unused param in activation checkpointing;
             # it's mainly used in the demo to encode spatial memories w/ consolidated masks)

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.