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

ultralytics/models/utils/loss.py

@@ -1,5 +1,9 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
+from __future__ import annotations
+
+from typing import Any
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -11,15 +15,14 @@ from .ops import HungarianMatcher
 
 
 class DETRLoss(nn.Module):
-    """
-    DETR (DEtection TRansformer) Loss class for calculating various loss components.
+    """DETR (DEtection TRansformer) Loss class for calculating various loss components.
 
-    This class computes classification loss, bounding box loss, GIoU loss, and optionally auxiliary losses for the
-    DETR object detection model.
+    This class computes classification loss, bounding box loss, GIoU loss, and optionally auxiliary losses for the DETR
+    object detection model.
 
     Attributes:
         nc (int): Number of classes.
-        loss_gain (dict): Coefficients for different loss components.
+        loss_gain (dict[str, float]): Coefficients for different loss components.
         aux_loss (bool): Whether to compute auxiliary losses.
         use_fl (bool): Whether to use FocalLoss.
         use_vfl (bool): Whether to use VarifocalLoss.
@@ -33,32 +36,31 @@ class DETRLoss(nn.Module):
 
     def __init__(
         self,
-        nc=80,
-        loss_gain=None,
-        aux_loss=True,
-        use_fl=True,
-        use_vfl=False,
-        use_uni_match=False,
-        uni_match_ind=0,
-        gamma=1.5,
-        alpha=0.25,
+        nc: int = 80,
+        loss_gain: dict[str, float] | None = None,
+        aux_loss: bool = True,
+        use_fl: bool = True,
+        use_vfl: bool = False,
+        use_uni_match: bool = False,
+        uni_match_ind: int = 0,
+        gamma: float = 1.5,
+        alpha: float = 0.25,
     ):
-        """
-        Initialize DETR loss function with customizable components and gains.
+        """Initialize DETR loss function with customizable components and gains.
 
         Uses default loss_gain if not provided. Initializes HungarianMatcher with preset cost gains. Supports auxiliary
         losses and various loss types.
 
         Args:
             nc (int): Number of classes.
-            loss_gain (dict): Coefficients for different loss components.
+            loss_gain (dict[str, float], optional): Coefficients for different loss components.
             aux_loss (bool): Whether to use auxiliary losses from each decoder layer.
             use_fl (bool): Whether to use FocalLoss.
             use_vfl (bool): Whether to use VarifocalLoss.
             use_uni_match (bool): Whether to use fixed layer for auxiliary branch label assignment.
             uni_match_ind (int): Index of fixed layer for uni_match.
             gamma (float): The focusing parameter that controls how much the loss focuses on hard-to-classify examples.
-            alpha (float | list): The balancing factor used to address class imbalance.
+            alpha (float): The balancing factor used to address class imbalance.
         """
         super().__init__()
 
@@ -75,19 +77,20 @@ class DETRLoss(nn.Module):
         self.uni_match_ind = uni_match_ind
         self.device = None
 
-    def _get_loss_class(self, pred_scores, targets, gt_scores, num_gts, postfix=""):
-        """
-        Compute classification loss based on predictions, target values, and ground truth scores.
+    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]:
+        """Compute classification loss based on predictions, target values, and ground truth scores.
 
         Args:
-            pred_scores (torch.Tensor): Predicted class scores with shape (batch_size, num_queries, num_classes).
-            targets (torch.Tensor): Target class indices with shape (batch_size, num_queries).
-            gt_scores (torch.Tensor): Ground truth confidence scores with shape (batch_size, num_queries).
+            pred_scores (torch.Tensor): Predicted class scores with shape (B, N, C).
+            targets (torch.Tensor): Target class indices with shape (B, N).
+            gt_scores (torch.Tensor): Ground truth confidence scores with shape (B, N).
             num_gts (int): Number of ground truth objects.
             postfix (str, optional): String to append to the loss name for identification in multi-loss scenarios.
 
         Returns:
-            loss_cls (torch.Tensor): Classification loss value.
+            (dict[str, torch.Tensor]): Dictionary containing classification loss value.
 
         Notes:
             The function supports different classification loss types:
@@ -115,22 +118,20 @@ class DETRLoss(nn.Module):
 
         return {name_class: loss_cls.squeeze() * self.loss_gain["class"]}
 
-    def _get_loss_bbox(self, pred_bboxes, gt_bboxes, postfix=""):
-        """
-        Compute bounding box and GIoU losses for predicted and ground truth bounding boxes.
+    def _get_loss_bbox(
+        self, pred_bboxes: torch.Tensor, gt_bboxes: torch.Tensor, postfix: str = ""
+    ) -> dict[str, torch.Tensor]:
+        """Compute bounding box and GIoU losses for predicted and ground truth bounding boxes.
 
         Args:
-            pred_bboxes (torch.Tensor): Predicted bounding boxes with shape (batch_size, num_queries, 4).
-            gt_bboxes (torch.Tensor): Ground truth bounding boxes with shape (N, 4), where N is the total
-                number of ground truth boxes.
-            postfix (str): String to append to the loss names for identification in multi-loss scenarios.
+            pred_bboxes (torch.Tensor): Predicted bounding boxes with shape (N, 4).
+            gt_bboxes (torch.Tensor): Ground truth bounding boxes with shape (N, 4).
+            postfix (str, optional): String to append to the loss names for identification in multi-loss scenarios.
 
         Returns:
-            loss (dict): Dictionary containing:
-                - loss_bbox{postfix} (torch.Tensor): L1 loss between predicted and ground truth boxes,
-                  scaled by the bbox loss gain.
-                - loss_giou{postfix} (torch.Tensor): GIoU loss between predicted and ground truth boxes,
-                  scaled by the giou loss gain.
+            (dict[str, torch.Tensor]): Dictionary containing:
+                - loss_bbox{postfix}: L1 loss between predicted and ground truth boxes, scaled by the bbox loss gain.
+                - loss_giou{postfix}: GIoU loss between predicted and ground truth boxes, scaled by the giou loss gain.
 
         Notes:
             If no ground truth boxes are provided (empty list), zero-valued tensors are returned for both losses.
@@ -184,32 +185,31 @@ class DETRLoss(nn.Module):
 
     def _get_loss_aux(
         self,
-        pred_bboxes,
-        pred_scores,
-        gt_bboxes,
-        gt_cls,
-        gt_groups,
-        match_indices=None,
-        postfix="",
-        masks=None,
-        gt_mask=None,
-    ):
-        """
-        Get auxiliary losses for intermediate decoder layers.
+        pred_bboxes: torch.Tensor,
+        pred_scores: torch.Tensor,
+        gt_bboxes: torch.Tensor,
+        gt_cls: torch.Tensor,
+        gt_groups: list[int],
+        match_indices: list[tuple] | None = None,
+        postfix: str = "",
+        masks: torch.Tensor | None = None,
+        gt_mask: torch.Tensor | None = None,
+    ) -> dict[str, torch.Tensor]:
+        """Get auxiliary losses for intermediate decoder layers.
 
         Args:
             pred_bboxes (torch.Tensor): Predicted bounding boxes from auxiliary layers.
             pred_scores (torch.Tensor): Predicted scores from auxiliary layers.
             gt_bboxes (torch.Tensor): Ground truth bounding boxes.
             gt_cls (torch.Tensor): Ground truth classes.
-            gt_groups (List[int]): Number of ground truths per image.
-            match_indices (List[tuple], optional): Pre-computed matching indices.
-            postfix (str): String to append to loss names.
+            gt_groups (list[int]): Number of ground truths per image.
+            match_indices (list[tuple], optional): Pre-computed matching indices.
+            postfix (str, optional): String to append to loss names.
             masks (torch.Tensor, optional): Predicted masks if using segmentation.
             gt_mask (torch.Tensor, optional): Ground truth masks if using segmentation.
 
         Returns:
-            (dict): Dictionary of auxiliary losses.
+            (dict[str, torch.Tensor]): Dictionary of auxiliary losses.
         """
         # NOTE: loss class, bbox, giou, mask, dice
         loss = torch.zeros(5 if masks is not None else 3, device=pred_bboxes.device)
@@ -255,32 +255,34 @@ class DETRLoss(nn.Module):
         return loss
 
     @staticmethod
-    def _get_index(match_indices):
-        """
-        Extract batch indices, source indices, and destination indices from match indices.
+    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.
 
         Args:
-            match_indices (List[tuple]): List of tuples containing matched indices.
+            match_indices (list[tuple]): List of tuples containing matched indices.
 
         Returns:
-            (tuple): Tuple containing (batch_idx, src_idx) and dst_idx.
+            batch_idx (tuple[torch.Tensor, torch.Tensor]): Tuple containing (batch_idx, src_idx).
+            dst_idx (torch.Tensor): Destination indices.
         """
         batch_idx = torch.cat([torch.full_like(src, i) for i, (src, _) in enumerate(match_indices)])
         src_idx = torch.cat([src for (src, _) in match_indices])
         dst_idx = torch.cat([dst for (_, dst) in match_indices])
         return (batch_idx, src_idx), dst_idx
 
-    def _get_assigned_bboxes(self, pred_bboxes, gt_bboxes, match_indices):
-        """
-        Assign predicted bounding boxes to ground truth bounding boxes based on match indices.
+    def _get_assigned_bboxes(
+        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.
 
         Args:
             pred_bboxes (torch.Tensor): Predicted bounding boxes.
             gt_bboxes (torch.Tensor): Ground truth bounding boxes.
-            match_indices (List[tuple]): List of tuples containing matched indices.
+            match_indices (list[tuple]): List of tuples containing matched indices.
 
         Returns:
-            (tuple): Tuple containing assigned predictions and ground truths.
+            pred_assigned (torch.Tensor): Assigned predicted bounding boxes.
+            gt_assigned (torch.Tensor): Assigned ground truth bounding boxes.
         """
         pred_assigned = torch.cat(
             [
@@ -298,32 +300,31 @@ class DETRLoss(nn.Module):
 
     def _get_loss(
         self,
-        pred_bboxes,
-        pred_scores,
-        gt_bboxes,
-        gt_cls,
-        gt_groups,
-        masks=None,
-        gt_mask=None,
-        postfix="",
-        match_indices=None,
-    ):
-        """
-        Calculate losses for a single prediction layer.
+        pred_bboxes: torch.Tensor,
+        pred_scores: torch.Tensor,
+        gt_bboxes: torch.Tensor,
+        gt_cls: torch.Tensor,
+        gt_groups: list[int],
+        masks: torch.Tensor | None = None,
+        gt_mask: torch.Tensor | None = None,
+        postfix: str = "",
+        match_indices: list[tuple] | None = None,
+    ) -> dict[str, torch.Tensor]:
+        """Calculate losses for a single prediction layer.
 
         Args:
             pred_bboxes (torch.Tensor): Predicted bounding boxes.
             pred_scores (torch.Tensor): Predicted class scores.
             gt_bboxes (torch.Tensor): Ground truth bounding boxes.
             gt_cls (torch.Tensor): Ground truth classes.
-            gt_groups (List[int]): Number of ground truths per image.
+            gt_groups (list[int]): Number of ground truths per image.
             masks (torch.Tensor, optional): Predicted masks if using segmentation.
             gt_mask (torch.Tensor, optional): Ground truth masks if using segmentation.
-            postfix (str): String to append to loss names.
-            match_indices (List[tuple], optional): Pre-computed matching indices.
+            postfix (str, optional): String to append to loss names.
+            match_indices (list[tuple], optional): Pre-computed matching indices.
 
         Returns:
-            (dict): Dictionary of losses.
+            (dict[str, torch.Tensor]): Dictionary of losses.
         """
         if match_indices is None:
             match_indices = self.matcher(
@@ -347,22 +348,25 @@ class DETRLoss(nn.Module):
             # **(self._get_loss_mask(masks, gt_mask, match_indices, postfix) if masks is not None and gt_mask is not None else {})
         }
 
-    def forward(self, pred_bboxes, pred_scores, batch, postfix="", **kwargs):
-        """
-        Calculate loss for predicted bounding boxes and scores.
+    def forward(
+        self,
+        pred_bboxes: torch.Tensor,
+        pred_scores: torch.Tensor,
+        batch: dict[str, Any],
+        postfix: str = "",
+        **kwargs: Any,
+    ) -> dict[str, torch.Tensor]:
+        """Calculate loss for predicted bounding boxes and scores.
 
         Args:
-            pred_bboxes (torch.Tensor): Predicted bounding boxes, shape [l, b, query, 4].
-            pred_scores (torch.Tensor): Predicted class scores, shape [l, b, query, num_classes].
-            batch (dict): Batch information containing:
-                cls (torch.Tensor): Ground truth classes, shape [num_gts].
-                bboxes (torch.Tensor): Ground truth bounding boxes, shape [num_gts, 4].
-                gt_groups (List[int]): Number of ground truths for each image in the batch.
-            postfix (str): Postfix for loss names.
+            pred_bboxes (torch.Tensor): Predicted bounding boxes, shape (L, B, N, 4).
+            pred_scores (torch.Tensor): Predicted class scores, shape (L, B, N, C).
+            batch (dict[str, Any]): Batch information containing cls, bboxes, and gt_groups.
+            postfix (str, optional): Postfix for loss names.
             **kwargs (Any): Additional arguments, may include 'match_indices'.
 
         Returns:
-            (dict): Computed losses, including main and auxiliary (if enabled).
+            (dict[str, torch.Tensor]): Computed losses, including main and auxiliary (if enabled).
 
         Notes:
             Uses last elements of pred_bboxes and pred_scores for main loss, and the rest for auxiliary losses if
@@ -387,26 +391,31 @@ class DETRLoss(nn.Module):
 
 
 class RTDETRDetectionLoss(DETRLoss):
-    """
-    Real-Time DeepTracker (RT-DETR) Detection Loss class that extends the DETRLoss.
+    """Real-Time DeepTracker (RT-DETR) Detection Loss class that extends the DETRLoss.
 
     This class computes the detection loss for the RT-DETR model, which includes the standard detection loss as well as
     an additional denoising training loss when provided with denoising metadata.
     """
 
-    def forward(self, preds, batch, dn_bboxes=None, dn_scores=None, dn_meta=None):
-        """
-        Forward pass to compute detection loss with optional denoising loss.
+    def forward(
+        self,
+        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.
 
         Args:
-            preds (tuple): Tuple containing predicted bounding boxes and scores.
-            batch (dict): Batch data containing ground truth information.
+            preds (tuple[torch.Tensor, torch.Tensor]): Tuple containing predicted bounding boxes and scores.
+            batch (dict[str, Any]): Batch data containing ground truth information.
             dn_bboxes (torch.Tensor, optional): Denoising bounding boxes.
             dn_scores (torch.Tensor, optional): Denoising scores.
-            dn_meta (dict, optional): Metadata for denoising.
+            dn_meta (dict[str, Any], optional): Metadata for denoising.
 
         Returns:
-            (dict): Dictionary containing total loss and denoising loss if applicable.
+            (dict[str, torch.Tensor]): Dictionary containing total loss and denoising loss if applicable.
         """
         pred_bboxes, pred_scores = preds
         total_loss = super().forward(pred_bboxes, pred_scores, batch)
@@ -429,17 +438,18 @@ class RTDETRDetectionLoss(DETRLoss):
         return total_loss
 
     @staticmethod
-    def get_dn_match_indices(dn_pos_idx, dn_num_group, gt_groups):
-        """
-        Get match indices for denoising.
+    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]]:
+        """Get match indices for denoising.
 
         Args:
-            dn_pos_idx (List[torch.Tensor]): List of tensors containing positive indices for denoising.
+            dn_pos_idx (list[torch.Tensor]): List of tensors containing positive indices for denoising.
             dn_num_group (int): Number of denoising groups.
-            gt_groups (List[int]): List of integers representing number of ground truths per image.
+            gt_groups (list[int]): List of integers representing number of ground truths per image.
 
         Returns:
-            (List[tuple]): List of tuples containing matched indices for denoising.
+            (list[tuple[torch.Tensor, torch.Tensor]]): List of tuples containing matched indices for denoising.
         """
         dn_match_indices = []
         idx_groups = torch.as_tensor([0, *gt_groups[:-1]]).cumsum_(0)
@@ -447,8 +457,9 @@ class RTDETRDetectionLoss(DETRLoss):
             if num_gt > 0:
                 gt_idx = torch.arange(end=num_gt, dtype=torch.long) + idx_groups[i]
                 gt_idx = gt_idx.repeat(dn_num_group)
-                assert len(dn_pos_idx[i]) == len(gt_idx), "Expected the same length, "
-                f"but got {len(dn_pos_idx[i])} and {len(gt_idx)} respectively."
+                assert len(dn_pos_idx[i]) == len(gt_idx), (
+                    f"Expected the same length, but got {len(dn_pos_idx[i])} and {len(gt_idx)} respectively."
+                )
                 dn_match_indices.append((dn_pos_idx[i], gt_idx))
             else:
                 dn_match_indices.append((torch.zeros([0], dtype=torch.long), torch.zeros([0], dtype=torch.long)))