# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license

from __future__ import annotations

import math
from typing import Any

import torch
import torch.nn as nn
import torch.nn.functional as F

from ultralytics.utils import DEFAULT_CFG
from ultralytics.utils.metrics import OKS_SIGMA, RLE_WEIGHT
from ultralytics.utils.ops import crop_mask, xywh2xyxy, xyxy2xywh
from ultralytics.utils.plotting_3d import (
    decode_cut_partial_side_edge_from_gt,
    decode_visible_face_edge_from_gt,
    select_gt_visible_faces,
)
from ultralytics.utils.tal import RotatedTaskAlignedAssigner, TaskAlignedAssigner, dist2bbox, dist2rbox, make_anchors
from ultralytics.utils.torch_utils import autocast

from .metrics import bbox_iou, probiou
from .tal import bbox2dist, rbox2dist


def _normalize_edge_depth_targets_to_model_space(depths: torch.Tensor, depth_scale: torch.Tensor | float) -> torch.Tensor:
    """Convert metric edge-depth targets back into the canonical ROI depth space used by the model."""
    if not isinstance(depths, torch.Tensor):
        depths = torch.as_tensor(depths)
    scale = depth_scale if isinstance(depth_scale, torch.Tensor) else torch.as_tensor(depth_scale, device=depths.device)
    scale = scale.to(device=depths.device, dtype=depths.dtype)
    if scale.numel() != 1 or not torch.isfinite(scale).all() or float(scale.abs().item()) <= 1e-12:
        return depths
    return depths / scale


class VarifocalLoss(nn.Module):
    r"""Varifocal loss by Zhang et al.

    Implements the Varifocal Loss function for addressing class imbalance in object detection by focusing on
    hard-to-classify examples and balancing positive/negative samples.

    .. math::
        \text{VFL}(p, q, y) = \begin{cases}
            -q \cdot (q \log(p) + (1-q) \log(1-p)), & y = 1 \\
            -\alpha \cdot p^\gamma \cdot \log(1-p), & y = 0
        \end{cases}

    Attributes:
        gamma (float): The focusing parameter that controls how much the loss focuses on hard-to-classify examples.
        alpha (float): The balancing factor used to address class imbalance.

    References:
        https://arxiv.org/abs/2008.13367
    """

    def __init__(self, gamma: float = 2.0, alpha: float = 0.75):
        r"""Initialize the VarifocalLoss class with focusing and balancing parameters.

        Args:
            gamma (float): Focusing parameter controlling emphasis on hard examples. Default: ``2.0``
            alpha (float): Balancing factor for class imbalance. Default: ``0.75``
        """
        super().__init__()
        self.gamma = gamma
        self.alpha = alpha

    def forward(self, pred_score: torch.Tensor, gt_score: torch.Tensor, label: torch.Tensor) -> torch.Tensor:
        r"""Compute varifocal loss between predictions and ground truth.

        Args:
            pred_score (torch.Tensor): Predicted scores (logits) of shape :math:`(B, N, C)`.
            gt_score (torch.Tensor): Ground truth quality scores of shape :math:`(B, N, C)`.
            label (torch.Tensor): Binary labels indicating positive/negative samples of shape :math:`(B, N, C)`.

        Returns:
            (torch.Tensor): Scalar varifocal loss value.
        """
        weight = self.alpha * pred_score.sigmoid().pow(self.gamma) * (1 - label) + gt_score * label
        with autocast(enabled=False):
            loss = (
                (F.binary_cross_entropy_with_logits(pred_score.float(), gt_score.float(), reduction="none") * weight)
                .mean(1)
                .sum()
            )
        return loss


class FocalLoss(nn.Module):
    """Wraps focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5).

    Implements the Focal Loss function for addressing class imbalance by down-weighting easy examples and focusing on
    hard negatives during training.

    Attributes:
        gamma (float): The focusing parameter that controls how much the loss focuses on hard-to-classify examples.
        alpha (torch.Tensor): The balancing factor used to address class imbalance.
    """

    def __init__(self, gamma: float = 1.5, alpha: float = 0.25):
        """Initialize FocalLoss class with focusing and balancing parameters."""
        super().__init__()
        self.gamma = gamma
        self.alpha = torch.tensor(alpha)

    def forward(self, pred: torch.Tensor, label: torch.Tensor) -> torch.Tensor:
        """Calculate focal loss with modulating factors for class imbalance."""
        loss = F.binary_cross_entropy_with_logits(pred, label, reduction="none")
        # p_t = torch.exp(-loss)
        # loss *= self.alpha * (1.000001 - p_t) ** self.gamma  # non-zero power for gradient stability

        # TF implementation https://github.com/tensorflow/addons/blob/v0.7.1/tensorflow_addons/losses/focal_loss.py
        pred_prob = pred.sigmoid()  # prob from logits
        p_t = label * pred_prob + (1 - label) * (1 - pred_prob)
        modulating_factor = (1.0 - p_t) ** self.gamma
        loss *= modulating_factor
        if (self.alpha > 0).any():
            self.alpha = self.alpha.to(device=pred.device, dtype=pred.dtype)
            alpha_factor = label * self.alpha + (1 - label) * (1 - self.alpha)
            loss *= alpha_factor
        return loss.mean(1).sum()


class DFLoss(nn.Module):
    """Criterion class for computing Distribution Focal Loss (DFL)."""

    def __init__(self, reg_max: int = 16) -> None:
        """Initialize the DFL module with regularization maximum."""
        super().__init__()
        self.reg_max = reg_max

    def __call__(self, pred_dist: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        """Return sum of left and right DFL losses from https://ieeexplore.ieee.org/document/9792391."""
        target = target.clamp_(0, self.reg_max - 1 - 0.01)
        tl = target.long()  # target left
        tr = tl + 1  # target right
        wl = tr - target  # weight left
        wr = 1 - wl  # weight right
        return (
            F.cross_entropy(pred_dist, tl.view(-1), reduction="none").view(tl.shape) * wl
            + F.cross_entropy(pred_dist, tr.view(-1), reduction="none").view(tl.shape) * wr
        ).mean(-1, keepdim=True)


class BboxLoss(nn.Module):
    """Criterion class for computing training losses for bounding boxes."""

    def __init__(self, reg_max: int = 16):
        """Initialize the BboxLoss module with regularization maximum and DFL settings."""
        super().__init__()
        self.dfl_loss = DFLoss(reg_max) if reg_max > 1 else None

    def forward(
        self,
        pred_dist: torch.Tensor,
        pred_bboxes: torch.Tensor,
        anchor_points: torch.Tensor,
        target_bboxes: torch.Tensor,
        target_scores: torch.Tensor,
        target_scores_sum: torch.Tensor,
        fg_mask: torch.Tensor,
        imgsz: torch.Tensor,
        stride: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Compute IoU and DFL losses for bounding boxes."""
        weight = target_scores.sum(-1)[fg_mask].unsqueeze(-1)
        iou = bbox_iou(pred_bboxes[fg_mask], target_bboxes[fg_mask], xywh=False, CIoU=True)
        loss_iou = ((1.0 - iou) * weight).sum() / target_scores_sum

        # DFL loss
        if self.dfl_loss:
            target_ltrb = bbox2dist(anchor_points, target_bboxes, self.dfl_loss.reg_max - 1)
            loss_dfl = self.dfl_loss(pred_dist[fg_mask].view(-1, self.dfl_loss.reg_max), target_ltrb[fg_mask]) * weight
            loss_dfl = loss_dfl.sum() / target_scores_sum
        else:
            target_ltrb = bbox2dist(anchor_points, target_bboxes)
            # normalize ltrb by image size
            target_ltrb = target_ltrb * stride
            target_ltrb[..., 0::2] /= imgsz[1]
            target_ltrb[..., 1::2] /= imgsz[0]
            pred_dist = pred_dist * stride
            pred_dist[..., 0::2] /= imgsz[1]
            pred_dist[..., 1::2] /= imgsz[0]
            loss_dfl = (
                F.l1_loss(pred_dist[fg_mask], target_ltrb[fg_mask], reduction="none").mean(-1, keepdim=True) * weight
            )
            loss_dfl = loss_dfl.sum() / target_scores_sum

        return loss_iou, loss_dfl


class RLELoss(nn.Module):
    """Residual Log-Likelihood Estimation Loss.

    Attributes:
        size_average (bool): Option to average the loss by the batch_size.
        use_target_weight (bool): Option to use weighted loss.
        residual (bool): Option to add L1 loss and let the flow learn the residual error distribution.

    References:
        https://arxiv.org/abs/2107.11291
        https://github.com/open-mmlab/mmpose/blob/main/mmpose/models/losses/regression_loss.py
    """

    def __init__(self, use_target_weight: bool = True, size_average: bool = True, residual: bool = True):
        """Initialize RLELoss with target weight and residual options.

        Args:
            use_target_weight (bool): Whether to use target weights for loss calculation.
            size_average (bool): Whether to average the loss over elements.
            residual (bool): Whether to include residual log-likelihood term.
        """
        super().__init__()
        self.size_average = size_average
        self.use_target_weight = use_target_weight
        self.residual = residual

    def forward(
        self, sigma: torch.Tensor, log_phi: torch.Tensor, error: torch.Tensor, target_weight: torch.Tensor = None
    ) -> torch.Tensor:
        """
        Args:
            sigma (torch.Tensor): Output sigma, shape (N, D).
            log_phi (torch.Tensor): Output log_phi, shape (N).
            error (torch.Tensor): Error, shape (N, D).
            target_weight (torch.Tensor): Weights across different joint types, shape (N).
        """
        log_sigma = torch.log(sigma)
        loss = log_sigma - log_phi.unsqueeze(1)

        if self.residual:
            loss += torch.log(sigma * 2) + torch.abs(error)

        if self.use_target_weight:
            assert target_weight is not None, "'target_weight' should not be None when 'use_target_weight' is True."
            if target_weight.dim() == 1:
                target_weight = target_weight.unsqueeze(1)
            loss *= target_weight

        if self.size_average:
            loss /= len(loss)

        return loss.sum()


class RotatedBboxLoss(BboxLoss):
    """Criterion class for computing training losses for rotated bounding boxes."""

    def __init__(self, reg_max: int):
        """Initialize the RotatedBboxLoss module with regularization maximum and DFL settings."""
        super().__init__(reg_max)

    def forward(
        self,
        pred_dist: torch.Tensor,
        pred_bboxes: torch.Tensor,
        anchor_points: torch.Tensor,
        target_bboxes: torch.Tensor,
        target_scores: torch.Tensor,
        target_scores_sum: torch.Tensor,
        fg_mask: torch.Tensor,
        imgsz: torch.Tensor,
        stride: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Compute IoU and DFL losses for rotated bounding boxes."""
        weight = target_scores.sum(-1)[fg_mask].unsqueeze(-1)
        iou = probiou(pred_bboxes[fg_mask], target_bboxes[fg_mask])
        loss_iou = ((1.0 - iou) * weight).sum() / target_scores_sum

        # DFL loss
        if self.dfl_loss:
            target_ltrb = rbox2dist(
                target_bboxes[..., :4], anchor_points, target_bboxes[..., 4:5], reg_max=self.dfl_loss.reg_max - 1
            )
            loss_dfl = self.dfl_loss(pred_dist[fg_mask].view(-1, self.dfl_loss.reg_max), target_ltrb[fg_mask]) * weight
            loss_dfl = loss_dfl.sum() / target_scores_sum
        else:
            target_ltrb = rbox2dist(target_bboxes[..., :4], anchor_points, target_bboxes[..., 4:5])
            target_ltrb = target_ltrb * stride
            target_ltrb[..., 0::2] /= imgsz[1]
            target_ltrb[..., 1::2] /= imgsz[0]
            pred_dist = pred_dist * stride
            pred_dist[..., 0::2] /= imgsz[1]
            pred_dist[..., 1::2] /= imgsz[0]
            loss_dfl = (
                F.l1_loss(pred_dist[fg_mask], target_ltrb[fg_mask], reduction="none").mean(-1, keepdim=True) * weight
            )
            loss_dfl = loss_dfl.sum() / target_scores_sum

        return loss_iou, loss_dfl


class MultiChannelDiceLoss(nn.Module):
    """Criterion class for computing multi-channel Dice losses."""

    def __init__(self, smooth: float = 1e-6, reduction: str = "mean"):
        """Initialize MultiChannelDiceLoss with smoothing and reduction options.

        Args:
            smooth (float): Smoothing factor to avoid division by zero.
            reduction (str): Reduction method ('mean', 'sum', or 'none').
        """
        super().__init__()
        self.smooth = smooth
        self.reduction = reduction

    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        """Calculate multi-channel Dice loss between predictions and targets."""
        assert pred.size() == target.size(), "the size of predict and target must be equal."

        pred = pred.sigmoid()
        intersection = (pred * target).sum(dim=(2, 3))
        union = pred.sum(dim=(2, 3)) + target.sum(dim=(2, 3))
        dice = (2.0 * intersection + self.smooth) / (union + self.smooth)
        dice_loss = 1.0 - dice
        dice_loss = dice_loss.mean(dim=1)

        if self.reduction == "mean":
            return dice_loss.mean()
        elif self.reduction == "sum":
            return dice_loss.sum()
        else:
            return dice_loss


class BCEDiceLoss(nn.Module):
    """Criterion class for computing combined BCE and Dice losses."""

    def __init__(self, weight_bce: float = 0.5, weight_dice: float = 0.5):
        """Initialize BCEDiceLoss with BCE and Dice weight factors.

        Args:
            weight_bce (float): Weight factor for BCE loss component.
            weight_dice (float): Weight factor for Dice loss component.
        """
        super().__init__()
        self.weight_bce = weight_bce
        self.weight_dice = weight_dice
        self.bce = nn.BCEWithLogitsLoss()
        self.dice = MultiChannelDiceLoss(smooth=1)

    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        """Calculate combined BCE and Dice loss between predictions and targets."""
        _, _, mask_h, mask_w = pred.shape
        if tuple(target.shape[-2:]) != (mask_h, mask_w):  # downsample to the same size as pred
            target = F.interpolate(target, (mask_h, mask_w), mode="nearest")
        return self.weight_bce * self.bce(pred, target) + self.weight_dice * self.dice(pred, target)


class KeypointLoss(nn.Module):
    """Criterion class for computing keypoint losses."""

    def __init__(self, sigmas: torch.Tensor) -> None:
        """Initialize the KeypointLoss class with keypoint sigmas."""
        super().__init__()
        self.sigmas = sigmas

    def forward(
        self, pred_kpts: torch.Tensor, gt_kpts: torch.Tensor, kpt_mask: torch.Tensor, area: torch.Tensor
    ) -> torch.Tensor:
        """Calculate keypoint loss factor and Euclidean distance loss for keypoints."""
        d = (pred_kpts[..., 0] - gt_kpts[..., 0]).pow(2) + (pred_kpts[..., 1] - gt_kpts[..., 1]).pow(2)
        kpt_loss_factor = kpt_mask.shape[1] / (torch.sum(kpt_mask != 0, dim=1) + 1e-9)
        # e = d / (2 * (area * self.sigmas) ** 2 + 1e-9)  # from formula
        e = d / ((2 * self.sigmas).pow(2) * (area + 1e-9) * 2)  # from cocoeval
        return (kpt_loss_factor.view(-1, 1) * ((1 - torch.exp(-e)) * kpt_mask)).mean()


class v8DetectionLoss:
    """Criterion class for computing training losses for YOLOv8 object detection."""

    def __init__(self, model, tal_topk: int = 10, tal_topk2: int | None = None):  # model must be de-paralleled
        """Initialize v8DetectionLoss with model parameters and task-aligned assignment settings."""
        device = next(model.parameters()).device  # get model device
        h = model.args  # hyperparameters

        m = model.model[-1]  # Detect() module
        self.bce = nn.BCEWithLogitsLoss(reduction="none")
        self.hyp = h
        self.stride = m.stride  # model strides
        self.nc = m.nc  # number of classes
        self.no = m.nc + m.reg_max * 4
        self.reg_max = m.reg_max
        self.device = device

        self.use_dfl = m.reg_max > 1

        self.assigner = TaskAlignedAssigner(
            topk=tal_topk,
            num_classes=self.nc,
            alpha=0.5,
            beta=6.0,
            stride=self.stride.tolist(),
            topk2=tal_topk2,
        )
        self.bbox_loss = BboxLoss(m.reg_max).to(device)
        self.proj = torch.arange(m.reg_max, dtype=torch.float, device=device)

    def _compute_binary_difficulty_loss(self, preds, batch, fg_mask, target_gt_idx, class_filter=None):
        """Compute binary difficulty loss on assigned foreground anchors.

        Raw difficulty levels 0-1 map to class 0, and levels 2-3 map to class 1.
        """
        device = fg_mask.device
        zero = torch.zeros(1, device=device)
        if "preds_diff" not in preds or not fg_mask.any() or "difficulty_levels" not in batch:
            return zero

        preds_diff = preds["preds_diff"]  # [B, 1, A]
        diff_all = (batch["difficulty_levels"].to(device).reshape(-1).long().clamp_(0, 3) >= 2).to(preds_diff.dtype)
        cls_all = batch["cls"].to(device).reshape(-1).long()
        batch_idx = batch["batch_idx"].to(device)
        batch_size = preds_diff.shape[0]
        gt_counts = torch.zeros(batch_size, dtype=torch.long, device=device)
        for i in range(batch_size):
            gt_counts[i] = (batch_idx == i).sum()
        gt_offsets = torch.zeros(batch_size + 1, dtype=torch.long, device=device)
        gt_offsets[1:] = gt_counts.cumsum(0)

        loss = torch.zeros(1, device=device)
        pos_count = 0
        for i in range(batch_size):
            fg_i = fg_mask[i]
            if not fg_i.any():
                continue
            gt_idx_i = target_gt_idx[i][fg_i].reshape(-1).long()
            gt_count_i = int(gt_counts[i].item())
            if gt_idx_i.numel() and ((gt_idx_i < 0).any() or (gt_idx_i >= gt_count_i).any()):
                raise RuntimeError(
                    f"Assigned GT index out of range for difficulty loss image {i}: "
                    f"valid [0, {max(gt_count_i - 1, 0)}], got min={int(gt_idx_i.min().item())}, "
                    f"max={int(gt_idx_i.max().item())}, num_pos={gt_idx_i.numel()}"
                )
            gt_abs_idx = gt_idx_i + int(gt_offsets[i].item())
            if class_filter is not None:
                class_mask = self._cls_mask(cls_all[gt_abs_idx], class_filter)
                if not class_mask.any():
                    continue
                gt_abs_idx = gt_abs_idx[class_mask]
                fg_indices = torch.nonzero(fg_i, as_tuple=False).squeeze(1)[class_mask]
            else:
                fg_indices = fg_i
            pred_i = preds_diff[i, 0, fg_indices]
            loss += self.bce_diff(pred_i, diff_all[gt_abs_idx])
            pos_count += pred_i.shape[0]

        return loss / pos_count if pos_count > 0 else zero

    def preprocess(self, targets: torch.Tensor, batch_size: int, scale_tensor: torch.Tensor) -> torch.Tensor:
        """Preprocess targets by converting to tensor format and scaling coordinates."""
        nl, ne = targets.shape
        if nl == 0:
            out = torch.zeros(batch_size, 0, ne - 1, device=self.device)
        else:
            i = targets[:, 0]  # image index
            _, counts = i.unique(return_counts=True)
            counts = counts.to(dtype=torch.int32)
            out = torch.zeros(batch_size, counts.max(), ne - 1, device=self.device)
            for j in range(batch_size):
                matches = i == j
                if n := matches.sum():
                    out[j, :n] = targets[matches, 1:]
            out[..., 1:5] = xywh2xyxy(out[..., 1:5].mul_(scale_tensor))
        return out

    def bbox_decode(self, anchor_points: torch.Tensor, pred_dist: torch.Tensor) -> torch.Tensor:
        """Decode predicted object bounding box coordinates from anchor points and distribution."""
        if self.use_dfl:
            b, a, c = pred_dist.shape  # batch, anchors, channels
            pred_dist = pred_dist.view(b, a, 4, c // 4).softmax(3).matmul(self.proj.type(pred_dist.dtype))
            # pred_dist = pred_dist.view(b, a, c // 4, 4).transpose(2,3).softmax(3).matmul(self.proj.type(pred_dist.dtype))
            # pred_dist = (pred_dist.view(b, a, c // 4, 4).softmax(2) * self.proj.type(pred_dist.dtype).view(1, 1, -1, 1)).sum(2)
        return dist2bbox(pred_dist, anchor_points, xywh=False)

    def _get_box_loss_inputs(
        self, target_scores: torch.Tensor, fg_mask: torch.Tensor, batch: dict[str, Any]
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Return box-loss targets, masking out virtual-camera samples when present."""
        camera_mode = batch.get("camera_mode")
        if not isinstance(camera_mode, (list, tuple)) or len(camera_mode) != target_scores.shape[0]:
            return target_scores, fg_mask

        virtual_mask = torch.tensor([mode == "virtual" for mode in camera_mode], device=self.device, dtype=torch.bool)
        if not virtual_mask.any():
            return target_scores, fg_mask

        box_target_scores = target_scores.clone()
        box_target_scores[virtual_mask] = 0
        box_fg_mask = fg_mask.clone()
        box_fg_mask[virtual_mask] = False
        return box_target_scores, box_fg_mask

    def get_assigned_targets_and_loss(self, preds: dict[str, torch.Tensor], batch: dict[str, Any]) -> tuple:
        """Calculate the sum of the loss for box, cls and dfl multiplied by batch size and return foreground mask and
        target indices.
        """
        loss = torch.zeros(3, device=self.device)  # box, cls, dfl
        pred_distri, pred_scores = (
            preds["boxes"].permute(0, 2, 1).contiguous(),
            preds["scores"].permute(0, 2, 1).contiguous(),
        )
        anchor_points, stride_tensor = make_anchors(preds["feats"], self.stride, 0.5)

        dtype = pred_scores.dtype
        batch_size = pred_scores.shape[0]
        imgsz = torch.tensor(preds["feats"][0].shape[2:], device=self.device, dtype=dtype) * self.stride[0]

        # Targets
        targets = torch.cat((batch["batch_idx"].view(-1, 1), batch["cls"].view(-1, 1), batch["bboxes"]), 1)
        targets = self.preprocess(targets.to(self.device), batch_size, scale_tensor=imgsz[[1, 0, 1, 0]])
        gt_labels, gt_bboxes = targets.split((1, 4), 2)  # cls, xyxy
        mask_gt = gt_bboxes.sum(2, keepdim=True).gt_(0.0)

        # Pboxes
        pred_bboxes = self.bbox_decode(anchor_points, pred_distri)  # xyxy, (b, h*w, 4)

        _, target_bboxes, target_scores, fg_mask, target_gt_idx = self.assigner(
            pred_scores.detach().sigmoid(),
            (pred_bboxes.detach() * stride_tensor).type(gt_bboxes.dtype),
            anchor_points * stride_tensor,
            gt_labels,
            gt_bboxes,
            mask_gt,
        )

        target_scores_sum = max(target_scores.sum(), 1)

        # Cls loss
        loss[1] = self.bce(pred_scores, target_scores.to(dtype)).sum() / target_scores_sum  # BCE

        # Bbox loss
        box_target_scores, box_fg_mask = self._get_box_loss_inputs(target_scores, fg_mask, batch)
        box_target_scores_sum = max(box_target_scores.sum(), 1)
        if box_fg_mask.sum():
            loss[0], loss[2] = self.bbox_loss(
                pred_distri,
                pred_bboxes,
                anchor_points,
                target_bboxes / stride_tensor,
                box_target_scores,
                box_target_scores_sum,
                box_fg_mask,
                imgsz,
                stride_tensor,
            )

        loss[0] *= self.hyp.box  # box gain
        loss[1] *= self.hyp.cls  # cls gain
        loss[2] *= self.hyp.dfl  # dfl gain
        return (
            (fg_mask, target_gt_idx, target_bboxes, anchor_points, stride_tensor),
            loss,
            loss.detach(),
        )  # loss(box, cls, dfl)

    def parse_output(
        self, preds: dict[str, torch.Tensor] | tuple[torch.Tensor, dict[str, torch.Tensor]]
    ) -> torch.Tensor:
        """Parse model predictions to extract features."""
        return preds[1] if isinstance(preds, tuple) else preds

    def __call__(
        self,
        preds: dict[str, torch.Tensor] | tuple[torch.Tensor, dict[str, torch.Tensor]],
        batch: dict[str, torch.Tensor],
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Calculate the sum of the loss for box, cls and dfl multiplied by batch size."""
        return self.loss(self.parse_output(preds), batch)

    def loss(self, preds: dict[str, torch.Tensor], batch: dict[str, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
        """Calculate detection loss using assigned targets."""
        batch_size = preds["boxes"].shape[0]
        loss, loss_detach = self.get_assigned_targets_and_loss(preds, batch)[1:]
        return loss * batch_size, loss_detach


class v8DetectionLossGround(v8DetectionLoss):
    """Criterion class for computing training losses for ground 2D detection.

    Uses standard YOLO loss without difficulty weighting. This class exists to maintain
    compatibility with the ground dataset format but applies the same loss as v8DetectionLoss.
    """

    def __init__(self, model, tal_topk: int = 10, tal_topk2: int | None = None):
        """Initialize ground 2D detection loss with binary difficulty classification."""
        super().__init__(model, tal_topk, tal_topk2)
        self.bce_diff = nn.BCEWithLogitsLoss(reduction="sum")

    def loss(self, preds: dict[str, torch.Tensor], batch: dict[str, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
        """Calculate detection loss plus binary difficulty classification loss."""
        batch_size = preds["boxes"].shape[0]
        (fg_mask, target_gt_idx, *_), det_loss, det_loss_detach = self.get_assigned_targets_and_loss(preds, batch)
        diff_loss = self._compute_binary_difficulty_loss(preds, batch, fg_mask, target_gt_idx)
        total_loss = det_loss * batch_size + (diff_loss * batch_size / det_loss.numel())
        all_items = torch.cat([det_loss_detach, diff_loss.detach().reshape(1)])
        return total_loss, all_items


class E2EGroundLoss:
    """Criterion class for computing training losses for end-to-end ground 2D detection."""

    def __init__(self, model, loss_fn=v8DetectionLossGround):
        """Initialize E2EGroundLoss with one-to-many and one-to-one detection losses using the provided model."""
        self.one2many = loss_fn(model, tal_topk=10)
        self.one2one = loss_fn(model, tal_topk=7, tal_topk2=1)
        self.updates = 0
        self.total = 1.0
        args = getattr(model, "args", None)
        self.o2m = float(getattr(args, "e2e_o2m_start", 0.8))
        self.o2o = self.total - self.o2m
        self.o2m_copy = self.o2m
        self.final_o2m = float(getattr(args, "e2e_o2m_final", 0.1))
        decay_epochs = getattr(args, "e2e_o2m_decay_epochs", None)
        default_decay_epochs = max(float(getattr(self.one2one.hyp, "epochs", 1)) - 1.0, 1.0)
        self.decay_epochs = max(float(decay_epochs), 1.0) if decay_epochs is not None else default_decay_epochs

    def __call__(self, preds: Any, batch: dict[str, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
        """Calculate the sum of the loss for box, cls and dfl multiplied by batch size."""
        preds = self.one2many.parse_output(preds)
        one2many, one2one = preds["one2many"], preds["one2one"]
        loss_one2many = self.one2many.loss(one2many, batch)
        loss_one2one = self.one2one.loss(one2one, batch)
        return loss_one2many[0] * self.o2m + loss_one2one[0] * self.o2o, loss_one2one[1]

    def update(self) -> None:
        """Update the weights for one-to-many and one-to-one losses based on the decay schedule."""
        self.updates += 1
        self.o2m = self.decay(self.updates)
        self.o2o = max(self.total - self.o2m, 0)

    def decay(self, x) -> float:
        """Calculate the decayed weight for one-to-many loss based on the current update step."""
        return max(1 - x / self.decay_epochs, 0) * (self.o2m_copy - self.final_o2m) + self.final_o2m


class v8Detection3DLoss(v8DetectionLoss):
    """Joint 2D+3D detection loss extending v8DetectionLoss.

    Adds 3D loss components (depth, UV, size, yaw, face visibility, cut classification)
    on top of standard 2D detection loss, using TAL assignment for 3D supervision.

    3D predictions: 41 channels = 4 faces × 6 + 17 whole-box.
    3D GT (labels_3d): 42 dims per object (columns 5-46 of 48-dim format).
    """

    def __init__(self, model, tal_topk=10, tal_topk2=None):
        """Initialize 3D detection loss."""
        super().__init__(model, tal_topk, tal_topk2)
        self.l1_loss = nn.L1Loss(reduction="sum")
        self.l1_loss_none = nn.L1Loss(reduction="none")
        self.bce_yaw = nn.BCEWithLogitsLoss(reduction="sum")
        self.ce_cut = nn.CrossEntropyLoss(reduction="sum")
        self.bce_diff = nn.BCEWithLogitsLoss(reduction="sum")
        self.norm_scales_3d = getattr(model, "norm_scales_3d", {})
        self.face_3d_classes = getattr(model, "face_3d_classes", set())
        self.complete_3d_classes = getattr(model, "complete_3d_classes", set())
        self.fake_3d_classes = getattr(model, "fake_3d_classes", set())
        self.loss_3d_weight = 0.0  # ramped during training
        self.edge_loss_gain = float(getattr(getattr(model, "args", None), "edge_aux_loss_gain", 1.0))
        self.face_visibility_score_thresh = float(
            getattr(self.hyp, "face_visibility_score_thresh", DEFAULT_CFG.face_visibility_score_thresh)
        )

    def loss(self, preds, batch):
        """Calculate 2D + 3D detection loss."""
        batch_size = preds["boxes"].shape[0]
        (fg_mask, target_gt_idx, target_bboxes, anchor_points, stride_tensor), det_loss, det_loss_detach = (
            self.get_assigned_targets_and_loss(preds, batch)
        )
        # Compute image size from feature maps for UV coordinate conversion
        imgsz = torch.tensor(preds["feats"][0].shape[2:], device=self.device, dtype=torch.float32) * self.stride[0]
        opt_3d_items, log_3d_items = self._compute_3d_loss(
            preds, batch, fg_mask, target_gt_idx, anchor_points, stride_tensor, imgsz
        )
        if self.fake_3d_classes and "preds_3d_fake" in preds:
            fake_opt_3d_items, _ = self._compute_3d_loss(
                preds,
                batch,
                fg_mask,
                target_gt_idx,
                anchor_points,
                stride_tensor,
                imgsz,
                preds_3d_key="preds_3d_fake",
                class_filter=self.fake_3d_classes,
                use_edge_loss=False,
            )
            opt_3d_items = opt_3d_items + fake_opt_3d_items
        diff_loss = self._compute_difficulty_loss(preds, batch, fg_mask, target_gt_idx)
        total_loss = (
            det_loss * batch_size
            + (diff_loss * batch_size / det_loss.numel())
            + opt_3d_items.sum() * self.loss_3d_weight * batch_size
        )
        all_items = torch.cat([det_loss_detach, diff_loss.detach().reshape(1), log_3d_items.detach()])
        return total_loss, all_items

    def _compute_difficulty_loss(self, preds, batch, fg_mask, target_gt_idx):
        """Compute binary difficulty classification loss on assigned foreground anchors."""
        classes_3d = set(self.face_3d_classes) | set(self.complete_3d_classes)
        return self._compute_binary_difficulty_loss(preds, batch, fg_mask, target_gt_idx, class_filter=classes_3d)

    def _compute_3d_loss(
        self,
        preds,
        batch,
        fg_mask,
        target_gt_idx,
        anchor_points,
        stride_tensor,
        imgsz,
        preds_3d_key="preds_3d",
        class_filter=None,
        use_edge_loss=True,
    ):
        """Compute optimization 3D terms and detached log diagnostics using TAL-assigned targets."""
        device = fg_mask.device
        dim3d = 6
        edge_face_dim = 15
        edge_point_dim = 3
        zero_opt = torch.zeros(9, device=device)
        zero_log = torch.zeros(12, device=device)

        if preds_3d_key not in preds or not fg_mask.any():
            return zero_opt, zero_log

        preds_3d = preds[preds_3d_key]  # [B, 41, A]
        # When edge auxiliary loss is disabled, skip the expensive GT edge decoding path entirely.
        preds_edge = preds.get("preds_edge") if use_edge_loss and self.edge_loss_gain > 0 else None
        batch_size = preds_3d.shape[0]
        labels_3d = batch.get("labels_3d")
        if labels_3d is None or len(labels_3d) == 0:
            return zero_opt, zero_log

        labels_3d = labels_3d.to(device)
        batch_idx = batch["batch_idx"].to(device)
        cls_all = batch["cls"].to(device)
        batch_calib = batch.get("calib")
        edge_faces_points_2d = batch.get("edge_faces_points_2d")
        edge_faces_depths = batch.get("edge_faces_depths")
        edge_faces_valid = batch.get("edge_faces_valid")
        edge_partial_points_2d = batch.get("edge_partial_points_2d")
        edge_partial_depths = batch.get("edge_partial_depths")
        edge_partial_face_type = batch.get("edge_partial_face_type")
        edge_partial_valid = batch.get("edge_partial_valid")
        if cls_all.dim() > 1:
            cls_all = cls_all.squeeze(-1)

        # Build per-image GT offset mapping
        gt_counts = torch.zeros(batch_size, dtype=torch.long, device=device)
        for i in range(batch_size):
            gt_counts[i] = (batch_idx == i).sum()
        gt_offsets = torch.zeros(batch_size + 1, dtype=torch.long, device=device)
        gt_offsets[1:] = gt_counts.cumsum(0)

        # Optimization accumulators
        lz3d_face = torch.zeros(1, device=device)
        luv_face = torch.zeros(1, device=device)
        lsize_face = torch.zeros(1, device=device)
        lfacecls = torch.zeros(1, device=device)
        lz3d = torch.zeros(1, device=device)
        luv = torch.zeros(1, device=device)
        lsize = torch.zeros(1, device=device)
        lyawcls = torch.zeros(1, device=device)
        lyawreg = torch.zeros(1, device=device)
        lcutcls = torch.zeros(1, device=device)
        ledge_uv = torch.zeros(1, device=device)
        ledge_z = torch.zeros(1, device=device)

        # Human-readable diagnostics in physical units where possible
        z3d_face_m_sum = torch.zeros(1, device=device)
        uv_face_px_sum = torch.zeros(1, device=device)
        size_face_m_sum = torch.zeros(1, device=device)
        z3d_whole_m_sum = torch.zeros(1, device=device)
        uv_whole_px_sum = torch.zeros(1, device=device)
        size_whole_m_sum = torch.zeros(1, device=device)
        edge_uv_px_sum = torch.zeros(1, device=device)
        edge_z_m_sum = torch.zeros(1, device=device)

        face_pos_cnt = face_size_cnt = face_vis_cnt = 0
        whole_cnt = whole_size_cnt = yaw_cls_cnt = cut_cls_cnt = 0
        yaw_reg_cnt = 0
        edge_point_cnt = 0

        for i in range(batch_size):
            fg_i = fg_mask[i]
            if not fg_i.any():
                continue

            gt_idx_i = target_gt_idx[i][fg_i]
            p3d_i = preds_3d[i, :, fg_i].T  # [num_pos, 41] already denormalized by Detect3D head
            pedge_i = preds_edge[i, :, fg_i].T if preds_edge is not None else None  # [num_pos, 60]
            anchor_pts_i = anchor_points[fg_i]  # [num_pos, 2]
            stride_i = stride_tensor[fg_i].reshape(-1, 1)  # [num_pos, 1]
            # Scale to convert normalized UV → grid coords: uv_grid = uv_norm * imgsz / stride
            uv_scale_i = imgsz[[1, 0]] / stride_i  # [num_pos, 2] (w_grid, h_grid)

            gt_indices_i, gt_3d_i, cls_i = self._gather_assigned_3d_targets(
                labels_3d, cls_all, gt_offsets, gt_counts, gt_idx_i, image_idx=i
            )
            gt_bboxes_i = batch["bboxes"][batch_idx == i].to(device)
            gt_bboxes_xyxy_i = xywh2xyxy(gt_bboxes_i) * imgsz[[1, 0, 1, 0]] if len(gt_bboxes_i) else gt_bboxes_i.new_zeros((0, 4))
            gt_start = int(gt_offsets[i].item())

            depth_scale_i = 1.0
            calib_i = None
            if batch_calib is not None and i < len(batch_calib):
                calib_i = batch_calib[i]
                if isinstance(calib_i, dict):
                    depth_scale_i = float(calib_i.get("depth_scale", 1.0))
            depth_scale_i = p3d_i.new_tensor(depth_scale_i)

            # labels_3d 42-dim layout (offset from 48-dim col 5):
            # [0-2]: x3d,y3d,z3d  [3-5]: l,h,w  [6]: rot_y  [7-8]: xc,yc  [9]: alpha
            # [10-17]: front(x3d,y3d,z3d,alpha,xc,yc,score,is_visible)
            # [18-25]: rear  [26-33]: left  [34-41]: right

            valid_3d = ~torch.isnan(gt_3d_i[:, 2])  # z3d not NaN
            is_face = self._cls_mask(cls_i, self.face_3d_classes)
            is_complete = self._cls_mask(cls_i, self.complete_3d_classes)
            if class_filter is not None:
                class_filter_mask = self._cls_mask(cls_i, class_filter)
                is_face &= class_filter_mask
                is_complete &= class_filter_mask
            is_any_3d = (is_face | is_complete) & valid_3d

            if not is_any_3d.any():
                continue

            # Split predictions
            p_front = p3d_i[:, :dim3d]
            p_rear = p3d_i[:, dim3d : dim3d * 2]
            p_left = p3d_i[:, dim3d * 2 : dim3d * 3]
            p_right = p3d_i[:, dim3d * 3 : dim3d * 4]
            p_whole = p3d_i[:, dim3d * 4 :]  # [num_pos, 17]

            edge_blocks = {
                "front": pedge_i[:, :edge_face_dim] if pedge_i is not None else None,
                "rear": pedge_i[:, edge_face_dim : edge_face_dim * 2] if pedge_i is not None else None,
                "left": pedge_i[:, edge_face_dim * 2 : edge_face_dim * 3] if pedge_i is not None else None,
                "right": pedge_i[:, edge_face_dim * 3 :] if pedge_i is not None else None,
            }
            edge_pred_stacked = (
                torch.stack([edge_blocks["front"], edge_blocks["rear"], edge_blocks["left"], edge_blocks["right"]], dim=1)
                .reshape(-1, 4, 5, edge_point_dim)
                if pedge_i is not None
                else None
            )
            edge_gt_cache = {}

            # --- Face losses (only for face_3d_classes) ---
            face_defs = [
                ("front", 0, p_front, 10),
                ("rear", 1, p_rear, 18),
                ("left", 2, p_left, 26),
                ("right", 3, p_right, 34),
            ]
            for face_name, face_type, p_face, gt_off in face_defs:
                face_is_visible = gt_3d_i[:, gt_off + 7]
                face_score = gt_3d_i[:, gt_off + 6]

                vis_mask = is_face & valid_3d & (face_is_visible == 1) & (face_score >= 0.0)
                if vis_mask.any():
                    lfacecls += self.l1_loss(p_face[vis_mask, 5], face_score[vis_mask])
                    face_vis_cnt += vis_mask.sum().item()

                pos_mask = vis_mask & (face_score >= self.face_visibility_score_thresh)
                if pos_mask.any():
                    pos_count = pos_mask.sum().item()

                    gt_z = gt_3d_i[pos_mask, gt_off + 2]
                    z_err_face = self.l1_loss_none(p_face[pos_mask, 0], gt_z)
                    lz3d_face += z_err_face.sum()
                    z3d_face_m_sum += (z_err_face * depth_scale_i).sum()
                    face_pos_cnt += pos_count

                    gt_uv = gt_3d_i[pos_mask, gt_off + 4 : gt_off + 6]
                    gt_uv_grid = gt_uv * uv_scale_i[pos_mask]
                    anchor_uv = anchor_pts_i[pos_mask]
                    uv_offsets = gt_uv_grid - anchor_uv
                    uv_err_face = self.l1_loss_none(p_face[pos_mask, 1:3], uv_offsets)
                    luv_face += uv_err_face.sum()
                    uv_face_px_sum += (uv_err_face * stride_i[pos_mask]).sum()

                    if face_name in ("front", "rear"):
                        gt_size = gt_3d_i[pos_mask][:, [4, 5]]
                    else:
                        gt_size = gt_3d_i[pos_mask][:, [3, 4]]
                    size_err_face = self.l1_loss_none(p_face[pos_mask, 3:5], gt_size)
                    lsize_face += size_err_face.sum()
                    size_face_m_sum += size_err_face.sum()
                    face_size_cnt += pos_count

                    if edge_blocks[face_name] is not None and calib_i is not None:
                        used_precomputed = False
                        if edge_faces_valid is not None and edge_faces_points_2d is not None and edge_faces_depths is not None:
                            edge_mask = pos_mask & edge_faces_valid[gt_indices_i, face_type]
                            if edge_mask.any():
                                gt_points_2d = edge_faces_points_2d[gt_indices_i[edge_mask], face_type].to(dtype=p3d_i.dtype)
                                gt_depths = edge_faces_depths[gt_indices_i[edge_mask], face_type].to(dtype=p3d_i.dtype)
                                gt_depths = _normalize_edge_depth_targets_to_model_space(gt_depths, depth_scale_i)
                                gt_points_grid = gt_points_2d / stride_i[edge_mask].unsqueeze(1)
                                gt_edge_offsets = gt_points_grid - anchor_pts_i[edge_mask].unsqueeze(1)
                                gt_edge_tensor = torch.cat((gt_edge_offsets, gt_depths.unsqueeze(-1)), dim=2)
                                edge_pred_valid = edge_blocks[face_name][edge_mask].reshape(-1, 5, edge_point_dim)
                                uv_err_edge = self.l1_loss_none(edge_pred_valid[:, :, :2], gt_edge_tensor[:, :, :2])
                                z_err_edge = self.l1_loss_none(edge_pred_valid[:, :, 2], gt_edge_tensor[:, :, 2])
                                ledge_uv += uv_err_edge.sum()
                                ledge_z += z_err_edge.sum()
                                edge_uv_px_sum += (uv_err_edge * stride_i[edge_mask].unsqueeze(1)).sum()
                                edge_z_m_sum += (z_err_edge * depth_scale_i).sum()
                                edge_point_cnt += int(edge_mask.sum().item()) * 5
                            used_precomputed = True

                        if not used_precomputed:
                            pos_indices = torch.nonzero(pos_mask, as_tuple=False).squeeze(1)
                            gt_edge_points = []
                            valid_edge_rows = []
                            for local_idx in pos_indices.tolist():
                                gt_abs_idx = int(gt_indices_i[local_idx].item())
                                cache_key = (gt_abs_idx, face_type)
                                if cache_key not in edge_gt_cache:
                                    gt_local_idx = gt_abs_idx - gt_start
                                    bbox_xyxy = None
                                    if 0 <= gt_local_idx < len(gt_bboxes_xyxy_i):
                                        bbox_xyxy = gt_bboxes_xyxy_i[gt_local_idx].detach().cpu().numpy()
                                    edge_gt_cache[cache_key] = decode_visible_face_edge_from_gt(
                                        labels_3d[gt_abs_idx].detach().cpu().numpy(),
                                        int(cls_all[gt_abs_idx].item()),
                                        calib_i,
                                        int(imgsz[1].item()),
                                        int(imgsz[0].item()),
                                        self.face_3d_classes,
                                        self.complete_3d_classes,
                                        face_type=face_type,
                                        score_thr=self.face_visibility_score_thresh,
                                        bbox_xyxy=bbox_xyxy,
                                    )
                                gt_edge = edge_gt_cache[cache_key]
                                if gt_edge is None:
                                    continue
                                gt_points_2d = torch.as_tensor(gt_edge["points_2d"], device=device, dtype=p3d_i.dtype)
                                gt_depths = torch.as_tensor(gt_edge["depths"], device=device, dtype=p3d_i.dtype)
                                gt_depths = _normalize_edge_depth_targets_to_model_space(gt_depths, depth_scale_i)
                                gt_points_grid = gt_points_2d / stride_i[local_idx]
                                gt_edge_offsets = gt_points_grid - anchor_pts_i[local_idx]
                                gt_edge_points.append(torch.cat((gt_edge_offsets, gt_depths.unsqueeze(1)), dim=1))
                                valid_edge_rows.append(local_idx)

                            if gt_edge_points:
                                gt_edge_tensor = torch.stack(gt_edge_points, dim=0)
                                valid_rows_tensor = torch.as_tensor(valid_edge_rows, device=device, dtype=torch.long)
                                edge_pred_valid = edge_blocks[face_name][valid_rows_tensor].reshape(-1, 5, edge_point_dim)
                                uv_err_edge = self.l1_loss_none(edge_pred_valid[:, :, :2], gt_edge_tensor[:, :, :2])
                                z_err_edge = self.l1_loss_none(edge_pred_valid[:, :, 2], gt_edge_tensor[:, :, 2])
                                ledge_uv += uv_err_edge.sum()
                                ledge_z += z_err_edge.sum()
                                edge_uv_px_sum += (uv_err_edge * stride_i[valid_rows_tensor].unsqueeze(1)).sum()
                                edge_z_m_sum += (z_err_edge * depth_scale_i).sum()
                                edge_point_cnt += len(valid_edge_rows) * 5

            def _is_face_cut(gt, off):
                """Check whether a face was invalidated by crop handling."""
                return torch.all(gt[:, off : off + 6] == -1, dim=1) & (gt[:, off + 7] <= 0)

            f_cut = _is_face_cut(gt_3d_i, 10)
            r_cut = _is_face_cut(gt_3d_i, 18)
            l_cut = _is_face_cut(gt_3d_i, 26)
            ri_cut = _is_face_cut(gt_3d_i, 34)
            is_cut = (r_cut & l_cut & ri_cut) | (f_cut & l_cut & ri_cut)

            whole_pos_mask = is_any_3d & ~is_cut
            if whole_pos_mask.any():
                whole_pos_count = whole_pos_mask.sum().item()
                gt_z_whole = gt_3d_i[whole_pos_mask, 2]
                z_err_whole = self.l1_loss_none(p_whole[whole_pos_mask, 0], gt_z_whole)
                lz3d += z_err_whole.sum()
                z3d_whole_m_sum += (z_err_whole * depth_scale_i).sum()
                whole_cnt += whole_pos_count

                gt_uv_whole = gt_3d_i[whole_pos_mask, 7:9]
                gt_uv_whole_grid = gt_uv_whole * uv_scale_i[whole_pos_mask]
                anchor_uv_w = anchor_pts_i[whole_pos_mask]
                uv_err_whole = self.l1_loss_none(p_whole[whole_pos_mask, 1:3], gt_uv_whole_grid - anchor_uv_w)
                luv += uv_err_whole.sum()
                uv_whole_px_sum += (uv_err_whole * stride_i[whole_pos_mask]).sum()

            whole_size_mask = is_any_3d
            if whole_size_mask.any():
                gt_lwh = gt_3d_i[whole_size_mask, 3:6]
                size_err_whole = self.l1_loss_none(p_whole[whole_size_mask, 3:6], gt_lwh)
                lsize += size_err_whole.sum()
                size_whole_m_sum += size_err_whole.sum()
                whole_size_cnt += whole_size_mask.sum().item()

            rot_mask = is_any_3d
            if rot_mask.any():
                rot_y = gt_3d_i[rot_mask, 6:7]
                delta_0 = rot_y
                delta_1 = rot_y - math.pi / 2
                delta_2 = rot_y + math.pi / 2
                ang_mask_t = (torch.abs(rot_y - math.pi) < torch.abs(rot_y + math.pi)).float()
                delta_3 = (rot_y - math.pi) * ang_mask_t + (rot_y + math.pi) * (1 - ang_mask_t)
                angles = torch.cat([delta_0, delta_1, delta_2, delta_3], dim=1)

                ang_cls = torch.clamp((math.pi * 0.5 - torch.abs(angles)) / (math.pi * 0.5), 0.0, 1.0)

                yaw_logits = p_whole[rot_mask, 6:10]
                lyawcls += self.bce_yaw(yaw_logits, ang_cls)
                yaw_cls_cnt += yaw_logits.shape[0] * 4

                angle_valid = torch.abs(angles) <= (math.pi / 2)
                bin_active = ang_cls > 0.1
                valid_yaw = angle_valid & bin_active
                yaw_reg_cnt += valid_yaw.sum().item()
                target_sin = torch.sin(angles)
                yaw_reg_loss = self.l1_loss_none(p_whole[rot_mask, 10:14], target_sin) * valid_yaw
                lyawreg += yaw_reg_loss.sum()

            cut_mask = is_face & valid_3d
            if cut_mask.any() and pedge_i is not None and calib_i is not None:
                used_precomputed = False
                if (
                    edge_pred_stacked is not None
                    and edge_partial_valid is not None
                    and edge_partial_points_2d is not None
                    and edge_partial_depths is not None
                    and edge_partial_face_type is not None
                ):
                    partial_mask = cut_mask & edge_partial_valid[gt_indices_i]
                    if partial_mask.any():
                        gt_points_2d = edge_partial_points_2d[gt_indices_i[partial_mask]].to(dtype=p3d_i.dtype)
                        gt_depths = edge_partial_depths[gt_indices_i[partial_mask]].to(dtype=p3d_i.dtype)
                        gt_depths = _normalize_edge_depth_targets_to_model_space(gt_depths, depth_scale_i)
                        face_rows = edge_partial_face_type[gt_indices_i[partial_mask]].to(dtype=torch.long)
                        gt_points_grid = gt_points_2d / stride_i[partial_mask].unsqueeze(1)
                        gt_edge_offsets = gt_points_grid - anchor_pts_i[partial_mask].unsqueeze(1)
                        gt_edge_tensor = torch.cat((gt_edge_offsets, gt_depths.unsqueeze(-1)), dim=2)
                        edge_pred_valid = edge_pred_stacked[partial_mask, face_rows]
                        uv_err_edge = self.l1_loss_none(edge_pred_valid[:, :, :2], gt_edge_tensor[:, :, :2])
                        z_err_edge = self.l1_loss_none(edge_pred_valid[:, :, 2], gt_edge_tensor[:, :, 2])
                        ledge_uv += uv_err_edge.sum()
                        ledge_z += z_err_edge.sum()
                        edge_uv_px_sum += (uv_err_edge * stride_i[partial_mask].unsqueeze(1)).sum()
                        edge_z_m_sum += (z_err_edge * depth_scale_i).sum()
                        edge_point_cnt += int(partial_mask.sum().item()) * 5
                    used_precomputed = True

                if not used_precomputed:
                    cut_indices = torch.nonzero(cut_mask, as_tuple=False).squeeze(1)
                    partial_gt_points = []
                    partial_valid_rows = []
                    partial_face_types = []
                    for local_idx in cut_indices.tolist():
                        gt_abs_idx = int(gt_indices_i[local_idx].item())
                        gt_local_idx = gt_abs_idx - gt_start
                        bbox_xyxy = None
                        if 0 <= gt_local_idx < len(gt_bboxes_xyxy_i):
                            bbox_xyxy = gt_bboxes_xyxy_i[gt_local_idx].detach().cpu().numpy()
                        partial_edge = decode_cut_partial_side_edge_from_gt(
                            labels_3d[gt_abs_idx].detach().cpu().numpy(),
                            int(cls_all[gt_abs_idx].item()),
                            calib_i,
                            int(imgsz[1].item()),
                            int(imgsz[0].item()),
                            self.face_3d_classes,
                            self.complete_3d_classes,
                            bbox_xyxy=bbox_xyxy,
                        )
                        if partial_edge is None:
                            continue
                        gt_points_2d = torch.as_tensor(partial_edge["points_2d"], device=device, dtype=p3d_i.dtype)
                        gt_depths = torch.as_tensor(partial_edge["depths"], device=device, dtype=p3d_i.dtype)
                        gt_depths = _normalize_edge_depth_targets_to_model_space(gt_depths, depth_scale_i)
                        gt_points_grid = gt_points_2d / stride_i[local_idx]
                        gt_edge_offsets = gt_points_grid - anchor_pts_i[local_idx]
                        partial_gt_points.append(torch.cat((gt_edge_offsets, gt_depths.unsqueeze(1)), dim=1))
                        partial_valid_rows.append(local_idx)
                        partial_face_types.append(int(partial_edge["face_type"]))

                    if partial_gt_points:
                        gt_edge_tensor = torch.stack(partial_gt_points, dim=0)
                        valid_rows_tensor = torch.as_tensor(partial_valid_rows, device=device, dtype=torch.long)
                        edge_pred_valid = torch.stack(
                            [edge_blocks[("front", "rear", "left", "right")[face_type]][row] for row, face_type in zip(valid_rows_tensor.tolist(), partial_face_types)],
                            dim=0,
                        ).reshape(-1, 5, edge_point_dim)
                        uv_err_edge = self.l1_loss_none(edge_pred_valid[:, :, :2], gt_edge_tensor[:, :, :2])
                        z_err_edge = self.l1_loss_none(edge_pred_valid[:, :, 2], gt_edge_tensor[:, :, 2])
                        ledge_uv += uv_err_edge.sum()
                        ledge_z += z_err_edge.sum()
                        edge_uv_px_sum += (uv_err_edge * stride_i[valid_rows_tensor].unsqueeze(1)).sum()
                        edge_z_m_sum += (z_err_edge * depth_scale_i).sum()
                        edge_point_cnt += len(partial_valid_rows) * 5

            if cut_mask.any():
                cut_label = torch.zeros(cut_mask.sum(), dtype=torch.long, device=device)
                r_c = r_cut[cut_mask]
                l_c = l_cut[cut_mask]
                ri_c = ri_cut[cut_mask]
                f_c = f_cut[cut_mask]
                cut_label[(r_c & l_c & ri_c)] = 1
                cut_label[(f_c & l_c & ri_c)] = 2
                lcutcls += self.ce_cut(p_whole[cut_mask, -3:], cut_label)
                cut_cls_cnt += cut_mask.sum().item()

        if face_pos_cnt > 0:
            lz3d_face /= face_pos_cnt
            luv_face /= face_pos_cnt
        if face_size_cnt > 0:
            lsize_face /= face_size_cnt
        if face_vis_cnt > 0:
            lfacecls /= face_vis_cnt
        if whole_cnt > 0:
            lz3d /= whole_cnt
            luv /= whole_cnt
            lsize /= whole_cnt
        if yaw_cls_cnt > 0:
            lyawcls /= yaw_cls_cnt
        if yaw_reg_cnt > 0:
            lyawreg /= yaw_reg_cnt
        if cut_cls_cnt > 0:
            lcutcls /= cut_cls_cnt
        if edge_point_cnt > 0:
            ledge_uv /= edge_point_cnt
            ledge_z /= edge_point_cnt

        z_scale = self.norm_scales_3d.get("z3d_scale", 1.0)
        size_scale = self.norm_scales_3d.get("size_scale", 1.0)

        loss_z = (lz3d + lz3d_face) / z_scale
        loss_uv = luv + luv_face
        loss_size = (lsize + lsize_face) / size_scale
        lyawreg_norm = lyawreg
        opt_items = torch.cat(
            [
                loss_z,
                loss_uv,
                loss_size,
                lyawcls,
                lyawreg_norm,
                lcutcls,
                lfacecls,
                ledge_uv * self.edge_loss_gain,
                ledge_z * self.edge_loss_gain,
            ]
        )

        whole_z_m = z3d_whole_m_sum / whole_cnt if whole_cnt > 0 else torch.zeros(1, device=device)
        whole_uv_px = uv_whole_px_sum / (whole_cnt * 2) if whole_cnt > 0 else torch.zeros(1, device=device)
        whole_size_m = size_whole_m_sum / (whole_size_cnt * 3) if whole_size_cnt > 0 else torch.zeros(1, device=device)
        face_z_m = z3d_face_m_sum / face_pos_cnt if face_pos_cnt > 0 else torch.zeros(1, device=device)
        face_uv_px = uv_face_px_sum / (face_pos_cnt * 2) if face_pos_cnt > 0 else torch.zeros(1, device=device)
        face_size_m = size_face_m_sum / (face_size_cnt * 2) if face_size_cnt > 0 else torch.zeros(1, device=device)
        yaw_deg = torch.asin(torch.clamp(lyawreg, -1.0, 1.0)) * (180.0 / math.pi)
        edge_uv_px = edge_uv_px_sum / (edge_point_cnt * 2) if edge_point_cnt > 0 else torch.zeros(1, device=device)
        edge_z_m = edge_z_m_sum / edge_point_cnt if edge_point_cnt > 0 else torch.zeros(1, device=device)
        log_items = torch.cat([
            whole_z_m,
            whole_uv_px,
            whole_size_m,
            lyawcls,
            yaw_deg,
            lcutcls,
            face_z_m,
            face_uv_px,
            face_size_m,
            lfacecls,
            edge_uv_px,
            edge_z_m,
        ])
        return opt_items, log_items


    @staticmethod
    def _gather_assigned_3d_targets(labels_3d, cls_all, gt_offsets, gt_counts, gt_idx_i, image_idx):
        """Gather flat 3D GT rows for a single image from local TAL assignments."""
        gt_idx_i = gt_idx_i.reshape(-1).to(dtype=torch.long)
        gt_count_i = int(gt_counts[image_idx].item())
        if gt_idx_i.numel() and ((gt_idx_i < 0).any() or (gt_idx_i >= gt_count_i).any()):
            raise RuntimeError(
                f"Assigned GT index out of range for image {image_idx}: "
                f"valid [0, {max(gt_count_i - 1, 0)}], got min={int(gt_idx_i.min().item())}, "
                f"max={int(gt_idx_i.max().item())}, num_pos={gt_idx_i.numel()}"
            )

        gt_start_i = int(gt_offsets[image_idx].item())
        gt_indices_i = gt_idx_i + gt_start_i
        gt_3d_i = labels_3d[gt_indices_i]
        cls_i = cls_all[gt_indices_i]
        return gt_indices_i, gt_3d_i, cls_i


    @staticmethod
    def _cls_mask(cls_tensor, cls_set):
        """Build boolean mask for targets whose class is in cls_set."""
        mask = torch.zeros_like(cls_tensor, dtype=torch.bool)
        for c in cls_set:
            mask |= cls_tensor == c
        return mask


class E2EGround3DLoss:
    """Criterion class for end-to-end ground 3D detection training."""

    def __init__(self, model, loss_fn=v8Detection3DLoss):
        """Initialize with one-to-many and one-to-one 3D detection losses."""
        self.one2many = loss_fn(model, tal_topk=10)
        self.one2one = loss_fn(model, tal_topk=7, tal_topk2=1)
        self.updates = 0
        self.total = 1.0
        args = getattr(model, "args", None)
        self.o2m = float(getattr(args, "e2e_o2m_start", 0.8))
        self.o2o = self.total - self.o2m
        self.o2m_copy = self.o2m
        self.final_o2m = float(getattr(args, "e2e_o2m_final", 0.1))
        decay_epochs = getattr(args, "e2e_o2m_decay_epochs", None)
        default_decay_epochs = max(float(getattr(self.one2one.hyp, "epochs", 1)) - 1.0, 1.0)
        self.decay_epochs = max(float(decay_epochs), 1.0) if decay_epochs is not None else default_decay_epochs

    def __call__(self, preds, batch):
        """Calculate the sum of the loss for box, cls, dfl, and 3D multiplied by batch size."""
        preds = self.one2many.parse_output(preds)
        one2many, one2one = preds["one2many"], preds["one2one"]
        loss_one2many = self.one2many.loss(one2many, batch)
        loss_one2one = self.one2one.loss(one2one, batch)
        return loss_one2many[0] * self.o2m + loss_one2one[0] * self.o2o, loss_one2one[1]

    def update(self):
        """Update the weights for one-to-many and one-to-one losses."""
        self.updates += 1
        self.o2m = self.decay(self.updates)
        self.o2o = max(self.total - self.o2m, 0)

    def decay(self, x):
        """Calculate the decayed weight for one-to-many loss."""
        return max(1 - x / self.decay_epochs, 0) * (self.o2m_copy - self.final_o2m) + self.final_o2m


class v8SegmentationLoss(v8DetectionLoss):
    """Criterion class for computing training losses for YOLOv8 segmentation."""

    def __init__(self, model, tal_topk: int = 10, tal_topk2: int | None = None):  # model must be de-paralleled
        """Initialize the v8SegmentationLoss class with model parameters and mask overlap setting."""
        super().__init__(model, tal_topk, tal_topk2)
        self.overlap = model.args.overlap_mask
        self.bcedice_loss = BCEDiceLoss(weight_bce=0.5, weight_dice=0.5)

    def loss(self, preds: dict[str, torch.Tensor], batch: dict[str, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
        """Calculate and return the combined loss for detection and segmentation."""
        pred_masks, proto = preds["mask_coefficient"].permute(0, 2, 1).contiguous(), preds["proto"]
        loss = torch.zeros(5, device=self.device)  # box, seg, cls, dfl, semseg
        if isinstance(proto, tuple) and len(proto) == 2:
            proto, pred_semseg = proto
        else:
            pred_semseg = None
        (fg_mask, target_gt_idx, target_bboxes, _, _), det_loss, _ = self.get_assigned_targets_and_loss(preds, batch)
        # NOTE: re-assign index for consistency for now. Need to be removed in the future.
        loss[0], loss[2], loss[3] = det_loss[0], det_loss[1], det_loss[2]

        batch_size, _, mask_h, mask_w = proto.shape  # batch size, number of masks, mask height, mask width
        if fg_mask.sum():
            # Masks loss
            masks = batch["masks"].to(self.device).float()
            if tuple(masks.shape[-2:]) != (mask_h, mask_w):  # downsample
                # masks = F.interpolate(masks[None], (mask_h, mask_w), mode="nearest")[0]
                proto = F.interpolate(proto, masks.shape[-2:], mode="bilinear", align_corners=False)

            imgsz = (
                torch.tensor(preds["feats"][0].shape[2:], device=self.device, dtype=pred_masks.dtype) * self.stride[0]
            )
            loss[1] = self.calculate_segmentation_loss(
                fg_mask,
                masks,
                target_gt_idx,
                target_bboxes,
                batch["batch_idx"].view(-1, 1),
                proto,
                pred_masks,
                imgsz,
            )
            if pred_semseg is not None:
                sem_masks = batch["sem_masks"].to(self.device)  # NxHxW
                sem_masks = F.one_hot(sem_masks.long(), num_classes=self.nc).permute(0, 3, 1, 2).float()  # NxCxHxW

                if self.overlap:
                    mask_zero = masks == 0  # NxHxW
                    sem_masks[mask_zero.unsqueeze(1).expand_as(sem_masks)] = 0
                else:
                    batch_idx = batch["batch_idx"].view(-1)  # [total_instances]
                    for i in range(batch_size):
                        instance_mask_i = masks[batch_idx == i]  # [num_instances_i, H, W]
                        if len(instance_mask_i) == 0:
                            continue
                        sem_masks[i, :, instance_mask_i.sum(dim=0) == 0] = 0

                loss[4] = self.bcedice_loss(pred_semseg, sem_masks)
                loss[4] *= self.hyp.box  # seg gain

        # WARNING: lines below prevent Multi-GPU DDP 'unused gradient' PyTorch errors, do not remove
        else:
            loss[1] += (proto * 0).sum() + (pred_masks * 0).sum()  # inf sums may lead to nan loss
            if pred_semseg is not None:
                loss[4] += (pred_semseg * 0).sum()

        loss[1] *= self.hyp.box  # seg gain
        return loss * batch_size, loss.detach()  # loss(box, seg, cls, dfl, semseg)

    @staticmethod
    def single_mask_loss(
        gt_mask: torch.Tensor, pred: torch.Tensor, proto: torch.Tensor, xyxy: torch.Tensor, area: torch.Tensor
    ) -> torch.Tensor:
        """Compute the instance segmentation loss for a single image.

        Args:
            gt_mask (torch.Tensor): Ground truth mask of shape (N, H, W), where N is the number of objects.
            pred (torch.Tensor): Predicted mask coefficients of shape (N, 32).
            proto (torch.Tensor): Prototype masks of shape (32, H, W).
            xyxy (torch.Tensor): Ground truth bounding boxes in xyxy format, normalized to [0, 1], of shape (N, 4).
            area (torch.Tensor): Area of each ground truth bounding box of shape (N,).

        Returns:
            (torch.Tensor): The calculated mask loss for a single image.

        Notes:
            The function uses the equation pred_mask = torch.einsum('in,nhw->ihw', pred, proto) to produce the
            predicted masks from the prototype masks and predicted mask coefficients.
        """
        pred_mask = torch.einsum("in,nhw->ihw", pred, proto)  # (n, 32) @ (32, 80, 80) -> (n, 80, 80)
        loss = F.binary_cross_entropy_with_logits(pred_mask, gt_mask, reduction="none")
        return (crop_mask(loss, xyxy).mean(dim=(1, 2)) / area).sum()

    def calculate_segmentation_loss(
        self,
        fg_mask: torch.Tensor,
        masks: torch.Tensor,
        target_gt_idx: torch.Tensor,
        target_bboxes: torch.Tensor,
        batch_idx: torch.Tensor,
        proto: torch.Tensor,
        pred_masks: torch.Tensor,
        imgsz: torch.Tensor,
    ) -> torch.Tensor:
        """Calculate the loss for instance segmentation.

        Args:
            fg_mask (torch.Tensor): A binary tensor of shape (BS, N_anchors) indicating which anchors are positive.
            masks (torch.Tensor): Ground truth masks of shape (BS, H, W) if `overlap` is False, otherwise (BS, ?, H, W).
            target_gt_idx (torch.Tensor): Indexes of ground truth objects for each anchor of shape (BS, N_anchors).
            target_bboxes (torch.Tensor): Ground truth bounding boxes for each anchor of shape (BS, N_anchors, 4).
            batch_idx (torch.Tensor): Batch indices of shape (N_labels_in_batch, 1).
            proto (torch.Tensor): Prototype masks of shape (BS, 32, H, W).
            pred_masks (torch.Tensor): Predicted masks for each anchor of shape (BS, N_anchors, 32).
            imgsz (torch.Tensor): Size of the input image as a tensor of shape (2), i.e., (H, W).

        Returns:
            (torch.Tensor): The calculated loss for instance segmentation.

        Notes:
            The batch loss can be computed for improved speed at higher memory usage.
            For example, pred_mask can be computed as follows:
                pred_mask = torch.einsum('in,nhw->ihw', pred, proto)  # (i, 32) @ (32, 160, 160) -> (i, 160, 160)
        """
        _, _, mask_h, mask_w = proto.shape
        loss = 0

        # Normalize to 0-1
        target_bboxes_normalized = target_bboxes / imgsz[[1, 0, 1, 0]]

        # Areas of target bboxes
        marea = xyxy2xywh(target_bboxes_normalized)[..., 2:].prod(2)

        # Normalize to mask size
        mxyxy = target_bboxes_normalized * torch.tensor([mask_w, mask_h, mask_w, mask_h], device=proto.device)

        for i, single_i in enumerate(zip(fg_mask, target_gt_idx, pred_masks, proto, mxyxy, marea, masks)):
            fg_mask_i, target_gt_idx_i, pred_masks_i, proto_i, mxyxy_i, marea_i, masks_i = single_i
            if fg_mask_i.any():
                mask_idx = target_gt_idx_i[fg_mask_i]
                if self.overlap:
                    gt_mask = masks_i == (mask_idx + 1).view(-1, 1, 1)
                    gt_mask = gt_mask.float()
                else:
                    gt_mask = masks[batch_idx.view(-1) == i][mask_idx]

                loss += self.single_mask_loss(
                    gt_mask, pred_masks_i[fg_mask_i], proto_i, mxyxy_i[fg_mask_i], marea_i[fg_mask_i]
                )

            # WARNING: lines below prevents Multi-GPU DDP 'unused gradient' PyTorch errors, do not remove
            else:
                loss += (proto * 0).sum() + (pred_masks * 0).sum()  # inf sums may lead to nan loss

        return loss / fg_mask.sum()


class v8PoseLoss(v8DetectionLoss):
    """Criterion class for computing training losses for YOLOv8 pose estimation."""

    def __init__(self, model, tal_topk: int = 10, tal_topk2: int = 10):  # model must be de-paralleled
        """Initialize v8PoseLoss with model parameters and keypoint-specific loss functions."""
        super().__init__(model, tal_topk, tal_topk2)
        self.kpt_shape = model.model[-1].kpt_shape
        self.bce_pose = nn.BCEWithLogitsLoss()
        is_pose = self.kpt_shape == [17, 3]
        nkpt = self.kpt_shape[0]  # number of keypoints
        sigmas = torch.from_numpy(OKS_SIGMA).to(self.device) if is_pose else torch.ones(nkpt, device=self.device) / nkpt
        self.keypoint_loss = KeypointLoss(sigmas=sigmas)

    def loss(self, preds: dict[str, torch.Tensor], batch: dict[str, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
        """Calculate the total loss and detach it for pose estimation."""
        pred_kpts = preds["kpts"].permute(0, 2, 1).contiguous()
        loss = torch.zeros(5, device=self.device)  # box, kpt_location, kpt_visibility, cls, dfl
        (fg_mask, target_gt_idx, target_bboxes, anchor_points, stride_tensor), det_loss, _ = (
            self.get_assigned_targets_and_loss(preds, batch)
        )
        # NOTE: re-assign index for consistency for now. Need to be removed in the future.
        loss[0], loss[3], loss[4] = det_loss[0], det_loss[1], det_loss[2]

        batch_size = pred_kpts.shape[0]
        imgsz = torch.tensor(preds["feats"][0].shape[2:], device=self.device, dtype=pred_kpts.dtype) * self.stride[0]

        # Pboxes
        pred_kpts = self.kpts_decode(anchor_points, pred_kpts.view(batch_size, -1, *self.kpt_shape))  # (b, h*w, 17, 3)

        # Keypoint loss
        if fg_mask.sum():
            keypoints = batch["keypoints"].to(self.device).float().clone()
            keypoints[..., 0] *= imgsz[1]
            keypoints[..., 1] *= imgsz[0]

            loss[1], loss[2] = self.calculate_keypoints_loss(
                fg_mask,
                target_gt_idx,
                keypoints,
                batch["batch_idx"].view(-1, 1),
                stride_tensor,
                target_bboxes,
                pred_kpts,
            )

        loss[1] *= self.hyp.pose  # pose gain
        loss[2] *= self.hyp.kobj  # kobj gain

        return loss * batch_size, loss.detach()  # loss(box, pose, kobj, cls, dfl)

    @staticmethod
    def kpts_decode(anchor_points: torch.Tensor, pred_kpts: torch.Tensor) -> torch.Tensor:
        """Decode predicted keypoints to image coordinates."""
        y = pred_kpts.clone()
        y[..., :2] *= 2.0
        y[..., 0] += anchor_points[:, [0]] - 0.5
        y[..., 1] += anchor_points[:, [1]] - 0.5
        return y

    def _select_target_keypoints(
        self,
        keypoints: torch.Tensor,
        batch_idx: torch.Tensor,
        target_gt_idx: torch.Tensor,
        masks: torch.Tensor,
    ) -> torch.Tensor:
        """Select target keypoints for each anchor based on batch index and target ground truth index.

        Args:
            keypoints (torch.Tensor): Ground truth keypoints, shape (N_kpts_in_batch, N_kpts_per_object, kpts_dim).
            batch_idx (torch.Tensor): Batch index tensor for keypoints, shape (N_kpts_in_batch, 1).
            target_gt_idx (torch.Tensor): Index tensor mapping anchors to ground truth objects, shape (BS, N_anchors).
            masks (torch.Tensor): Binary mask tensor indicating object presence, shape (BS, N_anchors).

        Returns:
            (torch.Tensor): Selected keypoints tensor, shape (BS, N_anchors, N_kpts_per_object, kpts_dim).
        """
        batch_idx = batch_idx.flatten()
        batch_size = len(masks)

        # Find the maximum number of keypoints in a single image
        max_kpts = torch.unique(batch_idx, return_counts=True)[1].max()

        # Create a tensor to hold batched keypoints
        batched_keypoints = torch.zeros(
            (batch_size, max_kpts, keypoints.shape[1], keypoints.shape[2]), device=keypoints.device
        )

        # TODO: any idea how to vectorize this?
        # Fill batched_keypoints with keypoints based on batch_idx
        for i in range(batch_size):
            keypoints_i = keypoints[batch_idx == i]
            batched_keypoints[i, : keypoints_i.shape[0]] = keypoints_i

        # Expand dimensions of target_gt_idx to match the shape of batched_keypoints
        target_gt_idx_expanded = target_gt_idx.unsqueeze(-1).unsqueeze(-1)

        # Use target_gt_idx_expanded to select keypoints from batched_keypoints
        selected_keypoints = batched_keypoints.gather(
            1, target_gt_idx_expanded.expand(-1, -1, keypoints.shape[1], keypoints.shape[2])
        )

        return selected_keypoints

    def calculate_keypoints_loss(
        self,
        masks: torch.Tensor,
        target_gt_idx: torch.Tensor,
        keypoints: torch.Tensor,
        batch_idx: torch.Tensor,
        stride_tensor: torch.Tensor,
        target_bboxes: torch.Tensor,
        pred_kpts: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Calculate the keypoints loss for the model.

        This function calculates the keypoints loss and keypoints object loss for a given batch. The keypoints loss is
        based on the difference between the predicted keypoints and ground truth keypoints. The keypoints object loss is
        a binary classification loss that classifies whether a keypoint is present or not.

        Args:
            masks (torch.Tensor): Binary mask tensor indicating object presence, shape (BS, N_anchors).
            target_gt_idx (torch.Tensor): Index tensor mapping anchors to ground truth objects, shape (BS, N_anchors).
            keypoints (torch.Tensor): Ground truth keypoints, shape (N_kpts_in_batch, N_kpts_per_object, kpts_dim).
            batch_idx (torch.Tensor): Batch index tensor for keypoints, shape (N_kpts_in_batch, 1).
            stride_tensor (torch.Tensor): Stride tensor for anchors, shape (N_anchors, 1).
            target_bboxes (torch.Tensor): Ground truth boxes in (x1, y1, x2, y2) format, shape (BS, N_anchors, 4).
            pred_kpts (torch.Tensor): Predicted keypoints, shape (BS, N_anchors, N_kpts_per_object, kpts_dim).

        Returns:
            kpts_loss (torch.Tensor): The keypoints loss.
            kpts_obj_loss (torch.Tensor): The keypoints object loss.
        """
        # Select target keypoints using helper method
        selected_keypoints = self._select_target_keypoints(keypoints, batch_idx, target_gt_idx, masks)

        # Divide coordinates by stride
        selected_keypoints[..., :2] /= stride_tensor.view(1, -1, 1, 1)

        kpts_loss = 0
        kpts_obj_loss = 0

        if masks.any():
            target_bboxes /= stride_tensor
            gt_kpt = selected_keypoints[masks]
            area = xyxy2xywh(target_bboxes[masks])[:, 2:].prod(1, keepdim=True)
            pred_kpt = pred_kpts[masks]
            kpt_mask = gt_kpt[..., 2] != 0 if gt_kpt.shape[-1] == 3 else torch.full_like(gt_kpt[..., 0], True)
            kpts_loss = self.keypoint_loss(pred_kpt, gt_kpt, kpt_mask, area)  # pose loss

            if pred_kpt.shape[-1] == 3:
                kpts_obj_loss = self.bce_pose(pred_kpt[..., 2], kpt_mask.float())  # keypoint obj loss

        return kpts_loss, kpts_obj_loss


class PoseLoss26(v8PoseLoss):
    """Criterion class for computing training losses for YOLOv8 pose estimation with RLE loss support."""

    def __init__(self, model, tal_topk: int = 10, tal_topk2: int | None = None):  # model must be de-paralleled
        """Initialize PoseLoss26 with model parameters and keypoint-specific loss functions including RLE loss."""
        super().__init__(model, tal_topk, tal_topk2)
        is_pose = self.kpt_shape == [17, 3]
        nkpt = self.kpt_shape[0]  # number of keypoints
        self.rle_loss = None
        self.flow_model = model.model[-1].flow_model if hasattr(model.model[-1], "flow_model") else None
        if self.flow_model is not None:
            self.rle_loss = RLELoss(use_target_weight=True).to(self.device)
            self.target_weights = (
                torch.from_numpy(RLE_WEIGHT).to(self.device) if is_pose else torch.ones(nkpt, device=self.device)
            )

    def loss(self, preds: dict[str, torch.Tensor], batch: dict[str, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
        """Calculate the total loss and detach it for pose estimation."""
        pred_kpts = preds["kpts"].permute(0, 2, 1).contiguous()
        loss = torch.zeros(
            6 if self.rle_loss else 5, device=self.device
        )  # box, kpt_location, kpt_visibility, cls, dfl[, rle]
        (fg_mask, target_gt_idx, target_bboxes, anchor_points, stride_tensor), det_loss, _ = (
            self.get_assigned_targets_and_loss(preds, batch)
        )
        # NOTE: re-assign index for consistency for now. Need to be removed in the future.
        loss[0], loss[3], loss[4] = det_loss[0], det_loss[1], det_loss[2]

        batch_size = pred_kpts.shape[0]
        imgsz = torch.tensor(preds["feats"][0].shape[2:], device=self.device, dtype=pred_kpts.dtype) * self.stride[0]

        pred_kpts = pred_kpts.view(batch_size, -1, *self.kpt_shape)  # (b, h*w, 17, 3)

        if self.rle_loss and preds.get("kpts_sigma", None) is not None:
            pred_sigma = preds["kpts_sigma"].permute(0, 2, 1).contiguous()
            pred_sigma = pred_sigma.view(batch_size, -1, self.kpt_shape[0], 2)  # (b, h*w, 17, 2)
            pred_kpts = torch.cat([pred_kpts, pred_sigma], dim=-1)  # (b, h*w, 17, 5)

        pred_kpts = self.kpts_decode(anchor_points, pred_kpts)

        # Keypoint loss
        if fg_mask.sum():
            keypoints = batch["keypoints"].to(self.device).float().clone()
            keypoints[..., 0] *= imgsz[1]
            keypoints[..., 1] *= imgsz[0]

            keypoints_loss = self.calculate_keypoints_loss(
                fg_mask,
                target_gt_idx,
                keypoints,
                batch["batch_idx"].view(-1, 1),
                stride_tensor,
                target_bboxes,
                pred_kpts,
            )
            loss[1] = keypoints_loss[0]
            loss[2] = keypoints_loss[1]
            if self.rle_loss is not None:
                loss[5] = keypoints_loss[2]

        loss[1] *= self.hyp.pose  # pose gain
        loss[2] *= self.hyp.kobj  # kobj gain
        if self.rle_loss is not None:
            loss[5] *= self.hyp.rle  # rle gain

        return loss * batch_size, loss.detach()  # loss(box, kpt_location, kpt_visibility, cls, dfl[, rle])

    @staticmethod
    def kpts_decode(anchor_points: torch.Tensor, pred_kpts: torch.Tensor) -> torch.Tensor:
        """Decode predicted keypoints to image coordinates."""
        y = pred_kpts.clone()
        y[..., 0] += anchor_points[:, [0]]
        y[..., 1] += anchor_points[:, [1]]
        return y

    def calculate_rle_loss(self, pred_kpt: torch.Tensor, gt_kpt: torch.Tensor, kpt_mask: torch.Tensor) -> torch.Tensor:
        """Calculate the RLE (Residual Log-likelihood Estimation) loss for keypoints.

        Args:
            pred_kpt (torch.Tensor): Predicted kpts with sigma, shape (N, num_keypoints, kpts_dim) where kpts_dim >= 4.
            gt_kpt (torch.Tensor): Ground truth keypoints, shape (N, num_keypoints, kpts_dim).
            kpt_mask (torch.Tensor): Mask for valid keypoints, shape (N, num_keypoints).

        Returns:
            (torch.Tensor): The RLE loss.
        """
        pred_kpt_visible = pred_kpt[kpt_mask]
        gt_kpt_visible = gt_kpt[kpt_mask]
        pred_coords = pred_kpt_visible[:, 0:2]
        pred_sigma = pred_kpt_visible[:, -2:]
        gt_coords = gt_kpt_visible[:, 0:2]

        target_weights = self.target_weights.unsqueeze(0).repeat(kpt_mask.shape[0], 1)
        target_weights = target_weights[kpt_mask]

        pred_sigma = pred_sigma.sigmoid()
        error = (pred_coords - gt_coords) / (pred_sigma + 1e-9)

        # Filter out NaN and Inf values to prevent MultivariateNormal validation errors
        valid_mask = ~(torch.isnan(error) | torch.isinf(error)).any(dim=-1)
        if not valid_mask.any():
            return torch.tensor(0.0, device=pred_kpt.device)

        error = error[valid_mask]
        error = error.clamp(-100, 100)  # Prevent numerical instability
        pred_sigma = pred_sigma[valid_mask]
        target_weights = target_weights[valid_mask]

        log_phi = self.flow_model.log_prob(error)

        return self.rle_loss(pred_sigma, log_phi, error, target_weights)

    def calculate_keypoints_loss(
        self,
        masks: torch.Tensor,
        target_gt_idx: torch.Tensor,
        keypoints: torch.Tensor,
        batch_idx: torch.Tensor,
        stride_tensor: torch.Tensor,
        target_bboxes: torch.Tensor,
        pred_kpts: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Calculate the keypoints loss for the model.

        This function calculates the keypoints loss and keypoints object loss for a given batch. The keypoints loss is
        based on the difference between the predicted keypoints and ground truth keypoints. The keypoints object loss is
        a binary classification loss that classifies whether a keypoint is present or not.

        Args:
            masks (torch.Tensor): Binary mask tensor indicating object presence, shape (BS, N_anchors).
            target_gt_idx (torch.Tensor): Index tensor mapping anchors to ground truth objects, shape (BS, N_anchors).
            keypoints (torch.Tensor): Ground truth keypoints, shape (N_kpts_in_batch, N_kpts_per_object, kpts_dim).
            batch_idx (torch.Tensor): Batch index tensor for keypoints, shape (N_kpts_in_batch, 1).
            stride_tensor (torch.Tensor): Stride tensor for anchors, shape (N_anchors, 1).
            target_bboxes (torch.Tensor): Ground truth boxes in (x1, y1, x2, y2) format, shape (BS, N_anchors, 4).
            pred_kpts (torch.Tensor): Predicted keypoints, shape (BS, N_anchors, N_kpts_per_object, kpts_dim).

        Returns:
            kpts_loss (torch.Tensor): The keypoints loss.
            kpts_obj_loss (torch.Tensor): The keypoints object loss.
            rle_loss (torch.Tensor): The RLE loss.
        """
        # Select target keypoints using inherited helper method
        selected_keypoints = self._select_target_keypoints(keypoints, batch_idx, target_gt_idx, masks)

        # Divide coordinates by stride
        selected_keypoints[..., :2] /= stride_tensor.view(1, -1, 1, 1)

        kpts_loss = 0
        kpts_obj_loss = 0
        rle_loss = 0

        if masks.any():
            target_bboxes /= stride_tensor
            gt_kpt = selected_keypoints[masks]
            area = xyxy2xywh(target_bboxes[masks])[:, 2:].prod(1, keepdim=True)
            pred_kpt = pred_kpts[masks]
            kpt_mask = gt_kpt[..., 2] != 0 if gt_kpt.shape[-1] == 3 else torch.full_like(gt_kpt[..., 0], True)
            kpts_loss = self.keypoint_loss(pred_kpt, gt_kpt, kpt_mask, area)  # pose loss

            if self.rle_loss is not None and (pred_kpt.shape[-1] == 4 or pred_kpt.shape[-1] == 5):
                rle_loss = self.calculate_rle_loss(pred_kpt, gt_kpt, kpt_mask)
                rle_loss = rle_loss.clamp(min=0)
            if pred_kpt.shape[-1] == 3 or pred_kpt.shape[-1] == 5:
                kpts_obj_loss = self.bce_pose(pred_kpt[..., 2], kpt_mask.float())  # keypoint obj loss

        return kpts_loss, kpts_obj_loss, rle_loss


class v8ClassificationLoss:
    """Criterion class for computing training losses for classification."""

    def __call__(self, preds: Any, batch: dict[str, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
        """Compute the classification loss between predictions and true labels."""
        preds = preds[1] if isinstance(preds, (list, tuple)) else preds
        loss = F.cross_entropy(preds, batch["cls"], reduction="mean")
        return loss, loss.detach()


class v8OBBLoss(v8DetectionLoss):
    """Calculates losses for object detection, classification, and box distribution in rotated YOLO models."""

    def __init__(self, model, tal_topk=10, tal_topk2: int | None = None):
        """Initialize v8OBBLoss with model, assigner, and rotated bbox loss; model must be de-paralleled."""
        super().__init__(model, tal_topk=tal_topk)
        self.assigner = RotatedTaskAlignedAssigner(
            topk=tal_topk,
            num_classes=self.nc,
            alpha=0.5,
            beta=6.0,
            stride=self.stride.tolist(),
            topk2=tal_topk2,
        )
        self.bbox_loss = RotatedBboxLoss(self.reg_max).to(self.device)

    def preprocess(self, targets: torch.Tensor, batch_size: int, scale_tensor: torch.Tensor) -> torch.Tensor:
        """Preprocess targets for oriented bounding box detection."""
        if targets.shape[0] == 0:
            out = torch.zeros(batch_size, 0, 6, device=self.device)
        else:
            i = targets[:, 0]  # image index
            _, counts = i.unique(return_counts=True)
            counts = counts.to(dtype=torch.int32)
            out = torch.zeros(batch_size, counts.max(), 6, device=self.device)
            for j in range(batch_size):
                matches = i == j
                if n := matches.sum():
                    bboxes = targets[matches, 2:]
                    bboxes[..., :4].mul_(scale_tensor)
                    out[j, :n] = torch.cat([targets[matches, 1:2], bboxes], dim=-1)
        return out

    def loss(self, preds: dict[str, torch.Tensor], batch: dict[str, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
        """Calculate and return the loss for oriented bounding box detection."""
        loss = torch.zeros(4, device=self.device)  # box, cls, dfl, angle
        pred_distri, pred_scores, pred_angle = (
            preds["boxes"].permute(0, 2, 1).contiguous(),
            preds["scores"].permute(0, 2, 1).contiguous(),
            preds["angle"].permute(0, 2, 1).contiguous(),
        )
        anchor_points, stride_tensor = make_anchors(preds["feats"], self.stride, 0.5)
        batch_size = pred_angle.shape[0]  # batch size

        dtype = pred_scores.dtype
        imgsz = torch.tensor(preds["feats"][0].shape[2:], device=self.device, dtype=dtype) * self.stride[0]

        # targets
        try:
            batch_idx = batch["batch_idx"].view(-1, 1)
            targets = torch.cat((batch_idx, batch["cls"].view(-1, 1), batch["bboxes"].view(-1, 5)), 1)
            rw, rh = targets[:, 4] * float(imgsz[1]), targets[:, 5] * float(imgsz[0])
            targets = targets[(rw >= 2) & (rh >= 2)]  # filter rboxes of tiny size to stabilize training
            targets = self.preprocess(targets.to(self.device), batch_size, scale_tensor=imgsz[[1, 0, 1, 0]])
            gt_labels, gt_bboxes = targets.split((1, 5), 2)  # cls, xywhr
            mask_gt = gt_bboxes.sum(2, keepdim=True).gt_(0.0)
        except RuntimeError as e:
            raise TypeError(
                "ERROR ❌ OBB dataset incorrectly formatted or not a OBB dataset.\n"
                "This error can occur when incorrectly training a 'OBB' model on a 'detect' dataset, "
                "i.e. 'yolo train model=yolo26n-obb.pt data=dota8.yaml'.\nVerify your dataset is a "
                "correctly formatted 'OBB' dataset using 'data=dota8.yaml' "
                "as an example.\nSee https://docs.ultralytics.com/datasets/obb/ for help."
            ) from e

        # Pboxes
        pred_bboxes = self.bbox_decode(anchor_points, pred_distri, pred_angle)  # xyxy, (b, h*w, 4)

        bboxes_for_assigner = pred_bboxes.clone().detach()
        # Only the first four elements need to be scaled
        bboxes_for_assigner[..., :4] *= stride_tensor
        _, target_bboxes, target_scores, fg_mask, _ = self.assigner(
            pred_scores.detach().sigmoid(),
            bboxes_for_assigner.type(gt_bboxes.dtype),
            anchor_points * stride_tensor,
            gt_labels,
            gt_bboxes,
            mask_gt,
        )

        target_scores_sum = max(target_scores.sum(), 1)

        # Cls loss
        # loss[1] = self.varifocal_loss(pred_scores, target_scores, target_labels) / target_scores_sum  # VFL way
        loss[1] = self.bce(pred_scores, target_scores.to(dtype)).sum() / target_scores_sum  # BCE

        # Bbox loss
        if fg_mask.sum():
            target_bboxes[..., :4] /= stride_tensor
            loss[0], loss[2] = self.bbox_loss(
                pred_distri,
                pred_bboxes,
                anchor_points,
                target_bboxes,
                target_scores,
                target_scores_sum,
                fg_mask,
                imgsz,
                stride_tensor,
            )
            weight = target_scores.sum(-1)[fg_mask]
            loss[3] = self.calculate_angle_loss(
                pred_bboxes, target_bboxes, fg_mask, weight, target_scores_sum
            )  # angle loss
        else:
            loss[0] += (pred_angle * 0).sum()

        loss[0] *= self.hyp.box  # box gain
        loss[1] *= self.hyp.cls  # cls gain
        loss[2] *= self.hyp.dfl  # dfl gain
        loss[3] *= self.hyp.angle  # angle gain

        return loss * batch_size, loss.detach()  # loss(box, cls, dfl, angle)

    def bbox_decode(
        self, anchor_points: torch.Tensor, pred_dist: torch.Tensor, pred_angle: torch.Tensor
    ) -> torch.Tensor:
        """Decode predicted object bounding box coordinates from anchor points and distribution.

        Args:
            anchor_points (torch.Tensor): Anchor points, (h*w, 2).
            pred_dist (torch.Tensor): Predicted rotated distance, (bs, h*w, 4).
            pred_angle (torch.Tensor): Predicted angle, (bs, h*w, 1).

        Returns:
            (torch.Tensor): Predicted rotated bounding boxes with angles, (bs, h*w, 5).
        """
        if self.use_dfl:
            b, a, c = pred_dist.shape  # batch, anchors, channels
            pred_dist = pred_dist.view(b, a, 4, c // 4).softmax(3).matmul(self.proj.type(pred_dist.dtype))
        return torch.cat((dist2rbox(pred_dist, pred_angle, anchor_points), pred_angle), dim=-1)

    def calculate_angle_loss(self, pred_bboxes, target_bboxes, fg_mask, weight, target_scores_sum, lambda_val=3):
        """Calculate oriented angle loss.

        Args:
            pred_bboxes (torch.Tensor): Predicted bounding boxes with shape [N, 5] (x, y, w, h, theta).
            target_bboxes (torch.Tensor): Target bounding boxes with shape [N, 5] (x, y, w, h, theta).
            fg_mask (torch.Tensor): Foreground mask indicating valid predictions.
            weight (torch.Tensor): Loss weights for each prediction.
            target_scores_sum (torch.Tensor): Sum of target scores for normalization.
            lambda_val (int): Controls the sensitivity to aspect ratio.

        Returns:
            (torch.Tensor): The calculated angle loss.
        """
        w_gt = target_bboxes[..., 2]
        h_gt = target_bboxes[..., 3]
        pred_theta = pred_bboxes[..., 4]
        target_theta = target_bboxes[..., 4]

        log_ar = torch.log((w_gt + 1e-9) / (h_gt + 1e-9))
        scale_weight = torch.exp(-(log_ar**2) / (lambda_val**2))

        delta_theta = pred_theta - target_theta
        delta_theta_wrapped = delta_theta - torch.round(delta_theta / math.pi) * math.pi
        ang_loss = torch.sin(2 * delta_theta_wrapped[fg_mask]) ** 2

        ang_loss = scale_weight[fg_mask] * ang_loss
        ang_loss = ang_loss * weight

        return ang_loss.sum() / target_scores_sum


class E2EDetectLoss:
    """Criterion class for computing training losses for end-to-end detection."""

    def __init__(self, model):
        """Initialize E2EDetectLoss with one-to-many and one-to-one detection losses using the provided model."""
        self.one2many = v8DetectionLoss(model, tal_topk=10)
        self.one2one = v8DetectionLoss(model, tal_topk=1)

    def __call__(self, preds: Any, batch: dict[str, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
        """Calculate the sum of the loss for box, cls and dfl multiplied by batch size."""
        preds = preds[1] if isinstance(preds, tuple) else preds
        one2many = preds["one2many"]
        loss_one2many = self.one2many(one2many, batch)
        one2one = preds["one2one"]
        loss_one2one = self.one2one(one2one, batch)
        return loss_one2many[0] + loss_one2one[0], loss_one2many[1] + loss_one2one[1]


class E2ELoss:
    """Criterion class for computing training losses for end-to-end detection."""

    def __init__(self, model, loss_fn=v8DetectionLoss):
        """Initialize E2ELoss with one-to-many and one-to-one detection losses using the provided model."""
        self.one2many = loss_fn(model, tal_topk=10)
        self.one2one = loss_fn(model, tal_topk=7, tal_topk2=1)
        self.updates = 0
        self.total = 1.0
        # init gain
        self.o2m = 0.8
        self.o2o = self.total - self.o2m
        self.o2m_copy = self.o2m
        # final gain
        self.final_o2m = 0.1

    def __call__(self, preds: Any, batch: dict[str, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
        """Calculate the sum of the loss for box, cls and dfl multiplied by batch size."""
        preds = self.one2many.parse_output(preds)
        one2many, one2one = preds["one2many"], preds["one2one"]
        loss_one2many = self.one2many.loss(one2many, batch)
        loss_one2one = self.one2one.loss(one2one, batch)
        return loss_one2many[0] * self.o2m + loss_one2one[0] * self.o2o, loss_one2one[1]

    def update(self) -> None:
        """Update the weights for one-to-many and one-to-one losses based on the decay schedule."""
        self.updates += 1
        self.o2m = self.decay(self.updates)
        self.o2o = max(self.total - self.o2m, 0)

    def decay(self, x) -> float:
        """Calculate the decayed weight for one-to-many loss based on the current update step."""
        return max(1 - x / max(self.one2one.hyp.epochs - 1, 1), 0) * (self.o2m_copy - self.final_o2m) + self.final_o2m


class TVPDetectLoss:
    """Criterion class for computing training losses for text-visual prompt detection."""

    def __init__(self, model, tal_topk=10, tal_topk2: int | None = None):
        """Initialize TVPDetectLoss with task-prompt and visual-prompt criteria using the provided model."""
        self.vp_criterion = v8DetectionLoss(model, tal_topk, tal_topk2)
        # NOTE: store following info as it's changeable in __call__
        self.hyp = self.vp_criterion.hyp
        self.ori_nc = self.vp_criterion.nc
        self.ori_no = self.vp_criterion.no
        self.ori_reg_max = self.vp_criterion.reg_max

    def parse_output(self, preds) -> dict[str, torch.Tensor]:
        """Parse model predictions to extract features."""
        return self.vp_criterion.parse_output(preds)

    def __call__(self, preds: Any, batch: dict[str, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
        """Calculate the loss for text-visual prompt detection."""
        return self.loss(self.parse_output(preds), batch)

    def loss(self, preds: dict[str, torch.Tensor], batch: dict[str, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
        """Calculate the loss for text-visual prompt detection."""
        if self.ori_nc == preds["scores"].shape[1]:
            loss = torch.zeros(3, device=self.vp_criterion.device, requires_grad=True)
            return loss, loss.detach()

        preds["scores"] = self._get_vp_features(preds)
        vp_loss = self.vp_criterion(preds, batch)
        box_loss = vp_loss[0][1]
        return box_loss, vp_loss[1]

    def _get_vp_features(self, preds: dict[str, torch.Tensor]) -> list[torch.Tensor]:
        """Extract visual-prompt features from the model output."""
        scores = preds["scores"]
        vnc = scores.shape[1]

        self.vp_criterion.nc = vnc
        self.vp_criterion.no = vnc + self.vp_criterion.reg_max * 4
        self.vp_criterion.assigner.num_classes = vnc
        return scores


class TVPSegmentLoss(TVPDetectLoss):
    """Criterion class for computing training losses for text-visual prompt segmentation."""

    def __init__(self, model, tal_topk=10):
        """Initialize TVPSegmentLoss with task-prompt and visual-prompt criteria using the provided model."""
        super().__init__(model)
        self.vp_criterion = v8SegmentationLoss(model, tal_topk)
        self.hyp = self.vp_criterion.hyp

    def __call__(self, preds: Any, batch: dict[str, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
        """Calculate the loss for text-visual prompt segmentation."""
        return self.loss(self.parse_output(preds), batch)

    def loss(self, preds: Any, batch: dict[str, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
        """Calculate the loss for text-visual prompt segmentation."""
        if self.ori_nc == preds["scores"].shape[1]:
            loss = torch.zeros(4, device=self.vp_criterion.device, requires_grad=True)
            return loss, loss.detach()

        preds["scores"] = self._get_vp_features(preds)
        vp_loss = self.vp_criterion(preds, batch)
        cls_loss = vp_loss[0][2]
        return cls_loss, vp_loss[1]