physics

factories.py

@@ -3,7 +3,7 @@ from typing import Any, List
 import deepinv as dinv
 import numpy as np
 import torch
-from deepinv.physics.generator import MotionBlurGenerator, SigmaGenerator
+from deepinv.physics.generator import MotionBlurGenerator, SigmaGenerator, GainGenerator
 from torchvision import transforms
 
 from datasets import Preprocessed_fastMRI, Preprocessed_LIDCIDRI, LsdirMiniDataset
@@ -127,7 +127,7 @@ class PhysicsWithGenerator(torch.nn.Module):
                                  "fixed_params": {"acceleration_factor": 4}}
         elif self.name == "CT":
             acceleration_factor = 10
-            img_h = 480
+            img_h = 512
             angles = int(img_h / acceleration_factor)
             # angles = torch.linspace(0, 180, steps=10)
             self.physics = dinv.physics.Tomography(
@@ -139,11 +139,11 @@ class PhysicsWithGenerator(torch.nn.Module):
                 noise_model=dinv.physics.GaussianNoise(sigma=1e-4),
                 max_iter=10,
             )
-            self.physics_generator = None
-            self.generator = SigmaGenerator(sigma_min=1e-5, sigma_max=1e-4, device=device_str)
-            self.saved_params = {"updatable_params": {"sigma": 1e-4},
+            self.physics_generator = SigmaGenerator(sigma_min=1e-4, sigma_max=1e-4, device=device_str)
+            self.generator = SigmaGenerator(sigma_min=1e-4, sigma_max=1e-4, device=device_str)
+            self.saved_params = {"updatable_params": {"sigma": 0.1},
                                  "updatable_params_converter": {"sigma": float},
-                                 "fixed_params": {"noise_sigma_min": 0.001, "noise_sigma_max": 0.2,
+                                 "fixed_params": {"noise_sigma_min": 1e-4, "noise_sigma_max": 1e-4,
                                                   "angles": angles, "max_iter": 10}}
 
     def display_saved_params(self) -> str:
@@ -235,6 +235,10 @@ class EvalModel(torch.nn.Module):
             self.model.eval()
 
     def forward(self, y: torch.Tensor, physics: torch.nn.Module) -> torch.Tensor:
+        physics.noise_model.sigma = torch.nn.Parameter(torch.tensor([1e-06]))
+        physics.noise_model.gain = torch.nn.Parameter(torch.tensor([1e-06]))
+        print('sigma = ', physics.noise_model.sigma)
+        print('gain = ', physics.noise_model.gain)
         return self.model(y, physics=physics)
 
 
@@ -401,7 +405,7 @@ class BaselineModel(torch.nn.Module):
             # Set the DPIR algorithm parameters
             sigma_float = physics.noise_model.sigma.item()  # sigma should be a single value
             lip_const = physics.compute_norm(physics.A_adjoint(y))
-            lamb, sigma_denoiser, stepsize, max_iter = self.get_DPIR_CT_params(sigma_float, max_iter=8,
+            lamb, sigma_denoiser, stepsize, max_iter = self.get_DPIR_CT_params(sigma_float, max_iter=1,  # for debugging speed
                                                                                lip_cons=lip_const.item())
             params_algo = {"stepsize": stepsize, "g_param": sigma_denoiser, "lambda": lamb}
             early_stop = False  # Do not stop algorithm with convergence criteria

model_factory.py

@@ -51,12 +51,12 @@ class ArtifactRemoval(nn.Module):
         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?
+            sigma = 1e-5  # 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?
+            gamma = 1e-5  # 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)