deepinv 0.1.0.dev0__py3-none-any.whl

deepinv/models/artifactremoval.py

@@ -0,0 +1,56 @@
+# import DeepInv
+import torch
+import torch.nn as nn
+
+
+class ArtifactRemoval(nn.Module):
+    r"""
+    Artifact removal architecture :math:`\phi(A^{\top}y)`.
+
+    The architecture is inspired by the FBPConvNet approach of https://arxiv.org/pdf/1611.03679
+    where a deep network :math:`\phi` is used to improve the linear reconstruction :math:`A^{\top}y`.
+
+    :param torch.nn.Module backbone_net: Base network :math:`\phi`, can be pretrained or not.
+    :param bool pinv: If ``True`` uses pseudo-inverse :math:`A^{\dagger}y` instead of the default transpose.
+    :param torch.device device: cpu or gpu.
+    """
+
+    def __init__(self, backbone_net, pinv=False, ckpt_path=None, device=None):
+        super(ArtifactRemoval, self).__init__()
+        self.pinv = pinv
+        self.backbone_net = backbone_net
+
+        if ckpt_path is not None:
+            self.backbone_net.load_state_dict(torch.load(ckpt_path), strict=True)
+            self.backbone_net.eval()
+
+        if type(self.backbone_net).__name__ == "UNetRes":
+            for _, v in self.backbone_net.named_parameters():
+                v.requires_grad = False
+            self.backbone_net = self.backbone_net.to(device)
+
+    def forward(self, y, physics, **kwargs):
+        r"""
+        Reconstructs a signal estimate from measurements y
+
+        :param torch.tensor y: measurements
+        :param deepinv.physics.Physics physics: forward operator
+        """
+        if isinstance(physics, nn.DataParallel):
+            physics = physics.module
+
+        y_in = physics.A_adjoint(y) if not self.pinv else physics.A_dagger(y)
+        if type(self.backbone_net).__name__ == "UNetRes":
+            noise_level_map = (
+                torch.FloatTensor(y_in.size(0), 1, y_in.size(2), y_in.size(3))
+                .fill_(kwargs["sigma"])
+                .to(y_in.dtype)
+            )
+            y_in = torch.cat((y_in, noise_level_map), 1)
+
+        if hasattr(physics.noise_model, "sigma"):
+            sigma = physics.noise_model.sigma
+        else:
+            sigma = None
+
+        return self.backbone_net(y_in, sigma)

deepinv/models/bm3d.py

@@ -0,0 +1,57 @@
+import numpy as np
+import torch
+import torch.nn as nn
+
+# Compat for optional dependency on BM3D
+try:
+    import bm3d
+except:
+    bm3d = ImportError("The bm3d package is not installed.")
+
+
+class BM3D(nn.Module):
+    """
+    BM3D denoiser.
+
+
+    This module wraps the BM3D denoiser from the `BM3D python package <https://pypi.org/project/bm3d/>`_.
+    The denoiser is applied sequentially to each noisy image in the batch.
+
+    The BM3D denoiser was introduced in "Image denoising by sparse 3D transform-domain collaborative filtering", by
+    Davob et al., IEEE Transactions on Image Processing (2007).
+
+
+    """
+
+    def __init__(self):
+        super().__init__()
+        if isinstance(bm3d, ImportError):
+            raise ImportError(
+                "BM3D denoiser not available. Please install the bm3d package with `pip install bm3d`."
+            ) from bm3d
+
+    def forward(self, x, sigma):
+        r"""
+        Run the denoiser on image with noise level :math:`\sigma`.
+
+        :param torch.Tensor x: noisy image
+        :param float sigma: noise level (not used)
+        """
+
+        out = torch.zeros_like(x)
+
+        for i in range(x.shape[0]):
+            out[i, :, :, :] = array2tensor(
+                bm3d.bm3d(tensor2array(x[i, :, :, :]), sigma)
+            )
+        return out
+
+
+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)