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

Code before:
        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
            alpha: weightage between cross entropy and logit constraint
            sigma: sigma of gaussian
        """

        super().__init__()

        if kernel_ops not in ["mean", "gaussian"]:
            raise ValueError("Kernel ops must be either mean or gaussian")

        if dim not in [2, 3]:
            raise ValueError("Supoorts 2d and 3d")

        if distance_type not in ["l1", "l2"]:
            raise ValueError("Distance type must be either L1 or L2")

        self.nc = classes
        self.dim = dim
        self.cross_entropy = nn.CrossEntropyLoss()
        self.distance_type = distance_type
        self.alpha = alpha
        self.ks = kernel_size
        self.svls_layer: Any

        if kernel_ops == "mean":
            self.svls_layer = MeanFilter(spatial_dims=dim, size=kernel_size)
            self.svls_layer.filter = self.svls_layer.filter / (kernel_size**dim)
        if kernel_ops == "gaussian":
            self.svls_layer = GaussianFilter(spatial_dims=dim, sigma=sigma)

    def get_constr_target(self, mask: torch.Tensor) -> torch.Tensor:
        """
        Converts the mask to one hot represenation and applies the spatial filter.
        """
        rmask: torch.Tensor

        if self.dim == 2:
            oh_labels = F.one_hot(mask.to(torch.int64), num_classes=self.nc).contiguous().permute(0, 3, 1, 2).float()
            rmask = self.svls_layer(oh_labels)


Review comment (suggestion):
```suggestion
            raise ValueError(f"Distance type must be either L1 or L2, got {distance_type}")
```

Code after:
        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
            alpha: weightage between cross entropy and logit constraint
            sigma: sigma of gaussian
        """

        super().__init__()

        if kernel_ops not in ["mean", "gaussian"]:
            raise ValueError("Kernel ops must be either mean or gaussian")

        if dim not in [2, 3]:
            raise ValueError(f"Support 2d and 3d, got dim={dim}.")

        if distance_type not in ["l1", "l2"]:
            raise ValueError(f"Distance type must be either L1 or L2, got {distance_type}")

        self.nc = classes
        self.dim = dim
        self.cross_entropy = nn.CrossEntropyLoss()
        self.distance_type = distance_type
        self.alpha = alpha
        self.ks = kernel_size
        self.svls_layer: Any

        if kernel_ops == "mean":
            self.svls_layer = MeanFilter(spatial_dims=dim, size=kernel_size)
            self.svls_layer.filter = self.svls_layer.filter / (kernel_size**dim)
        if kernel_ops == "gaussian":
            self.svls_layer = GaussianFilter(spatial_dims=dim, sigma=sigma)

    def get_constr_target(self, mask: torch.Tensor) -> torch.Tensor:
        """
        Converts the mask to one hot represenation and is smoothened with the selected spatial filter.

        Args:
            mask: the shape should be BH[WD].

        Returns:
            torch.Tensor: the shape would be BNH[WD], N being number of classes.
        """