Perceptual loss changes.

Author: Virginia Fernandez

monai/losses/perceptual.py

@@ -18,12 +18,16 @@
 
 from monai.utils import optional_import
 from monai.utils.enums import StrEnum
+from huggingface_hub import hf_hub_download
 
 LPIPS, _ = optional_import("lpips", name="LPIPS")
 torchvision, _ = optional_import("torchvision")
 
 
 class PercetualNetworkType(StrEnum):
+    """Types of neural networks that are supported by perceptua loss.
+    """
+
     alex = "alex"
     vgg = "vgg"
     squeeze = "squeeze"
@@ -108,9 +112,12 @@ def __init__(
 
         self.spatial_dims = spatial_dims
         self.perceptual_function: nn.Module
+
+        # If spatial_dims is 3, only MedicalNet supports 3D models, otherwise, spatial_dims=2 and fake_3D must be used.
         if spatial_dims == 3 and is_fake_3d is False:
             self.perceptual_function = MedicalNetPerceptualSimilarity(
-                net=network_type, verbose=False, channel_wise=channel_wise
+                net=network_type, verbose=False, channel_wise=channel_wise,
+                cache_dir=cache_dir
             )
         elif "radimagenet_" in network_type:
             self.perceptual_function = RadImageNetPerceptualSimilarity(net=network_type, verbose=False)
@@ -122,7 +129,9 @@ def __init__(
                 pretrained_state_dict_key=pretrained_state_dict_key,
             )
         else:
+            # VGG, AlexNet and SqueezeNet are independently handled by LPIPS.
             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
         self.channel_wise = channel_wise
@@ -194,7 +203,7 @@ class MedicalNetPerceptualSimilarity(nn.Module):
     """
     Component to perform the perceptual evaluation with the networks pretrained by Chen, et al. "Med3D: Transfer
     Learning for 3D Medical Image Analysis". This class uses torch Hub to download the networks from
-    "Warvito/MedicalNet-models".
+    "Project-MONAI/perceptual-models".
 
     Args:
         net: {``"medicalnet_resnet10_23datasets"``, ``"medicalnet_resnet50_23datasets"``}
@@ -205,11 +214,12 @@ class MedicalNetPerceptualSimilarity(nn.Module):
     """
 
     def __init__(
-        self, net: str = "medicalnet_resnet10_23datasets", verbose: bool = False, channel_wise: bool = False
+        self, net: str = "medicalnet_resnet_10_23datasets", verbose: bool = False, channel_wise: bool = False,
+        cache_dir: str | None = None,
     ) -> None:
         super().__init__()
         torch.hub._validate_not_a_forked_repo = lambda a, b, c: True
-        self.model = torch.hub.load("warvito/MedicalNet-models", model=net, verbose=verbose)
+        self.model = torch.hub.load("Project-MONAI/perceptual-models:main", model=net, verbose=verbose, cache_dir=cache_dir)
         self.eval()
 
         self.channel_wise = channel_wise
@@ -287,17 +297,20 @@ class RadImageNetPerceptualSimilarity(nn.Module):
     """
     Component to perform the perceptual evaluation with the networks pretrained on RadImagenet (pretrained by Mei, et
     al. "RadImageNet: An Open Radiologic Deep Learning Research Dataset for Effective Transfer Learning"). This class
-    uses torch Hub to download the networks from "Warvito/radimagenet-models".
+    uses torch Hub to download the networks from "Project-MONAI/perceptual-models".
 
     Args:
         net: {``"radimagenet_resnet50"``}
             Specifies the network architecture to use. Defaults to ``"radimagenet_resnet50"``.
         verbose: if false, mute messages from torch Hub load function.
     """
 
-    def __init__(self, net: str = "radimagenet_resnet50", verbose: bool = False) -> None:
+    def __init__(self, net: str = "radimagenet_resnet50", 
+                 verbose: bool = False,
+                 cache_dir: str | None = None) -> None:
         super().__init__()
-        self.model = torch.hub.load("Warvito/radimagenet-models", model=net, verbose=verbose)
+        self.model = torch.hub.load("Project-MONAI/perceptual-models", model=net, verbose=verbose,
+                                    cache_dir=cache_dir)
         self.eval()
 
         for param in self.parameters():