deepinv 0.1.0.dev0__py3-none-any.whl

deepinv/physics/remote_sensing.py

@@ -0,0 +1,123 @@
+from deepinv.physics.noise import GaussianNoise
+from deepinv.physics.forward import LinearPhysics
+from deepinv.physics.blur import Downsampling
+from deepinv.physics.range import Decolorize
+from deepinv.utils import TensorList
+
+
+class Pansharpen(LinearPhysics):
+    """
+    Pansharpening forward operator.
+
+    The measurements consist of a high resolution grayscale image and a low resolution RGB image, and
+    are represented using :class:`deepinv.utils.TensorList`, where the first element is the RGB image and the second
+    element is the grayscale image.
+
+    By default, the downsampling is done with a gaussian filter with standard deviation equal to the downsampling,
+    however, the user can provide a custom downsampling filter.
+
+    It is possible to assign a different noise model to the RGB and grayscale images.
+
+    Example usage:
+
+    ::
+
+        import deepinv
+        import torch
+
+        x = torch.randn(1, 3, 256, 256)
+        physics = deepinv.physics.Pansharpen(img_size=x.shape[1:], device=x.device)
+
+        y = physics(x) # returns a TensorList with the RGB and grayscale images
+
+        x_adj = physics.A_adjoint(y)
+        x_pinv = physics.A_dagger(y)
+
+        deepinv.utils.plot([y[0], y[1], x_adj, x_pinv, x], titles=['low res color', 'high res gray',
+                                                                  'A_adjoint', 'A_dagger', 'x'])
+
+    :param tuple[int] img_size: size of the input image.
+    :param int factor: downsampling factor.
+    :param torch.nn.Module noise_color: noise model for the RGB image.
+    :param torch.nn.Module noise_gray: noise model for the grayscale image.
+    :param torch.Tensor, str, NoneType filter: Downsampling filter. It can be 'gaussian', 'bilinear' or 'bicubic' or a
+        custom ``torch.Tensor`` filter. If ``None``, no filtering is applied.
+    :param str padding: options are ``'valid'``, ``'circular'``, ``'replicate'`` and ``'reflect'``.
+        If ``padding='valid'`` the blurred output is smaller than the image (no padding)
+        otherwise the blurred output has the same size as the image.
+    """
+
+    def __init__(
+        self,
+        img_size,
+        factor=4,
+        noise_color=GaussianNoise(sigma=0.0),
+        noise_gray=GaussianNoise(sigma=0.05),
+        filter="gaussian",
+        device="cpu",
+        padding="circular",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.downsampling = Downsampling(
+            img_size=img_size,
+            factor=factor,
+            filter=filter,
+            device=device,
+            padding=padding,
+        )
+
+        self.noise_color = noise_color
+        self.noise_gray = noise_gray
+        self.colorize = Decolorize()
+
+    def A(self, x):
+        return TensorList([self.downsampling(x), self.colorize(x)])
+
+    def A_adjoint(self, y):
+        return self.downsampling.A_adjoint(y[0]) + self.colorize.A_adjoint(y[1])
+
+    def forward(self, x):
+        return TensorList(
+            [self.noise_color(self.downsampling(x)), self.noise_gray(self.colorize(x))]
+        )
+
+
+# test code
+# if __name__ == "__main__":
+#     device = "cuda:0"
+#     import torch
+#     import torchvision
+#     import deepinv
+#
+#     device = "cuda:0"
+#
+#     x = torchvision.io.read_image("../../datasets/celeba/img_align_celeba/085307.jpg")
+#     x = x.unsqueeze(0).float().to(device) / 255
+#     x = torchvision.transforms.Resize((160, 180))(x)
+#
+#     class Toy(LinearPhysics):
+#         def __init__(self, **kwargs):
+#             super().__init__(**kwargs)
+#             self.A = lambda x: x * 2
+#             self.A_adjoint = lambda x: x * 2
+#
+#     sigma_noise = 0.1
+#     kernel = torch.zeros((1, 1, 15, 15), device=device)
+#     kernel[:, :, 7, :] = 1 / 15
+#     # physics = deepinv.physics.BlurFFT(img_size=x.shape[1:], filter=kernel, device=device)
+#     physics = Pansharpen(factor=8, img_size=x.shape[1:], device=device)
+#
+#     y = physics(x)
+#
+#     xhat2 = physics.A_adjoint(y)
+#     xhat1 = physics.A_dagger(y)
+#
+#     physics.compute_norm(x)
+#     physics.adjointness_test(x)
+#
+#     deepinv.utils.plot(
+#         [y[0], y[1], xhat2, xhat1, x],
+#         titles=["low res color", "high res gray", "A_adjoint", "A_dagger", "x"],
+#     )

deepinv/physics/singlepixel.py

@@ -0,0 +1,218 @@
+from deepinv.physics.forward import DecomposablePhysics
+import torch
+import numpy as np
+
+
+def hadamard_1d(u, normalize=True):
+    """
+    Multiply H_n @ u where H_n is the Hadamard matrix of dimension n x n.
+    n must be a power of 2.
+
+    Parameters:
+        u: Tensor of shape (..., n)
+        normalize: if True, divide the result by 2^{m/2} where m = log_2(n).
+    Returns:
+        product: Tensor of shape (..., n)
+    """
+    n = u.shape[-1]
+    m = int(np.log2(n))
+    assert n == 1 << m, "n must be a power of 2"
+    x = u[..., np.newaxis]
+    for d in range(m)[::-1]:
+        x = torch.cat(
+            (x[..., ::2, :] + x[..., 1::2, :], x[..., ::2, :] - x[..., 1::2, :]), dim=-1
+        )
+    return x.squeeze(-2) / 2 ** (m / 2) if normalize else x.squeeze(-2)
+
+
+def hadamard_2d(x):
+    """
+    Computes 2 dimensional Hadamard transform using 1 dimensional transform.
+    """
+    out = hadamard_1d(hadamard_1d(x).transpose(-1, -2)).transpose(-1, -2)
+    return out
+
+
+class SinglePixelCamera(DecomposablePhysics):
+    r"""
+    Single pixel imaging camera.
+
+    Linear imaging operator with binary entries.
+
+    If ``fast=False``, the operator uses a 2D subsampled hadamard transform, which keeps the first :math:`m` modes
+    according to the `sequency ordering <https://en.wikipedia.org/wiki/Walsh_matrix#Sequency_ordering>`_.
+    In this case, the images should have a size which is a power of 2.
+
+    If ``fast=False``, the operator is a random iid binary matrix with equal probability of :math:`1/\sqrt{m}` or
+    :math:`-1/\sqrt{m}`.
+
+    Both options allow for an efficient singular value decomposition (see :meth:`deepinv.physics.DecomposablePhysics`)
+    The operator is always applied independently across channels.
+
+    It is recommended to use ``fast=True`` for image sizes bigger than 32 x 32, since the forward computation with
+    ``fast=False`` has an :math:`O(mn)` complexity, whereas with ``fast=True`` it has an :math:`O(n \log n)` complexity.
+
+    An existing operator can be loaded from a saved ``.pth`` file via ``self.load_state_dict(save_path)``,
+    in a similar fashion to :meth:`torch.nn.Module`.
+
+    :param int m: number of single pixel measurements per acquisition.
+    :param tuple img_shape: shape (C, H, W) of images.
+    :param bool fast: The operator is iid binary if false, otherwise A is a 2D subsampled hadamard transform.
+    :param str device: Device to store the forward matrix.
+
+    """
+
+    def __init__(
+        self, m, img_shape, fast=True, device="cpu", dtype=torch.float32, **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.name = f"spcamera_m{m}"
+        self.img_shape = img_shape
+        self.fast = fast
+        self.device = device
+
+        if self.fast:
+            C, H, W = img_shape
+            mi = min(int(np.sqrt(m)), H)
+            mj = min(m - mi, W)
+
+            revi = get_permutation_list(H)[:mi]
+            revj = get_permutation_list(W)[:mj]
+
+            assert H == 1 << int(np.log2(H)), "image height must be a power of 2"
+            assert W == 1 << int(np.log2(W)), "image width must be a power of 2"
+
+            mask = torch.zeros(img_shape).unsqueeze(0)
+            for i in range(len(revi)):
+                for j in range(len(revj)):
+                    mask[0, :, revi[i], revj[j]] = 1
+
+            mask = mask.to(device)
+            self.mask = torch.nn.Parameter(mask, requires_grad=False)
+
+        else:
+            n = int(np.prod(img_shape[1:]))
+            A = torch.ones((m, n), device=device)
+            A[torch.randn_like(A) > 0.5] = -1.0
+            A /= np.sqrt(m)  # normalize
+            u, mask, vh = torch.linalg.svd(A, full_matrices=False)
+
+            self.mask = mask.to(device).unsqueeze(0).type(dtype)
+            self.vh = vh.to(device).type(dtype)
+            self.u = u.to(device).type(dtype)
+
+            self.u = torch.nn.Parameter(self.u, requires_grad=False)
+            self.vh = torch.nn.Parameter(self.vh, requires_grad=False)
+            self.mask = torch.nn.Parameter(self.mask, requires_grad=False)
+
+    def V_adjoint(self, x):
+        if self.fast:
+            y = hadamard_2d(x)
+        else:
+            N, C = x.shape[0], self.img_shape[0]
+            x = x.reshape(N, C, -1)
+            y = torch.einsum("ijk, mk->ijm", x, self.vh)
+        return y
+
+    def V(self, y):
+        if self.fast:
+            x = hadamard_2d(y)
+        else:
+            N = y.shape[0]
+            C, H, W = self.img_shape[0], self.img_shape[1], self.img_shape[2]
+            x = torch.einsum("ijk, km->ijm", y, self.vh)
+            x = x.reshape(N, C, H, W)
+        return x
+
+    def U_adjoint(self, x):
+        if self.fast:
+            out = x
+        else:
+            out = torch.einsum("ijk, km->ijm", x, self.u)
+        return out
+
+    def U(self, x):
+        if self.fast:
+            out = x
+        else:
+            out = torch.einsum("ijk, mk->ijm", x, self.u)
+        return out
+
+
+def gray_decode(n):
+    m = n >> 1
+    while m:
+        n ^= m
+        m >>= 1
+    return n
+
+
+def reverse(n, numbits):
+    return sum(1 << (numbits - 1 - i) for i in range(numbits) if n >> i & 1)
+
+
+def get_permutation_list(n):
+    rev = np.zeros((n), dtype=int)
+    for l in range(n):
+        rev[l] = reverse(l, np.log2(n).astype(int))
+
+    rev2 = np.zeros_like(rev)
+    for l in range(n):
+        rev2[l] = rev[gray_decode(l)]
+
+    return rev2
+
+
+# test code
+# if __name__ == "__main__":
+#     import matplotlib.pyplot as plt
+#     import deepinv as dinv
+#     import torchvision
+#
+#     device = "cuda:0"
+#     x = torchvision.io.read_image("../../datasets/celeba/img_align_celeba/085307.jpg")
+#     x = x.unsqueeze(0).float().to(device) / 255
+#     x = torchvision.transforms.Resize((16, 8))(x)
+#
+#     m = 20
+#     physics = SinglePixelCamera(m, (3, 16, 8), fast=False, device=device)
+#
+#     y = physics(x)
+#
+#     xhat = physics.A_adjoint(y)
+#
+#     dinv.utils.plot([x, xhat])
+#
+#     print(physics.adjointness_test(x))
+#     print(physics.compute_norm(x))
+#     # mi = min(int(np.sqrt(m)), x.shape[-2])
+#     # mj = min(m - mi, x.shape[-2])
+#     #
+#     # revi = get_permutation_list(x.shape[-2])[:mi]
+#     # revj = get_permutation_list(x.shape[-1])[:mj]
+#     #
+#     # mask = torch.zeros_like(x)
+#     # for i in range(len(revi)):
+#     #     for j in range(len(revj)):
+#     #         mask[0, :, revi[i], revj[j]] = 1
+#     #
+#     # # generate low pass hadamard mask
+#     # f = hadamard_2d(x)
+#     # f = f * mask
+#     # out = hadamard_2d(f)
+#     #
+#     # dinv.utils.plot_batch([x, out, f])
+#     #
+#     # rev = get_permutation_list(8)
+#     # imgs = []
+#     # for i in range(8):
+#     #     y = torch.zeros((8, 1, 8, 8), device=dinv.device)
+#     #     for j in range(8):
+#     #         x = torch.zeros((8, 8), device=dinv.device)
+#     #         x[rev[i], rev[j]] = 1
+#     #         x = hadamard_2d(x)
+#     #         y[j, 0, :, :] = x
+#     #
+#     #     imgs.append(y)
+#     #
+#     # dinv.utils.plot_batch(imgs)

deepinv/physics/tomography.py

@@ -0,0 +1,321 @@
+import torch
+import numpy as np
+from torch import nn
+import torch.nn.functional as F
+from deepinv.physics.forward import LinearPhysics
+
+if torch.__version__ > "1.2.0":
+    affine_grid = lambda theta, size: F.affine_grid(theta, size, align_corners=True)
+    grid_sample = lambda input, grid, mode="bilinear": F.grid_sample(
+        input, grid, align_corners=True, mode=mode
+    )
+else:
+    affine_grid = F.affine_grid
+    grid_sample = F.grid_sample
+
+# constants
+PI = 4 * torch.ones(1).atan()
+SQRT2 = (2 * torch.ones(1)).sqrt()
+
+
+def fftfreq(n):
+    val = 1.0 / n
+    results = torch.zeros(n)
+    N = (n - 1) // 2 + 1
+    p1 = torch.arange(0, N)
+    results[:N] = p1
+    p2 = torch.arange(-(n // 2), 0)
+    results[N:] = p2
+    return results * val
+
+
+def deg2rad(x):
+    return x * 4 * torch.ones(1, device=x.device, dtype=x.dtype).atan() / 180
+
+
+class AbstractFilter(nn.Module):
+    def __init__(self, device="cpu", dtype=torch.float):
+        super().__init__()
+        self.device = device
+        self.dtype = dtype
+
+    def forward(self, x):
+        input_size = x.shape[2]
+        projection_size_padded = max(
+            64, int(2 ** (2 * torch.tensor(input_size)).float().log2().ceil())
+        )
+        pad_width = projection_size_padded - input_size
+        padded_tensor = F.pad(x, (0, 0, 0, pad_width))
+        f = self._get_fourier_filter(padded_tensor.shape[2]).to(x.device)
+        fourier_filter = self.create_filter(f)
+        fourier_filter = fourier_filter.unsqueeze(-2)
+
+        projection = (
+            torch.view_as_real(torch.fft.fft(padded_tensor.transpose(2, 3))).transpose(
+                2, 3
+            )
+            * fourier_filter
+        )
+        result = torch.view_as_real(
+            torch.fft.ifft(torch.view_as_complex(projection).transpose(2, 3))
+        )[..., 0]
+        result = result.transpose(2, 3)[:, :, :input_size, :]
+
+        return result
+
+    def _get_fourier_filter(self, size):
+        n = torch.cat(
+            [torch.arange(1, size / 2 + 1, 2), torch.arange(size / 2 - 1, 0, -2)]
+        )
+
+        f = torch.zeros(size, dtype=self.dtype, device=self.device)
+        f[0] = 0.25
+        f[1::2] = -1 / (PI * n) ** 2
+
+        fourier_filter = torch.view_as_real(torch.fft.fft(f, dim=-1))
+        fourier_filter[:, 1] = fourier_filter[:, 0]
+
+        return 2 * fourier_filter
+
+    def create_filter(self, f):
+        raise NotImplementedError
+
+
+class RampFilter(AbstractFilter):
+    def __init__(self, **kwargs):
+        super(RampFilter, self).__init__(**kwargs)
+
+    def create_filter(self, f):
+        return f
+
+
+class Radon(nn.Module):
+    def __init__(
+        self,
+        in_size=None,
+        theta=None,
+        circle=False,
+        dtype=torch.float,
+        device=torch.device("cpu"),
+    ):
+        super().__init__()
+        self.circle = circle
+        self.theta = theta
+        if theta is None:
+            self.theta = torch.arange(180)
+        self.dtype = dtype
+        self.all_grids = None
+        if in_size is not None:
+            self.all_grids = self._create_grids(self.theta, in_size, circle).to(device)
+
+    def forward(self, x):
+        N, C, W, H = x.shape
+        assert W == H, "Input image must be square"
+
+        if (
+            self.all_grids is None
+        ):  # if in_size was not given, we have to create the grid online.
+            self.all_grids = self._create_grids(
+                self.theta, W, self.circle, device=x.device
+            )
+
+        if not self.circle:
+            diagonal = SQRT2 * W
+            pad = int((diagonal - W).ceil())
+            new_center = (W + pad) // 2
+            old_center = W // 2
+            pad_before = new_center - old_center
+            pad_width = (pad_before, pad - pad_before)
+            x = F.pad(x, (pad_width[0], pad_width[1], pad_width[0], pad_width[1]))
+
+        N, C, W, _ = x.shape
+        out = torch.zeros(N, C, W, len(self.theta), device=x.device, dtype=self.dtype)
+
+        for i in range(len(self.theta)):
+            rotated = grid_sample(x, self.all_grids[i].repeat(N, 1, 1, 1).to(x.device))
+            out[..., i] = rotated.sum(2)
+        return out
+
+    def _create_grids(self, angles, grid_size, circle, device="cpu"):
+        if not circle:
+            grid_size = int((SQRT2 * grid_size).ceil())
+        all_grids = []
+        for theta in angles:
+            theta = deg2rad(theta)
+            R = torch.tensor(
+                [[[theta.cos(), theta.sin(), 0], [-theta.sin(), theta.cos(), 0]]],
+                dtype=self.dtype,
+                device=device,
+            )
+            all_grids.append(affine_grid(R, torch.Size([1, 1, grid_size, grid_size])))
+        return torch.stack(all_grids)
+
+
+class IRadon(nn.Module):
+    def __init__(
+        self,
+        in_size=None,
+        theta=None,
+        circle=False,
+        use_filter=True,
+        out_size=None,
+        dtype=torch.float,
+        device=torch.device("cpu"),
+    ):
+        super().__init__()
+        self.circle = circle
+        self.device = device
+        self.theta = theta if theta is not None else torch.arange(180).to(self.device)
+        self.out_size = out_size
+        self.in_size = in_size
+        self.dtype = dtype
+        self.ygrid, self.xgrid, self.all_grids = None, None, None
+        if in_size is not None:
+            self.ygrid, self.xgrid = self._create_yxgrid(in_size, circle)
+            self.all_grids = self._create_grids(self.theta, in_size, circle).to(
+                self.device
+            )
+        self.filter = (
+            RampFilter(dtype=self.dtype, device=self.device)
+            if use_filter
+            else lambda x: x
+        )
+
+    def forward(self, x, filtering=True):
+        it_size = x.shape[2]
+        ch_size = x.shape[1]
+
+        if self.in_size is None:
+            self.in_size = (
+                int((it_size / SQRT2).floor()) if not self.circle else it_size
+            )
+        # if None in [self.ygrid, self.xgrid, self.all_grids]:
+        if self.ygrid is None or self.xgrid is None or self.all_grids is None:
+            self.ygrid, self.xgrid = self._create_yxgrid(self.in_size, self.circle)
+            self.all_grids = self._create_grids(self.theta, self.in_size, self.circle)
+
+        x = self.filter(x) if filtering else x
+
+        reco = torch.zeros(
+            x.shape[0], ch_size, it_size, it_size, device=self.device, dtype=self.dtype
+        )
+        for i_theta in range(len(self.theta)):
+            reco += grid_sample(
+                x, self.all_grids[i_theta].repeat(reco.shape[0], 1, 1, 1)
+            )
+
+        if not self.circle:
+            W = self.in_size
+            diagonal = it_size
+            pad = int(torch.tensor(diagonal - W, dtype=torch.float).ceil())
+            new_center = (W + pad) // 2
+            old_center = W // 2
+            pad_before = new_center - old_center
+            pad_width = (pad_before, pad - pad_before)
+            reco = F.pad(
+                reco, (-pad_width[0], -pad_width[1], -pad_width[0], -pad_width[1])
+            )
+
+        if self.circle:
+            reconstruction_circle = (self.xgrid**2 + self.ygrid**2) <= 1
+            reconstruction_circle = reconstruction_circle.repeat(
+                x.shape[0], ch_size, 1, 1
+            )
+            reco[~reconstruction_circle] = 0.0
+
+        reco = reco * PI.item() / (2 * len(self.theta))
+
+        if self.out_size is not None:
+            pad = (self.out_size - self.in_size) // 2
+            reco = F.pad(reco, (pad, pad, pad, pad))
+
+        return reco
+
+    def _create_yxgrid(self, in_size, circle):
+        if not circle:
+            in_size = int((SQRT2 * in_size).ceil())
+        unitrange = torch.linspace(-1, 1, in_size, dtype=self.dtype, device=self.device)
+        return torch.meshgrid(unitrange, unitrange, indexing="ij")
+
+    def _XYtoT(self, theta):
+        T = self.xgrid * (deg2rad(theta)).cos() - self.ygrid * (deg2rad(theta)).sin()
+        return T
+
+    def _create_grids(self, angles, grid_size, circle):
+        if not circle:
+            grid_size = int((SQRT2 * grid_size).ceil())
+        all_grids = []
+        for i_theta in range(len(angles)):
+            X = (
+                torch.ones(grid_size, dtype=self.dtype, device=self.device)
+                .view(-1, 1)
+                .repeat(1, grid_size)
+                * i_theta
+                * 2.0
+                / (len(angles) - 1)
+                - 1.0
+            )
+            Y = self._XYtoT(angles[i_theta])
+            all_grids.append(
+                torch.cat((X.unsqueeze(-1), Y.unsqueeze(-1)), dim=-1).unsqueeze(0)
+            )
+        return torch.stack(all_grids)
+
+
+class Tomography(LinearPhysics):
+    r"""
+    (Computed) Tomography operator.
+
+    The Radon transform is the integral transform which takes a square image :math:`x` defined on the plane to a function
+    :math:`y=Rx` defined on the (two-dimensional) space of lines in the plane, whose value at a particular line is equal
+    to the line integral of the function over that line.
+
+    .. note::
+
+        The pseudo-inverse is computed using the filtered back-projection algorithm with a Ramp filter.
+        This is not the exact linear pseudo-inverse of the Radon transform, but it is a good approximation which is
+        robust to noise.
+
+    .. warning::
+
+        The adjoint operator has small numerical errors due to interpolation.
+
+    :param int img_width: width/height of the square image input.
+    :param int, torch.tensor angles: If the type is ``int``, the angles are sampled uniformly between 0 and 360 degrees.
+        If the type is ``torch.tensor``, the angles are the ones provided (e.g., ``torch.linspace(0, 180, steps=10)``).
+    :param bool circle: If ``True`` both forward and backward projection will be restricted to pixels inside a circle
+        inscribed in the square image.
+    :param str device: gpu or cpu.
+    """
+
+    def __init__(
+        self,
+        img_width,
+        angles,
+        circle=False,
+        device=torch.device("cpu"),
+        dtype=torch.float,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        if isinstance(angles, int) or isinstance(angles, float):
+            theta = torch.linspace(0, 180, steps=angles + 1, device=device)[:-1]
+        else:
+            theta = angles.to(device)
+
+        self.radon = Radon(
+            img_width, theta, circle=circle, device=device, dtype=dtype
+        ).to(device)
+        self.iradon = IRadon(
+            img_width, theta, circle=circle, device=device, dtype=dtype
+        ).to(device)
+
+    def A(self, x):
+        return self.radon(x)
+
+    def A_dagger(self, y):
+        return self.iradon(y)
+
+    def A_adjoint(self, y):
+        return self.iradon(y, filtering=False)

deepinv/sampling/__init__.py

@@ -0,0 +1,2 @@
+from .langevin import MonteCarlo, ULA, SKRock
+from .diffusion import DDRM, DiffusionSampler, DiffPIR, DPS