losses/dual_loss.py

import torch
from torch import nn, Tensor
import torch.nn.functional as F
from typing import List, Tuple, Dict

from .dm_loss import DMLoss
from .multiscale_mae import MultiscaleMAE
from .utils import _reshape_density



class DualLoss(nn.Module):
    def __init__(
        self,
        input_size: int,
        block_size: int,
        bins: List[Tuple[float, float]],
        bin_centers: List[float],
        cls_loss: str = "ce",
        reg_loss: str = "dm",
        weight_tv: float = 0.01,
        weight_cls: float = 0.1,
        weight_reg: float = 0.1,
        numItermax: int = 100,
        regularization: float = 10.0,
        scales: List[int] = [1, 2, 4],
        min_scale_weight: float = 0.25,
        max_scale_weight: float = 0.75,
        alpha: float = 0.5,
    ) -> None:
        super().__init__()
        assert len(bins) == len(bin_centers) >= 2, f"Expected bins and bin_centers to have at least 2 elements, got {len(bins)} and {len(bin_centers)}"
        assert all([len(b) == 2 for b in bins]), f"Expected all bins to be of length 2, got {bins}"
        assert all(b[0] <= p <= b[1] for b, p in zip(bins, bin_centers)), f"Expected bin_centers to be within the range of the corresponding bin, got {bins} and {bin_centers}"
        assert cls_loss in ["ce", "mae", "mse", "none"], f"Expected cls_loss to be one of ['ce', 'mae', 'mse', 'none'], got {cls_loss}"
        assert reg_loss in ["dm", "msmae", "mae", "mse", "none"], f"Expected reg_loss to be one of ['dm', 'msmae', 'mae', 'mse', 'none'], got {reg_loss}"
        assert not (cls_loss == "none" and reg_loss == "none"), "Expected at least one of cls_loss and reg_loss to be provided"
        assert weight_cls is None or weight_cls >= 0, f"Expected weight_cls to be non-negative, got {weight_cls}"
        assert weight_reg is None or weight_reg >= 0, f"Expected weight_reg to be non-negative, got {weight_reg}"
        assert weight_tv is None or weight_tv >= 0, f"Expected weight_tv to be non-negative, got {weight_tv}"
        assert min_scale_weight is None or max_scale_weight is None or max_scale_weight >= min_scale_weight > 0, f"Expected max_scale_weight to be greater than or equal to min_scale_weight, got {min_scale_weight} and {max_scale_weight}"
        assert alpha is None or 1 > alpha > 0, f"Expected alpha to be between 0 and 1, got {alpha}"

        if reg_loss == "dm":
            assert numItermax is not None and numItermax > 0, f"Expected numItermax to be a positive integer, got {numItermax}"
            assert regularization is not None and regularization > 0, f"Expected regularization to be a positive float, got {regularization}"
            assert weight_tv is not None and weight_tv >= 0, f"Expected weight_tv to be non-negative, got {weight_tv}"
        else:
            weight_tv, numItermax, regularization = None, None, None
        
        if reg_loss == "msmae":
            assert isinstance(scales, (list, tuple)) and len(scales) > 0 and all(isinstance(s, int) and s > 0 for s in scales), f"Expected scales to be a list of positive integers, got {scales}"
            assert max_scale_weight >= min_scale_weight > 0, f"Expected max_scale_weight to be greater than or equal to min_scale_weight, got {min_scale_weight} and {max_scale_weight}"
            assert 1 > alpha > 0, f"Expected alpha to be between 0 and 1, got {alpha}"
        else:
            scales = None
            min_scale_weight, max_scale_weight = None, None
            alpha = None

        weight_cls = weight_cls if weight_cls is not None else 0
        weight_reg = weight_reg if weight_reg is not None else 0

        self.input_size, self.block_size = input_size, block_size
        self.num_blocks_h, self.num_blocks_w = input_size // block_size, input_size // block_size
        self.bins, self.bin_centers, self.num_bins = bins, bin_centers, len(bins)
        self.cls_loss, self.reg_loss = cls_loss, reg_loss
        self.weight_cls, self.weight_reg = weight_cls, weight_reg
        self.numItermax, self.regularization = numItermax, regularization
        self.weight_tv = weight_tv
        self.scales = scales
        self.min_scale_weight, self.max_scale_weight = min_scale_weight, max_scale_weight

        if cls_loss == "ce":
            self.cls_loss_fn = nn.CrossEntropyLoss(reduction="none")
            self.weight_cls = 1.0
        elif cls_loss == "mae":
            self.cls_loss_fn = nn.L1Loss(reduction="none")
            self.weight_cls = weight_cls
        elif cls_loss == "mse":
            self.cls_loss_fn = nn.MSELoss(reduction="none")
            self.weight_cls = weight_cls
        else:  # cls_loss == "none"
            self.cls_loss_fn = None
            self.weight_cls = 0

        if reg_loss == "dm":
            self.reg_loss_fn = DMLoss(
                input_size=input_size,
                block_size=block_size,
                numItermax=numItermax,
                regularization=regularization,
                weight_ot=weight_reg,
                weight_tv=weight_tv,
                weight_cnt=0,  # Calculate the count loss separately
            )
            self.weight_reg = 1.0
        elif reg_loss == "msmae":
            self.reg_loss_fn = MultiscaleMAE(scales=scales, weights=None, min_scale_weight=min_scale_weight, max_scale_weight=max_scale_weight, alpha=alpha)
            self.weight_reg = 1.0
        elif reg_loss == "mae":
            self.reg_loss_fn = nn.L1Loss(reduction="none")
            self.weight_reg = weight_reg
        elif reg_loss == "mse":
            self.reg_loss_fn = nn.MSELoss(reduction="none")
            self.weight_reg = weight_reg
        else:
            self.reg_loss_fn = None
            self.weight_reg = 0
        
        self.cnt_loss_fn = nn.L1Loss(reduction="none")
        
    def _bin_count(self, density_map: Tensor) -> Tensor:
        class_map = torch.zeros_like(density_map, dtype=torch.long)
        for idx, (low, high) in enumerate(self.bins):
            mask = (density_map >= low) & (density_map <= high)
            class_map[mask] = idx
        return class_map.squeeze(1)  # remove channel dimension

    def forward(
        self,
        pred_logit_map: Tensor,
        pred_den_map: Tensor,
        gt_den_map: Tensor,
        gt_points: List[Tensor]
    ) -> Tuple[Tensor, Dict[str, Tensor]]:
        B = pred_logit_map.shape[0]
        assert pred_logit_map.shape == (B, self.num_bins, self.num_blocks_h, self.num_blocks_w), f"Expected pred_logit_map to have shape {B, self.num_bins, self.num_blocks_h, self.num_blocks_w}, got {pred_logit_map.shape}"
        if gt_den_map.shape[-2:] != (self.num_blocks_h, self.num_blocks_w):
            assert gt_den_map.shape[-2:] == (self.input_size, self.input_size), f"Expected gt_den_map to have shape {B, 1, self.input_size, self.input_size}, got {gt_den_map.shape}"
            gt_den_map = _reshape_density(gt_den_map, block_size=self.block_size)
        assert pred_den_map.shape == gt_den_map.shape == (B, 1, self.num_blocks_h, self.num_blocks_w), f"Expected pred_den_map and gt_den_map to have shape (B, 1, H, W), got {pred_den_map.shape} and {gt_den_map.shape}"
        assert len(gt_points) == B, f"Expected gt_points to have length B, got {len(gt_points)}"
        
        loss_info = {}

        if self.weight_cls > 0:
            gt_class_map = self._bin_count(gt_den_map)
            if self.cls_loss == "ce":
                cls_loss = self.cls_loss_fn(pred_logit_map, gt_class_map).sum(dim=(-1, -2)).mean()
                loss_info["cls_ce_loss"] = cls_loss.detach()
            else:  # self.cls_loss in ["mae", "mse"]
                gt_prob_map = F.one_hot(gt_class_map, num_classes=self.num_bins).float()  # B, H, W -> B, H, W, N
                gt_prob_map = gt_prob_map.permute(0, 3, 1, 2)  # B, H, W, N -> B, N, H, W
                pred_prob_map = pred_logit_map.softmax(dim=1)
                cls_loss = self.cls_loss_fn(pred_prob_map, gt_prob_map).sum(dim=(-1, -2)).mean()
                loss_info[f"cls_{self.cls_loss}_loss"] = cls_loss.detach()
        else:
            cls_loss = 0

        if self.weight_reg > 0:
            if self.reg_loss == "dm":
                reg_loss, reg_loss_info = self.reg_loss_fn(
                    pred_den_map=pred_den_map,
                    gt_den_map=gt_den_map,
                    gt_points=gt_points,
                )
                loss_info.update({f"reg_{k}": v for k, v in reg_loss_info.items()})
            elif self.reg_loss == "msmae":
                reg_loss, reg_loss_info = self.reg_loss_fn(pred_den_map, gt_den_map)
                loss_info.update({f"reg_{k}": v for k, v in reg_loss_info.items()})
            else:  # self.reg_loss in ["mae", "mse"]
                reg_loss = self.reg_loss_fn(pred_den_map, gt_den_map).sum(dim=(-1, -2)).mean()
                loss_info[f"reg_{self.reg_loss}_loss"] = reg_loss.detach()
        else:
            reg_loss = 0

        gt_cnt = torch.tensor([len(p) for p in gt_points], dtype=torch.float32, device=pred_den_map.device)
        cnt_loss = self.cnt_loss_fn(pred_den_map.sum(dim=(1, 2, 3)), gt_cnt).mean()
        loss_info["cnt_loss"] = cnt_loss.detach()

        total_loss = self.weight_cls * cls_loss + self.weight_reg * reg_loss + cnt_loss
        loss_info["total_loss"] = total_loss.detach()
        loss_info = dict(sorted(loss_info.items()))  # sort by key for nicer printing

        return total_loss, loss_info