| import torch |
| from torch import nn, Tensor |
| from einops import rearrange |
| from typing import List, Tuple |
| from .utils import _reshape_density, _bin_count |
|
|
| EPS = 1e-8 |
|
|
|
|
| class ZIPoissonNLL(nn.Module): |
| def __init__( |
| self, |
| reduction: str = "mean", |
| ) -> None: |
| super().__init__() |
| assert reduction in ["none", "mean", "sum"], f"Expected reduction to be one of ['none', 'mean', 'sum'], got {reduction}." |
| self.reduction = reduction |
|
|
| def forward( |
| self, |
| logit_pi_maps: Tensor, |
| lambda_maps: Tensor, |
| gt_den_maps: Tensor, |
| ) -> Tensor: |
| assert len(logit_pi_maps.shape) == len(lambda_maps.shape) == len(gt_den_maps.shape) == 4, f"Expected 4D (B, C, H, W) tensor, got {logit_pi_maps.shape}, {lambda_maps.shape}, and {gt_den_maps.shape}" |
| B, _, H, W = lambda_maps.shape |
| assert logit_pi_maps.shape == (B, 2, H, W), f"Expected logit_pi_maps to have shape (B, 2, H, W), got {logit_pi_maps.shape}" |
| assert lambda_maps.shape == (B, 1, H, W), f"Expected lambda_maps to have shape (B, 1, H, W), got {lambda_maps.shape}" |
| if gt_den_maps.shape[2:] != (H, W): |
| gt_h, gt_w = gt_den_maps.shape[-2], gt_den_maps.shape[-1] |
| assert gt_h % H == 0 and gt_w % W == 0 and gt_h // H == gt_w // W, f"Expected the spatial dimension of gt_den_maps to be a multiple of that of lambda_maps, got {gt_den_maps.shape} and {lambda_maps.shape}" |
| gt_den_maps = _reshape_density(gt_den_maps, block_size=gt_h // H) |
| assert gt_den_maps.shape == (B, 1, H, W), f"Expected gt_den_maps to have shape (B, 1, H, W), got {gt_den_maps.shape}" |
|
|
| pi_maps = logit_pi_maps.softmax(dim=1) |
| zero_indices = (gt_den_maps == 0).float() |
| zero_loss = -torch.log(pi_maps[:, 0:1] + pi_maps[:, 1:] * torch.exp(-lambda_maps) + EPS) * zero_indices |
|
|
| poisson_log_p = gt_den_maps * torch.log(lambda_maps + EPS) - lambda_maps |
| nonzero_loss = (-torch.log(pi_maps[:, 1:] + EPS) - poisson_log_p) * (1.0 - zero_indices) |
|
|
| loss = (zero_loss + nonzero_loss).sum(dim=(-1, -2)) |
| if self.reduction == "mean": |
| loss = loss.mean() |
| elif self.reduction == "sum": |
| loss = loss.sum() |
|
|
| return loss, {"zipnll": loss.detach()} |
|
|
|
|
| class ZICrossEntropy(nn.Module): |
| def __init__( |
| self, |
| bins: List[Tuple[int, int]], |
| reduction: str = "mean", |
| ) -> None: |
| super().__init__() |
| assert all([low <= high for low, high in bins]), f"Expected bins to be a list of tuples (low, high) where low <= high, got {bins}" |
| assert reduction in ["mean", "sum"], f"Expected reduction to be one of ['none', 'mean', 'sum'], got {reduction}." |
|
|
| self.bins = bins |
| self.reduction = reduction |
| self.ce_loss_fn = nn.CrossEntropyLoss(reduction="none") |
|
|
| def forward( |
| self, |
| logit_maps: Tensor, |
| gt_den_maps: Tensor, |
| ) -> Tensor: |
| assert len(logit_maps.shape) == len(gt_den_maps.shape) == 4, f"Expected 4D (B, C, H, W) tensor, got {logit_maps.shape} and {gt_den_maps.shape}" |
| B, _, H, W = logit_maps.shape |
| assert logit_maps.shape[0] == B and logit_maps.shape[2:] == (H, W), f"Expected logit_maps to have shape (B, C, H, W), got {logit_maps.shape}" |
| if gt_den_maps.shape[2:] != (H, W): |
| gt_h, gt_w = gt_den_maps.shape[-2], gt_den_maps.shape[-1] |
| assert gt_h % H == 0 and gt_w % W == 0 and gt_h // H == gt_w // W, f"Expected the spatial dimension of gt_den_maps to be a multiple of that of logit_maps, got {gt_den_maps.shape} and {logit_maps.shape}" |
| gt_den_maps = _reshape_density(gt_den_maps, block_size=gt_h // H) |
| assert gt_den_maps.shape == (B, 1, H, W), f"Expected gt_den_maps to have shape (B, 1, H, W), got {gt_den_maps.shape}" |
|
|
| gt_class_maps = _bin_count(gt_den_maps, bins=self.bins) |
| gt_class_maps = rearrange(gt_class_maps, "B H W -> B (H W)") |
| logit_maps = rearrange(logit_maps, "B C H W -> B (H W) C") |
|
|
| loss = 0.0 |
| for idx in range(gt_class_maps.shape[0]): |
| gt_class_map, logit_map = gt_class_maps[idx], logit_maps[idx] |
| mask = gt_class_map > 0 |
| |
| gt_class_map = gt_class_map[mask] - 1 |
| logit_map = logit_map[mask] |
| loss += self.ce_loss_fn(logit_map, gt_class_map).sum() |
|
|
| if self.reduction == "mean": |
| loss /= gt_class_maps.shape[0] |
| |
| return loss, {"cls_zice": loss.detach()} |
| |
