Repository: Project-MONAI/MONAI
File: monai/losses/nacl_loss.py (Python)

Code before:
#     http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the 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 torch.nn.modules.loss import _Loss

from monai.utils import pytorch_after


def get_mean_kernel_2d(ksize: int = 3) -> torch.Tensor:
    mean_kernel = torch.ones([ksize, ksize]) / (ksize**2)
    return mean_kernel


def get_mean_kernel_3d(ksize: int = 3) -> torch.Tensor:
    mean_kernel = torch.ones([ksize, ksize, ksize]) / (ksize**3)
    return mean_kernel


def get_gaussian_kernel_2d(ksize: int = 3, sigma: float = 1.0) -> torch.Tensor:
    x_grid = torch.arange(ksize).repeat(ksize).view(ksize, ksize)
    y_grid = x_grid.t()
    xy_grid = torch.stack([x_grid, y_grid], dim=-1).float()
    mean = (ksize - 1) / 2.0
    variance = sigma**2.0
    gaussian_kernel: torch.Tensor = (1.0 / (2.0 * math.pi * variance + 1e-16)) * torch.exp(
        -torch.sum((xy_grid - mean) ** 2.0, dim=-1) / (2 * variance + 1e-16)
    )
    gaussian_kernel = gaussian_kernel / torch.sum(gaussian_kernel)
    return gaussian_kernel


def get_gaussian_kernel_3d(ksize: int = 3, sigma: float = 1.0) -> torch.Tensor:
    x_coord = torch.arange(ksize)
    x_grid_2d = x_coord.repeat(ksize).view(ksize, ksize)
    x_grid = x_coord.repeat(ksize * ksize).view(ksize, ksize, ksize)
    y_grid_2d = x_grid_2d.t()
    y_grid = y_grid_2d.repeat(ksize, 1).view(ksize, ksize, ksize)
    z_grid = y_grid_2d.repeat(1, ksize).view(ksize, ksize, ksize)
    xyz_grid = torch.stack([x_grid, y_grid, z_grid], dim=-1).float()

Review comment (suggestion):
Perhaps we can combine these two to one function and move it under this [folder](https://github.com/Project-MONAI/MONAI/blob/dev/monai/networks/layers/convutils.py) to make it to be a more common used utils. What do you think?

cc @ericspod @Nic-Ma 

Code after:
#     http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

from typing import Any

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.modules.loss import _Loss

from monai.networks.layers import GaussianFilter, MeanFilter


class NACLLoss(_Loss):
    """
    Neighbor-Aware Calibration Loss (NACL) is primarily developed for developing calibrated models in image segmentation.
    NACL computes standard cross-entropy loss with a linear penalty that enforces the logit distributions
    to match a soft class proportion of surrounding pixel.

    Murugesan, Balamurali, et al.
    "Trust your neighbours: Penalty-based constraints for model calibration."
    International Conference on Medical Image Computing and Computer-Assisted Intervention, MICCAI 2023.
    https://arxiv.org/abs/2303.06268

    Murugesan, Balamurali, et al.
    "Neighbor-Aware Calibration of Segmentation Networks with Penalty-Based Constraints."
    https://arxiv.org/abs/2401.14487
    """

    def __init__(
        self,
        classes: int,
        dim: int,
        kernel_size: int = 3,
        kernel_ops: str = "mean",
        distance_type: str = "l1",
        alpha: float = 0.1,
        sigma: float = 1.0,
    ) -> None:
        """
        Args:
            classes: number of classes
            dim: dimension of data (supports 2d and 3d)
            kernel_size: size of the spatial kernel
            distance_type: l1/l2 distance between spatial kernel and predicted logits