Stylegan implementation

gandlf_synth/data/datasets.py

@@ -44,10 +44,8 @@ def _prepare_multichannel_image(self, index):
             for channel_column in channel_columns
         ]
         tio_scalar_image = tio.ScalarImage(channel_file_paths)
-
         if self.transforms:
             tio_scalar_image = self.transforms(tio_scalar_image)
-        # TODO think if this is valid
         image = tio_scalar_image.data.squeeze(-1).float()
         return image
 

gandlf_synth/losses/__init__.py

@@ -3,12 +3,15 @@
 from torch.nn import Module
 from torch.nn import MSELoss, L1Loss, CrossEntropyLoss, BCEWithLogitsLoss, BCELoss
 
+from gandlf_synth.losses.custom_losses import PlainMeanLoss
+
 global_losses_dict = {
     "cel": CrossEntropyLoss,
     "l1": L1Loss,
     "bce": BCELoss,
     "bcelogits": BCEWithLogitsLoss,
     "mse": MSELoss,
+    "plain_mean": PlainMeanLoss,
 }
 
 WEIGHT_KEYS = ["weight", "pos_weight"]

gandlf_synth/losses/custom_losses.py

@@ -0,0 +1,13 @@
+import torch
+from torch.nn import Module
+
+
+class PlainMeanLoss(Module):
+    def __init__(self):
+        """
+        A simple mean loss that averages the input tensor.
+        """
+        super().__init__()
+
+    def forward(self, x):
+        return torch.mean(x)