Bugfix

.gitignore

@@ -1,2 +1,2 @@
-.ipynb
-__pycache__
+.ipynb_checkpoints/
+__pycache__/

app.py

@@ -14,6 +14,9 @@ from torchvision import transforms
 from evals import PhysicsWithGenerator, EvalModel, BaselineModel, EvalDataset, Metric
 
 
+DEVICE_STR = 'cuda'
+
+
 ### Gradio Utils
 def generate_imgs(dataset: EvalDataset, idx: int,
                   model: EvalModel, baseline: BaselineModel,
@@ -152,8 +155,8 @@ with gr.Blocks(title=title, css=custom_css) as interface:
     # Loading things
     model_a_placeholder = gr.State(lambda: get_eval_model_on_DEVICE_STR("unext_emb_physics_config_C", ""))  # lambda expression to instanciate a callable in a gr.State
     model_b_placeholder = gr.State(lambda: get_baseline_model_on_DEVICE_STR("DRUNET"))  # lambda expression to instanciate a callable in a gr.State
-    dataset_placeholder = gr.State(get_dataset_on_DEVICE_STR("DIV2K_valid_HR"))
-    physics_placeholder = gr.State(lambda: get_physics_generator_on_DEVICE_STR("Denoising"))  # lambda expression to instanciate a callable in a gr.State
+    dataset_placeholder = gr.State(get_dataset_on_DEVICE_STR("Natural"))
+    physics_placeholder = gr.State(lambda: get_physics_generator_on_DEVICE_STR("MotionBlur_easy"))  # lambda expression to instanciate a callable in a gr.State
     metrics_placeholder = gr.State(get_list_metrics_on_DEVICE_STR(["PSNR"]))
 
     @gr.render(inputs=[model_a_placeholder, model_b_placeholder, dataset_placeholder, physics_placeholder, metrics_placeholder])
@@ -265,5 +268,4 @@ with gr.Blocks(title=title, css=custom_css) as interface:
                                          metrics_placeholder], 
                                  outputs=[idx_slider, clean, y_image, model_a_out, model_b_out, physics_params, y_metrics, out_a_metric, out_b_metric])
 
-if __name__ == "__main__":
-    interface.launch()
+interface.launch()

evals.py

@@ -486,8 +486,6 @@ class BaselineModel(torch.nn.Module):
             x_adj = physics.A_adjoint(y)
             output = output[..., :x_adj.size(-2), :x_adj.size(-1)]
             return output
-        elif 'UNROLLED_DPIR' in self.name:
-            return self.model(y, physics=physics)
         else:
             return self.model(y)
 
@@ -504,19 +502,19 @@ class EvalDataset(torch.utils.data.Dataset):
         if self.name not in self.all_datasets:
             raise ValueError(f"{self.name} is unavailable.")
         if self.name == 'Natural':
-            self.root = 'datasets/LSDIR_samples'
+            self.root = 'img_samples/LSDIR_samples'
             self.transform = transforms.Compose([transforms.ToTensor()])
             self.dataset = dinv.datasets.LsdirHR(root=self.root,
                                                  download=False,
                                                  transform=self.transform)
         elif self.name == 'MRI':
-            self.root = 'datasets/FastMRI_samples'
+            self.root = 'img_samples/FastMRI_samples'
             self.transform = transforms.CenterCrop((640, 320))  # , pad_if_needed=True)
             self.dataset = Preprocessed_fastMRI(root=self.root,
                                                 transform=self.transform,
                                                 preprocess=False)
         elif self.name == "CT":
-            self.root = 'datasets/LIDC_IDRI_samples'
+            self.root = 'img_samples/LIDC_IDRI_samples'
             self.transform = None
             self.dataset = Preprocessed_LIDCIDRI(root=self.root,
                                                  transform=self.transform)

models/blocks.py

@@ -7,7 +7,7 @@ import torch.nn.functional as F
 from deepinv.models.unet import BFBatchNorm2d
 from deepinv.physics.blur import gaussian_blur
 from deepinv.physics.functional import conv2d
-from deepinv.utils.tensorlist import TensorList
+from deepinv.utils import TensorList
 
 from timm.models.layers import trunc_normal_, DropPath
 

utils.py

@@ -9,8 +9,16 @@ from physics.multiscale import Pad
 
 
 class ArtifactRemoval(nn.Module):
+    r"""
+    Artifact removal architecture :math:`\phi(A^{\top}y)`.
+
+    This differs from the dinv.models.ArtifactRemoval in that it allows to forward the physics.
+
+    In the end we should not use this for unext !!
+    """
+
     def __init__(self, backbone_net, pinv=False, ckpt_path=None, device=None, fm_mode=False):
-        super().__init__()
+        super(ArtifactRemoval, self).__init__()
         self.pinv = pinv
         self.backbone_net = backbone_net
         self.fm_mode = fm_mode
@@ -24,7 +32,14 @@ class ArtifactRemoval(nn.Module):
                 v.requires_grad = False
             self.backbone_net = self.backbone_net.to(device)
 
+
     def forward_basic(self, y=None, physics=None, x_in=None, t=None, **kwargs):
+        r"""
+        Reconstructs a signal estimate from measurements y
+
+        :param torch.tensor y: measurements
+        :param deepinv.physics.Physics physics: forward operator
+        """
         if physics is None:
             physics = dinv.physics.Denoising(noise_model=dinv.physics.GaussianNoise(sigma=0.), device=y.device)
 
@@ -35,8 +50,15 @@ class ArtifactRemoval(nn.Module):
 
         x_in = physics.A_adjoint(y) if not self.pinv else physics.A_dagger(y)
 
-        sigma = getattr(physics.noise_model, "sigma", 1e-3)
-        gamma = getattr(physics.noise_model, "gain", 1e-3)
+        if hasattr(physics.noise_model, "sigma"):
+            sigma = physics.noise_model.sigma
+        else:
+            sigma = 1e-3  # WARNING: this is a default value that we may not want to use?
+
+        if hasattr(physics.noise_model, "gain"):
+            gamma = physics.noise_model.gain
+        else:
+            gamma = 1e-3  # WARNING: this is a default value that we may not want to use?
 
         out = self.backbone_net(x_in, physics=physics, y=y, sigma=sigma, gamma=gamma, t=t)
 
@@ -45,14 +67,18 @@ class ArtifactRemoval(nn.Module):
 
         return out
 
-    def forward(self, y=None, physics=None, x_in=None, **kwargs):
-        return self.forward_basic(physics=physics, y=y, x_in=x_in, **kwargs)
+    def forward(self,  y=None, physics=None, x_in=None, **kwargs):
+        if 'unext' in type(self.backbone_net).__name__.lower():
+            return self.forward_basic(physics=physics, y=y, x_in=x_in, **kwargs)
+        else:
+            return self.backbone_net(physics=physics, y=y, **kwargs)
 
 
 def get_model(
     model_name="unext_emb_physics_config_C",
     device="cpu",
     in_channels=[1, 2, 3],
+    grayscale=False,
     conv_type="base",
     pool_type="base",
     layer_scale_init_value=1e-6,
@@ -65,6 +91,7 @@ def get_model(
     antialias="gaussian",
     nc_base=64,
     cond_type="base",
+    blind=False,
     pretrained_pth=None,
     weight_tied=True,
     N=4,
@@ -73,41 +100,31 @@ def get_model(
     relu_in_encoding=False,
     skip_in_encoding=True,
 ):
+    """
+    Load the model.
+
+    :param str model_name: name of the model
+    :param str device: device
+    :param bool grayscale: if True, the model is trained on grayscale images
+    :param bool train: if True, the model is trained
+    :return: model
+    """
     model_name = model_name.lower()
-    nc = [nc_base * 2**i for i in range(4)]
 
     if model_name == "pdnet":
         return get_PDNet_architecture(in_channels=in_channels, out_channels=in_channels, device=device)
 
-    elif model_name == "unrolled_dpir":
-        model = UNeXt(
-            in_channels=in_channels,
-            out_channels=in_channels,
-            device=device,
-            conv_type=conv_type,
-            pool_type=pool_type,
-            layer_scale_init_value=layer_scale_init_value,
-            init_type=init_type,
-            gain_init_conv=gain_init_conv,
-            gain_init_linear=gain_init_linear,
-            drop_prob=drop_prob,
-            replk=replk,
-            mult_fact=mult_fact,
-            antialias=antialias,
-            nc=nc,
-            cond_type=cond_type,
-            emb_physics=False,
-            config=None,
-            pretrained_pth=pretrained_pth,
-        ).to(device)
-        model = get_unrolled_architecture(model=model, weight_tied=weight_tied, device=device)
-        return ArtifactRemoval(model, pinv=True, device=device)
-
     elif model_name == "unext_emb_physics_config_c":
+        n_chan = [1, 2, 3] # 6 for old head grayscale, complex and color = 1 + 2 + 3
+        residual = True if "residual" in model_name else False
+        nc = [nc_base * 2**i for i in range(4)]
+
+
         model = UNeXt(
             in_channels=in_channels,
             out_channels=in_channels,
             device=device,
+            residual=residual,
             conv_type=conv_type,
             pool_type=pool_type,
             layer_scale_init_value=layer_scale_init_value,