[Fix] PSNR computation

app.py

@@ -52,6 +52,7 @@ def generate_imgs_from_user(image,
         x = torch.cat((x, torch.zeros_like(x)), dim=1)
 
     return generate_imgs(x, physics, use_gen, baseline, model, metrics)
+
 def generate_imgs_from_dataset(dataset: EvalDataset, idx: int,
                                physics: PhysicsWithGenerator, use_gen: bool,
                                baseline: BaselineModel, model: EvalModel,
@@ -108,26 +109,24 @@ def generate_imgs(x: torch.Tensor,
         if out_baseline.shape != out.shape:
             out_baseline = out_baseline[..., w_1:w_2, h_1:h_2]
 
+    ### Process y when y shape is different from x shape
+    if physics.name == 'MRI' and physics.name == 'CT':
+        y_plot = physics.physics.prox_l2(physics.physics.A_adjoint(y), y, 1e4)
+    else:
+        y_plot = y.clone()
+
     ### Metrics
     metrics_y = ""
     metrics_out = ""
     metrics_out_baseline = ""
     for metric in metrics:
-        if y.shape == x.shape:
-            metrics_y += f"{metric.name} = {metric(y, x).item():.4f}" + "\n"
+        #if y.shape == x.shape:
+        metrics_y += f"{metric.name} = {metric(y_plot, x).item():.4f}" + "\n"
         metrics_out += f"{metric.name} = {metric(out, x).item():.4f}" + "\n"
         metrics_out_baseline += f"{metric.name} = {metric(out_baseline, x).item():.4f}" + "\n"
     metrics_out += f"Inference time = {ram_time:.3f}s"
     metrics_out_baseline += f"Inference time = {dpir_time:.3f}s"
 
-    ### Process y when y shape is different from x shape
-    if physics.name == "MRI":
-        y_plot = physics.physics.prox_l2(physics.physics.A_adjoint(y), y, 1e4)
-    elif physics.name == "CT":
-        y_plot = physics.physics.A_adjoint(y)
-    else:
-        y_plot = y.clone()
-
     ### Processing images for plotting :
     #     - clip value outside of [0,1]
     #     - shape (1, C, H, W) -> (C, H, W)

factories.py

@@ -140,7 +140,7 @@ class PhysicsWithGenerator(torch.nn.Module):
                 max_iter=10,
             )
             self.physics_generator = None
-            self.generator = SigmaGenerator(sigma_min=0.001, sigma_max=0.2, device=device_str)
+            self.generator = SigmaGenerator(sigma_min=1e-5, sigma_max=1e-4, device=device_str)
             self.saved_params = {"updatable_params": {"sigma": 1e-4},
                                  "updatable_params_converter": {"sigma": float},
                                  "fixed_params": {"noise_sigma_min": 0.001, "noise_sigma_max": 0.2,