ENH adding new tasks (MRI, demoisaicing)

benchmark_utils/custom_models.py

@@ -1,22 +1,16 @@
 from deepinv.models import UNet
 
-class CustomUNet(UNet):
-    def __init__(self, in_channels, out_channels, scales=3, batch_norm=False, is_mri=False):
-        self.name = "CustomUNet"
+class MRIUNet(UNet):
+    def __init__(self, in_channels, out_channels, scales=3, batch_norm=False):
+        self.name = "MRIUNet"
         self.in_channels = in_channels
-        self.is_mri = is_mri
 
         super().__init__(in_channels=in_channels, out_channels=out_channels, scales=scales, batch_norm=batch_norm)
 
     def forward(self, x, sigma=None, **kwargs):
-        if self.is_mri:
-            # Reshape for MRI specific processing
-            x = x.reshape(1, self.in_channels, x.shape[3], x.shape[4])
-
+        # Reshape for MRI specific processing
+        x = x.reshape(1, self.in_channels, x.shape[3], x.shape[4])
+
         x = super().forward(x, sigma=sigma, **kwargs)
-
-        if self.is_mri:
-            # Reshape for when MRI
-            x = x.reshape(2, 4, x.shape[2], x.shape[3])
 
         return x

benchmark_utils/fastmri_dataset.py

@@ -30,6 +30,6 @@ def __getitem__(self, idx):
 
         # We add an imaginary part of zeros
         x = torch.cat([x, torch.zeros_like(x)], dim=0)
-        y = y.reshape(8, y.shape[2], y.shape[3])
+        #y = y.reshape(8, y.shape[2], y.shape[3])
 
         return x, y

objective.py

@@ -82,7 +82,7 @@ def evaluate_result(self, model, model_name, device):
                     model(y_i[None], self.physics) for y_i in y
                 ])
             else:
-                x_hat = model(y)
+                x_hat = model(y, self.physics)
 
             if (self.dataset_name == 'FastMRI'):
                 transform = torchvision.transforms.Compose(
@@ -185,4 +185,7 @@ def get_objective(self):
         # for `Solver.set_objective`. This defines the
         # benchmark's API for passing the objective to the solver.
         # It is customizable for each benchmark.
-        return dict(train_dataset=self.train_dataset, physics=self.physics, image_size=self.image_size)
+        return dict(train_dataset=self.train_dataset,
+                    physics=self.physics,
+                    image_size=self.image_size,
+                    dataset_name=self.dataset_name,)

solvers/diffpir.py

@@ -16,7 +16,7 @@ class Solver(BaseSolver):
 
     requirements = []
 
-    def set_objective(self, train_dataset, physics, image_size):
+    def set_objective(self, train_dataset, physics, image_size, dataset_name):
         batch_size = 2
         self.train_dataloader = DataLoader(
             train_dataset, batch_size=batch_size, shuffle=False

solvers/dpir.py

@@ -23,7 +23,7 @@ class Solver(BaseSolver):
 
     requirements = []
 
-    def set_objective(self, train_dataset, physics, image_size):
+    def set_objective(self, train_dataset, physics, image_size, dataset_name):
         batch_size = 2
         self.train_dataloader = DataLoader(
             train_dataset, batch_size=batch_size, shuffle=False
@@ -33,6 +33,7 @@ def set_objective(self, train_dataset, physics, image_size):
         )
         self.physics = physics
         self.image_size = image_size
+        self.dataset_name = dataset_name
 
     def run(self, n_iter):
         best_sigma = 0
@@ -50,25 +51,27 @@ def run(self, n_iter):
         for sigma in np.linspace(0.01, 0.1, 10):
             model = model_class(sigma=sigma, device=self.device)
 
+            psnr = []
+
             for x, y in self.train_dataloader:
                 x, y = x.to(self.device), y.to(self.device)
 
-            x_hat = model(y)
-            
-            if (self.dataset_name == 'FastMRI'):
-                transform = torchvision.transforms.Compose(
-                    [
-                        torchvision.transforms.CenterCrop(x.shape[-2:]),
-                        dinv.metric.functional.complex_abs,
-                    ]
-                )
-
-                CustomPSNR.transform = transform
-                
-                psnr.append(CustomPSNR()(x_hat, x))
-            else:
-                psnr.append(dinv.metric.PSNR()(x_hat, x))
-                
+                x_hat = model(y, self.physics)
+
+                if (self.dataset_name == 'FastMRI'):
+                    transform = torchvision.transforms.Compose(
+                        [
+                            torchvision.transforms.CenterCrop(x.shape[-2:]),
+                            dinv.metric.functional.complex_abs,
+                        ]
+                    )
+
+                    CustomPSNR.transform = transform
+
+                    psnr.append(CustomPSNR()(x_hat, x))
+                else:
+                    psnr.append(dinv.metric.PSNR()(x_hat, x))
+
             psnr = torch.mean(torch.cat(psnr)).item()
 
             #results = dinv.test(

solvers/u-net.py

@@ -7,7 +7,7 @@
     import deepinv as dinv
     import torchvision
     from benchmark_utils.metrics import CustomMSE, CustomPSNR
-    from benchmark_utils.custom_models import CustomUNet
+    from benchmark_utils.custom_models import MRIUNet
 
 
 class Solver(BaseSolver):
@@ -22,7 +22,7 @@ class Solver(BaseSolver):
 
     requirements = []
 
-    def set_objective(self, train_dataset, physics, image_size):
+    def set_objective(self, train_dataset, physics, image_size, dataset_name):
         batch_size = 1
         self.train_dataloader = DataLoader(
             train_dataset, batch_size=batch_size, shuffle=False
@@ -32,15 +32,23 @@ def set_objective(self, train_dataset, physics, image_size):
         )
         self.physics = physics.to(self.device)
         self.image_size = image_size
+        self.dataset_name = dataset_name
 
     def run(self, n_iter):
         epochs = 4
 
         x, y = next(iter(self.train_dataloader))
 
-        model = dinv.models.UNet(
-            in_channels=y.shape[1], out_channels=x.shape[1], scales=3, batch_norm=False
-        ).to(self.device)
+        if self.dataset_name == 'FastMRI':
+            model = MRIUNet(
+                in_channels=y.shape[1] * y.shape[2], out_channels=x.shape[1], scales=3,
+                batch_norm=False
+            ).to(self.device)
+        else:
+            model = dinv.models.UNet(
+                in_channels=y.shape[1], out_channels=x.shape[1], scales=3,
+                batch_norm=False
+            ).to(self.device)
 
         verbose = True  # print training information
         wandb_vis = False  # plot curves and images in Weight&Bias