gradio demo

.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.pth.tar filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,2 @@
+.ipynb
+__pycache__

README.md

@@ -0,0 +1,13 @@
+---
+title: Denoising
+emoji: 💻
+colorFrom: blue
+colorTo: yellow
+sdk: gradio
+sdk_version: 4.19.2
+app_file: app.py
+pinned: false
+license: bsd-3-clause
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,269 @@
+import json
+import os
+import random
+from functools import partial
+from pathlib import Path
+from typing import List
+
+import deepinv as dinv
+import gradio as gr
+import torch
+from PIL import Image
+from torchvision import transforms
+
+from evals import PhysicsWithGenerator, EvalModel, BaselineModel, EvalDataset, Metric
+
+
+### Gradio Utils
+def generate_imgs(dataset: EvalDataset, idx: int, 
+                  model: EvalModel, baseline: BaselineModel, 
+                  physics: PhysicsWithGenerator, use_gen: bool, 
+                  metrics: List[Metric]):
+    ### Load 1 image
+    x = dataset[idx]    # shape : (3, 256, 256)
+    x = x.unsqueeze(0)  # shape : (1, 3, 256, 256)
+
+    with torch.no_grad():
+        ### Compute y
+        y = physics(x, use_gen)  # possible reduction in img shape due to Blurring
+
+        ### Compute x_hat
+        out = model(y=y, physics=physics.physics)
+        out_baseline = baseline(y=y, physics=physics.physics)
+
+        ### Process tensors before metric computation
+        if "Blur" in physics.name:
+            w_1, w_2 = (x.shape[2] - y.shape[2]) // 2, (x.shape[2] + y.shape[2]) // 2
+            h_1, h_2 = (x.shape[3] - y.shape[3]) // 2, (x.shape[3] + y.shape[3]) // 2
+
+            x = x[..., w_1:w_2, h_1:h_2]
+            out = out[..., w_1:w_2, h_1:h_2]
+            if out_baseline.shape != out.shape:
+                out_baseline = out_baseline[..., w_1:w_2, h_1:h_2]
+
+        ### 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"
+            metrics_out += f"{metric.name} = {metric(out, x).item():.4f}" + "\n"
+            metrics_out_baseline += f"{metric.name} = {metric(out_baseline, x).item():.4f}" + "\n"
+
+    ### Process y when y shape is different from x shape
+    if physics.name == "MRI" or "SR" in physics.name:
+        y_plot = physics.physics.prox_l2(physics.physics.A_adjoint(y), y, 1e4)
+    else:
+        y_plot = y.clone()
+
+    ### Processing images for plotting :
+    #     - clip value outside of [0,1]
+    #     - shape (1, C, H, W) -> (C, H, W)
+    #     - torch.Tensor object -> Pil object
+    process_img = partial(dinv.utils.plotting.preprocess_img, rescale_mode="clip")
+    to_pil = transforms.ToPILImage()
+    x = to_pil(process_img(x)[0].to('cpu'))
+    y = to_pil(process_img(y_plot)[0].to('cpu'))
+    out = to_pil(process_img(out)[0].to('cpu'))
+    out_baseline = to_pil(process_img(out_baseline)[0].to('cpu'))
+
+
+    return x, y, out, out_baseline, physics.display_saved_params(), metrics_y, metrics_out, metrics_out_baseline
+
+def update_random_idx_and_generate_imgs(dataset: EvalDataset, 
+                                        model: EvalModel, 
+                                        baseline: BaselineModel, 
+                                        physics: PhysicsWithGenerator,
+                                        use_gen: bool,
+                                        metrics: List[Metric]):
+    idx = random.randint(0, len(dataset))
+    x, y, out, out_baseline, saved_params_str, metrics_y, metrics_out, metrics_out_baseline = generate_imgs(dataset, 
+                                                                                                            idx, 
+                                                                                                            model, 
+                                                                                                            baseline, 
+                                                                                                            physics, 
+                                                                                                            use_gen, 
+                                                                                                            metrics)
+    return idx, x, y, out, out_baseline, saved_params_str, metrics_y, metrics_out, metrics_out_baseline
+
+def save_imgs(dataset: EvalDataset, idx: int, physics: PhysicsWithGenerator,
+              model_a: EvalModel | BaselineModel, model_b: EvalModel | BaselineModel,
+              x: Image.Image, y: Image.Image,
+              out_a: Image.Image, out_b: Image.Image,
+              y_metrics_str: str,
+              out_a_metric_str : str, out_b_metric_str: str) -> None:
+
+    ### PROCESSES STR
+    physics_params_str = ""
+    for param_name, param_value in physics.saved_params["updatable_params"].items():
+        physics_params_str += f"{param_name}_{param_value}-"
+    physics_params_str = physics_params_str[:-1] if physics_params_str.endswith("-") else physics_params_str
+    y_metrics_str = y_metrics_str.replace(" = ", "_").replace("\n", "-")
+    y_metrics_str = y_metrics_str[:-1] if y_metrics_str.endswith("-") else y_metrics_str
+    out_a_metric_str = out_a_metric_str.replace(" = ", "_").replace("\n", "-")
+    out_a_metric_str = out_a_metric_str[:-1] if out_a_metric_str.endswith("-") else out_a_metric_str
+    out_b_metric_str = out_b_metric_str.replace(" = ", "_").replace("\n", "-")
+    out_b_metric_str = out_b_metric_str[:-1] if out_b_metric_str.endswith("-") else out_b_metric_str
+
+    save_path = SAVE_IMG_DIR / f"{dataset.name}+{idx}+{physics.name}+{physics_params_str}+{y_metrics_str}+{model_a.name}+{out_a_metric_str}+{model_b.name}+{out_b_metric_str}.png"
+    titles = [f"{dataset.name}[{idx}]",
+              f"y = {physics.name}(x)",
+              f"{model_a.name}",
+              f"{model_b.name}"]
+
+    # Pil object -> torch.Tensor
+    to_tensor = transforms.ToTensor()
+    x = to_tensor(x)
+    y = to_tensor(y)
+    out_a = to_tensor(out_a)
+    out_b = to_tensor(out_b)
+
+    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_dataset_on_DEVICE_STR = partial(EvalDataset, device_str=DEVICE_STR)
+get_physics_generator_on_DEVICE_STR = partial(PhysicsWithGenerator, device_str=DEVICE_STR)
+
+def get_model(model_name, ckpt_pth):
+    if model_name in BaselineModel.all_baselines:
+        return get_baseline_model_on_DEVICE_STR(model_name)
+    else:
+        return get_eval_model_on_DEVICE_STR(model_name, ckpt_pth)
+
+
+### Gradio Blocks interface
+
+# Define custom CSS
+custom_css = """
+.fixed-textbox textarea {
+    height: 90px !important;  /* Adjust height to fit exactly 4 lines */
+    overflow: scroll;         /* Add a scroll bar if necessary */
+    resize: none;             /* User can resize vertically the textbox */
+}
+"""
+
+title = "Inverse problem playground"  # displayed on gradio tab and in the gradio page
+with gr.Blocks(title=title, css=custom_css) as interface:
+    gr.Markdown("## " + title)
+
+    # 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
+    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):
+        ### 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]
+
+        # Components: Inputs/Outputs + Load EvalDataset/PhysicsWithGenerator/EvalModel/BaselineModel
+        with gr.Row():
+            with gr.Column():
+                with gr.Row():
+                    with gr.Column():
+                        clean = gr.Image(label=f"{dataset_name} IMAGE", interactive=False)
+                        physics_params = gr.Textbox(label="Physics parameters", elem_classes=["fixed-textbox"], value=physics.display_saved_params())
+                    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()), 
+                                               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")
+                    
+            with gr.Column():
+                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)
+                    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)
+                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
+        with gr.Row():
+            with gr.Column():
+                choose_metrics = gr.CheckboxGroup(choices=Metric.all_metrics,
+                                                  value=metric_names,
+                                                  label="Choose metrics you are interested")
+            with gr.Column():
+                load_button = gr.Button("Load images...")
+                load_random_button = gr.Button("Load randomly...")
+        save_button = gr.Button("Save images...")
+
+        ### 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,
+                              inputs=choose_physics,
+                              outputs=physics_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, 
+                                  model_a_placeholder,
+                                  model_b_placeholder,
+                                  physics_placeholder,
+                                  use_generator_button,
+                                  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, 
+                                         model_a_placeholder,
+                                         model_b_placeholder,
+                                         physics_placeholder,
+                                         use_generator_button,
+                                         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()

datasets.py

@@ -0,0 +1,84 @@
+from pathlib import Path
+from typing import Callable, Optional
+import os
+
+import torch
+from torch.utils.data import Dataset
+
+
+class Preprocessed_fastMRI(torch.utils.data.Dataset):
+    """FastMRI from preprocessed data for faster lading."""
+
+    def __init__(
+        self,
+        root: str,
+        transform: Optional[Callable] = None,
+        preprocess: bool = False,
+    ) -> None:
+        self.root = root
+        self.transform = transform
+        self.preprocess = preprocess
+
+        # should contain all the information to load a data sample from the storage
+        self.sample_identifiers = []
+
+        # append all filenames in self.root ending with .pt
+        for root, _, files in os.walk(self.root):
+            for file in files:
+                if file.endswith(".pt"):
+                    self.sample_identifiers.append(file)
+
+    def __len__(self) -> int:
+        return len(self.sample_identifiers)
+
+    def __getitem__(self, idx: int):
+        fname = self.sample_identifiers[idx]
+
+        tensor = torch.load(os.path.join(self.root, fname), weights_only=True)
+        img = tensor['data'].float()
+
+        if self.transform is not None:
+            img = self.transform(img)
+
+        if not self.preprocess:
+            return img
+
+        else:
+            # remove extension and prefix from filename
+            fname = Path(fname).stem
+            return img, fname
+
+
+class Preprocessed_LIDCIDRI(torch.utils.data.Dataset):
+    """FastMRI from preprocessed data for faster lading."""
+
+    def __init__(
+        self,
+        root: str,
+        transform: Optional[Callable] = None,
+    ) -> None:
+        self.root = root
+        self.transform = transform
+
+        # should contain all the information to load a data sample from the storage
+        self.sample_identifiers = []
+
+        # append all filenames in self.root ending with .pt
+        for root, _, files in os.walk(self.root):
+            for file in files:
+                if file.endswith(".pt"):
+                    self.sample_identifiers.append(file)
+
+    def __len__(self) -> int:
+        return len(self.sample_identifiers)
+
+    def __getitem__(self, idx: int):
+        fname = self.sample_identifiers[idx]
+
+        tensor = torch.load(os.path.join(self.root, fname), weights_only=True)
+        img = tensor['data'].float()
+
+        if self.transform is not None:
+            img = self.transform(img)
+
+        img = img.unsqueeze(0)  # add channel dim

evals.py

@@ -0,0 +1,564 @@
+from typing import Any, List
+
+import deepinv as dinv
+import numpy as np
+import torch
+from deepinv.physics.generator import MotionBlurGenerator, SigmaGenerator
+from torchvision import transforms
+
+from datasets import Preprocessed_fastMRI, Preprocessed_LIDCIDRI
+from utils import get_model
+
+DEFAULT_MODEL_PARAMS = {
+    "in_channels": [1, 2, 3],
+    "grayscale": False,
+    "conv_type": "base",
+    "pool_type": "base",
+    "layer_scale_init_value": 1e-6,
+    "init_type": "ortho",
+    "gain_init_conv": 1.0,
+    "gain_init_linear": 1.0,
+    "drop_prob": 0.0,
+    "replk": False,
+    "mult_fact": 4,
+    "antialias": "gaussian",
+    "nc_base": 64,
+    "cond_type": "base",
+    "blind": False,
+    "pretrained_pth": None,
+    "N": 2,
+    "c_mult": 2,
+    "depth_encoding": 2,
+    "relu_in_encoding": False,
+    "skip_in_encoding": True
+}
+
+
+class PhysicsWithGenerator(torch.nn.Module):
+    """Interface between Physics, Generator and Gradio."""
+    all_physics = ["MotionBlur_easy", "MotionBlur_medium", "MotionBlur_hard", "GaussianBlur",
+                   "MRI", "CT"]
+
+    def __init__(self, physics_name: str, device_str: str = "cpu") -> None:
+        super().__init__()
+
+        self.name = physics_name
+        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.generator = self.physics_generator + SigmaGenerator(sigma_min=0.01, sigma_max=0.01, device=device_str)
+            self.saved_params = {"updatable_params": {"sigma": 0.05},
+                                 "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.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,
+                                                  "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.generator = self.physics_generator + SigmaGenerator(sigma_min=0.1, sigma_max=0.1, device=device_str)
+            self.saved_params = {"updatable_params": {"sigma": 0.05},
+                                 "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":
+            psf_size = 31
+            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), num_channels=1,
+                                                           device=device_str)
+            self.generator = self.physics_generator + self.sigma_generator
+            self.saved_params = {"updatable_params": {"sigma": 0.05},
+                                 "updatable_params_converter": {"sigma": float},
+                                 "fixed_params": {"noise_sigma_min": 0.001, "noise_sigma_max": 0.2,
+                                                  "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_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},
+                                 "updatable_params_converter": {"sigma": float},
+                                 "fixed_params": {"noise_sigma_min": 0.001, "noise_sigma_max": 0.2,
+                                                  "acceleration_factor": 4}}
+        elif self.name == "CT":
+            acceleration_factor = 10
+            img_h = 480
+            angles = int(img_h / acceleration_factor)
+            # angles = torch.linspace(0, 180, steps=10)
+            self.physics = dinv.physics.Tomography(
+                img_width=img_h,
+                angles=angles,
+                circle=False,
+                normalize=True,
+                device=device_str,
+                noise_model=dinv.physics.GaussianNoise(sigma=1e-4),
+                max_iter=10,
+            )
+            self.physics_generator = None
+            self.generator = self.sigma_generator
+            self.saved_params = {"updatable_params": {"sigma": 0.1},
+                                 "updatable_params_converter": {"sigma": float},
+                                 "fixed_params": {"noise_sigma_min": 0.001, "noise_sigma_max": 0.,
+                                                  "angles": angles, "max_iter": 10}}
+
+    def display_saved_params(self) -> str:
+        """Printable version of saved_params."""
+        updatable_params_str = "Updatable parameters:\n"
+        for param_name, param_value in self.saved_params["updatable_params"].items():
+            updatable_params_str += f"\t\t{param_name} = {param_value}" + "\n"
+
+        fixed_params_str = "Fixed parameters:\n"
+        for param_name, param_value in self.saved_params["fixed_params"].items():
+            fixed_params_str += f"\t\t{param_name} = {param_value}" + "\n"
+
+        return updatable_params_str + fixed_params_str
+
+    def _update_save_params(self, key: str, value: Any) -> None:
+        """Update value of an existing key in save_params."""
+        if key in list(self.saved_params["updatable_params"].keys()):
+            if type(value) == str:  # it may be only a str representation
+                # type: str -> ???
+                value = self.saved_params["updatable_params_converter"][key](value)
+            elif isinstance(value, torch.Tensor):
+                value = value.item()  # type: torch.Tensor -> float
+                value = float(f"{value:.4f}")  # keeps only 4 significant digits
+            self.saved_params["updatable_params"][key] = value
+
+    def update_and_display_params(self, key, value) -> str:
+        """_update_save_params + update physics with saved_params + display_saved_params"""
+        self._update_save_params(key, value)
+
+        if self.name == "Denoising":
+            self.physics.noise_model.update_parameters(**self.saved_params["updatable_params"])
+        else:
+            self.physics.update_parameters(**self.saved_params["updatable_params"])
+
+        return self.display_saved_params()
+
+    def update_saved_params_and_physics(self, **kwargs) -> None:
+        """Update save_params and update physics."""
+        for key, value in kwargs.items():
+            self._update_save_params(key, value)
+
+        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"):
+            use_gen = True
+        elif self.name in ["MRI"] and not hasattr(self.physics, "mask"):
+            use_gen = True
+
+        if use_gen:
+            kwargs = self.generator.step(batch_size=x.shape[0])  # generate a set of params for each sample
+            self.update_saved_params_and_physics(**kwargs)
+
+        return self.physics(x)
+
+
+class EvalModel(torch.nn.Module):
+    """Eval model.
+
+    Is there a difference with BaselineModel ?
+        -> BaselineModel should be models that are already trained and will have fixed weights.
+        -> Eval model will change depending on differents checkpoints.
+    """
+    all_models = ["unext_emb_physics_config_C"]
+
+    def __init__(self, model_name: str, ckpt_pth: str = "", device_str: str = "cpu") -> None:
+        """Load the model we want to evaluate."""
+        super().__init__()
+        self.base_name = model_name
+        self.ckpt_pth = ckpt_pth
+        self.name = self.base_name
+        if self.base_name not in self.all_models:
+            raise ValueError(f"{self.base_name} is unavailable.")
+        if self.base_name == "unext_emb_physics_config_C":
+            if self.ckpt_pth == "":
+                self.ckpt_pth = "ckpt/ram_ckp_10.pth.tar"
+            self.model = get_model(model_name=self.base_name,
+                                   device='cpu',
+                                   **DEFAULT_MODEL_PARAMS)
+
+            # load model checkpoint
+            state_dict = torch.load(self.ckpt_pth, map_location=lambda storage, loc: storage)[
+                'state_dict']  # load on cpu
+            self.model.load_state_dict(state_dict)
+            self.model.to(device_str)
+            self.model.eval()
+
+            # add epoch in the model name
+            epoch = torch.load(self.ckpt_pth, map_location=lambda storage, loc: storage)['epoch']
+            self.name = self.name + f"+{epoch}"
+
+    def forward(self, y: torch.Tensor, physics: torch.nn.Module) -> torch.Tensor:
+        return self.model(y, physics=physics)
+
+
+class BaselineModel(torch.nn.Module):
+    """Baseline model.
+
+    Is there a difference with EvalModel ?
+        -> 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"]
+
+    def __init__(self, model_name: str, device_str: str = "cpu") -> None:
+        super().__init__()
+        self.base_name = model_name
+        self.ckpt_pth = ""
+        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"
+            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
+
+            # Specify the denoising prior
+            self.prior = dinv.optim.prior.PnP(denoiser=drunet)
+        elif self.name == "DPIR_MRI":
+            class ComplexDenoiser(torch.nn.Module):
+                def __init__(self, denoiser):
+                    super().__init__()
+                    self.denoiser = denoiser
+
+                def forward(self, x, sigma):
+                    noisy_batch = torch.cat((x[:, 0:1, ...], x[:, 1:2, ...]), 0)
+                    input_min = noisy_batch.min()
+                    denoised_batch = self.denoiser(noisy_batch - input_min, sigma)
+                    denoised_batch = denoised_batch + input_min
+                    denoised = torch.cat((denoised_batch[0:1, ...], denoised_batch[1:2, ...]), 1)
+                    return denoised
+
+            # Load PnP denoiser backbone
+            n_channels = 1
+            ckpt_pth = "ckpt/drunet_gray.pth"
+            drunet = dinv.models.DRUNet(in_channels=n_channels, out_channels=n_channels, device=device_str,
+                                        pretrained=ckpt_pth)
+            complex_drunet = ComplexDenoiser(drunet)
+            complex_drunet.eval()
+
+            # Specify the denoising prior
+            self.prior = dinv.optim.prior.PnP(denoiser=complex_drunet)
+        elif self.name == "DPIR_CT":
+            class CTDenoiser(torch.nn.Module):
+                def __init__(self, denoiser):
+                    super().__init__()
+                    self.denoiser = denoiser
+
+                def forward(self, x, sigma):
+                    x = x - x.min()
+                    denoised = self.denoiser(x, sigma)
+                    denoised = denoised + x.min()
+                    return denoised
+
+            # Load PnP denoiser backbone
+            n_channels = 1
+            ckpt_pth = "ckpt/drunet_gray.pth"
+            drunet = dinv.models.DRUNet(in_channels=n_channels, out_channels=n_channels, device=device_str,
+                                        pretrained=ckpt_pth)
+            ct_drunet = CTDenoiser(drunet)
+            ct_drunet.eval()
+
+            # Specify the denoising prior
+            self.prior = dinv.optim.prior.PnP(denoiser=ct_drunet)
+
+    def circular_roll(self, tensor, p_h, p_w):
+        return tensor.roll(shifts=(p_h, p_w), dims=(-2, -1))
+
+    def get_DPIR_params(self, noise_level_img, max_iter=8):
+        r"""
+        Default parameters for the DPIR Plug-and-Play algorithm.
+
+        :param float noise_level_img: Noise level of the input image.
+        :return: tuple(list with denoiser noise level per iteration, list with stepsize per iteration, iterations).
+        """
+        max_iter = 8
+        s1 = 49.0 / 255.0
+        s2 = max(noise_level_img, 0.01)
+        sigma_denoiser = np.logspace(np.log10(s1), np.log10(s2), max_iter).astype(
+            np.float32
+        )
+        stepsize = (sigma_denoiser / max(0.01, noise_level_img)) ** 2
+        lamb = 1 / 0.23
+        return list(sigma_denoiser), list(lamb * stepsize)
+
+    def get_DPIR_MRI_params(self, noise_level_img: float, max_iter: int = 8):
+        r"""
+        Default parameters for the DPIR Plug-and-Play algorithm.
+
+        :param float noise_level_img: Noise level of the input image.
+        """
+        s1 = 49.0 / 255.0
+        s2 = noise_level_img
+        sigma_denoiser = np.logspace(np.log10(s1), np.log10(s2), max_iter).astype(
+            np.float32
+        )
+        stepsize = (sigma_denoiser / max(0.01, noise_level_img)) ** 2
+        lamb = 1.
+        return lamb, list(sigma_denoiser), list(stepsize), max_iter
+
+    def get_DPIR_CT_params(self, noise_level_img: float, max_iter: int = 8, lip_cons: float = 1.0):
+        r"""
+        Default parameters for the DPIR Plug-and-Play algorithm.
+
+        :param float noise_level_img: Noise level of the input image.
+        """
+        s1 = 49.0 / 255.0 * lip_cons
+        s2 = noise_level_img
+        sigma_denoiser = np.logspace(np.log10(s1), np.log10(s2), max_iter).astype(
+            np.float32
+        )
+        stepsize = (sigma_denoiser / max(0.01, noise_level_img)) ** 2  #
+        lamb = 1.
+        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":
+            # Set the DPIR algorithm parameters
+            sigma_float = physics.noise_model.sigma.item()  # sigma should be a single value
+            max_iter = 8
+
+            sigma_denoiser, stepsize = self.get_DPIR_params(sigma_float, max_iter=max_iter)
+            params_algo = {"stepsize": stepsize, "g_param": sigma_denoiser}
+            early_stop = False  # Do not stop algorithm with convergence criteria
+
+            # instantiate DPIR
+            model = dinv.optim.optim_builder(
+                iteration="HQS",
+                prior=self.prior,
+                data_fidelity=dinv.optim.data_fidelity.L2(),
+                early_stop=early_stop,
+                max_iter=max_iter,
+                verbose=True,
+                params_algo=params_algo,
+            )
+            return model(y, physics=physics)
+        elif self.name == "DPIR_MRI":
+            sigma_float = max(physics.noise_model.sigma.item(), 0.015)  # sigma should be a single value
+            lamb, sigma_denoiser, stepsize, max_iter = self.get_DPIR_MRI_params(sigma_float, max_iter=16)
+            stepsize = [stepsize[0]] * max_iter
+            params_algo = {"stepsize": stepsize, "g_param": sigma_denoiser, "lambda": lamb}
+            early_stop = False  # Do not stop algorithm with convergence criteria
+
+            # Instantiate the algorithm class to solve the IP
+            model = dinv.optim.optim_builder(
+                iteration="HQS",
+                prior=self.prior,
+                data_fidelity=dinv.optim.data_fidelity.L2(),
+                early_stop=early_stop,
+                max_iter=max_iter,
+                verbose=True,
+                params_algo=params_algo,
+            )
+            return model(y, physics=physics)
+        elif 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))
+            lamb, sigma_denoiser, stepsize, max_iter = self.get_DPIR_CT_params(sigma_float, max_iter=8,
+                                                                               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
+
+            def custom_init(y, physic_op):
+                x_init = physic_op.prox_l2(physic_op.A_adjoint(y), y, gamma=1e4)
+                return {"est": (x_init, x_init)}
+
+            # Instantiate the algorithm class to solve the IP
+            algo = dinv.optim.optim_builder(
+                iteration="HQS",
+                prior=self.prior,
+                data_fidelity=dinv.optim.data_fidelity.L2(),
+                early_stop=early_stop,
+                max_iter=max_iter,
+                verbose=True,
+                params_algo=params_algo,
+                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
+        elif 'UNROLLED_DPIR' in self.name:
+            return self.model(y, physics=physics)
+        else:
+            return self.model(y)
+
+
+class EvalDataset(torch.utils.data.Dataset):
+    """
+    We expect that images are 480x480.
+    """
+    all_datasets = ["Natural", "MRI", "CT"]
+
+    def __init__(self, dataset_name: str, device_str: str = "cpu") -> None:
+        self.name = dataset_name
+        self.device_str = device_str
+        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.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.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.transform = None
+            self.dataset = Preprocessed_LIDCIDRI(root=self.root,
+                                                 transform=self.transform)
+
+    def __len__(self) -> int:
+        return len(self.dataset)
+
+    def __getitem__(self, idx: int) -> torch.Tensor:
+        return self.dataset[idx].to(self.device_str)
+
+
+class Metric():
+    """Metrics and utilities."""
+    all_metrics = ["PSNR", "SSIM", "LPIPS"]
+
+    def __init__(self, metric_name: str, device_str: str = "cpu") -> None:
+        self.name = metric_name
+        if self.name not in self.all_metrics:
+            raise ValueError(f"{self.name} is unavailable.")
+        elif self.name == "PSNR":
+            self.metric = dinv.loss.metric.PSNR()
+        elif self.name == "SSIM":
+            self.metric = dinv.loss.metric.SSIM()
+        elif self.name == "LPIPS":
+            self.metric = dinv.loss.metric.LPIPS(device=device_str)
+
+    def __call__(self, x_net: torch.Tensor, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
+        # it may happen that x_net and x do not have the same size, in which case we take the minimum size of both
+        if x_net.shape[-1] != x.shape[-1]:
+            min_size = min(x_net.shape[-1], x.shape[-1])
+            x_net_crop = x_net[..., x_net.shape[-2] // 2 - min_size // 2: x_net.shape[-2] // 2 + min_size // 2,
+                         x_net.shape[-1] // 2 - min_size // 2: x_net.shape[-1] // 2 + min_size // 2]
+            x_crop = x[..., x_net.shape[-2] // 2 - min_size // 2: x_net.shape[-2] // 2 + min_size // 2,
+                     x_net.shape[-1] // 2 - min_size // 2: x_net.shape[-1] // 2 + min_size // 2]
+        else:
+            x_net_crop = x_net
+            x_crop = x
+        return self.metric(x_net_crop, x_crop)
+
+    @classmethod
+    def get_list_metrics(cls, metric_names: List[str], device_str: str = "cpu") -> List["Metric"]:
+        l = []
+        for metric_name in metric_names:
+            l.append(cls(metric_name, device_str=device_str))
+        return l