remove model choices

app.py

@@ -20,7 +20,7 @@ DEVICE_STR = 'cuda'
 ### Gradio Utils
 def generate_imgs(dataset: EvalDataset, idx: int,
                   model: EvalModel, baseline: BaselineModel,
-                  physics: PhysicsWithGenerator, use_gen: bool, 
+                  physics: PhysicsWithGenerator, use_gen: bool,
                   metrics: List[Metric]):
     ### Load 1 image
     x = dataset[idx]    # shape : (3, 256, 256)
@@ -80,7 +80,7 @@ def update_random_idx_and_generate_imgs(dataset: EvalDataset,
                                         physics: PhysicsWithGenerator,
                                         use_gen: bool,
                                         metrics: List[Metric]):
-    idx = random.randint(0, len(dataset))
+    idx = random.randint(0, len(dataset)-1)
     x, y, out, out_baseline, saved_params_str, metrics_y, metrics_out, metrics_out_baseline = generate_imgs(dataset, 
                                                                                                             idx, 
                                                                                                             model, 
@@ -125,10 +125,19 @@ def save_imgs(dataset: EvalDataset, idx: int, physics: PhysicsWithGenerator,
     dinv.utils.plot([x, y, out_a, out_b], titles=titles, show=False, save_fn=save_path)
 
 get_list_metrics_on_DEVICE_STR = partial(Metric.get_list_metrics, device_str=DEVICE_STR)
-get_baseline_model_on_DEVICE_STR = partial(BaselineModel, device_str=DEVICE_STR)
 get_eval_model_on_DEVICE_STR = partial(EvalModel, device_str=DEVICE_STR)
+get_baseline_model_on_DEVICE_STR = partial(BaselineModel, device_str=DEVICE_STR)
 get_dataset_on_DEVICE_STR = partial(EvalDataset, device_str=DEVICE_STR)
-get_physics_generator_on_DEVICE_STR = partial(PhysicsWithGenerator, device_str=DEVICE_STR)
+get_physics_on_DEVICE_STR = partial(PhysicsWithGenerator, device_str=DEVICE_STR)
+
+def get_physics(physics_name):
+    if physics_name == 'MRI':
+        baseline = get_baseline_model_on_DEVICE_STR('DPIR_MRI')
+    elif physics_name == 'CT':
+        baseline = get_baseline_model_on_DEVICE_STR('DPIR_CT')
+    else:
+        baseline = get_baseline_model_on_DEVICE_STR('DPIR')
+    return get_physics_on_DEVICE_STR(physics_name), baseline
 
 def get_model(model_name, ckpt_pth):
     if model_name in BaselineModel.all_baselines:
@@ -154,18 +163,14 @@ 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
+    model_b_placeholder = gr.State(lambda: get_baseline_model_on_DEVICE_STR("DPIR"))  # 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
+    physics_placeholder = gr.State(lambda: get_physics_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])
-    def dynamic_layout(model_a, model_b, dataset, physics, metrics):
+    @gr.render(inputs=[dataset_placeholder, physics_placeholder, metrics_placeholder])
+    def dynamic_layout(dataset, physics, metrics):
         ### LAYOUT
-        model_a_name = model_a.base_name
-        model_a_full_name = model_a.name
-        model_b_name = model_b.base_name
-        model_b_full_name = model_b.name
         dataset_name = dataset.name
         physics_name = physics.name
         metric_names = [metric.name for metric in metrics]
@@ -180,87 +185,65 @@ with gr.Blocks(title=title, css=custom_css) as interface:
                     with gr.Column():
                         y_image = gr.Image(label=f"{physics_name} IMAGE", interactive=False)
                         y_metrics = gr.Textbox(label="Metrics(y, x)", elem_classes=["fixed-textbox"],)
-                with gr.Row():
-                    choose_dataset = gr.Radio(choices=EvalDataset.all_datasets,
-                                              label="List of EvalDataset",
-                                              value=dataset_name,
-                                              scale=2)
-                    idx_slider = gr.Slider(minimum=0, maximum=len(dataset)-1, step=1, label="Sample index", scale=1)
-        
+
                 choose_physics = gr.Radio(choices=PhysicsWithGenerator.all_physics,
                                           label="List of PhysicsWithGenerator",
                                           value=physics_name)
                 with gr.Row():
-                    key_selector = gr.Dropdown(choices=list(physics.saved_params["updatable_params"].keys()), 
+                    key_selector = gr.Dropdown(choices=list(physics.saved_params["updatable_params"].keys()),
                                                label="Updatable Parameter Key",
                                                scale=2)
                     value_text = gr.Textbox(label="Update Value", scale=2)
-                    with gr.Column(scale=1):
-                        update_button = gr.Button("Update Param")
-                        use_generator_button = gr.Checkbox(label="Use param generator")
-                    
+                    update_button = gr.Button("Manually update parameter value", scale=1)
+
             with gr.Column():
-                with gr.Row():      
+                with gr.Row():
                     with gr.Column():
-                        model_a_out = gr.Image(label=f"{model_a_full_name} OUTPUT", interactive=False)
-                        out_a_metric = gr.Textbox(label="Metrics(model_a(y), x)", elem_classes=["fixed-textbox"])
-                        load_model_a = gr.Button("Load model A...", scale=1)
+                        model_a_out = gr.Image(label="RAM OUTPUT", interactive=False)
+                        out_a_metric = gr.Textbox(label="Metrics(RAM(y, physics), x)", elem_classes=["fixed-textbox"])
                     with gr.Column():
-                        model_b_out = gr.Image(label=f"{model_b_full_name} OUTPUT", interactive=False)
-                        out_b_metric = gr.Textbox(label="Metrics(model_b(y), x)", elem_classes=["fixed-textbox"])
-                        load_model_b = gr.Button("Load model B...", scale=1)
-                with gr.Row():
-                    choose_model_a = gr.Dropdown(choices=EvalModel.all_models + BaselineModel.all_baselines,
-                                                 label="List of Model A",
-                                                 value=model_a_name,
-                                                 scale=2)
-                    path_a_str = gr.Textbox(value=model_a.ckpt_pth, label="Checkpoint path", scale=3)
+                        model_b_out = gr.Image(label="DPIR OUTPUT", interactive=False)
+                        out_b_metric = gr.Textbox(label="Metrics(DPIR(y, physics), x)", elem_classes=["fixed-textbox"])
                 with gr.Row():
-                    choose_model_b = gr.Dropdown(choices=EvalModel.all_models + BaselineModel.all_baselines,
-                                                 label="List of Model B",
-                                                 value=model_b_name,
-                                                 scale=2)
-                    path_b_str = gr.Textbox(value=model_b.ckpt_pth, label="Checkpoint path", scale=3)
-        
-        # Components: Load Metric + Load/Save Buttons
+                    choose_dataset = gr.Radio(choices=EvalDataset.all_datasets,
+                                              label="List of EvalDataset",
+                                              value=dataset_name,
+                                              scale=2)
+                    idx_slider = gr.Slider(minimum=0, maximum=len(dataset)-1, step=1, label="Sample index", scale=1)
+
+        # Components: Load Metric + Load image Buttons
         with gr.Row():
-            with gr.Column():
+            with gr.Column(scale=2):
                 choose_metrics = gr.CheckboxGroup(choices=Metric.all_metrics,
                                                   value=metric_names,
                                                   label="Choose metrics you are interested")
-            with gr.Column():
+            use_generator_button = gr.Checkbox(label="Generate valid physics parameters", scale=1)
+            with gr.Column(scale=1):
                 load_button = gr.Button("Load images...")
-                load_random_button = gr.Button("Load randomly...")
-        save_button = gr.Button("Save images...")
+	                load_random_button = gr.Button("Load randomly...")
 
         ### Event listeners
         choose_dataset.change(fn=get_dataset_on_DEVICE_STR,
                               inputs=choose_dataset,
                               outputs=dataset_placeholder)
-        choose_physics.change(fn=get_physics_generator_on_DEVICE_STR,
+        choose_physics.change(fn=get_physics,
                               inputs=choose_physics,
-                              outputs=physics_placeholder)
+                              outputs=[physics_placeholder, model_b_placeholder])
         update_button.click(fn=physics.update_and_display_params, inputs=[key_selector, value_text], outputs=physics_params)
-        load_model_a.click(fn=get_model,
-                           inputs=[choose_model_a, path_a_str],
-                           outputs=model_a_placeholder)
-        load_model_b.click(fn=get_model,
-                           inputs=[choose_model_b, path_b_str],
-                           outputs=model_b_placeholder)
         choose_metrics.change(fn=get_list_metrics_on_DEVICE_STR,
                               inputs=choose_metrics,
                               outputs=metrics_placeholder)
-        load_button.click(fn=generate_imgs, 
-                          inputs=[dataset_placeholder, 
-                                  idx_slider, 
+        load_button.click(fn=generate_imgs,
+                          inputs=[dataset_placeholder,
+                                  idx_slider,
                                   model_a_placeholder,
                                   model_b_placeholder,
                                   physics_placeholder,
                                   use_generator_button,
-                                  metrics_placeholder], 
+                                  metrics_placeholder],
                           outputs=[clean, y_image, model_a_out, model_b_out, physics_params, y_metrics, out_a_metric, out_b_metric])
-        load_random_button.click(fn=update_random_idx_and_generate_imgs, 
-                                 inputs=[dataset_placeholder, 
+        load_random_button.click(fn=update_random_idx_and_generate_imgs,
+                                 inputs=[dataset_placeholder,
                                          model_a_placeholder,
                                          model_b_placeholder,
                                          physics_placeholder,

datasets.py

@@ -93,7 +93,7 @@ class LsdirMiniDataset(torch.utils.data.Dataset):
 	transform: Optional[Callable] = None,
     ) -> None:
         self.root = root
-        self.image_files = [f for f in os.listdir(self.root) if f.lower().endswith(('.png', '.JPEG'))]
+        self.image_files = [f for f in os.listdir(self.root) if f.lower().endswith(('.png', '.jpeg'))]
         self.transform = transform
 
     def __len__(self) -> int:

evals.py

@@ -47,22 +47,21 @@ class PhysicsWithGenerator(torch.nn.Module):
         if self.name not in self.all_physics:
             raise ValueError(f"{self.name} is unavailable.")
 
-        self.sigma_generator = SigmaGenerator(sigma_min=0.001, sigma_max=0.2, device=device_str)
         if self.name == "MotionBlur_easy":
             psf_size = 31
-            self.physics = dinv.physics.Blur(noise_model=dinv.physics.GaussianNoise(sigma=.01), padding="valid",
-                                             device=device_str)
-            self.physics_generator = MotionBlurGenerator((psf_size, psf_size), l=0.1, sigma=0.1, device=device_str) + SigmaGenerator(sigma_min=0.01, sigma_max=0.01, device=device_str)
+            self.physics = dinv.physics.Blur(noise_model=dinv.physics.GaussianNoise(sigma=.01),
+                                             padding="valid", device=device_str)
+            self.physics_generator = MotionBlurGenerator((psf_size, psf_size), l=0.1, sigma=0.1, device=device_str)
             self.generator = self.physics_generator + SigmaGenerator(sigma_min=0.01, sigma_max=0.01, device=device_str)
-            self.saved_params = {"updatable_params": {"sigma": 0.05},
+            self.saved_params = {"updatable_params": {"sigma": 0.01},
                                  "updatable_params_converter": {"sigma": float},
                                  "fixed_params": {"noise_sigma_min": 0.01, "noise_sigma_max": 0.01,
                                                   "psf_size": 31, "motion_gen_l": 0.1, "motion_gen_s": 0.1}}
         elif self.name == "MotionBlur_medium":
             psf_size = 31
-            self.physics = dinv.physics.Blur(noise_model=dinv.physics.GaussianNoise(sigma=.05), padding="valid",
-                                             device=device_str)
-            self.physics_generator = MotionBlurGenerator((psf_size, psf_size), l=0.6, sigma=0.5, device=device_str) + SigmaGenerator(sigma_min=0.05, sigma_max=0.05, device=device_str)
+            self.physics = dinv.physics.Blur(noise_model=dinv.physics.GaussianNoise(sigma=.05),
+                                             padding="valid", device=device_str)
+            self.physics_generator = MotionBlurGenerator((psf_size, psf_size), l=0.6, sigma=0.5, device=device_str)
             self.generator = self.physics_generator + SigmaGenerator(sigma_min=0.05, sigma_max=0.05, device=device_str)
             self.saved_params = {"updatable_params": {"sigma": 0.05},
                                  "updatable_params_converter": {"sigma": float},
@@ -70,62 +69,61 @@ class PhysicsWithGenerator(torch.nn.Module):
                                                   "psf_size": 31, "motion_gen_l": 0.6, "motion_gen_s": 0.5}}
         elif self.name == "MotionBlur_hard":
             psf_size = 31
-            self.physics = dinv.physics.Blur(noise_model=dinv.physics.GaussianNoise(sigma=.1), padding="valid",
-                                             device=device_str)
-            self.physics_generator = MotionBlurGenerator((psf_size, psf_size), l=1.2, sigma=1.0, device=device_str) + SigmaGenerator(sigma_min=0.1, sigma_max=0.1, device=device_str)
+            self.physics = dinv.physics.Blur(noise_model=dinv.physics.GaussianNoise(sigma=.1),
+                                             padding="valid", device=device_str)
+            self.physics_generator = MotionBlurGenerator((psf_size, psf_size), l=1.2, sigma=1.0, device=device_str)
             self.generator = self.physics_generator + SigmaGenerator(sigma_min=0.1, sigma_max=0.1, device=device_str)
-            self.saved_params = {"updatable_params": {"sigma": 0.05},
+            self.saved_params = {"updatable_params": {"sigma": 0.1},
                                  "updatable_params_converter": {"sigma": float},
                                  "fixed_params": {"noise_sigma_min": 0.1, "noise_sigma_max": 0.1,
                                                   "psf_size": 31, "motion_gen_l": 1.2, "motion_gen_s": 1.0}}
         elif self.name == "GaussianBlur_easy":
             psf_size = 31
-            self.physics = dinv.physics.Blur(noise_model=dinv.physics.GaussianNoise(sigma=0.01), padding="valid",
-                                             device=device_str)
+            self.physics = dinv.physics.Blur(noise_model=dinv.physics.GaussianNoise(sigma=0.01),
+                                             padding="valid", device=device_str)
             self.physics_generator = GaussianBlurGenerator(psf_size=(psf_size, psf_size),
                                                            sigma_min=1.0, sigma_max=1.0,
                                                            num_channels=1,
-                                                           device=device_str) + SigmaGenerator(sigma_min=0.01, sigma_max=0.01, device=device_str)
-            self.generator = self.physics_generator
-            self.saved_params = {"updatable_params": {},
-                                 "updatable_params_converter": {},
-                                 "fixed_params": {"noise_sigma": 0.01, "blur_sigma": 1.0,
-                                                  "psf_size": 31, "num_channels": 1}}
+                                                           device=device_str)
+            self.generator = self.physics_generator + SigmaGenerator(sigma_min=0.01, sigma_max=0.01, device=device_str)
+            self.saved_params = {"updatable_params": {"sigma": 0.01},
+                                 "updatable_params_converter": {"sigma": float},
+                                 "fixed_params": {"noise_sigma_min": 0.01, "noise_sigma_max": 0.01,
+                                                  "blur_sigma": 1.0, "psf_size": 31, "num_channels": 1}}
         elif self.name == "GaussianBlur_medium":
             psf_size = 31
-            self.physics = dinv.physics.Blur(noise_model=dinv.physics.GaussianNoise(sigma=0.05), padding="valid",
-                                             device=device_str)
+            self.physics = dinv.physics.Blur(noise_model=dinv.physics.GaussianNoise(sigma=0.05),
+                                             padding="valid", device=device_str)
             self.physics_generator = GaussianBlurGenerator(psf_size=(psf_size, psf_size),
                                                            sigma_min=2.0, sigma_max=2.0,
                                                            num_channels=1,
-                                                           device=device_str) + SigmaGenerator(sigma_min=0.05, sigma_max=0.05, device=device_str)
-            self.generator = self.physics_generator
-            self.saved_params = {"updatable_params": {},
-                                 "updatable_params_converter": {},
-                                 "fixed_params": {"noise_sigma": 0.05, "blur_sigma": 2.0,
-                                                  "psf_size": 31, "num_channels": 1}}
+                                                           device=device_str)
+            self.generator = self.physics_generator + SigmaGenerator(sigma_min=0.05, sigma_max=0.05, device=device_str)
+            self.saved_params = {"updatable_params": {"sigma": 0.05},
+                                 "updatable_params_converter": {"sigma": float},
+                                 "fixed_params": {"noise_sigma_min": 0.05, "noise_sigma_max": 0.05,
+                                                  "blur_sigma": 2.0, "psf_size": 31, "num_channels": 1}}
         elif self.name == "GaussianBlur_hard":
             psf_size = 31
-            self.physics = dinv.physics.Blur(noise_model=dinv.physics.GaussianNoise(sigma=0.05), padding="valid",
-                                             device=device_str)
+            self.physics = dinv.physics.Blur(noise_model=dinv.physics.GaussianNoise(sigma=0.05),
+                                             padding="valid", device=device_str)
             self.physics_generator = GaussianBlurGenerator(psf_size=(psf_size, psf_size),
                                                            sigma_min=4.0, sigma_max=4.0,
                                                            num_channels=1,
-                                                           device=device_str) + SigmaGenerator(sigma_min=0.1, sigma_max=0.1, device=device_str)
-            self.generator = self.physics_generator
-            self.saved_params = {"updatable_params": {},
-                                 "updatable_params_converter": {},
-                                 "fixed_params": {"noise_sigma": 0.1, "blur_sigma": 4.0,
-                                                  "psf_size": 31, "num_channels": 1}}
+                                                           device=device_str)
+            self.generator = self.physics_generator + SigmaGenerator(sigma_min=0.05, sigma_max=0.05, device=device_str)
+            self.saved_params = {"updatable_params": {"sigma": 0.05},
+                                 "updatable_params_converter": {"sigma": float},
+                                 "fixed_params": {"noise_sigma_min": 0.05, "noise_sigma_max": 0.05,
+                                                  "blur_sigma": 4.0, "psf_size": 31, "num_channels": 1}}
         elif self.name == "MRI":
-            self.physics = dinv.physics.MRI(img_size=(640, 320), noise_model=dinv.physics.GaussianNoise(sigma=.01),
-                                            device=device_str)
+            self.physics = dinv.physics.MRI(noise_model=dinv.physics.GaussianNoise(sigma=.01),
+                                            img_size=(640, 320), device=device_str)
             self.physics_generator = dinv.physics.generator.RandomMaskGenerator((2, 640, 320), acceleration_factor=4)
-            self.generator = self.physics_generator  # + self.sigma_generator
-            self.saved_params = {"updatable_params": {"sigma": 0.05},
+            self.generator = self.physics_generator
+            self.saved_params = {"updatable_params": {"sigma": 0.01},
                                  "updatable_params_converter": {"sigma": float},
-                                 "fixed_params": {"noise_sigma_min": 0.001, "noise_sigma_max": 0.2,
-                                                  "acceleration_factor": 4}}
+                                 "fixed_params": {"acceleration_factor": 4}}
         elif self.name == "CT":
             acceleration_factor = 10
             img_h = 480
@@ -141,10 +139,10 @@ class PhysicsWithGenerator(torch.nn.Module):
                 max_iter=10,
             )
             self.physics_generator = None
-            self.generator = self.sigma_generator
-            self.saved_params = {"updatable_params": {"sigma": 0.1},
+            self.generator = SigmaGenerator(sigma_min=0.001, sigma_max=0.2, 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.,
+                                 "fixed_params": {"noise_sigma_min": 0.001, "noise_sigma_max": 0.2,
                                                   "angles": angles, "max_iter": 10}}
 
     def display_saved_params(self) -> str:
@@ -189,7 +187,7 @@ class PhysicsWithGenerator(torch.nn.Module):
         self.physics.update(**kwargs)
 
     def forward(self, x: torch.Tensor, use_gen: bool) -> torch.Tensor:
-        if self.name in ["MotionBlur_easy", "MotionBlur_medium", "MotionBlur_hard", "GaussianBlur"] and not hasattr(self.physics, "filter"):
+        if self.name in ["MotionBlur_easy", "MotionBlur_medium", "MotionBlur_hard", "GaussianBlur_easy", "GaussianBlur_medium", "GaussianBlur_hard"] and not hasattr(self.physics, "filter"):
             use_gen = True
         elif self.name in ["MRI"] and not hasattr(self.physics, "mask"):
             use_gen = True
@@ -247,8 +245,7 @@ class BaselineModel(torch.nn.Module):
         -> BaselineModel should be models that are already trained and will have fixed weights.
         -> Eval model will change depending on differents checkpoints.
     """
-    all_baselines = ["DRUNET", "PnP-PGD-DRUNET", "SWINIRx2", "SWINIRx4", "DPIR",
-                     "DPIR_MRI", "DPIR_CT", "PDNET"]
+    all_baselines = ["DPIR", "DPIR_MRI", "DPIR_CT"]
 
     def __init__(self, model_name: str, device_str: str = "cpu") -> None:
         super().__init__()
@@ -257,58 +254,6 @@ class BaselineModel(torch.nn.Module):
         self.name = self.base_name
         if self.name not in self.all_baselines:
             raise ValueError(f"{self.name} is unavailable.")
-        elif self.name == "DRUNET":
-            n_channels = 3
-            ckpt_pth = "ckpt/drunet_deepinv_color_finetune_22k.pth"
-            self.model = dinv.models.DRUNet(in_channels=n_channels,
-                                            out_channels=n_channels,
-                                            device=device_str,
-                                            pretrained=ckpt_pth)
-            self.model.eval()  # Set the model to evaluation mode
-        elif self.name == 'PDNET':
-            ckpt_pth = "ckpt/pdnet.pth.tar"
-            self.model = get_model(model_name='pdnet',
-                                   device=device_str)
-            self.model.eval()
-            self.model.load_state_dict(torch.load(ckpt_pth, map_location=lambda storage, loc: storage)['state_dict'])
-        elif self.name == "SWINIRx2":
-            n_channels = 3
-            scale = 2
-            ckpt_pth = "ckpt/001_classicalSR_DF2K_s64w8_SwinIR-M_x2.pth"
-            upsampler = 'nearest+conv' if 'realSR' in ckpt_pth else 'pixelshuffle'
-            self.model = dinv.models.SwinIR(upscale=scale, in_chans=n_channels, img_size=64, window_size=8,
-                                            img_range=1., depths=[6, 6, 6, 6, 6, 6], embed_dim=180,
-                                            num_heads=[6, 6, 6, 6, 6, 6],
-                                            mlp_ratio=2, upsampler=upsampler, resi_connection='1conv',
-                                            pretrained=ckpt_pth)
-            self.model.to(device_str)
-            self.model.eval()  # Set the model to evaluation mode
-        elif self.name == "SWINIRx4":
-            n_channels = 3
-            scale = 4
-            ckpt_pth = "ckpt/001_classicalSR_DF2K_s64w8_SwinIR-M_x4.pth"
-            upsampler = 'nearest+conv' if 'realSR' in ckpt_pth else 'pixelshuffle'
-            self.model = dinv.models.SwinIR(upscale=scale, in_chans=n_channels, img_size=64, window_size=8,
-                                            img_range=1., depths=[6, 6, 6, 6, 6, 6], embed_dim=180,
-                                            num_heads=[6, 6, 6, 6, 6, 6],
-                                            mlp_ratio=2, upsampler=upsampler, resi_connection='1conv',
-                                            pretrained=ckpt_pth)
-            self.model.to(device_str)
-            self.model.eval()  # Set the model to evaluation mode
-
-        elif self.name == "PnP-PGD-DRUNET":
-            n_channels = 3
-            ckpt_pth = "ckpt/drunet_deepinv_color_finetune_22k.pth"
-            drunet = dinv.models.DRUNet(in_channels=n_channels,
-                                        out_channels=n_channels,
-                                        device=device_str,
-                                        pretrained=ckpt_pth)
-            drunet.eval()  # Set the model to evaluation mode
-            self.model = dinv.optim.optim_builder(iteration="PGD",
-                                                  prior=dinv.optim.PnP(drunet).to(device_str),
-                                                  data_fidelity=dinv.optim.L2(),
-                                                  max_iter=20,
-                                                  params_algo={'stepsize': 1., 'g_param': .05})
         elif self.name == "DPIR":
             n_channels = 3
             ckpt_pth = "ckpt/drunet_deepinv_color_finetune_22k.pth"
@@ -418,11 +363,7 @@ class BaselineModel(torch.nn.Module):
         return lamb, list(sigma_denoiser), list(stepsize), max_iter
 
     def forward(self, y: torch.Tensor, physics: torch.nn.Module) -> torch.Tensor:
-        if self.name == "DRUNET":
-            return self.model(y, sigma=physics.noise_model.sigma)
-        elif self.name == "PnP-PGD-DRUNET":
-            return self.model(y, physics=physics)
-        elif self.name == "DPIR":
+        if self.name == "DPIR":
             # Set the DPIR algorithm parameters
             sigma_float = physics.noise_model.sigma.item()  # sigma should be a single value
             max_iter = 8
@@ -460,7 +401,7 @@ class BaselineModel(torch.nn.Module):
                 params_algo=params_algo,
             )
             return model(y, physics=physics)
-        elif self.name == "DPIR_CT":
+        else self.name == "DPIR_CT":
             # 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))
@@ -485,42 +426,10 @@ class BaselineModel(torch.nn.Module):
                 custom_init=custom_init
             )
             return algo(y, physics=physics)
-        elif self.name == 'SWINIRx4':
-            window_size = 8
-            scale = 4
-            _, _, h_old, w_old = y.size()
-            h_pad = (h_old // window_size + 1) * window_size - h_old
-            w_pad = (w_old // window_size + 1) * window_size - w_old
-            img_lq = torch.cat([y, torch.flip(y, [2])], 2)[:, :, :h_old + h_pad, :]
-            img_lq = torch.cat([img_lq, torch.flip(img_lq, [3])], 3)[:, :, :, :w_old + w_pad]
-            output = self.model(img_lq)
-            output = output[..., :h_old * scale, :w_old * scale]
-            output = self.circular_roll(output, -2, -2)
-            # check shape of adjoint
-            x_adj = physics.A_adjoint(y)
-            output = output[..., :x_adj.size(-2), :x_adj.size(-1)]
-            return output
-        elif self.name == 'SWINIRx2':
-            window_size = 8
-            scale = 2
-            _, _, h_old, w_old = y.size()
-            h_pad = (h_old // window_size + 1) * window_size - h_old
-            w_pad = (w_old // window_size + 1) * window_size - w_old
-            img_lq = torch.cat([y, torch.flip(y, [2])], 2)[:, :, :h_old + h_pad, :]
-            img_lq = torch.cat([img_lq, torch.flip(img_lq, [3])], 3)[:, :, :, :w_old + w_pad]
-            output = self.model(img_lq)
-            output = output[..., :h_old * scale, :w_old * scale]
-            output = self.circular_roll(output, -1, -1)
-            # check shape of adjoint
-            x_adj = physics.A_adjoint(y)
-            output = output[..., :x_adj.size(-2), :x_adj.size(-1)]
-            return output
-        else:
-            return self.model(y)
 
 
 class EvalDataset(torch.utils.data.Dataset):
-    """"""
+
     all_datasets = ["Natural", "MRI", "CT"]
 
     def __init__(self, dataset_name: str, device_str: str = "cpu") -> None: