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

Code before:
            raise ValueError(
                "MedicalNet networks are only compatible with ``spatial_dims=3``."
                "Argument is_fake_3d must be set to False."
            )

        if channelwise and "medicalnet_" not in network_type:
            raise ValueError("Channelwise loss is only compatible with MedicalNet networks.")

        if network_type.lower() not in list(PercetualNetworkType):
            raise ValueError(
                "Unrecognised criterion entered for Adversarial Loss. Must be one in: %s"
                % ", ".join(PercetualNetworkType)
            )

        if cache_dir:
            torch.hub.set_dir(cache_dir)
            # raise a warning that this may change the default cache dir for all torch.hub calls
            warnings.warn(
                f"Setting cache_dir to {cache_dir}, this may change the default cache dir for all torch.hub calls."
            )

        self.spatial_dims = spatial_dims
        self.perceptual_function: nn.Module
        if spatial_dims == 3 and is_fake_3d is False:
            self.perceptual_function = MedicalNetPerceptualSimilarity(
                net=network_type, verbose=False, channelwise=channelwise
            )
        elif "radimagenet_" in network_type:
            self.perceptual_function = RadImageNetPerceptualSimilarity(net=network_type, verbose=False)
        elif network_type == "resnet50":
            self.perceptual_function = TorchvisionModelPerceptualSimilarity(
                net=network_type,
                pretrained=pretrained,
                pretrained_path=pretrained_path,
                pretrained_state_dict_key=pretrained_state_dict_key,
            )
        else:
            self.perceptual_function = LPIPS(pretrained=pretrained, net=network_type, verbose=False)
        self.is_fake_3d = is_fake_3d
        self.fake_3d_ratio = fake_3d_ratio

    def _calculate_axis_loss(self, input: torch.Tensor, target: torch.Tensor, spatial_axis: int) -> torch.Tensor:
        """
        Calculate perceptual loss in one of the axis used in the 2.5D approach. After the slices of one spatial axis
        is transformed into different instances in the batch, we compute the loss using the 2D approach.

        Args:
            input: input 5D tensor. BNHWD
            target: target 5D tensor. BNHWD
            spatial_axis: spatial axis to obtain the 2D slices.
        """

Review comment (suggestion):
```suggestion
                net=network_type, verbose=False, channel_wise=channel_wise
```

Code after:
            raise NotImplementedError("Perceptual loss is implemented only in 2D and 3D.")

        if (spatial_dims == 2 or is_fake_3d) and "medicalnet_" in network_type:
            raise ValueError(
                "MedicalNet networks are only compatible with ``spatial_dims=3``."
                "Argument is_fake_3d must be set to False."
            )

        if channel_wise and "medicalnet_" not in network_type:
            raise ValueError("Channel-wise loss is only compatible with MedicalNet networks.")

        if network_type.lower() not in list(PercetualNetworkType):
            raise ValueError(
                "Unrecognised criterion entered for Adversarial Loss. Must be one in: %s"
                % ", ".join(PercetualNetworkType)
            )

        if cache_dir:
            torch.hub.set_dir(cache_dir)
            # raise a warning that this may change the default cache dir for all torch.hub calls
            warnings.warn(
                f"Setting cache_dir to {cache_dir}, this may change the default cache dir for all torch.hub calls."
            )

        self.spatial_dims = spatial_dims
        self.perceptual_function: nn.Module
        if spatial_dims == 3 and is_fake_3d is False:
            self.perceptual_function = MedicalNetPerceptualSimilarity(
                net=network_type, verbose=False, channel_wise=channel_wise
            )
        elif "radimagenet_" in network_type:
            self.perceptual_function = RadImageNetPerceptualSimilarity(net=network_type, verbose=False)
        elif network_type == "resnet50":
            self.perceptual_function = TorchvisionModelPerceptualSimilarity(
                net=network_type,
                pretrained=pretrained,
                pretrained_path=pretrained_path,
                pretrained_state_dict_key=pretrained_state_dict_key,
            )
        else:
            self.perceptual_function = LPIPS(pretrained=pretrained, net=network_type, verbose=False)
        self.is_fake_3d = is_fake_3d
        self.fake_3d_ratio = fake_3d_ratio

    def _calculate_axis_loss(self, input: torch.Tensor, target: torch.Tensor, spatial_axis: int) -> torch.Tensor:
        """
        Calculate perceptual loss in one of the axis used in the 2.5D approach. After the slices of one spatial axis
        is transformed into different instances in the batch, we compute the loss using the 2D approach.

        Args:
            input: input 5D tensor. BNHWD