factories.py

from typing import Any, List

import deepinv as dinv
import numpy as np
import torch
from deepinv.physics.generator import MotionBlurGenerator, SigmaGenerator, GainGenerator
from torchvision import transforms

from datasets import Preprocessed_fastMRI, Preprocessed_LIDCIDRI, LsdirMiniDataset
from model_factory import get_model
from physics.blur_generator import GaussianBlurGenerator


class PhysicsWithGenerator(torch.nn.Module):
    """Interface between Physics, Generator and Gradio."""
    all_physics = ["Inpainting", "MotionBlur_medium", "MotionBlur_hard",
                   "GaussianBlur_easy", "GaussianBlur_medium", "GaussianBlur_hard",
                   "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.")

        if self.name == "MotionBlur_medium":
            psf_size = 31
            self.physics = dinv.physics.Blur(noise_model=dinv.physics.GaussianNoise(sigma=.05).to(device_str),
                                             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},
                                 "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).to(device_str),
                                             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.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).to(device_str),
                                             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)
            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).to(device_str),
                                             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)
            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).to(device_str),
                                             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)
            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 == "Inpainting":
            sigma = 0.05
            split_ratio = 0.3
            pixelwise = True
            self.physics = dinv.physics.Inpainting(tensor_size=(3, 256, 256), mask=split_ratio,
                                                   noise_model=dinv.physics.GaussianNoise(sigma=sigma),
                                                   device=device_str)
            self.physics_generator = dinv.physics.generator.BernoulliSplittingMaskGenerator((3, 256, 256),
                                                                                              split_ratio=split_ratio, pixelwise=pixelwise,
                                                                                              random_split_ratio=False, device=device_str)
            self.generator = dinv.physics.generator.BernoulliSplittingMaskGenerator((3, 256, 256),
                                                                                              split_ratio=split_ratio, pixelwise=pixelwise,
                                                                                              random_split_ratio=False, device=device_str)

            self.saved_params = {"updatable_params": {},
                                 "updatable_params_converter": {"sigma": float},
                                 "fixed_params": {"sigma": sigma}}
        elif self.name == "MRI":
            self.physics = dinv.physics.MRI(noise_model=dinv.physics.GaussianNoise(sigma=.01).to(device_str),
                                            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.saved_params = {"updatable_params": {"sigma": 0.01},
                                 "updatable_params_converter": {"sigma": float},
                                 "fixed_params": {"acceleration_factor": 4}}
        elif self.name == "CT":
            acceleration_factor = 10
            img_h = 512
            angles = torch.linspace(0, 180, steps=int(img_h / acceleration_factor))
            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-3).to(device_str),
                max_iter=10,
            )
            self.physics_generator = SigmaGenerator(sigma_min=1e-3, sigma_max=1e-3, device=device_str)
            self.generator = SigmaGenerator(sigma_min=1e-3, sigma_max=1e-3, device=device_str)
            self.saved_params = {"updatable_params": {"sigma": 0.1},
                                 "updatable_params_converter": {"sigma": float},
                                 "fixed_params": {"noise_sigma_min": 1e-3, "noise_sigma_max": 1e-3,
                                                  "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 value != "" and key in list(self.saved_params["updatable_params"].keys()):
            if type(value) == str:  # it may be only a str representation
                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: str, value: Any) -> 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_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

        if use_gen:
            if self.name == 'MRI':  # RandomMaskGenerator deoends on image size
                _, _, h, w = x.shape
                self.generator = dinv.physics.generator.RandomMaskGenerator((2, h, w), acceleration_factor=4)

            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.
    """

    def __init__(self, device_str: str = "cpu") -> None:
        """Load the model we want to evaluate."""
        super().__init__()
        self.name = 'RAM'
        self.model = get_model()
        self.model.to(device_str)
        self.model.eval()

    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 = ["DPIR", "DPIR_MRI", "DPIR_CT"]

    def __init__(self, model_name: str, device_str: str = "cpu") -> None:
        super().__init__()
        self.name = model_name
        self.ckpt_pth = ""
        if self.name not in self.all_baselines:
            raise ValueError(f"{self.name} is unavailable.")
        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 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 == "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=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

            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)


class EvalDataset(torch.utils.data.Dataset):

    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 = 'img_samples/LSDIR_samples'
            self.transform = transforms.Compose([transforms.ToTensor()])
            self.dataset = LsdirMiniDataset(root=self.root,
                                            transform=self.transform)
        elif self.name == 'MRI':
            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 = 'img_samples/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, crop=31, **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[..., crop:-crop, crop:-crop], x_crop[..., crop:-crop, crop:-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