deepinv 0.1.0.dev0__py3-none-any.whl

deepinv/models/equivariant.py

@@ -0,0 +1,135 @@
+import numpy as np
+
+import torch
+
+
+class EquivariantDenoiser(torch.nn.Module):
+    r"""
+    Turns the input denoiser into an equivariant denoiser with respect to geometric transforms.
+
+    Recall that a denoiser is equivariant with respect to a group of transformations if it commutes with the action of
+    the group. More precisely, let :math:`\mathcal{G}` be a group of transformations :math:`\{T_g\}_{g\in \mathcal{G}}`
+    and :math:`\denoisername` a denoiser. Then, :math:`\denoisername` is equivariant with respect to :math:`\mathcal{G}`
+    if :math:`\denoisername(T_g(x)) = T_g(\denoisername(x))` for any image :math:`x` and any :math:`g\in \mathcal{G}`.
+
+    The denoiser can be turned into an equivariant denoiser by averaging over the group of transforms, i.e.
+
+    .. math::
+        \operatorname{D}^{\text{eq}}_{\sigma}(x) = \frac{1}{|\mathcal{G}|}\sum_{g\in \mathcal{G}} T_g^{-1}(\operatorname{D}_{\sigma}(T_g(x))).
+
+    Otherwise, as proposed in <https://arxiv.org/abs/2312.01831>`_, a Monte-Carlo approximation can be obtained by
+    sampling :math:`g \sim \mathcal{G}` at random and applying
+
+    .. math::
+        \operatorname{D}^{\text{MC}}_{\sigma}(x) = T_g^{-1}(\operatorname{D}_{\sigma}(T_g(x))).
+
+
+    :param callable denoiser: Denoiser :math:`\operatorname{D}_{\sigma}`.
+    :param str transform: type of geometric transformation. Can be either 'rotations', 'flips' or 'rotoflips'.
+        If 'rotations', the group of transformations contains the 4 rotations by multiples of 90 degrees; if 'flips',
+        the group of transformations contains the 2 horizontal and vertical flips; if 'rotoflips', the group of
+        transformations contains the 8 rotations and flips.
+    :param bool random: if True, the denoiser is applied to a randomly transformed version of the input image.
+        If False, the denoiser is applied to the average of all the transformed images, turning the denoiser into an
+        equivariant denoiser with respect to the chosen group of transformations. Otherwise, it is a Monte-Carlo
+        approximation of an equivariant denoiser.
+    """
+
+    def __init__(self, denoiser, transform="rotations", random=True):
+        super().__init__()
+        self.denoiser = denoiser
+        self.rotations = True if "rot" in transform else False
+        self.flips = True if "flip" in transform else False
+        self.random = random
+
+    def forward(self, x, sigma):
+        r"""
+        Applies the denoiser to the input image with the appropriate transformation.
+
+        :param torch.Tensor x: input image.
+        :param float sigma: noise level.
+        :return: denoised image.
+        """
+        return denoise_rotate(
+            self.denoiser,
+            x,
+            sigma,
+            rotations=self.rotations,
+            flips=self.flips,
+            random=self.random,
+        )
+
+
+def denoise_rotate(
+    denoiser,
+    image,
+    sigma,
+    rotations=True,
+    flips=False,
+    random=True,
+):
+    r"""
+    Applies a geometric transform (rotations and/or flips) to the input image, denoises it with the denoiser and
+    transform back the result. The output is either the average of all the transformed images (if random=False) or a
+    randomly transformed version of the denoised image (if random=True).
+
+    :param callable denoiser: Denoiser :math:`\operatorname{D}_{\sigma}`.
+    :param torch.Tensor image: input image.
+    :param float sigma: noise level.
+    :param bool rotations: if True, rotations are applied to the input image.
+    :param bool flips: if True, flips are applied to the input image.
+    :param bool random: if True, the denoiser is applied to a randomly transformed version of the input image.
+    :return: denoised image.
+    """
+    if random:
+        if rotations:
+            idx = np.random.randint(8) if flips else np.random.randint(4)
+        elif flips:
+            idx = np.random.choice([4, 6])
+        denoised = denoise_rotate_flip_fn(denoiser, image, sigma, idx)
+    else:
+        if rotations:
+            list_idx = list(range(8)) if flips else list(range(4))
+        elif flips:
+            list_idx = [4, 6]
+        denoised = torch.zeros_like(image)
+        for idx in list_idx:
+            denoised = denoised + denoise_rotate_flip_fn(denoiser, image, sigma, idx)
+        denoised = denoised / len(list_idx)
+    return denoised
+
+
+def denoise_rotate_flip_fn(denoiser, x, sigma_den, idx):
+    if idx == 0:
+        out = denoiser(x, sigma_den)
+    elif idx == 1:
+        out = rot3(denoiser(rot1(x), sigma_den))
+    elif idx == 2:
+        out = rot2(denoiser(rot2(x), sigma_den))
+    elif idx == 3:
+        out = rot1(denoiser(rot3(x), sigma_den))
+    elif idx == 4:
+        out = hflip(denoiser(hflip(x), sigma_den))
+    elif idx == 5:
+        out = hflip(rot3(denoiser(rot1(hflip(x)), sigma_den)))
+    elif idx == 6:
+        out = hflip(rot2(denoiser(rot2(hflip(x)), sigma_den)))
+    elif idx == 7:
+        out = hflip(rot1(denoiser(rot3(hflip(x)), sigma_den)))
+    return out
+
+
+def hflip(x):
+    return torch.flip(x, dims=[-1])
+
+
+def rot1(x):
+    return torch.rot90(x, k=1, dims=[-2, -1])
+
+
+def rot2(x):
+    return torch.rot90(x, k=2, dims=[-2, -1])
+
+
+def rot3(x):
+    return torch.rot90(x, k=3, dims=[-2, -1])

deepinv/models/median.py

@@ -0,0 +1,51 @@
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.modules.utils import _pair, _quadruple
+
+# code adapted from https://gist.github.com/rwightman/f2d3849281624be7c0f11c85c87c1598
+
+
+class MedianFilter(nn.Module):
+    r"""
+    Median filter.
+
+
+    :param int kernel_size: size of pooling kernel, int or 2-tuple
+    :param padding: pool padding, int or 4-tuple (l, r, t, b) as in pytorch F.pad
+    :param same: override padding and enforce same padding, boolean
+    """
+
+    def __init__(self, kernel_size=9, padding=0, same=True):
+        super(MedianFilter, self).__init__()
+        self.k = _pair(kernel_size)
+        self.stride = _pair(1)
+        self.padding = _quadruple(padding)  # convert to l, r, t, b
+        self.same = same
+
+    def _padding(self, x):
+        if self.same:
+            ih, iw = x.size()[2:]
+            if ih % self.stride[0] == 0:
+                ph = max(self.k[0] - self.stride[0], 0)
+            else:
+                ph = max(self.k[0] - (ih % self.stride[0]), 0)
+            if iw % self.stride[1] == 0:
+                pw = max(self.k[1] - self.stride[1], 0)
+            else:
+                pw = max(self.k[1] - (iw % self.stride[1]), 0)
+            pl = pw // 2
+            pr = pw - pl
+            pt = ph // 2
+            pb = ph - pt
+            padding = (pl, pr, pt, pb)
+        else:
+            padding = self.padding
+        return padding
+
+    def forward(self, x, sigma=None):
+        # using existing pytorch functions and tensor ops so that we get autograd,
+        # would likely be more efficient to implement from scratch at C/Cuda level
+        x = F.pad(x, self._padding(x), mode="reflect")
+        x = x.unfold(2, self.k[0], self.stride[0]).unfold(3, self.k[1], self.stride[1])
+        x = x.contiguous().view(x.size()[:4] + (-1,)).median(dim=-1)[0]
+        return x

deepinv/models/scunet.py

@@ -0,0 +1,490 @@
+# Code taken from https://github.com/cszn/SCUNet/blob/main/models/network_scunet.py
+import torch
+import torch.nn as nn
+import numpy as np
+from einops import rearrange
+from einops.layers.torch import Rearrange
+from .utils import get_weights_url
+
+# Compatibility with optional dependency on timm
+try:
+    import timm
+    from timm.models.layers import trunc_normal_, DropPath
+except ImportError as e:
+    timm = e
+
+
+class WMSA(nn.Module):
+    """Self-attention module in Swin Transformer"""
+
+    def __init__(self, input_dim, output_dim, head_dim, window_size, type):
+        if isinstance(timm, ImportError):
+            raise ImportError(
+                "timm is needed to use the SCUNet class. Please install it with `pip install timm`"
+            ) from timm
+        super(WMSA, self).__init__()
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+        self.head_dim = head_dim
+        self.scale = self.head_dim**-0.5
+        self.n_heads = input_dim // head_dim
+        self.window_size = window_size
+        self.type = type
+        self.embedding_layer = nn.Linear(self.input_dim, 3 * self.input_dim, bias=True)
+
+        # self.relative_position_params = nn.Parameter(torch.zeros(self.n_heads, 2 * window_size - 1, 2 * window_size -1))
+        self.relative_position_params = nn.Parameter(
+            torch.zeros((2 * window_size - 1) * (2 * window_size - 1), self.n_heads)
+        )
+
+        self.linear = nn.Linear(self.input_dim, self.output_dim)
+
+        trunc_normal_(self.relative_position_params, std=0.02)
+        self.relative_position_params = torch.nn.Parameter(
+            self.relative_position_params.view(
+                2 * window_size - 1, 2 * window_size - 1, self.n_heads
+            )
+            .transpose(1, 2)
+            .transpose(0, 1)
+        )
+
+    def generate_mask(self, h, w, p, shift):
+        """generating the mask of SW-MSA
+        Args:
+            shift: shift parameters in CyclicShift.
+        Returns:
+            attn_mask: should be (1 1 w p p),
+        """
+        # supporting sqaure.
+        attn_mask = torch.zeros(
+            h,
+            w,
+            p,
+            p,
+            p,
+            p,
+            dtype=torch.bool,
+            device=self.relative_position_params.device,
+        )
+        if self.type == "W":
+            return attn_mask
+
+        s = p - shift
+        attn_mask[-1, :, :s, :, s:, :] = True
+        attn_mask[-1, :, s:, :, :s, :] = True
+        attn_mask[:, -1, :, :s, :, s:] = True
+        attn_mask[:, -1, :, s:, :, :s] = True
+        attn_mask = rearrange(
+            attn_mask, "w1 w2 p1 p2 p3 p4 -> 1 1 (w1 w2) (p1 p2) (p3 p4)"
+        )
+        return attn_mask
+
+    def forward(self, x):
+        """Forward pass of Window Multi-head Self-attention module.
+        Args:
+            x: input tensor with shape of [b h w c];
+            attn_mask: attention mask, fill -inf where the value is True;
+        Returns:
+            output: tensor shape [b h w c]
+        """
+        if self.type != "W":
+            x = torch.roll(
+                x,
+                shifts=(-(self.window_size // 2), -(self.window_size // 2)),
+                dims=(1, 2),
+            )
+        x = rearrange(
+            x,
+            "b (w1 p1) (w2 p2) c -> b w1 w2 p1 p2 c",
+            p1=self.window_size,
+            p2=self.window_size,
+        )
+        h_windows = x.size(1)
+        w_windows = x.size(2)
+        # sqaure validation
+        # assert h_windows == w_windows
+
+        x = rearrange(
+            x,
+            "b w1 w2 p1 p2 c -> b (w1 w2) (p1 p2) c",
+            p1=self.window_size,
+            p2=self.window_size,
+        )
+        qkv = self.embedding_layer(x)
+        q, k, v = rearrange(
+            qkv, "b nw np (threeh c) -> threeh b nw np c", c=self.head_dim
+        ).chunk(3, dim=0)
+        sim = torch.einsum("hbwpc,hbwqc->hbwpq", q, k) * self.scale
+        # Adding learnable relative embedding
+        sim = sim + rearrange(self.relative_embedding(), "h p q -> h 1 1 p q")
+        # Using Attn Mask to distinguish different subwindows.
+        if self.type != "W":
+            attn_mask = self.generate_mask(
+                h_windows, w_windows, self.window_size, shift=self.window_size // 2
+            )
+            sim = sim.masked_fill_(attn_mask, float("-inf"))
+
+        probs = nn.functional.softmax(sim, dim=-1)
+        output = torch.einsum("hbwij,hbwjc->hbwic", probs, v)
+        output = rearrange(output, "h b w p c -> b w p (h c)")
+        output = self.linear(output)
+        output = rearrange(
+            output,
+            "b (w1 w2) (p1 p2) c -> b (w1 p1) (w2 p2) c",
+            w1=h_windows,
+            p1=self.window_size,
+        )
+
+        if self.type != "W":
+            output = torch.roll(
+                output,
+                shifts=(self.window_size // 2, self.window_size // 2),
+                dims=(1, 2),
+            )
+        return output
+
+    def relative_embedding(self):
+        cord = torch.tensor(
+            np.array(
+                [
+                    [i, j]
+                    for i in range(self.window_size)
+                    for j in range(self.window_size)
+                ]
+            )
+        )
+        relation = cord[:, None, :] - cord[None, :, :] + self.window_size - 1
+        # negative is allowed
+        return self.relative_position_params[
+            :, relation[:, :, 0].long(), relation[:, :, 1].long()
+        ]
+
+
+class Block(nn.Module):
+    def __init__(
+        self,
+        input_dim,
+        output_dim,
+        head_dim,
+        window_size,
+        drop_path,
+        type="W",
+        input_resolution=None,
+    ):
+        """SwinTransformer Block"""
+        super(Block, self).__init__()
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+        assert type in ["W", "SW"]
+        self.type = type
+        if input_resolution <= window_size:
+            self.type = "W"
+
+        # print(
+        #    "Block Initial Type: {}, drop_path_rate:{:.6f}".format(self.type, drop_path)
+        # )
+        self.ln1 = nn.LayerNorm(input_dim)
+        self.msa = WMSA(input_dim, input_dim, head_dim, window_size, self.type)
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+        self.ln2 = nn.LayerNorm(input_dim)
+        self.mlp = nn.Sequential(
+            nn.Linear(input_dim, 4 * input_dim),
+            nn.GELU(),
+            nn.Linear(4 * input_dim, output_dim),
+        )
+
+    def forward(self, x):
+        x = x + self.drop_path(self.msa(self.ln1(x)))
+        x = x + self.drop_path(self.mlp(self.ln2(x)))
+        return x
+
+
+class ConvTransBlock(nn.Module):
+    def __init__(
+        self,
+        conv_dim,
+        trans_dim,
+        head_dim,
+        window_size,
+        drop_path,
+        type="W",
+        input_resolution=None,
+    ):
+        """SwinTransformer and Conv Block"""
+        super(ConvTransBlock, self).__init__()
+        self.conv_dim = conv_dim
+        self.trans_dim = trans_dim
+        self.head_dim = head_dim
+        self.window_size = window_size
+        self.drop_path = drop_path
+        self.type = type
+        self.input_resolution = input_resolution
+
+        assert self.type in ["W", "SW"]
+        if self.input_resolution <= self.window_size:
+            self.type = "W"
+
+        self.trans_block = Block(
+            self.trans_dim,
+            self.trans_dim,
+            self.head_dim,
+            self.window_size,
+            self.drop_path,
+            self.type,
+            self.input_resolution,
+        )
+        self.conv1_1 = nn.Conv2d(
+            self.conv_dim + self.trans_dim,
+            self.conv_dim + self.trans_dim,
+            1,
+            1,
+            0,
+            bias=True,
+        )
+        self.conv1_2 = nn.Conv2d(
+            self.conv_dim + self.trans_dim,
+            self.conv_dim + self.trans_dim,
+            1,
+            1,
+            0,
+            bias=True,
+        )
+
+        self.conv_block = nn.Sequential(
+            nn.Conv2d(self.conv_dim, self.conv_dim, 3, 1, 1, bias=False),
+            nn.ReLU(True),
+            nn.Conv2d(self.conv_dim, self.conv_dim, 3, 1, 1, bias=False),
+        )
+
+    def forward(self, x):
+        conv_x, trans_x = torch.split(
+            self.conv1_1(x), (self.conv_dim, self.trans_dim), dim=1
+        )
+        conv_x = self.conv_block(conv_x) + conv_x
+        trans_x = Rearrange("b c h w -> b h w c")(trans_x)
+        trans_x = self.trans_block(trans_x)
+        trans_x = Rearrange("b h w c -> b c h w")(trans_x)
+        res = self.conv1_2(torch.cat((conv_x, trans_x), dim=1))
+        x = x + res
+
+        return x
+
+
+class SCUNet(nn.Module):
+    r"""
+    SCUNet denoising network.
+
+    The Swin-Conv-UNet (SCUNet) denoising was introduced in `Practical Blind Denoising via Swin-Conv-UNet and
+    Data Synthesis <https://arxiv.org/abs/2203.13278>`_.
+
+    :param int in_nc: number of input channels. Default: 3.
+    :param list config: number of layers in each stage. Default: [4, 4, 4, 4, 4, 4, 4].
+    :param int dim: number of channels in each layer. Default: 64.
+    :param float drop_path_rate: drop path per sample rate (stochastic depth) for each layer. Default: 0.0.
+    :param int input_resolution: input resolution. Default: 256.
+    :param bool 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 the default architecture).
+        Finally, ``pretrained`` can also be set as a path to the user's own pretrained weights. Default: 'download'.
+        See :ref:`pretrained-weights <pretrained-weights>` for more details.
+    :param bool train: training or testing mode. Default: False.
+    :param str device: gpu or cpu. Default: 'cpu'.
+    ....
+    """
+
+    def __init__(
+        self,
+        in_nc=3,
+        config=[4, 4, 4, 4, 4, 4, 4],
+        dim=64,
+        drop_path_rate=0.0,
+        input_resolution=256,
+        pretrained="download",
+        train=False,
+        device="cpu",
+    ):
+        super(SCUNet, self).__init__()
+        self.config = config
+        self.dim = dim
+        self.head_dim = 32
+        self.window_size = 8
+
+        # drop path rate for each layer
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(config))]
+
+        self.m_head = [nn.Conv2d(in_nc, dim, 3, 1, 1, bias=False)]
+
+        begin = 0
+        self.m_down1 = [
+            ConvTransBlock(
+                dim // 2,
+                dim // 2,
+                self.head_dim,
+                self.window_size,
+                dpr[i + begin],
+                "W" if not i % 2 else "SW",
+                input_resolution,
+            )
+            for i in range(config[0])
+        ] + [nn.Conv2d(dim, 2 * dim, 2, 2, 0, bias=False)]
+
+        begin += config[0]
+        self.m_down2 = [
+            ConvTransBlock(
+                dim,
+                dim,
+                self.head_dim,
+                self.window_size,
+                dpr[i + begin],
+                "W" if not i % 2 else "SW",
+                input_resolution // 2,
+            )
+            for i in range(config[1])
+        ] + [nn.Conv2d(2 * dim, 4 * dim, 2, 2, 0, bias=False)]
+
+        begin += config[1]
+        self.m_down3 = [
+            ConvTransBlock(
+                2 * dim,
+                2 * dim,
+                self.head_dim,
+                self.window_size,
+                dpr[i + begin],
+                "W" if not i % 2 else "SW",
+                input_resolution // 4,
+            )
+            for i in range(config[2])
+        ] + [nn.Conv2d(4 * dim, 8 * dim, 2, 2, 0, bias=False)]
+
+        begin += config[2]
+        self.m_body = [
+            ConvTransBlock(
+                4 * dim,
+                4 * dim,
+                self.head_dim,
+                self.window_size,
+                dpr[i + begin],
+                "W" if not i % 2 else "SW",
+                input_resolution // 8,
+            )
+            for i in range(config[3])
+        ]
+
+        begin += config[3]
+        self.m_up3 = [nn.ConvTranspose2d(8 * dim, 4 * dim, 2, 2, 0, bias=False)] + [
+            ConvTransBlock(
+                2 * dim,
+                2 * dim,
+                self.head_dim,
+                self.window_size,
+                dpr[i + begin],
+                "W" if not i % 2 else "SW",
+                input_resolution // 4,
+            )
+            for i in range(config[4])
+        ]
+
+        begin += config[4]
+        self.m_up2 = [nn.ConvTranspose2d(4 * dim, 2 * dim, 2, 2, 0, bias=False)] + [
+            ConvTransBlock(
+                dim,
+                dim,
+                self.head_dim,
+                self.window_size,
+                dpr[i + begin],
+                "W" if not i % 2 else "SW",
+                input_resolution // 2,
+            )
+            for i in range(config[5])
+        ]
+
+        begin += config[5]
+        self.m_up1 = [nn.ConvTranspose2d(2 * dim, dim, 2, 2, 0, bias=False)] + [
+            ConvTransBlock(
+                dim // 2,
+                dim // 2,
+                self.head_dim,
+                self.window_size,
+                dpr[i + begin],
+                "W" if not i % 2 else "SW",
+                input_resolution,
+            )
+            for i in range(config[6])
+        ]
+
+        self.m_tail = [nn.Conv2d(dim, in_nc, 3, 1, 1, bias=False)]
+
+        self.m_head = nn.Sequential(*self.m_head)
+        self.m_down1 = nn.Sequential(*self.m_down1)
+        self.m_down2 = nn.Sequential(*self.m_down2)
+        self.m_down3 = nn.Sequential(*self.m_down3)
+        self.m_body = nn.Sequential(*self.m_body)
+        self.m_up3 = nn.Sequential(*self.m_up3)
+        self.m_up2 = nn.Sequential(*self.m_up2)
+        self.m_up1 = nn.Sequential(*self.m_up1)
+        self.m_tail = nn.Sequential(*self.m_tail)
+        # self.apply(self._init_weights)
+
+        if pretrained is not None:
+            if pretrained == "download":
+                name = "scunet_color_real_psnr.pth"
+                url = get_weights_url(model_name="scunet", file_name=name)
+                ckpt_drunet = torch.hub.load_state_dict_from_url(
+                    url, map_location=lambda storage, loc: storage, file_name=name
+                )
+            else:
+                ckpt_drunet = torch.load(
+                    pretrained, map_location=lambda storage, loc: storage
+                )
+
+            self.load_state_dict(ckpt_drunet, strict=True)
+
+        if not train:
+            self.eval()
+            for _, v in self.named_parameters():
+                v.requires_grad = False
+
+        if device is not None:
+            self.to(device)
+
+    def forward_scunet(self, x0):
+        h, w = x0.size()[-2:]
+        paddingBottom = int(np.ceil(h / 64) * 64 - h)
+        paddingRight = int(np.ceil(w / 64) * 64 - w)
+        x0 = nn.ReplicationPad2d((0, paddingRight, 0, paddingBottom))(x0)
+
+        x1 = self.m_head(x0)
+        x2 = self.m_down1(x1)
+        x3 = self.m_down2(x2)
+        x4 = self.m_down3(x3)
+        x = self.m_body(x4)
+        x = self.m_up3(x + x4)
+        x = self.m_up2(x + x3)
+        x = self.m_up1(x + x2)
+        x = self.m_tail(x + x1)
+
+        x = x[..., :h, :w]
+
+        return x
+
+    def forward(self, x, sigma):  # This is a blind model: sigma is not used
+        den = self.forward_scunet(x)
+        return den
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+
+# if __name__ == '__main__':
+#     # torch.cuda.empty_cache()
+#     net = SCUNet(pretrained='download', device='cpu', train=False)
+#
+#     x = torch.randn((2, 3, 64, 128))
+#     x = net(x)
+#     print(x.shape)