deepinv 0.1.0.dev0__py3-none-any.whl

deepinv/models/dip.py

@@ -0,0 +1,214 @@
+import torch
+import torch.nn as nn
+import numpy as np
+from deepinv.loss import MCLoss
+from tqdm import tqdm
+
+
+def add_module(self, module):
+    self.add_module(str(len(self) + 1), module)
+
+
+torch.nn.Module.add = add_module
+
+
+class ConvDecoder(nn.Module):
+    r"""
+    Convolutional decoder network.
+
+    The architecture was introduced in `"Accelerated MRI with Un-trained Neural Networks" <https://arxiv.org/abs/2007.02471>`_,
+    and it is well suited as a deep image prior (see :class:`deepinv.models.DeepImagePrior`).
+
+
+    :param tuple img_shape: shape of the output image.
+    :param tuple in_size: size of the input vector.
+    :param int layers: number of layers in the network.
+    :param int channels: number of channels in the network.
+
+    """
+
+    #  Code adapted from https://github.com/MLI-lab/ConvDecoder/tree/master by Darestani and Heckel.
+    def __init__(self, img_shape, in_size=(4, 4), layers=7, channels=256):
+        super(ConvDecoder, self).__init__()
+
+        out_size = img_shape[1:]
+        output_channels = img_shape[0]
+
+        # parameter setup
+        kernel_size = 3
+        strides = [1] * (layers - 1)
+
+        # compute up-sampling factor from one layer to another
+        scale_x, scale_y = (
+            (out_size[0] / in_size[0]) ** (1.0 / (layers - 1)),
+            (out_size[1] / in_size[1]) ** (1.0 / (layers - 1)),
+        )
+        hidden_size = [
+            (
+                int(np.ceil(scale_x**n * in_size[0])),
+                int(np.ceil(scale_y**n * in_size[1])),
+            )
+            for n in range(1, (layers - 1))
+        ] + [out_size]
+
+        # hidden layers
+        self.net = nn.Sequential()
+        for i in range(layers - 1):
+            self.net.add(nn.Upsample(size=hidden_size[i], mode="nearest"))
+            conv = nn.Conv2d(
+                channels,
+                channels,
+                kernel_size,
+                strides[i],
+                padding=(kernel_size - 1) // 2,
+                bias=True,
+            )
+            self.net.add(conv)
+            self.net.add(nn.ReLU())
+            self.net.add(nn.BatchNorm2d(channels, affine=True))
+        # final layer
+        self.net.add(
+            nn.Conv2d(
+                channels,
+                channels,
+                kernel_size,
+                strides[i],
+                padding=(kernel_size - 1) // 2,
+                bias=True,
+            )
+        )
+        self.net.add(nn.ReLU())
+        self.net.add(nn.BatchNorm2d(channels, affine=True))
+        self.net.add(nn.Conv2d(channels, output_channels, 1, 1, padding=0, bias=True))
+
+    def forward(self, x, scale_out=1):
+        return self.net(x) * scale_out
+
+
+class DeepImagePrior(torch.nn.Module):
+    r"""
+
+    Deep Image Prior reconstruction.
+
+    This method is based on the paper `"Deep Image Prior" by Ulyanov et al. (2018)
+    <https://arxiv.org/abs/1711.10925>`_, and reconstructs
+    an image by minimizing the loss function
+
+    .. math::
+
+        \min_{\theta}  \|y-Af_{\theta}(z)\|^2
+
+    where :math:`z` is a random input and :math:`f_{\theta}` is a convolutional decoder network with parameters
+    :math:`\theta`. The minimization should be stopped early to avoid overfitting. The method uses the Adam
+    optimizer.
+
+    .. note::
+
+        This method only works with certain convolutional decoder networks. We recommend using the
+        network :class:`deepinv.models.ConvDecoder`.
+
+
+    .. note::
+
+        The number of iterations and learning rate are set to the values used in the original paper. However, these
+        values may not be optimal for all problems. We recommend experimenting with different values.
+
+    :param torch.nn.Module generator: Convolutional decoder network.
+    :param list, tuple input_size: Size (C,H,W) of the input noise vector :math:`z`.
+    :param int iterations: Number of optimization iterations.
+    :param float learning_rate: Learning rate of the Adam optimizer.
+    :param bool verbose: If ``True``, print progress.
+    :param bool re_init: If ``True``, re-initialize the network parameters before each reconstruction.
+
+    """
+
+    def __init__(
+        self,
+        generator,
+        input_size,
+        iterations=2500,
+        learning_rate=1e-2,
+        verbose=False,
+        re_init=False,
+    ):
+        super().__init__()
+        self.generator = generator
+        self.max_iter = int(iterations)
+        self.lr = learning_rate
+        self.loss = MCLoss()
+        self.verbose = verbose
+        self.re_init = re_init
+        self.input_size = input_size
+
+    def forward(self, y, physics):
+        r"""
+        Reconstruct an image from the measurement :math:`y`. The reconstruction is performed by solving a minimiza
+        problem.
+
+        .. warning::
+
+            The optimization is run for every test batch. Thus, this method can be slow when tested on a large
+            number of test batches.
+
+        :param torch.Tensor y: Measurement.
+        :param torch.Tensor physics: Physics model.
+        """
+        if self.re_init:
+            for layer in self.generator.children():
+                if hasattr(layer, "reset_parameters"):
+                    layer.reset_parameters()
+
+        self.generator.requires_grad_(True)
+        z = torch.randn(self.input_size, device=y.device).unsqueeze(0)
+        optimizer = torch.optim.Adam(self.generator.parameters(), lr=self.lr)
+
+        for it in tqdm(range(self.max_iter), disable=(not self.verbose)):
+            x = self.generator(z)
+            error = self.loss(y, x, physics)
+            optimizer.zero_grad()
+            error.backward()
+            optimizer.step()
+
+        return self.generator(z)
+
+
+# test code
+# if __name__ == "__main__":
+#     device = "cuda:0"
+#     import torchvision
+#     import deepinv as dinv
+#
+#     device = dinv.utils.get_freer_gpu()
+#
+#     x = torchvision.io.read_image("../../datasets/celeba/img_align_celeba/085307.jpg")
+#     x = x.unsqueeze(0).float().to(device) / 255
+#     x = torchvision.transforms.Resize((128, 128))(x)
+#
+#     physics = dinv.physics.Inpainting(
+#         tensor_size=x.shape[1:],
+#         device=device,
+#         noise_model=dinv.physics.GaussianNoise(sigma=0.05),
+#     )
+#
+#     y = physics(x)
+#
+#     iterations = 1000
+#     lr = 1e-2
+#     channels = 256
+#     in_size = [8, 8]
+#     backbone = ConvDecoder(
+#         img_shape=x.shape[1:], in_size=in_size, channels=channels
+#     ).to(device)
+#
+#     model = DeepImagePrior(
+#         backbone,
+#         learning_rate=lr,
+#         re_init=True,
+#         iterations=iterations,
+#         verbose=True,
+#         input_size=[channels] + in_size,
+#     ).to(device)
+#
+#     x_hat = model(y, physics)
+#
+#     dinv.utils.plot([x, y, x_hat], titles=["GT", "Meas.", "Recon."])

deepinv/models/dncnn.py

@@ -0,0 +1,131 @@
+import torch.nn as nn
+import torch
+from .utils import get_weights_url
+import math
+
+
+class DnCNN(nn.Module):
+    r"""
+    DnCNN convolutional denoiser.
+
+    The architecture was introduced by Zhang et al. in https://arxiv.org/abs/1608.03981 and is composed of a series of
+    convolutional layers with ReLU activation functions. The number of layers can be specified by the user. Unlike the
+    original paper, this implementation does not include batch normalization layers.
+
+    The network can be initialized with pretrained weights, which can be downloaded from an online repository. The
+    pretrained weights are trained with the default parameters of the network, i.e. 20 layers, 64 channels and biases.
+
+    :param int in_channels: input image channels
+    :param int out_channels: output image channels
+    :param int depth: number of convolutional layers
+    :param bool bias: use bias in the convolutional layers
+    :param int nf: number of channels per convolutional layer
+    :param str, None pretrained: use a pretrained network. If ``pretrained=None``, the weights will be initialized at random
+        using Pytorch's default initialization. If ``pretrained='download'``, the weights will be downloaded from an
+        online repository (only available for architecture with depth 20, 64 channels and biases).
+        It is possible to download weights trained via the regularization method in https://epubs.siam.org/doi/abs/10.1137/20M1387961
+        using ``pretrained='download_lipschitz'``.
+        Finally, ``pretrained`` can also be set as a path to the user's own pretrained weights.
+        See :ref:`pretrained-weights <pretrained-weights>` for more details.
+    :param bool train: training or testing mode
+    :param str device: gpu or cpu
+    """
+
+    def __init__(
+        self,
+        in_channels=3,
+        out_channels=3,
+        depth=20,
+        bias=True,
+        nf=64,
+        pretrained="download",
+        train=False,
+        device="cpu",
+    ):
+        super(DnCNN, self).__init__()
+
+        self.depth = depth
+
+        self.in_conv = nn.Conv2d(
+            in_channels, nf, kernel_size=3, stride=1, padding=1, bias=bias
+        )
+        self.conv_list = nn.ModuleList(
+            [
+                nn.Conv2d(nf, nf, kernel_size=3, stride=1, padding=1, bias=bias)
+                for _ in range(self.depth - 2)
+            ]
+        )
+        self.out_conv = nn.Conv2d(
+            nf, out_channels, kernel_size=3, stride=1, padding=1, bias=bias
+        )
+
+        self.nl_list = nn.ModuleList([nn.ReLU() for _ in range(self.depth - 1)])
+
+        # if pretrain and ckpt_path is not None:
+        #    self.load_state_dict(torch.load(ckpt_path, map_location=lambda storage, loc: storage), strict=True)
+
+        if pretrained is not None:
+            if pretrained.startswith("download"):
+                name = ""
+                if bias and depth == 20:
+                    if pretrained == "download_lipschitz":
+                        if in_channels == 3 and out_channels == 3:
+                            name = "dncnn_sigma2_lipschitz_color.pth"
+                        elif in_channels == 1 and out_channels == 1:
+                            name = "dncnn_sigma2_lipschitz_gray.pth"
+                    else:
+                        if in_channels == 3 and out_channels == 3:
+                            name = "dncnn_sigma2_color.pth"
+                        elif in_channels == 1 and out_channels == 1:
+                            name = "dncnn_sigma2_gray.pth"
+
+                if name == "":
+                    raise Exception(
+                        "No pretrained weights were found online that match the chosen architecture"
+                    )
+                url = get_weights_url(model_name="dncnn", file_name=name)
+                ckpt = torch.hub.load_state_dict_from_url(
+                    url, map_location=lambda storage, loc: storage, file_name=name
+                )
+            else:
+                ckpt = torch.load(pretrained, map_location=lambda storage, loc: storage)
+            self.load_state_dict(ckpt, strict=True)
+
+        if not train:
+            self.eval()
+            for _, v in self.named_parameters():
+                v.requires_grad = False
+        else:
+            self.apply(weights_init_kaiming)
+
+        if device is not None:
+            self.to(device)
+
+    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)
+        """
+        x1 = self.in_conv(x)
+        x1 = self.nl_list[0](x1)
+
+        for i in range(self.depth - 2):
+            x_l = self.conv_list[i](x1)
+            x1 = self.nl_list[i + 1](x_l)
+
+        return self.out_conv(x1) + x
+
+
+def weights_init_kaiming(m):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        nn.init.kaiming_normal_(m.weight.data, a=0, mode="fan_in")
+    elif classname.find("Linear") != -1:
+        nn.init.kaiming_normal_(m.weight.data, a=0, mode="fan_in")
+    elif classname.find("BatchNorm") != -1:
+        m.weight.data.normal_(mean=0, std=math.sqrt(2.0 / 9.0 / 64.0)).clamp_(
+            -0.025, 0.025
+        )
+        nn.init.constant(m.bias.data, 0.0)