DEBUG:rasterio.env:New GDAL environment <rasterio._env.GDALEnv object at 0x3421667c0> created
DEBUG:rasterio._env:GDAL_DATA found in environment.
DEBUG:rasterio._env:PROJ_DATA found in environment.
DEBUG:rasterio._env:Started GDALEnv: self=<rasterio._env.GDALEnv object at 0x3421667c0>.
DEBUG:rasterio.env:Entered env context: <rasterio.env.Env object at 0x342166910>
DEBUG:rasterio._io:Output nodata value read from file: -32768.0
DEBUG:rasterio._io:Output nodata values: [-32768.0]
DEBUG:rasterio._io:all_valid: False
DEBUG:rasterio._io:mask_flags: ([<MaskFlags.nodata: 8>],)
DEBUG:rasterio._io:Jump straight to _read()
DEBUG:rasterio._io:Window: Window(col_off=14336, row_off=4096, width=2048, height=2048)
DEBUG:rasterio._io:IO window xoff=14336.0 yoff=4096.0 width=2048.0 height=2048.0
DEBUG:rasterio.env:Exiting env context: <rasterio.env.Env object at 0x342166910>
DEBUG:rasterio.env:Cleared existing <rasterio._env.GDALEnv object at 0x3421667c0> options
DEBUG:rasterio._env:Stopped GDALEnv <rasterio._env.GDALEnv object at 0x3421667c0>.
DEBUG:rasterio.env:Exiting outermost env
DEBUG:rasterio.env:Exited env context: <rasterio.env.Env object at 0x342166910>
Using encoder feature at index 2
[DEBUG] Channel 0 shape: (2048, 2048)
[DEBUG] Channel 1 shape: (2048, 2048)
x.shape after concat at stage 1: torch.Size([4, 576, 256, 256])
(2007.0, 1786.0, 0.0) [<class 'float'>, <class 'float'>, <class 'float'>]
(2044.0, 1831.0, 0.0) [<class 'float'>, <class 'float'>, <class 'float'>]
(1856.0, 1345.0, 1.9375784397125244) [<class 'float'>, <class 'float'>, <class 'float'>]
[DEBUG] Decoder stage 2
Using encoder feature at index 1
x.shape after concat at stage 2: torch.Size([4, 424, 512, 512])
[DEBUG] Decoder stage 3
Using encoder feature at index 0
encoder_feature.shape after resizing at stage 3: torch.Size([4, 24, 1024, 1024])
x.shape after concat at stage 3: torch.Size([4, 280, 1024, 1024])
[DEBUG] feature_maps in FPNFuse: type=<class 'list'>
[DEBUG] Valid feature_maps[0] shape: torch.Size([4, 512, 128, 128])
[DEBUG] Valid feature_maps[1] shape: torch.Size([4, 384, 256, 256])
[DEBUG] Valid feature_maps[2] shape: torch.Size([4, 256, 512, 512])
[DEBUG] Valid feature_maps[3] shape: torch.Size([4, 128, 1024, 1024])
yolo_output shape: torch.Size([4, 21, 128, 128])
objectness_map shape: torch.Size([4, 1, 256, 256])
vessel_map shape: torch.Size([4, 1, 256, 256])
fishing_map shape: torch.Size([4, 1, 256, 256])
size_map shape: torch.Size([4, 1, 256, 256])
offset_map shape: torch.Size([4, 2, 256, 256])
yolo_output shape: torch.Size([4, 21, 128, 128])
classification_output shape: torch.Size([4, 2, 128, 128])
yolo_output upsampled to: torch.Size([4, 21, 256, 256])
classification_output upsampled to: torch.Size([4, 2, 256, 256])
[2024-09-30 18:20:48,826][__main__][ERROR] - Expected keys 'CENTERNET_OUTPUT_YOLO' and 'CENTERNET_TARGET_YOLO' not found in model outputs or encoded targets.
DEBUG:rasterio.env:Entering env context: <rasterio.env.Env object at 0x349885460>
DEBUG:rasterio.env:Starting outermost env
DEBUG:rasterio.env:No GDAL environment exists
DEBUG:rasterio.env:New GDAL environment <rasterio._env.GDALEnv object at 0x3498b4e80> created
DEBUG:rasterio._env:GDAL_DATA found in environment.
DEBUG:rasterio._env:PROJ_DATA found in environment.
DEBUG:rasterio._env:Started GDALEnv: self=<rasterio._env.GDALEnv object at 0x3498b4e80>.
DEBUG:rasterio.env:Entered env context: <rasterio.env.Env object at 0x349885460>
DEBUG:rasterio._base:Sharing flag: 0
DEBUG:rasterio._base:Nodata success: 1, Nodata value: -32768.000000
DEBUG:rasterio._base:Dataset <open DatasetReader name='/Volumes/Crucial/xview3_dataset/full_dataset/validation/4a97701b4bd81bf7v/VH_dB.tif' mode='r'> is started.
DEBUG:rasterio.env:Exiting env context: <rasterio.env.Env object at 0x349885460>
DEBUG:rasterio.env:Cleared existing <rasterio._env.GDALEnv object at 0x3498b4e80> options
DEBUG:rasterio._env:Stopped GDALEnv <rasterio._env.GDALEnv object at 0x3498b4e80>.
DEBUG:rasterio.env:Exiting outermost env

here're some files to review to fix the issues

#circlenet_loss.py

from functools import partial
from typing import Dict, List, Tuple

import torch
from pytorch_toolbelt.losses import SoftBCEWithLogitsLoss, BinaryBiTemperedLogisticLoss, BinaryFocalLoss, BalancedBCEWithLogitsLoss
from torch import nn, Tensor
from torch.nn import functional as F

from .functional import binary_shannon_entropy_loss, objectness_reduced_focal_loss, objectness_mse_loss
from ..constants import (
    CENTERNET_OUTPUT_SIZE,
    CENTERNET_OUTPUT_OFFSET,
    CENTERNET_OUTPUT_OBJECTNESS_MAP,
    CENTERNET_TARGET_SIZE,
    CENTERNET_TARGET_OFFSET,
    CENTERNET_TARGET_FISHING_MAP,
    CENTERNET_OUTPUT_FISHING_MAP,
    CENTERNET_OUTPUT_VESSEL_MAP,
    CENTERNET_TARGET_VESSEL_MAP,
    CENTERNET_TARGET_OBJECTNESS_MAP,
    CENTERNET_TARGET_YOLO,
    CENTERNET_OUTPUT_YOLO
)
from ...constants import IGNORE_LABEL, TARGET_SAMPLE_WEIGHT

__all__ = ["CircleNetLoss"]
from xview3.utils import TorchDevice
device = TorchDevice.choose_torch_device()

class CircleNetLoss(nn.Module):
    def __init__(self, objectness_loss: str = "rfl", classifier_loss: str = "bce", size_loss: str = "mse", offset_loss: str = "mse", regularize_ignored: bool = False, regularization_weight: float = 1.0, ignore_index: int = IGNORE_LABEL, num_classes: int = 20, anchors: list = [[10, 13], [16, 30], [33, 23]]):
        super().__init__()
        self.ignore_index = ignore_index  # Ensure ignore_index is set
        self.objectness_loss = partial(objectness_mse_loss, ignore_index=ignore_index) if objectness_loss == "mse" else partial(objectness_reduced_focal_loss, ignore_index=ignore_index)
        self.classifier_loss = SoftBCEWithLogitsLoss(ignore_index=ignore_index, reduction="none")
        self.size_loss = mse_loss if size_loss == "mse" else smooth_l1_loss
        self.offset_loss = mse_loss if offset_loss == "mse" else l1_loss
        self.regularize_ignored = regularize_ignored
        self.regularization_weight = regularization_weight

        if objectness_loss == "rfl":
            self.objectness_loss = partial(objectness_reduced_focal_loss, ignore_index=self.ignore_index, alpha=2, beta=4, eps=0)
        elif objectness_loss == "mse":
            self.objectness_loss = partial(objectness_mse_loss, ignore_index=self.ignore_index)
        else:
            raise KeyError(objectness_loss)

        if classifier_loss == "bce":
            self.classifier_loss = SoftBCEWithLogitsLoss(ignore_index=ignore_index, reduction="none")
        elif classifier_loss == "bbce":
            self.classifier_loss = BalancedBCEWithLogitsLoss(ignore_index=ignore_index, reduction="none")
        elif classifier_loss == "soft_bce":
            self.classifier_loss = SoftBCEWithLogitsLoss(ignore_index=ignore_index, reduction="none", smooth_factor=0.05)
        elif classifier_loss == "bitempered":
            self.classifier_loss = BinaryBiTemperedLogisticLoss(t1=0.5, t2=4.0, ignore_index=ignore_index, reduction="none")
        elif classifier_loss == "focal":
            self.classifier_loss = BinaryFocalLoss(ignore_index=ignore_index, alpha=None, reduction="none")
        else:
            raise KeyError(classifier_loss)

        if size_loss == "huber":
            self.size_loss = huber_loss
        elif size_loss == "mse":
            self.size_loss = mse_loss
        elif size_loss == "smooth_l1":
            self.size_loss = smooth_l1_loss
        elif size_loss == "rel_error":
            self.size_loss = relative_length_error
        else:
            raise KeyError(size_loss)

        if offset_loss == "mse":
            self.offset_loss = mse_loss
        elif offset_loss == "l1":
            self.offset_loss = l1_loss
        elif offset_loss == "smooth_l1":
            self.offset_loss = smooth_l1_loss
        elif offset_loss == "huber":
            self.offset_loss = huber_loss
        elif offset_loss is None:
            self.offset_loss = None
        else:
            raise KeyError(offset_loss)

        # Initialize YOLO loss module
        self.yolo_loss_module = YOLOLoss(num_classes=num_classes, anchors=anchors)

    def forward(self, **kwargs) -> Dict[str, Tensor]:
        gt_objectness: Tensor = kwargs[CENTERNET_TARGET_OBJECTNESS_MAP]
        gt_is_vessel: Tensor = kwargs[CENTERNET_TARGET_VESSEL_MAP]
        gt_is_fishing: Tensor = kwargs[CENTERNET_TARGET_FISHING_MAP]
        gt_size: Tensor = kwargs[CENTERNET_TARGET_SIZE]
        gt_offset: Tensor = kwargs.get(CENTERNET_TARGET_OFFSET, torch.zeros_like(gt_size))
        pred_size: Tensor = kwargs[CENTERNET_OUTPUT_SIZE]

        if pred_size.shape[-2:] != gt_size.shape[-2:]:
            pred_size = F.interpolate(pred_size, size=gt_size.shape[-2:], mode='bilinear', align_corners=False)

        # Initialize sample_weights with the correct shape
        batch_size = gt_objectness.size(0)
        sample_weights = kwargs.get(
            TARGET_SAMPLE_WEIGHT,
            torch.ones(batch_size, 1, gt_size.shape[2], gt_size.shape[3], device=gt_objectness.device)
        )

        pred_objectness: Tensor = kwargs.get(CENTERNET_OUTPUT_OBJECTNESS_MAP, torch.zeros_like(gt_objectness))
        pred_is_vessel: Tensor = kwargs.get(CENTERNET_OUTPUT_VESSEL_MAP, torch.zeros_like(gt_is_vessel))
        pred_is_fishing: Tensor = kwargs.get(CENTERNET_OUTPUT_FISHING_MAP, torch.zeros_like(gt_is_fishing))
        pred_size: Tensor = kwargs.get(CENTERNET_OUTPUT_SIZE, torch.zeros_like(gt_size))
        pred_offset: Tensor = kwargs.get(CENTERNET_OUTPUT_OFFSET, torch.zeros_like(gt_offset))

        # Ensure pred_is_vessel has the same shape as gt_is_vessel by upsampling if necessary
        if pred_is_vessel.shape[-2:] != gt_is_vessel.shape[-2:]:
            pred_is_vessel = F.interpolate(pred_is_vessel, size=gt_is_vessel.shape[-2:], mode='bilinear', align_corners=False)

        # Ensure pred_is_fishing has the same shape as gt_is_fishing by upsampling if necessary
        if pred_is_fishing.shape[-2:] != gt_is_fishing.shape[-2:]:
            pred_is_fishing = F.interpolate(pred_is_fishing, size=gt_is_fishing.shape[-2:], mode='bilinear', align_corners=False)

        if pred_objectness.shape[-2:] != gt_objectness.shape[-2:]:
            pred_objectness = F.interpolate(pred_objectness, size=gt_objectness.shape[-2:], mode='bilinear', align_corners=False)

        # Mask and loss calculation
        peaks_mask = gt_objectness.eq(1)
        true_vessels_mask = gt_is_vessel.eq(1)
        num_objects = peaks_mask.sum(dtype=torch.float32).clamp_min(1.0)
        num_vessels = (true_vessels_mask & peaks_mask).sum(dtype=torch.float32).clamp_min(1.0)

        uncertain_objectness_mask = gt_objectness.eq(self.ignore_index)
        certain_objectness_mask = ~uncertain_objectness_mask

        # Compute classification losses
        vessel_loss = self.classifier_loss(pred_is_vessel, gt_is_vessel) * gt_objectness * certain_objectness_mask
        fishing_loss = self.classifier_loss(pred_is_fishing, gt_is_fishing) * gt_objectness * certain_objectness_mask
        pos_objectness_loss, neg_objectness_loss = self.objectness_loss(pred_objectness, gt_objectness, pos_mask=peaks_mask)

        if self.regularize_ignored:
            uncertain_vessel_mask = gt_is_vessel.eq(self.ignore_index)
            uncertain_fishing_mask = (gt_is_fishing.eq(self.ignore_index) & gt_is_vessel.eq(1)) | uncertain_vessel_mask
            # Plus regularization loss for ignored values in the objectness heatmap
            obj_reg_loss = binary_shannon_entropy_loss(pred_objectness) * uncertain_objectness_mask

            vsl_reg_loss = binary_shannon_entropy_loss(pred_is_vessel) * uncertain_vessel_mask

            fsh_reg_loss = binary_shannon_entropy_loss(pred_is_fishing) * uncertain_fishing_mask

            obj_reg_loss = (
                obj_reg_loss.sum(dtype=torch.float32)
                * self.regularization_weight
                / uncertain_objectness_mask.sum(dtype=torch.float32).clamp_min(1.0)
            )

            vsl_reg_loss = (
                vsl_reg_loss.sum(dtype=torch.float32)
                * self.regularization_weight
                / uncertain_vessel_mask.sum(dtype=torch.float32).clamp_min(1.0)
            )
            fsh_reg_loss = (
                fsh_reg_loss.sum(dtype=torch.float32)
                * self.regularization_weight
                / uncertain_fishing_mask.sum(dtype=torch.float32).clamp_min(1.0)
            )
        else:
            obj_reg_loss = 0
            vsl_reg_loss = 0
            fsh_reg_loss = 0

        valid_size_mask = peaks_mask * (gt_size > 0)
        valid_size_mask = gt_size > 0  # Size loss only for positive lengths
        num_objects_with_length = valid_size_mask.sum(dtype=torch.float32).clamp_min(1.0)

        # Reshape size_loss to match the shape of sample_weights
        size_loss = self.size_loss(pred_size, gt_size, valid_size_mask)
        size_loss = size_loss.view_as(sample_weights)

        # Ensure sample_weights has the correct shape
        height, width = vessel_loss.shape[2], vessel_loss.shape[3]
        sample_weights = sample_weights.expand(-1, -1, height, width)

        print(f"vessel_loss shape: {vessel_loss.shape}")
        print(f"sample_weights shape: {sample_weights.shape}")
        print(f"size_loss shape: {size_loss.shape}")

        weighted_vessel_loss = sample_weights * vessel_loss
        weighted_fishing_loss = sample_weights * fishing_loss
        weighted_size_loss = sample_weights * size_loss

        vessel_loss_reduced = weighted_vessel_loss.sum(dim=[2, 3])  # Shape: [batch_size, 1]
        fishing_loss_reduced = weighted_fishing_loss.sum(dim=[2, 3])  # Shape: [batch_size, 1]
        size_loss_reduced = weighted_size_loss.sum(dim=[2, 3])  # Shape: [batch_size, 1]

        print(f"weighted_vessel_loss shape: {weighted_vessel_loss.shape}")
        print(f"weighted_fishing_loss shape: {weighted_fishing_loss.shape}")
        print(f"weighted_size_loss shape: {weighted_size_loss.shape}")

        objectness_loss = neg_objectness_loss + pos_objectness_loss
        reduced_obj_loss = objectness_loss.sum(dim=[2, 3])  # Shape: [batch_size, 1]

        # Ensure reduced_obj_loss has the same spatial dimensions as sample_weights
        reduced_obj_loss = reduced_obj_loss.unsqueeze(-1).unsqueeze(-1).expand_as(sample_weights)

        # Now calculate mean_obj_loss
        mean_obj_loss = (sample_weights * reduced_obj_loss).sum() / num_objects + obj_reg_loss

        mean_vessel_loss = weighted_vessel_loss.sum() / num_objects + vsl_reg_loss
        mean_fishing_loss = weighted_fishing_loss.sum() / num_vessels + fsh_reg_loss
        mean_size_loss = weighted_size_loss.sum() / num_objects_with_length

        # Assume yolo_output and yolo_targets are obtained correctly at the call site
        yolo_output = kwargs.get(CENTERNET_OUTPUT_YOLO)
        if isinstance(yolo_output, dict):
            yolo_output = yolo_output['output']  # Assuming 'output' key contains the tensor
        elif isinstance(yolo_output, list):
            yolo_output = torch.cat(yolo_output, dim=1)  # Assuming list contains tensors to be concatenated along channel axis

        yolo_targets = kwargs.get(CENTERNET_TARGET_YOLO)
        if isinstance(yolo_targets, dict):
            yolo_targets = yolo_targets['output']  # Assuming 'output' key contains the tensor
        elif isinstance(yolo_targets, list):
            yolo_targets = torch.cat(yolo_targets, dim=1)  # Assuming list contains tensors to be concatenated along channel axis

        if not isinstance(yolo_output, Tensor) or not isinstance(yolo_targets, Tensor):
            raise ValueError("YOLO outputs and targets must be tensors.")

        # The following line will now work, given the tensors are passed correctly
        if yolo_output.shape[1] != yolo_targets.shape[1]:
            raise ValueError("Channel mismatch: predictions and targets must have the same number of channels.")
        # Calculate YOLO loss using the YOLO loss module
        yolo_loss = self.yolo_loss_module(yolo_output, yolo_targets)

        losses = {
            "objectness": mean_obj_loss.type_as(pred_objectness),
            "vessel": mean_vessel_loss.type_as(pred_objectness),
            "fishing": mean_fishing_loss.type_as(pred_objectness),
            "size": mean_size_loss.type_as(pred_objectness),
            "vsl_reg_loss": vsl_reg_loss,
            "fsh_reg_loss": fsh_reg_loss,
            "obj_reg_loss": obj_reg_loss,
            "num_objects": num_objects,
            "num_vessels": num_vessels,
            "num_objects_with_length": num_objects_with_length,
            "yolo_loss": yolo_loss.type_as(pred_objectness),
        }

        if self.offset_loss is not None:
            offset_loss = self.offset_loss(pred_offset, gt_offset, peaks_mask).sum(
                dim=1, keepdim=True
            )
            mean_off_loss = offset_loss.sum(dtype=torch.float32) / num_objects
            losses["offset"] = mean_off_loss.type_as(pred_objectness)

        return losses


def mse_loss(prediction, target, pos_mask):
    """MSE loss with a mask for valid objects."""
    prediction = F.interpolate(prediction, size=target.shape[-2:], mode='bilinear', align_corners=False)
    loss = F.mse_loss(prediction, target, reduction="none") * pos_mask
    return loss


def l1_loss(prediction, target, pos_mask):
    loss = F.l1_loss(prediction, target, reduction="none") * pos_mask
    return loss


def smooth_l1_loss(prediction, target, pos_mask):
    loss = F.smooth_l1_loss(prediction, target, reduction="none") * pos_mask
    return loss


def huber_loss(prediction, target, pos_mask):
    loss = F.huber_loss(prediction, target, reduction="none", delta=1) * pos_mask
    return loss


def relative_length_error(prediction, target, pos_mask, eps=1e-4):
    loss = pos_mask * F.smooth_l1_loss(prediction, target, reduction="none") / (target + eps)
    return loss


class YOLOLoss(nn.Module):
    def __init__(self, num_classes: int, anchors: list):
        super(YOLOLoss, self).__init__()
        self.num_classes = num_classes
        self.anchors = anchors
        self.mse_loss = nn.MSELoss(reduction='sum')  # Consider using 'mean' if preferred
        self.bce_loss = nn.BCELoss(reduction='sum')  # Consider using 'mean' if preferred

    def calculate_iou(self, box1: torch.Tensor, box2: torch.Tensor) -> torch.Tensor:
        """Calculate Intersection Over Union (IoU) between two sets of boxes."""
        x1y1 = torch.max(box1[..., :2], box2[..., :2])
        x2y2 = torch.min(box1[..., 2:], box2[..., 2:])
        intersection = torch.clamp((x2y2 - x1y1), min=0).prod(-1)
        area_box1 = (box1[..., 2] - box1[..., 0]) * (box1[..., 3] - box1[..., 1])
        area_box2 = (box2[..., 2] - box2[..., 0]) * (box2[..., 3] - box2[..., 1])
        union = area_box1 + area_box2 - intersection
        return intersection / union

    def forward(self, yolo_pred: torch.Tensor, yolo_targets: torch.Tensor) -> torch.Tensor:
        # Ensure predictions and targets have the same number of channels
        if yolo_pred.shape[1] != yolo_targets.shape[1]:
            raise ValueError(f"Channel mismatch: predictions have {yolo_pred.shape[1]} channels, targets have {yolo_targets.shape[1]} channels.")
        
        # Ensure spatial dimensions match
        if yolo_pred.shape[-2:] != yolo_targets.shape[-2:]:
            yolo_pred = F.interpolate(yolo_pred, size=yolo_targets.shape[-2:], mode='bilinear', align_corners=False)

        # Extract bounding boxes, confidence, and class predictions from YOLO output
        bbox_pred = yolo_pred[:, 0:4, :, :]  # [batch, 4, H, W]
        conf_pred = yolo_pred[:, 4:5, :, :]  # [batch, 1, H, W]
        class_pred = yolo_pred[:, 5:, :, :]  # [batch, num_classes, H, W]

        bbox_targets = yolo_targets[:, 0:4, :, :]  # [batch, 4, H, W]
        conf_targets = yolo_targets[:, 4:5, :, :]  # [batch, 1, H, W]
        class_targets = yolo_targets[:, 5:, :, :]  # [batch, num_classes, H, W]

        # Compute losses: bbox (MSE), confidence (BCE), and class (BCE or CrossEntropy)
        bbox_loss = self.mse_loss(bbox_pred, bbox_targets)
        conf_loss = self.bce_loss(conf_pred, conf_targets)
        class_loss = self.bce_loss(class_pred, class_targets)

        # Combine losses with appropriate weighting
        total_loss = bbox_loss + conf_loss + class_loss

        return total_loss

#functional.py

from typing import Tuple

import torch
import torch.nn.functional as F
from pytorch_toolbelt.losses.functional import focal_loss_with_logits, softmax_focal_loss_with_logits
from torch import Tensor
from typing import Tuple

__all__ = [
    "objectness_mse_loss",
    "objectness_reduced_focal_loss",
    "reduced_focal_loss_per_image",
    "reduced_focal_loss",
    "squeeze_heatmap",
    "shannon_entropy_loss",
    "binary_shannon_entropy_loss",
]

def objectness_mse_loss(pred: Tensor, gt: Tensor, pos_mask, ignore_index) -> Tuple[Tensor, Tensor]:
    neg_mask = ~pos_mask
    ign_mask = gt.eq(ignore_index)

    pt = pred.sigmoid()

    loss = F.mse_loss(pt, gt, reduction="none")
    pos_loss: Tensor = loss * pos_mask
    neg_loss: Tensor = loss * neg_mask
    return torch.masked_fill(pos_loss, ign_mask, 0), torch.masked_fill(neg_loss, ign_mask, 0)


def objectness_reduced_focal_loss(pred: Tensor, gt: Tensor, pos_mask, ignore_index, alpha: float, beta: float, eps: float = 1e-4):
    neg_mask = ~pos_mask
    ign_mask = gt.eq(ignore_index)

    pt = pred.sigmoid().clamp(eps, 1 - eps)

    pos_loss: Tensor = -torch.pow(1 - pt, alpha) * F.logsigmoid(pred) * pos_mask
    neg_loss: Tensor = -torch.pow(1.0 - gt, beta) * torch.pow(pt, alpha) * F.logsigmoid(-pred) * neg_mask

    return torch.masked_fill(pos_loss, ign_mask, 0), torch.masked_fill(neg_loss, ign_mask, 0)


def squeeze_heatmap(heatmap: Tensor, keepdim: bool = True) -> Tensor:
    """

    Args:
        heatmap: [B,C,H,W]
        keepdim:

    Returns:
        Binary tensor of shape [B, 1, H, W], where True values corresponds to peaks in the heatmap
    """
    pos_mask = heatmap.eq(1).any(dim=1, keepdim=keepdim)
    return pos_mask


def reduced_focal_loss(pred: Tensor, gt: Tensor, alpha: float, beta: float, eps: float = 1e-4) -> Tensor:
    pos_inds = gt.eq(1.0)
    num_pos = pos_inds.sum()

    pos_mask = pos_inds
    neg_mask = ~pos_mask

    pt = pred.sigmoid().clamp(eps, 1 - eps)

    pos_loss: Tensor = torch.pow(1 - pt, alpha) * torch.nn.functional.logsigmoid(pred) * pos_mask
    neg_loss: Tensor = torch.pow(1.0 - gt, beta) * torch.pow(pt, alpha) * torch.nn.functional.logsigmoid(-pred) * neg_mask

    pos_loss = -pos_loss.sum(dtype=torch.float32)
    neg_loss = -neg_loss.sum(dtype=torch.float32)
    loss = (neg_loss + pos_loss) / num_pos.clamp_min(1)
    return loss


def reduced_focal_loss_per_image(pred: Tensor, gt: Tensor, alpha: float, beta: float, eps: float = 1e-4) -> Tensor:
    pos_inds = gt.eq(1.0)
    num_pos = pos_inds.sum(dtype=torch.float32, dim=(1, 2, 3))  # [B]

    pos_mask = pos_inds
    neg_mask = ~pos_mask

    pt = pred.sigmoid().clamp(eps, 1 - eps)

    pos_loss: Tensor = torch.pow(1 - pt, alpha) * torch.nn.functional.logsigmoid(pred) * pos_mask
    neg_loss: Tensor = torch.pow(1.0 - gt, beta) * torch.pow(pt, alpha) * torch.nn.functional.logsigmoid(-pred) * neg_mask

    pos_loss = -pos_loss.sum(dtype=torch.float32, dim=(1, 2, 3))
    neg_loss = -neg_loss.sum(dtype=torch.float32, dim=(1, 2, 3))
    loss = (neg_loss + pos_loss) / num_pos.clamp_min(1)
    return loss.mean()


def balanced_binary_cross_entropy_with_logits(outputs: Tensor, targets: Tensor, reduction: str = "mean", gamma=1.0) -> Tensor:
    targets = targets.float()
    outputs = outputs.float()

    one_minus_beta: Tensor = targets.mean()
    beta = 1.0 - one_minus_beta

    pos_term = beta * targets.pow(gamma) * torch.nn.functional.logsigmoid(outputs)
    neg_term = one_minus_beta * (1 - targets).pow(gamma) * torch.nn.functional.logsigmoid(-outputs)

    loss = -(pos_term + neg_term)

    if reduction == "mean":
        loss = loss.mean()

    if reduction == "sum":
        loss = loss.sum()

    return loss


def centernet_classification_bce_loss(prediction: Tensor, target_classes: Tensor) -> Tensor:
    # If there is no label, we ignore this pixel
    ignore_mask = target_classes.sum(dim=1, keepdim=True).eq(0)

    loss = F.binary_cross_entropy_with_logits(prediction, target_classes, reduction="none")
    keep_mask = (~ignore_mask).type_as(loss)
    num_pos = keep_mask.sum(dtype=torch.float32).clamp_min(1)
    loss = (loss * keep_mask).sum(dtype=torch.float32)
    return loss / num_pos


def centernet_classification_ce_loss(prediction: Tensor, target_classes: Tensor, weight_mask: Tensor) -> Tensor:
    # If two classes overlap or there is no label, we ignore this pixel
    target_labels = target_classes.argmax(dim=1)

    num_pos = weight_mask.sum().clamp_min(1)
    loss = F.cross_entropy(prediction, target_labels, reduction="none") * weight_mask
    return loss.sum() / num_pos


def cnet_cls_focal_heatmap(prediction: Tensor, target: Tensor, eps: float = 1e-6) -> Tensor:
    pos_mask = squeeze_heatmap(target)
    num_pos = pos_mask.sum().clamp_min(eps)

    loss = focal_loss_with_logits(prediction, target, alpha=None, gamma=2, reduction="none")
    loss_sum = (loss * pos_mask).sum()
    loss_mean = loss_sum / num_pos
    return loss_mean


def cnet_cls_softmax_focal_heatmap(prediction: Tensor, target: Tensor, eps: float = 1e-6) -> Tensor:
    target_labels = torch.argmax(target, dim=1)

    pos_mask = squeeze_heatmap(target)
    num_pos = pos_mask.sum().clamp_min(eps)

    loss = softmax_focal_loss_with_logits(prediction, target_labels, gamma=2, reduction="none")
    loss_sum = (loss * pos_mask).sum()
    loss_mean = loss_sum / num_pos
    return loss_mean


def shannon_entropy_loss(x, dim=1):
    sm = x.softmax(dim=dim)
    lsm = F.log_softmax(x, dim=dim)
    entropy = (sm * lsm).sum(dim=dim, keepdim=True)
    return -entropy


def binary_shannon_entropy_loss(x):
    pt = x.sigmoid()
    entropy = F.logsigmoid(x) * pt + F.logsigmoid(-x) * (1 - pt)
    return -entropy



#multilabel_trimmed_loss.py

from functools import partial
from typing import Dict, Tuple

import torch
from pytorch_toolbelt.losses import SoftBCEWithLogitsLoss, BinaryBiTemperedLogisticLoss, BinaryFocalLoss, BalancedBCEWithLogitsLoss
from torch import nn, Tensor
from torch.nn import functional as F

from .functional import binary_shannon_entropy_loss
from ..constants import (
    CENTERNET_OUTPUT_SIZE,
    CENTERNET_OUTPUT_OFFSET,
    CENTERNET_OUTPUT_OBJECTNESS_MAP,
    CENTERNET_TARGET_SIZE,
    CENTERNET_TARGET_OFFSET,
    CENTERNET_TARGET_FISHING_MAP,
    CENTERNET_OUTPUT_FISHING_MAP,
    CENTERNET_OUTPUT_VESSEL_MAP,
    CENTERNET_TARGET_VESSEL_MAP,
    CENTERNET_TARGET_OBJECTNESS_MAP,
    CENTERNET_OUTPUT_YOLO,
    CENTERNET_TARGET_YOLO
)
from ...constants import IGNORE_LABEL

__all__ = ["MultilabelCenterTrimmedLoss"]


class MultilabelCenterTrimmedLoss(nn.Module):
    fraction: int
    samplewise: bool
    regularize_ignored: bool

    def __init__(
        self,
        fraction: int,
        samplewise: bool = True,
        objectness_loss: str = "rfl",
        classifier_loss="bce",
        size_loss: str = "mse",
        offset_loss: str = "mse",
        regularize_ignored: bool = False,
        regularization_weight: float = 1.0,
        ignore_index: int = IGNORE_LABEL,
        num_classes: int = 20,
        anchors: list = [[10, 13], [16, 30], [33, 23]],
    ):
        super().__init__()
        self.fraction = fraction
        self.objectness_loss = objectness_loss
        self.samplewise = samplewise
        self.regularize_ignored = regularize_ignored
        self.regularization_weight = regularization_weight
        self.ignore_index = ignore_index

        if objectness_loss == "rfl":
            self.objectness_loss = partial(self.reduced_focal_loss, alpha=2, beta=4, eps=0)
        elif objectness_loss == "l1":
            self.objectness_loss = self.objectness_l1_loss
        elif objectness_loss == "mse":
            self.objectness_loss = self.objectness_mse_loss
        elif objectness_loss == "focal_mse":
            self.objectness_loss = self.objectness_focal_mse_loss
        else:
            raise KeyError(objectness_loss)

        if classifier_loss == "bce":
            self.classifier_loss = SoftBCEWithLogitsLoss(ignore_index=ignore_index, reduction="none")
        elif classifier_loss == "bbce":
            self.classifier_loss = BalancedBCEWithLogitsLoss(ignore_index=ignore_index, reduction="none")
        elif classifier_loss == "soft_bce":
            self.classifier_loss = SoftBCEWithLogitsLoss(ignore_index=ignore_index, reduction="none", smooth_factor=0.05)
        elif classifier_loss == "bitempered":
            self.classifier_loss = BinaryBiTemperedLogisticLoss(t1=0.5, t2=4.0, ignore_index=ignore_index, reduction="none")
        elif classifier_loss == "focal":
            self.classifier_loss = BinaryFocalLoss(ignore_index=ignore_index, alpha=None, reduction="none")
        else:
            raise KeyError(classifier_loss)

        if size_loss == "huber":
            self.size_loss = self.huber_loss
        elif size_loss == "mse":
            self.size_loss = self.mse_loss
        elif size_loss == "smooth_l1":
            self.size_loss = self.smooth_l1_loss
        elif size_loss == "rel_error":
            self.size_loss = self.relative_length_error
        else:
            raise KeyError(size_loss)

        if offset_loss == "mse":
            self.offset_loss = self.mse_loss
        elif offset_loss == "l1":
            self.offset_loss = self.l1_loss
        elif offset_loss == "smooth_l1":
            self.offset_loss = self.smooth_l1_loss
        elif offset_loss == "huber":
            self.offset_loss = self.huber_loss
        elif offset_loss is None:
            self.offset_loss = None
        else:
            raise KeyError(offset_loss)

        # Initialize YOLO loss module
        self.yolo_loss_module = YOLOLoss(num_classes=num_classes, anchors=anchors)

    def objectness_mse_loss(self, pred: Tensor, gt: Tensor, pos_mask) -> Tuple[Tensor, Tensor]:
        ign_mask = gt.eq(self.ignore_index)
        pt = pred.sigmoid()
        neg_mask = ~pos_mask
        loss = F.mse_loss(pt, gt, reduction="none")
        pos_loss: Tensor = loss * pos_mask
        neg_loss: Tensor = loss * neg_mask
        return torch.masked_fill(pos_loss, ign_mask, 0), torch.masked_fill(neg_loss, ign_mask, 0)

    def objectness_focal_mse_loss(self, pred: Tensor, gt: Tensor, pos_mask: Tensor, alpha=2.0, beta=2.0) -> Tuple[Tensor, Tensor]:
        pt = pred.sigmoid()
        neg_mask = ~pos_mask
        loss = F.mse_loss(pt, gt, reduction="none")
        pos_loss: Tensor = loss * pos_mask
        neg_loss: Tensor = loss * neg_mask * torch.pow(pt, alpha) * torch.pow(1.0 - gt, beta)
        return pos_loss, neg_loss

    def objectness_l1_loss(self, pred: Tensor, gt: Tensor, pos_mask):
        pt = pred.sigmoid()
        ign_mask = gt.eq(self.ignore_index)
        neg_mask = ~pos_mask
        loss = F.l1_loss(pt, gt, reduction="none")
        pos_loss: Tensor = loss * pos_mask
        neg_loss: Tensor = (loss * neg_mask) ** 4
        return torch.masked_fill(pos_loss, ign_mask, 0), torch.masked_fill(neg_loss, ign_mask, 0)

    def reduced_focal_loss(self, pred: Tensor, gt: Tensor, pos_mask, alpha: float, beta: float, eps: float = 1e-4):
        neg_mask = ~pos_mask
        ign_mask = gt.eq(self.ignore_index)

        pt = pred.sigmoid().clamp(eps, 1 - eps)

        pos_loss: Tensor = -torch.pow(1 - pt, alpha) * F.logsigmoid(pred) * pos_mask
        neg_loss: Tensor = -torch.pow(1.0 - gt, beta) * torch.pow(pt, alpha) * F.logsigmoid(-pred) * neg_mask

        return torch.masked_fill(pos_loss, ign_mask, 0), torch.masked_fill(neg_loss, ign_mask, 0)

    def mse_loss(self, prediction, target, pos_mask):
        loss = F.mse_loss(prediction, target, reduction="none") * pos_mask
        return loss

    def l1_loss(self, prediction, target, pos_mask):
        loss = F.l1_loss(prediction, target, reduction="none") * pos_mask
        return loss

    def smooth_l1_loss(self, prediction, target, pos_mask):
        loss = F.smooth_l1_loss(prediction, target, reduction="none") * pos_mask
        return loss

    def huber_loss(self, prediction, target, pos_mask):
        loss = F.huber_loss(prediction, target, reduction="none", delta=1) * pos_mask
        return loss

    def relative_length_error(self, prediction, target, pos_mask, eps=1e-4):
        loss = pos_mask * F.smooth_l1_loss(prediction, target, reduction="none") / (target + eps)
        return loss

    # @torch.cuda.amp.autocast(False)
    def forward(self, **kwargs) -> Dict[str, Tensor]:
        gt_objectness: Tensor = kwargs[CENTERNET_TARGET_OBJECTNESS_MAP]
        gt_is_vessel: Tensor = kwargs[CENTERNET_TARGET_VESSEL_MAP]
        gt_is_fishing: Tensor = kwargs[CENTERNET_TARGET_FISHING_MAP]
        gt_size: Tensor = kwargs[CENTERNET_TARGET_SIZE]
        gt_yolo: Tensor = kwargs(CENTERNET_TARGET_YOLO)
        if gt_yolo is None:
            raise ValueError(f"Missing key '{CENTERNET_TARGET_YOLO}' in input dictionary")

        pred_objectness: Tensor = kwargs[CENTERNET_OUTPUT_OBJECTNESS_MAP]
        pred_is_vessel: Tensor = kwargs[CENTERNET_OUTPUT_VESSEL_MAP]
        pred_is_fishing: Tensor = kwargs[CENTERNET_OUTPUT_FISHING_MAP]
        pred_size: Tensor = kwargs[CENTERNET_OUTPUT_SIZE]
        pred_yolo: Tensor = kwargs(CENTERNET_OUTPUT_YOLO, torch.zeros_like(pred_objectness))  # Default to zeros if not present

        peaks_mask = gt_objectness.eq(1)
        true_vessels_mask = gt_is_vessel.eq(1)

        num_objects = peaks_mask.sum(dtype=torch.float32).clamp_min(1.0)
        num_vessels = (true_vessels_mask & peaks_mask).sum(dtype=torch.float32).clamp_min(1.0)

        uncertain_objectness_mask = gt_objectness.eq(self.ignore_index)
        certain_objectness_mask = ~uncertain_objectness_mask

        # Returns unreduced losses
        pos_objectness_loss, neg_objectness_loss = self.objectness_loss(pred_objectness, gt_objectness, pos_mask=peaks_mask)

        # Weight classification loss by objectness (This make sense for target encoding=center).
        # This way pixels closer to the center of the keypoint will have bigger influence than ones on the periphery
        vessel_loss = self.classifier_loss(pred_is_vessel, gt_is_vessel) * gt_objectness * certain_objectness_mask
        fishing_loss = self.classifier_loss(pred_is_fishing, gt_is_fishing) * gt_objectness * certain_objectness_mask
        yolo_loss = self.yolo_loss_module(pred_yolo, gt_yolo)

        if self.regularize_ignored:
            uncertain_vessel_mask = gt_is_vessel.eq(self.ignore_index)
            uncertain_fishing_mask = gt_is_fishing.eq(self.ignore_index)
            # Plus regularization loss for ignored values in the objectness heatmap
            objectness_reg_loss = binary_shannon_entropy_loss(pred_objectness) * uncertain_objectness_mask

            vsl_reg_loss = binary_shannon_entropy_loss(pred_is_vessel) * certain_objectness_mask * uncertain_vessel_mask * gt_objectness
            fsh_reg_loss = (
                binary_shannon_entropy_loss(pred_is_fishing)
                * certain_objectness_mask
                * uncertain_fishing_mask
                * true_vessels_mask
                * gt_objectness
            )

            pos_objectness_loss += objectness_reg_loss * self.regularization_weight
            vessel_loss += vsl_reg_loss * self.regularization_weight
            fishing_loss += fsh_reg_loss * self.regularization_weight

        valid_size_mask = peaks_mask * (gt_size > 0)  # Size loss only for positive lengths
        num_objects_with_length = valid_size_mask.sum(dtype=torch.float32).clamp_min(1.0)
        size_loss = self.size_loss(pred_size, gt_size, valid_size_mask)

        if self.offset_loss is not None:
            offset_loss = self.offset_loss(kwargs[CENTERNET_OUTPUT_OFFSET], kwargs[CENTERNET_TARGET_OFFSET], peaks_mask).sum(
                dim=1, keepdim=True
            )
        else:
            offset_loss = torch.zeros_like(size_loss)

        if self.fraction > 0:
            if self.samplewise:
                (pos_objectness_loss, neg_objectness_loss, vessel_loss, fishing_loss, offset_loss, size_loss,) = samplewise_trimming(
                    pred_objectness,
                    pred_is_vessel,
                    pred_is_fishing,
                    pos_objectness_loss,
                    neg_objectness_loss,
                    vessel_loss,
                    fishing_loss,
                    offset_loss,
                    size_loss,
                )
            else:
                (pos_objectness_loss, neg_objectness_loss, vessel_loss, fishing_loss, offset_loss, size_loss,) = batchwise_trimming(
                    pred_objectness,
                    pred_is_vessel,
                    pred_is_fishing,
                    pos_objectness_loss,
                    neg_objectness_loss,
                    vessel_loss,
                    fishing_loss,
                    offset_loss,
                    size_loss,
                )

        mean_obj_loss = (neg_objectness_loss.sum(dtype=torch.float32) + pos_objectness_loss.sum(dtype=torch.float32)) / num_objects
        mean_vessel_loss = (vessel_loss.sum(dtype=torch.float32)) / num_objects
        mean_fishing_loss = (fishing_loss.sum(dtype=torch.float32)) / num_vessels
        mean_siz_loss = size_loss.sum(dtype=torch.float32) / num_objects_with_length

        losses = {
            "objectness": mean_obj_loss.type_as(pred_objectness),
            "vessel": mean_vessel_loss.type_as(pred_objectness),
            "fishing": mean_fishing_loss.type_as(pred_objectness),
            "size": mean_siz_loss.type_as(pred_objectness),
            "yolo_loss": yolo_loss.type_as(pred_objectness),  # Add YOLO loss to the losses dictionary
        }
        if self.offset_loss is not None:
            mean_off_loss = offset_loss.sum(dtype=torch.float32) / num_objects
            losses["offset"] = mean_off_loss.type_as(pred_objectness)
        return losses


def samplewise_trimming(
    pred_objectness,
    pred_is_vessel,
    pred_is_fishing,
    pos_objectness_loss,
    neg_objectness_loss,
    vessel_loss,
    fishing_loss,
    offset_loss,
    size_loss,
    num_elements_to_ignore,
):
    pred_objectness = pred_objectness.flatten(2)  # [B, 1, N]
    pred_is_vessel = pred_is_vessel.flatten(2)  # [B, 1, N]
    pred_is_fishing = pred_is_fishing.flatten(2)  # [B, 1, N]
    pos_objectness_loss = pos_objectness_loss.flatten(2)  # [B, 1, N]
    neg_objectness_loss = neg_objectness_loss.flatten(2)  # [B, 1, N]
    vessel_loss = vessel_loss.flatten(2)  # [B, 1, N]
    fishing_loss = fishing_loss.flatten(2)  # [B, 1, N]
    offset_loss = offset_loss.flatten(2)  # [B, 1, N]
    size_loss = size_loss.flatten(2)  # [B, 1, N]

    # Here we ignore top-N negative losses for negative objectness loss (E.g target is predicted, but groundtruth is negative).
    # Since some objects may be missing in gt, and we don't want to include those wrong false-positives in the training signal
    # And use keep indexes to propagate ignored pixels in other losses
    _, indices = torch.topk(neg_objectness_loss, k=num_elements_to_ignore, dim=2, largest=True, sorted=False)

    non_ignored_mask = torch.ones_like(neg_objectness_loss).scatter_(2, indices, 0)

    # Instead of just setting loss to zeros in those hard example predictions, we minimize entropy loss
    # for objectness & classification to push the predictions to either 0 or 1
    neg_objectness_loss = neg_objectness_loss * non_ignored_mask + binary_shannon_entropy_loss(pred_objectness) * (1 - non_ignored_mask)
    vessel_loss = vessel_loss * non_ignored_mask + binary_shannon_entropy_loss(pred_is_vessel) * (1 - non_ignored_mask)
    fishing_loss = fishing_loss * non_ignored_mask + binary_shannon_entropy_loss(pred_is_fishing) * (1 - non_ignored_mask)
    # For offset & size there is no regularization, so we exclude corresponding pixels from the grad. contribution
    offset_loss = offset_loss * non_ignored_mask
    size_loss = size_loss * non_ignored_mask

    return (
        pos_objectness_loss,
        neg_objectness_loss,
        vessel_loss,
        fishing_loss,
        offset_loss,
        size_loss,
    )


def batchwise_trimming(
    pred_objectness,
    pred_is_vessel,
    pred_is_fishing,
    pos_objectness_loss,
    neg_objectness_loss,
    vessel_loss,
    fishing_loss,
    offset_loss,
    size_loss,
    num_elements_to_ignore,
):
    pred_objectness = pred_objectness.permute(1, 0, 2, 3).flatten(1)  # [1, BxN]
    pred_is_vessel = pred_is_vessel.permute(1, 0, 2, 3).flatten(1)  # [C, BxN]
    pred_is_fishing = pred_is_fishing.permute(1, 0, 2, 3).flatten(1)  # [C, BxN]
    pos_objectness_loss = pos_objectness_loss.permute(1, 0, 2, 3).flatten(1)  # [1, BxN]
    neg_objectness_loss = neg_objectness_loss.permute(1, 0, 2, 3).flatten(1)  # [1, BxN]
    vessel_loss = vessel_loss.permute(1, 0, 2, 3).flatten(1)  # [1, BxN]
    fishing_loss = fishing_loss.permute(1, 0, 2, 3).flatten(1)  # [1, BxN]
    offset_loss = offset_loss.permute(1, 0, 2, 3).flatten(1)  # [1, BxN]
    size_loss = size_loss.permute(1, 0, 2, 3).flatten(1)  # [1, BxN]

    # Here we ignore top-N negative losses for negative objectness loss (E.g target is predicted, but groundtruth is negative).
    # Since some objects may be missing in gt, and we don't want to include those wrong false-positives in the training signal
    # And use keep indexes to propagate ignored pixels in other losses

    _, indices = torch.topk(neg_objectness_loss, k=num_elements_to_ignore, dim=1, largest=True, sorted=False)

    non_ignored_mask = torch.ones_like(neg_objectness_loss).scatter_(1, indices, 0)

    # Instead of just setting loss to zeros in those hard example predictions, we minimize entropy loss
    # for objectness & classification to push the predictions to either 0 or 1
    neg_objectness_loss = neg_objectness_loss * non_ignored_mask + binary_shannon_entropy_loss(pred_objectness) * (1 - non_ignored_mask)
    vessel_loss = vessel_loss * non_ignored_mask + binary_shannon_entropy_loss(pred_is_vessel) * (1 - non_ignored_mask)
    fishing_loss = fishing_loss * non_ignored_mask + binary_shannon_entropy_loss(pred_is_fishing) * (1 - non_ignored_mask)
    # For offset & size there is no regularization, so we exclude corresponding pixels from the grad. contribution
    offset_loss = offset_loss * non_ignored_mask
    size_loss = size_loss * non_ignored_mask

    return (
        pos_objectness_loss,
        neg_objectness_loss,
        vessel_loss,
        fishing_loss,
        offset_loss,
        size_loss,
    )


class YOLOLoss(nn.Module):
    def __init__(self, num_classes: int, anchors: list):
        super(YOLOLoss, self).__init__()
        self.num_classes = num_classes
        self.anchors = anchors
        self.mse_loss = nn.MSELoss(reduction='sum')  # Consider using 'mean' if preferred
        self.bce_loss = nn.BCELoss(reduction='sum')  # Consider using 'mean' if preferred

    def calculate_iou(self, box1: torch.Tensor, box2: torch.Tensor) -> torch.Tensor:
        """Calculate Intersection Over Union (IoU) between two sets of boxes."""
        x1y1 = torch.max(box1[..., :2], box2[..., :2])
        x2y2 = torch.min(box1[..., 2:], box2[..., 2:])
        intersection = torch.clamp((x2y2 - x1y1), min=0).prod(-1)
        area_box1 = (box1[..., 2] - box1[..., 0]) * (box1[..., 3] - box1[..., 1])
        area_box2 = (box2[..., 2] - box2[..., 0]) * (box2[..., 3] - box2[..., 1])
        union = area_box1 + area_box2 - intersection
        return intersection / union

    def forward(self, yolo_pred: torch.Tensor, yolo_targets: torch.Tensor) -> torch.Tensor:
        # Ensure predictions and targets have the same number of channels
        if yolo_pred.shape[1] != yolo_targets.shape[1]:
            raise ValueError(f"Channel mismatch: predictions have {yolo_pred.shape[1]} channels, targets have {yolo_targets.shape[1]} channels.")
        
        # Ensure spatial dimensions match
        if yolo_pred.shape[-2:] != yolo_targets.shape[-2:]:
            yolo_pred = F.interpolate(yolo_pred, size=yolo_targets.shape[-2:], mode='bilinear', align_corners=False)

        # Extract bounding boxes, confidence, and class predictions from YOLO output
        bbox_pred = yolo_pred[:, 0:4, :, :]  # [batch, 4, H, W]
        conf_pred = yolo_pred[:, 4:5, :, :]  # [batch, 1, H, W]
        class_pred = yolo_pred[:, 5:, :, :]  # [batch, num_classes, H, W]

        bbox_targets = yolo_targets[:, 0:4, :, :]  # [batch, 4, H, W]
        conf_targets = yolo_targets[:, 4:5, :, :]  # [batch, 1, H, W]
        class_targets = yolo_targets[:, 5:, :, :]  # [batch, num_classes, H, W]

        # Compute losses: bbox (MSE), confidence (BCE), and class (BCE or CrossEntropy)
        bbox_loss = self.mse_loss(bbox_pred, bbox_targets)
        conf_loss = self.bce_loss(conf_pred, conf_targets)
        class_loss = self.bce_loss(class_pred, class_targets)

        # Combine losses with appropriate weighting
        total_loss = bbox_loss + conf_loss + class_loss

        return total_loss