deepinv 0.1.0.dev0__py3-none-any.whl

deepinv/models/unet.py

@@ -0,0 +1,337 @@
+import torch
+import torch.nn as nn
+from .drunet import test_pad
+
+
+class BFBatchNorm2d(nn.BatchNorm2d):
+    r"""
+    From Mohan et al.
+
+    "Robust And Interpretable Blind Image Denoising Via Bias-Free Convolutional Neural Networks"
+    S. Mohan, Z. Kadkhodaie, E. P. Simoncelli, C. Fernandez-Granda
+    Int'l. Conf. on Learning Representations (ICLR), Apr 2020.
+    """
+
+    def __init__(
+        self, num_features, eps=1e-5, momentum=0.1, use_bias=False, affine=True
+    ):
+        super(BFBatchNorm2d, self).__init__(num_features, eps, momentum)
+        self.use_bias = use_bias
+        self.affine = affine
+
+    def forward(self, x):
+        self._check_input_dim(x)
+        y = x.transpose(0, 1)
+        return_shape = y.shape
+        y = y.contiguous().view(x.size(1), -1)
+        if self.use_bias:
+            mu = y.mean(dim=1)
+        sigma2 = y.var(dim=1)
+        if self.training is not True:
+            if self.use_bias:
+                y = y - self.running_mean.view(-1, 1)
+            y = y / (self.running_var.view(-1, 1) ** 0.5 + self.eps)
+        else:
+            if self.track_running_stats is True:
+                with torch.no_grad():
+                    if self.use_bias:
+                        self.running_mean = (
+                            1 - self.momentum
+                        ) * self.running_mean + self.momentum * mu
+                    self.running_var = (
+                        1 - self.momentum
+                    ) * self.running_var + self.momentum * sigma2
+            if self.use_bias:
+                y = y - mu.view(-1, 1)
+            y = y / (sigma2.view(-1, 1) ** 0.5 + self.eps)
+        if self.affine:
+            y = self.weight.view(-1, 1) * y
+            if self.use_bias:
+                y += self.bias.view(-1, 1)
+
+        return y.view(return_shape).transpose(0, 1)
+
+
+class UNet(nn.Module):
+    r"""
+    U-Net convolutional denoiser.
+
+    This network is a fully convolutional denoiser based on the U-Net architecture. The number of downsample steps
+    can be controlled with the `scales` parameter. The number of trainable parameters increases with the number of
+    scales.
+
+    :param int in_channels: input image channels
+    :param int out_channels: output image channels
+    :param bool residual: use a skip-connection between output and output.
+    :param bool circular_padding: circular padding for the convolutional layers.
+    :param bool cat: use skip-connections between intermediate levels.
+    :param bool bias: use learnable biases.
+    :param int scales: Number of downsampling steps used in the U-Net. The options are 2,3,4 and 5.
+        The number of trainable parameters increases with the scale.
+    """
+
+    def __init__(
+        self,
+        in_channels=1,
+        out_channels=1,
+        residual=True,
+        circular_padding=False,
+        cat=True,
+        bias=True,
+        batch_norm=True,
+        scales=4,
+    ):
+        super(UNet, self).__init__()
+        self.name = "unet"
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        self.residual = residual
+        self.cat = cat
+        self.compact = scales
+        self.Maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
+
+        def conv_block(ch_in, ch_out):
+            if batch_norm:
+                return nn.Sequential(
+                    nn.Conv2d(
+                        ch_in,
+                        ch_out,
+                        kernel_size=3,
+                        stride=1,
+                        padding=1,
+                        bias=bias,
+                        padding_mode="circular" if circular_padding else "zeros",
+                    ),
+                    BFBatchNorm2d(ch_out, use_bias=bias),
+                    nn.ReLU(inplace=True),
+                    nn.Conv2d(
+                        ch_out, ch_out, kernel_size=3, stride=1, padding=1, bias=bias
+                    ),
+                    BFBatchNorm2d(ch_out, use_bias=bias),
+                    nn.ReLU(inplace=True),
+                )
+            else:
+                return nn.Sequential(
+                    nn.Conv2d(
+                        ch_in,
+                        ch_out,
+                        kernel_size=3,
+                        stride=1,
+                        padding=1,
+                        bias=bias,
+                        padding_mode="circular" if circular_padding else "zeros",
+                    ),
+                    nn.ReLU(inplace=True),
+                    nn.Conv2d(
+                        ch_out, ch_out, kernel_size=3, stride=1, padding=1, bias=bias
+                    ),
+                    nn.ReLU(inplace=True),
+                )
+
+        def up_conv(ch_in, ch_out):
+            if batch_norm:
+                return nn.Sequential(
+                    nn.Upsample(scale_factor=2),
+                    nn.Conv2d(
+                        ch_in, ch_out, kernel_size=3, stride=1, padding=1, bias=bias
+                    ),
+                    BFBatchNorm2d(ch_out, use_bias=bias),
+                    nn.ReLU(inplace=True),
+                )
+            else:
+                return nn.Sequential(
+                    nn.Upsample(scale_factor=2),
+                    nn.Conv2d(
+                        ch_in, ch_out, kernel_size=3, stride=1, padding=1, bias=bias
+                    ),
+                    nn.ReLU(inplace=True),
+                )
+
+        self.Conv1 = conv_block(ch_in=in_channels, ch_out=64)
+        self.Conv2 = conv_block(ch_in=64, ch_out=128)
+        self.Conv3 = (
+            conv_block(ch_in=128, ch_out=256) if self.compact in [3, 4, 5] else None
+        )
+        self.Conv4 = (
+            conv_block(ch_in=256, ch_out=512) if self.compact in [4, 5] else None
+        )
+        self.Conv5 = conv_block(ch_in=512, ch_out=1024) if self.compact in [5] else None
+
+        self.Up5 = up_conv(ch_in=1024, ch_out=512) if self.compact in [5] else None
+        self.Up_conv5 = (
+            conv_block(ch_in=1024, ch_out=512) if self.compact in [5] else None
+        )
+
+        self.Up4 = up_conv(ch_in=512, ch_out=256) if self.compact in [4, 5] else None
+        self.Up_conv4 = (
+            conv_block(ch_in=512, ch_out=256) if self.compact in [4, 5] else None
+        )
+
+        self.Up3 = up_conv(ch_in=256, ch_out=128) if self.compact in [3, 4, 5] else None
+        self.Up_conv3 = (
+            conv_block(ch_in=256, ch_out=128) if self.compact in [3, 4, 5] else None
+        )
+
+        self.Up2 = up_conv(ch_in=128, ch_out=64)
+        self.Up_conv2 = conv_block(ch_in=128, ch_out=64)
+
+        self.Conv_1x1 = nn.Conv2d(
+            in_channels=64,
+            out_channels=out_channels,
+            bias=bias,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+        )
+
+        if self.compact == 5:
+            self._forward = self.forward_standard
+        if self.compact == 4:
+            self._forward = self.forward_compact4
+        if self.compact == 3:
+            self._forward = self.forward_compact3
+        if self.compact == 2:
+            self._forward = self.forward_compact2
+
+    def forward(self, x, sigma=None):
+        r"""
+        Run the denoiser on noisy image. The noise level is not used in this denoiser.
+
+        :param torch.Tensor x: noisy image.
+        :param float sigma: noise level (not used).
+        """
+
+        factor = 2 ** (self.compact - 1)
+        if x.size(2) % factor == 0 and x.size(3) % factor == 0:
+            return self._forward(x)
+        else:
+            return test_pad(self._forward, x, modulo=factor)
+
+    def forward_standard(self, x):
+        # encoding path
+        cat_dim = 1
+        input = x
+        x1 = self.Conv1(input)
+
+        x2 = self.Maxpool(x1)
+        x2 = self.Conv2(x2)
+
+        x3 = self.Maxpool(x2)
+        x3 = self.Conv3(x3)
+
+        x4 = self.Maxpool(x3)
+        x4 = self.Conv4(x4)
+
+        x5 = self.Maxpool(x4)
+        x5 = self.Conv5(x5)
+
+        # decoding + concat path
+        d5 = self.Up5(x5)
+        if self.cat:
+            d5 = torch.cat((x4, d5), dim=cat_dim)
+            d5 = self.Up_conv5(d5)
+
+        d4 = self.Up4(d5)
+        if self.cat:
+            d4 = torch.cat((x3, d4), dim=cat_dim)
+            d4 = self.Up_conv4(d4)
+
+        d3 = self.Up3(d4)
+        if self.cat:
+            d3 = torch.cat((x2, d3), dim=cat_dim)
+            d3 = self.Up_conv3(d3)
+
+        d2 = self.Up2(d3)
+        if self.cat:
+            d2 = torch.cat((x1, d2), dim=cat_dim)
+            d2 = self.Up_conv2(d2)
+
+        d1 = self.Conv_1x1(d2)
+
+        out = d1 + x if self.residual and self.in_channels == self.out_channels else d1
+        return out
+
+    def forward_compact4(self, x):
+        # def forward_compact4(self, x):
+        # encoding path
+        cat_dim = 1
+        input = x
+
+        x1 = self.Conv1(input)  # 1->64
+
+        x2 = self.Maxpool(x1)
+        x2 = self.Conv2(x2)  # 64->128
+
+        x3 = self.Maxpool(x2)
+        x3 = self.Conv3(x3)  # 128->256
+
+        x4 = self.Maxpool(x3)
+        x4 = self.Conv4(x4)  # 256->512
+
+        d4 = self.Up4(x4)  # 512->256
+        if self.cat:
+            d4 = torch.cat((x3, d4), dim=cat_dim)
+            d4 = self.Up_conv4(d4)
+
+        d3 = self.Up3(d4)  # 256->128
+        if self.cat:
+            d3 = torch.cat((x2, d3), dim=cat_dim)
+            d3 = self.Up_conv3(d3)
+
+        d2 = self.Up2(d3)  # 128->64
+        if self.cat:
+            d2 = torch.cat((x1, d2), dim=cat_dim)
+            d2 = self.Up_conv2(d2)
+
+        d1 = self.Conv_1x1(d2)
+
+        out = d1 + x if self.residual and self.in_channels == self.out_channels else d1
+        return out
+
+    def forward_compact3(self, x):
+        # encoding path
+        cat_dim = 1
+        input = x
+        x1 = self.Conv1(input)
+
+        x2 = self.Maxpool(x1)
+        x2 = self.Conv2(x2)
+
+        x3 = self.Maxpool(x2)
+        x3 = self.Conv3(x3)
+
+        d3 = self.Up3(x3)
+        if self.cat:
+            d3 = torch.cat((x2, d3), dim=cat_dim)
+            d3 = self.Up_conv3(d3)
+
+        d2 = self.Up2(d3)
+        if self.cat:
+            d2 = torch.cat((x1, d2), dim=cat_dim)
+            d2 = self.Up_conv2(d2)
+
+        d1 = self.Conv_1x1(d2)
+
+        out = d1 + x if self.residual and self.in_channels == self.out_channels else d1
+        return out
+
+    def forward_compact2(self, x):
+        # encoding path
+        cat_dim = 1
+        input = x
+        x1 = self.Conv1(input)
+
+        x2 = self.Maxpool(x1)
+        x2 = self.Conv2(x2)
+
+        d2 = self.Up2(x2)
+        if self.cat:
+            d2 = torch.cat((x1, d2), dim=cat_dim)
+            d2 = self.Up_conv2(d2)
+
+        d1 = self.Conv_1x1(d2)
+
+        out = d1 + x if self.residual and self.in_channels == self.out_channels else d1
+        return out

deepinv/models/utils.py

@@ -0,0 +1,22 @@
+import torch
+import numpy as np
+
+
+def tensor2array(img):
+    img = img.cpu().detach().numpy()
+    img = np.transpose(img, (1, 2, 0))
+    return img
+
+
+def array2tensor(img):
+    return torch.from_numpy(img).permute(2, 0, 1)
+
+
+def get_weights_url(model_name, file_name):
+    return (
+        "https://huggingface.co/deepinv/"
+        + model_name
+        + "/resolve/main/"
+        + file_name
+        + "?download=true"
+    )

deepinv/models/wavdict.py

@@ -0,0 +1,231 @@
+import torch
+import torch.nn as nn
+
+try:
+    import pytorch_wavelets
+except:
+    pytorch_wavelets = ImportError("The pytorch_wavelets package is not installed.")
+
+
+class WaveletPrior(nn.Module):
+    r"""
+    Wavelet denoising with the :math:`\ell_1` norm.
+
+
+    This denoiser is defined as the solution to the optimization problem:
+
+    .. math::
+
+        \underset{x}{\arg\min} \;  \|x-y\|^2 + \lambda \|\Psi x\|_n
+
+    where :math:`\Psi` is an orthonormal wavelet transform, :math:`\lambda>0` is a hyperparameter, and where
+    :math:`\|\cdot\|_n` is either the :math:`\ell_1` norm (``non_linearity="soft"``) or
+    the :math:`\ell_0` norm (``non_linearity="hard"``). A variant of the :math:`\ell_0` norm is also available
+    (``non_linearity="topk"``), where the thresholding is done by keeping the :math:`k` largest coefficients
+    in each wavelet subband and setting the others to zero.
+
+    The solution is available in closed-form, thus the denoiser is cheap to compute.
+
+    :param int level: decomposition level of the wavelet transform
+    :param str wv: mother wavelet (follows the `PyWavelets convention
+        <https://pywavelets.readthedocs.io/en/latest/ref/wavelets.html>`_) (default: "db8")
+    :param str device: cpu or gpu
+    :param str non_linearity: ``"soft"``, ``"hard"`` or ``"topk"`` thresholding (default: ``"soft"``).
+        If ``"topk"``, only the top-k wavelet coefficients are kept.
+    """
+
+    def __init__(self, level=3, wv="db8", device="cpu", non_linearity="soft"):
+        if isinstance(pytorch_wavelets, ImportError):
+            raise ImportError(
+                "pytorch_wavelets is needed to use the WaveletPrior class. "
+                "It should be installed with `pip install "
+                "git+https://github.com/fbcotter/pytorch_wavelets.git`"
+            ) from pytorch_wavelets
+        super().__init__()
+        self.level = level
+        self.dwt = pytorch_wavelets.DWTForward(J=self.level, wave=wv).to(device)
+        self.iwt = pytorch_wavelets.DWTInverse(wave=wv).to(device)
+        self.device = device
+        self.non_linearity = non_linearity
+
+    def get_ths_map(self, ths):
+        if isinstance(ths, float) or isinstance(ths, int):
+            ths_map = ths
+        elif len(ths.shape) == 0 or ths.shape[0] == 1:
+            ths_map = ths.to(self.device)
+        else:
+            ths_map = (
+                ths.unsqueeze(0)
+                .unsqueeze(0)
+                .unsqueeze(-1)
+                .unsqueeze(-1)
+                .to(self.device)
+            )
+        return ths_map
+
+    def prox_l1(self, x, ths=0.1):
+        r"""
+        Soft thresholding of the wavelet coefficients.
+
+        :param torch.Tensor x: wavelet coefficients.
+        :param float, torch.Tensor ths: threshold.
+        """
+        ths_map = self.get_ths_map(ths)
+        return torch.maximum(
+            torch.tensor([0], device=x.device).type(x.dtype), x - ths_map
+        ) + torch.minimum(torch.tensor([0], device=x.device).type(x.dtype), x + ths_map)
+
+    def prox_l0(self, x, ths=0.1):
+        r"""
+        Hard thresholding of the wavelet coefficients.
+
+        :param torch.Tensor x: wavelet coefficients.
+        :param float, torch.Tensor ths: threshold.
+        """
+        if isinstance(ths, float):
+            ths_map = ths
+        else:
+            ths_map = self.get_ths_map(ths)
+            ths_map = ths_map.repeat(
+                1, 1, 1, x.shape[-2], x.shape[-1]
+            )  # Reshaping to image wavelet shape
+        out = x.clone()
+        out[abs(out) < ths_map] = 0
+        return out
+
+    def hard_threshold_topk(self, x, ths=0.1):
+        r"""
+        Hard thresholding of the wavelet coefficients by keeping only the top-k coefficients and setting the others to
+        0.
+
+        :param torch.Tensor x: wavelet coefficients.
+        :param float, int ths: top k coefficients to keep. If ``float``, it is interpreted as a proportion of the total
+            number of coefficients. If ``int``, it is interpreted as the number of coefficients to keep.
+        """
+        if isinstance(ths, float):
+            k = int(ths * x.shape[-2] * x.shape[-1])
+        else:
+            k = int(ths)
+
+        # Reshape arrays to 2D and initialize output to 0
+        x_flat = x.view(x.shape[0], -1)
+        out = torch.zeros_like(x_flat)
+
+        topk_indices_flat = torch.topk(abs(x_flat), k, dim=-1)[1]
+
+        # Convert the flattened indices to the original indices of x
+        batch_indices = (
+            torch.arange(x.shape[0], device=x.device).unsqueeze(1).repeat(1, k)
+        )
+        topk_indices = torch.stack([batch_indices, topk_indices_flat], dim=-1)
+
+        # Set output's top-k elements to values from original x
+        out[tuple(topk_indices.view(-1, 2).t())] = x_flat[
+            tuple(topk_indices.view(-1, 2).t())
+        ]
+        return torch.reshape(out, x.shape)
+
+    def forward(self, x, ths=0.1):
+        r"""
+        Run the model on a noisy image.
+
+        :param torch.Tensor x: noisy image.
+        :param int, float, torch.Tensor ths: thresholding parameter.
+            If ``non_linearity`` equals ``"soft"`` or ``"hard"``, ``ths`` serves as a (soft or hard)
+            thresholding parameter for the wavelet coefficients. If ``non_linearity`` equals ``"topk"``,
+            ``ths`` can indicate the number of wavelet coefficients
+            that are kept (if ``int``) or the proportion of coefficients that are kept (if ``float``).
+
+        """
+        h, w = x.size()[-2:]
+        padding_bottom = h % 2
+        padding_right = w % 2
+        x = torch.nn.ReplicationPad2d((0, padding_right, 0, padding_bottom))(x)
+
+        coeffs = self.dwt(x)
+        for l in range(self.level):
+            ths_cur = (
+                ths
+                if (
+                    isinstance(ths, int)
+                    or isinstance(ths, float)
+                    or len(ths.shape) == 0
+                    or ths.shape[0] == 1
+                )
+                else ths[l]
+            )
+            if self.non_linearity == "soft":
+                coeffs[1][l] = self.prox_l1(coeffs[1][l], ths_cur)
+            elif self.non_linearity == "hard":
+                coeffs[1][l] = self.prox_l0(coeffs[1][l], ths_cur)
+            elif self.non_linearity == "topk":
+                coeffs[1][l] = self.hard_threshold_topk(coeffs[1][l], ths_cur)
+        y = self.iwt(coeffs)
+
+        y = y[..., :h, :w]
+        return y
+
+
+class WaveletDict(nn.Module):
+    r"""
+    Overcomplete Wavelet denoising with the :math:`\ell_1` norm.
+
+    This denoiser is defined as the solution to the optimization problem:
+
+    .. math::
+
+        \underset{x}{\arg\min} \;  \|x-y\|^2 + \lambda \|\Psi x\|_n
+
+    where :math:`\Psi` is an overcomplete wavelet transform, composed of 2 or more wavelets, i.e.,
+    :math:`\Psi=[\Psi_1,\Psi_2,\dots,\Psi_L]`, :math:`\lambda>0` is a hyperparameter, and where
+    :math:`\|\cdot\|_n` is either the :math:`\ell_1` norm (``non_linearity="soft"``),
+    the :math:`\ell_0` norm (``non_linearity="hard"``) or a variant of the :math:`\ell_0` norm
+    (``non_linearity="topk"``) where only the top-k coefficients are kept; see :meth:`deepinv.models.WaveletPrior` for
+    more details.
+
+    The solution is not available in closed-form, thus the denoiser runs an optimization algorithm for each test image.
+
+    :param int level: decomposition level of the wavelet transform.
+    :param list[str] wv: list of mother wavelets. The names of the wavelets can be found in `here
+        <https://wavelets.pybytes.com/>`_. (default: ["db8", "db4"]).
+    :param str device: cpu or gpu.
+    :param int max_iter: number of iterations of the optimization algorithm (default: 10).
+    :param str non_linearity: "soft", "hard" or "topk" thresholding (default: "soft")
+    """
+
+    def __init__(
+        self, level=3, list_wv=["db8", "db4"], max_iter=10, non_linearity="soft"
+    ):
+        super().__init__()
+        self.level = level
+        self.list_prox = nn.ModuleList(
+            [
+                WaveletPrior(level=level, wv=wv, non_linearity=non_linearity)
+                for wv in list_wv
+            ]
+        )
+        self.max_iter = max_iter
+
+    def forward(self, y, ths=0.1):
+        r"""
+        Run the model on a noisy image.
+
+        :param torch.Tensor y: noisy image.
+        :param float, torch.Tensor ths: noise level.
+        """
+        z_p = y.repeat(len(self.list_prox), 1, 1, 1, 1)
+        p_p = torch.zeros_like(z_p)
+        x = p_p.clone()
+        for it in range(self.max_iter):
+            x_prev = x.clone()
+            for p in range(len(self.list_prox)):
+                p_p[p, ...] = self.list_prox[p](z_p[p, ...], ths)
+            x = torch.mean(p_p.clone(), axis=0)
+            for p in range(len(self.list_prox)):
+                z_p[p, ...] = x + z_p[p, ...].clone() - p_p[p, ...]
+            rel_crit = torch.linalg.norm((x - x_prev).flatten()) / torch.linalg.norm(
+                x.flatten() + 1e-6
+            )
+            if rel_crit < 1e-3:
+                break
+        return x

deepinv/optim/__init__.py

@@ -0,0 +1,5 @@
+from .data_fidelity import DataFidelity, L2, L1, IndicatorL2, PoissonLikelihood
+from .optimizers import BaseOptim, optim_builder
+from .fixed_point import FixedPoint
+from .prior import Prior, ScorePrior, Tikhonov, PnP, RED, L1Prior
+from .optim_iterators.optim_iterator import OptimIterator