deepinv 0.1.0.dev0__py3-none-any.whl

deepinv/physics/lidar.py

@@ -0,0 +1,123 @@
+import torch
+from deepinv.physics.forward import Physics
+from deepinv.physics.noise import PoissonNoise
+
+
+class SinglePhotonLidar(Physics):
+    r"""
+    Single photon lidar operator for depth ranging.
+
+    See https://ieeexplore.ieee.org/abstract/document/9127841 for a review of this imaging method.
+
+    The forward operator is given by
+
+    .. math::
+        y_{i,j,t} = \mathcal{P}(h(t-d_{i,j}) r_{i,j} + b_{i,j})
+
+    where :math:`\mathcal{P}` is the Poisson noise model, :math:`h(t)` is a Gaussian impulse response function at
+    time :math:`t`, :math:`d_{i,j}` is the depth of the scene at pixel :math:`(i,j)`,
+    :math:`r_{i,j}` is the intensity of the scene at pixel :math:`(i,j)` and :math:`b_{i,j}` is the background noise
+    at pixel :math:`(i,j)`.
+
+    For a pixel grid of size (H,W) and batch size B, the signals have size (B, 3, H, W), where the first channel
+    contains the depth of the scene :math:`d`, the second channel contains the intensity of the scene :math:`r` and
+    the third channel contains the per pixel background noise levels :math:`b`.
+
+    :param float sigma: Standard deviation of the Gaussian impulse response function.
+    :param int bins: Number of histogram bins per pixel.
+    :param str device: Device to use (gpu or cpu).
+    """
+
+    def __init__(self, sigma=1.0, bins=50, device="cpu"):
+        super().__init__()
+
+        self.T = bins
+        self.grid = torch.meshgrid(torch.arange(bins), indexing="ij")[0].to(device)
+        self.sigma = torch.nn.Parameter(
+            torch.tensor(sigma, device=device), requires_grad=False
+        )
+        self.noise_model = PoissonNoise()
+
+        h = ((self.grid - 3 * sigma) / self.sigma).pow(2)
+        h = torch.exp(-h / 2.0)
+        h = h[: int(6 * sigma)]
+        h = h / h.sum()
+        self.irf = h.unsqueeze(0).unsqueeze(0)  # set impulse response function
+        self.grid = self.grid.unsqueeze(0).unsqueeze(2).unsqueeze(3)
+
+    def A(self, x):
+        r"""
+        Applies the forward operator.
+
+        Input is of size (B, 3, H, W) and output is of size (B, bins, H, W)
+
+        :param torch.tensor x: tensor containing the depth, intensity and background noise levels.
+        """
+
+        h = ((self.grid - x[:, 0, :, :]) / self.sigma).pow(2)
+        h = torch.exp(-h / 2.0)
+        h = h / h.sum(dim=1, keepdim=True)
+        y = x[:, 1, :, :] * h + x[:, 2, :, :]
+        return y
+
+    def A_dagger(self, y):
+        r"""
+        Applies Matched filtering to find the peaks.
+
+        Input is of size (B, bins, H, W), output of size (B, 3, H, W).
+
+        :param torch.tensor y: measurements
+        """
+        B, T, H, W = y.shape
+
+        # reshape to (B*H*W, 1, T)
+        y = y.permute(0, 2, 3, 1).reshape(B * H * W, 1, T)
+
+        # Apply irf using convolution
+        x = torch.nn.functional.conv1d(y, self.irf, padding="same")
+
+        # Find peak value in each channel
+        _, x = torch.max(x, dim=-1, keepdim=True)
+        x = x.type(torch.float32)
+        offset = self.irf.shape[-1] // 2
+        x -= 3 * self.sigma - offset - 0.5
+
+        mask = torch.ones_like(y)
+        grid = self.grid.squeeze(-1).squeeze(-1)  # (1, T)
+
+        mask *= (x - 4 * self.sigma) < grid
+        mask *= (x + 4 * self.sigma) > grid
+
+        b = (y * (1 - mask)).sum(dim=-1, keepdim=True)
+        r = y.sum(dim=-1, keepdim=True) - b
+        b /= T
+
+        x = torch.stack([x, r, b], dim=-1)
+
+        # reshape to (B, 3, H, W)
+        x = x.reshape(B, H, W, 3).permute(0, 3, 1, 2)
+
+        return x
+
+
+# if __name__ == "__main__":
+#     import matplotlib.pyplot as plt
+#     import deepinv as dinv
+#
+#     bins = 40
+#     device = "cuda:0"
+#     physics = SinglePhotonLidar(bins=bins, device=device)
+#
+#     x = torch.ones((1, 3, 2, 4), device=device)
+#     x[:, 0, :, :] *= bins / 2
+#     x[:, 1, :, :] *= 300
+#     x[:, 2, :, :] *= 1
+#
+#     y = physics(x)
+#     xhat = physics.A_dagger(y)
+#
+#     y0 = y[0, :, 0, 0].detach().cpu().numpy()
+#     plt.plot(y0)
+#     plt.show()
+#
+#     print(f"MSE {dinv.utils.cal_mse(x, xhat)}")

deepinv/physics/mri.py

@@ -0,0 +1,329 @@
+import numpy as np
+import torch
+import torch.fft
+from typing import List, Optional
+from deepinv.physics.forward import DecomposablePhysics
+
+
+class MRI(DecomposablePhysics):
+    r"""
+    Single-coil accelerated magnetic resonance imaging.
+
+    The linear operator operates in 2D slices and is defined as
+
+    .. math::
+
+        y = SFx
+
+    where :math:`S` applies a mask (subsampling operator), and :math:`F` is the 2D discrete Fourier Transform.
+    This operator has a simple singular value decomposition, so it inherits the structure of
+    :meth:`deepinv.physics.DecomposablePhysics` and thus have a fast pseudo-inverse and prox operators.
+
+    The complex images :math:`x` and measurements :math:`y` should be of size (B, 2, H, W) where the first channel corresponds to the real part
+    and the second channel corresponds to the imaginary part.
+
+    :param torch.tensor mask: the mask values should be binary.
+        The mask size should be of the form (H,W) where H is the image height and W is the image width.
+    :param torch.device device: cpu or gpu.
+    """
+
+    def __init__(
+        self,
+        mask=None,
+        image_size=(320, 320),
+        acceleration_factor=4,
+        device="cpu",
+        seed=None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.device = device
+        self.image_size = image_size
+
+        if mask is not None:
+            mask = mask.to(device).unsqueeze(0).unsqueeze(0)
+        else:
+            mask = (
+                self.sample_mask(
+                    image_size=image_size,
+                    acceleration_factor=acceleration_factor,
+                    seed=seed,
+                )
+                .unsqueeze(0)
+                .unsqueeze(0)
+            )
+
+        self.mask = torch.nn.Parameter(
+            torch.cat([mask, mask], dim=1), requires_grad=False
+        )
+
+    def reset(self, **kwargs):
+        r"""
+        Resets the physics, i.e. re-samples a new mask and new noise realization (if any).
+        """
+        super().reset(**kwargs)
+        mask = (
+            self.sample_mask(image_size=self.image_size, **kwargs)
+            .unsqueeze(0)
+            .unsqueeze(0)
+        )
+
+        self.mask = torch.nn.Parameter(
+            torch.cat([mask, mask], dim=1), requires_grad=False
+        )
+
+    def V_adjoint(self, x):  # (B, 2, H, W) -> (B, H, W, 2)
+        y = fft2c_new(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
+        return y
+
+    def U(self, x):
+        return x[:, self.mask.squeeze(0) > 0]
+
+    def U_adjoint(self, x):
+        _, c, h, w = self.mask.shape
+        out = torch.zeros((x.shape[0], c, h, w), device=x.device)
+        out[:, self.mask.squeeze(0) > 0] = x
+        return out
+
+    def V(self, x):  # (B, 2, H, W) -> (B, H, W, 2)
+        x = x.permute(0, 2, 3, 1)
+        return ifft2c_new(x).permute(0, 3, 1, 2)
+
+    def sample_mask(self, image_size=(320, 320), acceleration_factor=4, seed=None):
+        r"""
+        Create a mask of vertical lines.
+
+        :param tuple image_size: image size.
+        :param int acceleration_factor: acceleration factor.
+        :param int seed: random seed.
+        :return: mask of size (H, W) with values in {0, 1}.
+        """
+        if seed is not None:
+            np.random.seed(seed)
+        if acceleration_factor == 4:
+            central_lines_percent = 0.08
+            num_lines_center = int(central_lines_percent * image_size[-1])
+            side_lines_percent = 0.25 - central_lines_percent
+            num_lines_side = int(side_lines_percent * image_size[-1])
+        if acceleration_factor == 8:
+            central_lines_percent = 0.04
+            num_lines_center = int(central_lines_percent * image_size[-1])
+            side_lines_percent = 0.125 - central_lines_percent
+            num_lines_side = int(side_lines_percent * image_size[-1])
+        mask = torch.zeros(image_size)
+        center_line_indices = torch.linspace(
+            image_size[0] // 2 - num_lines_center // 2,
+            image_size[0] // 2 + num_lines_center // 2 + 1,
+            steps=50,
+            dtype=torch.long,
+        )
+        mask[:, center_line_indices] = 1
+        random_line_indices = np.random.choice(
+            image_size[0], size=(num_lines_side // 2,), replace=False
+        )
+        mask[:, random_line_indices] = 1
+        return mask.float().to(self.device)
+
+
+#
+# reference: https://github.com/facebookresearch/fastMRI/blob/main/fastmri/fftc.py
+def fft2c_new(data: torch.Tensor, norm: str = "ortho") -> torch.Tensor:
+    r"""
+    Apply centered 2 dimensional Fast Fourier Transform.
+    :param torch.tensor data: Complex valued input data containing at least 3 dimensions:
+        dimensions -2 & -1 are spatial dimensions and dimension -3 has size
+        2. All other dimensions are assumed to be batch dimensions.
+    :param bool norm: Normalization mode. See ``torch.fft.fft``.
+    :return: (torch.tensor) the FFT of the input.
+    """
+    if not data.shape[-1] == 2:
+        raise ValueError("Tensor does not have separate complex dim.")
+
+    data = ifftshift(data, dim=[-3, -2])
+    data = torch.view_as_real(
+        torch.fft.fftn(  # type: ignore
+            torch.view_as_complex(data), dim=(-2, -1), norm=norm
+        )
+    )
+    data = fftshift(data, dim=[-3, -2])
+
+    return data
+
+
+def ifft2c_new(data: torch.Tensor, norm: str = "ortho") -> torch.Tensor:
+    """
+    Apply centered 2-dimensional Inverse Fast Fourier Transform.
+    Args:
+        data: Complex valued input data containing at least 3 dimensions:
+            dimensions -2 & -1 are spatial dimensions and dimension -3 has size
+            2. All other dimensions are assumed to be batch dimensions.
+        norm: Normalization mode. See ``torch.fft.ifft``.
+    Returns:
+        The IFFT of the input.
+    """
+    if not data.shape[-1] == 2:
+        raise ValueError("Tensor does not have separate complex dim.")
+
+    data = ifftshift(data, dim=[-3, -2])
+    data = torch.view_as_real(
+        torch.fft.ifftn(  # type: ignore
+            torch.view_as_complex(data), dim=(-2, -1), norm=norm
+        )
+    )
+    data = fftshift(data, dim=[-3, -2])
+
+    return data
+
+
+# Helper functions
+def roll_one_dim(x: torch.Tensor, shift: int, dim: int) -> torch.Tensor:
+    """
+    Similar to roll but for only one dim.
+    Args:
+        x: A PyTorch tensor.
+        shift: Amount to roll.
+        dim: Which dimension to roll.
+    Returns:
+        Rolled version of x.
+    """
+    shift = shift % x.size(dim)
+    if shift == 0:
+        return x
+
+    left = x.narrow(dim, 0, x.size(dim) - shift)
+    right = x.narrow(dim, x.size(dim) - shift, shift)
+
+    return torch.cat((right, left), dim=dim)
+
+
+def roll(x: torch.Tensor, shift: List[int], dim: List[int]) -> torch.Tensor:
+    """
+    Similar to np.roll but applies to PyTorch Tensors.
+    Args:
+        x: A PyTorch tensor.
+        shift: Amount to roll.
+        dim: Which dimension to roll.
+    Returns:
+        Rolled version of x.
+    """
+    if len(shift) != len(dim):
+        raise ValueError("len(shift) must match len(dim)")
+
+    for s, d in zip(shift, dim):
+        x = roll_one_dim(x, s, d)
+
+    return x
+
+
+def fftshift(x: torch.Tensor, dim: Optional[List[int]] = None) -> torch.Tensor:
+    """
+    Similar to np.fft.fftshift but applies to PyTorch Tensors
+    Args:
+        x: A PyTorch tensor.
+        dim: Which dimension to fftshift.
+    Returns:
+        fftshifted version of x.
+    """
+    if dim is None:
+        # this weird code is necessary for toch.jit.script typing
+        dim = [0] * (x.dim())
+        for i in range(1, x.dim()):
+            dim[i] = i
+
+    # also necessary for torch.jit.script
+    shift = [0] * len(dim)
+    for i, dim_num in enumerate(dim):
+        shift[i] = x.shape[dim_num] // 2
+
+    return roll(x, shift, dim)
+
+
+def ifftshift(x: torch.Tensor, dim: Optional[List[int]] = None) -> torch.Tensor:
+    """
+    Similar to np.fft.ifftshift but applies to PyTorch Tensors
+    Args:
+        x: A PyTorch tensor.
+        dim: Which dimension to ifftshift.
+    Returns:
+        ifftshifted version of x.
+    """
+    if dim is None:
+        # this weird code is necessary for toch.jit.script typing
+        dim = [0] * (x.dim())
+        for i in range(1, x.dim()):
+            dim[i] = i
+
+    # also necessary for torch.jit.script
+    shift = [0] * len(dim)
+    for i, dim_num in enumerate(dim):
+        shift[i] = (x.shape[dim_num] + 1) // 2
+
+    return roll(x, shift, dim)
+
+
+# if __name__ == "__main__":
+#     # deepinv test
+#     from deepinv.tests.test_physics import (
+#         test_operators_norm,
+#         test_operators_adjointness,
+#         test_pseudo_inverse,
+#         device,
+#     )
+#     import deepinv as dinv
+#     from fastmri.data import subsample
+#
+#     imsize = (25, 32)
+#     # Create a mask function
+#     mask_func = subsample.RandomMaskFunc(center_fractions=[0.08], accelerations=[4])
+#     m = mask_func.sample_mask((imsize[1], imsize[0]), offset=None)
+#
+#     # mask = torch.ones((imsize[0], 1)) * (m[0] + m[1]).permute(1, 0)
+#     mask = torch.ones(imsize)
+#     mask[mask > 1] = 1
+#
+#     sigma = 0.1
+#     # physics = MRI(mask=mask, device=dinv.device)
+#     physics = dinv.physics.Denoising()
+#     physics.noise_model = dinv.physics.GaussianNoise(sigma)
+#
+#     # choose a reconstruction architecture
+#     backbone = dinv.models.MedianFilter()
+#
+#     class denoiser(torch.nn.Module):
+#         def __init__(self):
+#             super().__init__()
+#
+#         def forward(self, x, sigma=None):
+#             return x
+#
+#     f = dinv.models.ArtifactRemoval(backbone)
+#
+#     batch_size = 1
+#
+#     for tau in np.logspace(-5, 3, 1):
+#         x = torch.ones((batch_size, 2) + imsize, device=dinv.device)
+#         y = physics(x)
+#
+#         # choose training losses
+#         loss = dinv.loss.SureGaussianLoss(sigma, tau=tau)
+#         x_net = f(y, physics)
+#         mse = dinv.metric.mse()(physics.A(x), physics.A(x_net))
+#         sure = loss(y, x_net, physics, f)
+#
+#         print(f"tau:{tau:.2e}  mse: {mse:.2e}, sure: {sure:.2e}")
+#         rel_error = (sure - mse).abs() / mse
+#         print(f"rel_error: {rel_error:.2e}")
+#
+#     d = physics.A_adjoint(y)
+#     dinv.utils.plot([d.sum(1).unsqueeze(1), x.sum(1).unsqueeze(1)])
+#
+#     print("adjoint test....")
+#     test_operators_adjointness(
+#         "MRI", (2, 320, 320), dinv.device
+#     )  # pass, tensor(0., device='cuda:0')
+#     print("norm test....")
+#     test_operators_norm("MRI", (2, 320, 320), dinv.device)  # pass
+#     print("pinv test....")
+#     test_pseudo_inverse("MRI", (2, 320, 320), dinv.device)  # pass
+#
+#     print("pass all...")

deepinv/physics/noise.py

@@ -0,0 +1,180 @@
+import torch
+
+
+class GaussianNoise(torch.nn.Module):
+    r"""
+
+    Gaussian noise :math:`y=z+\epsilon` where :math:`\epsilon\sim \mathcal{N}(0,I\sigma^2)`.
+
+    It can be added to a physics operator in its construction or by setting the ``noise_model``
+    attribute of the physics operator.
+
+
+    ::
+
+        >>> from deepinv.physics import Denoising, GaussianNoise
+        >>> import torch
+        >>> physics = Denoising()
+        >>> physics.noise_model = GaussianNoise()
+        >>> x = torch.rand(1, 1, 2, 2)
+        >>> y = physics(x)
+
+    :param float sigma: Standard deviation of the noise.
+
+    """
+
+    def __init__(self, sigma=0.1):
+        super().__init__()
+        self.sigma = torch.nn.Parameter(torch.tensor(sigma), requires_grad=False)
+
+    def forward(self, x):
+        r"""
+        Adds the noise to measurements x
+
+        :param torch.Tensor x: measurements
+        :returns: noisy measurements
+        """
+        return x + torch.randn_like(x) * self.sigma
+
+
+class UniformGaussianNoise(torch.nn.Module):
+    r"""
+    Gaussian noise :math:`y=z+\epsilon` where
+    :math:`\epsilon\sim \mathcal{N}(0,I\sigma^2)` and
+    :math:`\sigma \sim\mathcal{U}(\sigma_{\text{min}}, \sigma_{\text{max}})`
+
+    It can be added to a physics operator in its construction or by setting:
+
+    ::
+
+        >>> physics.noise_model = UniformGaussianNoise()
+
+
+    :param float sigma_min: minimum standard deviation of the noise.
+    :param float sigma_max: maximum standard deviation of the noise.
+    :param float, torch.Tensor sigma: standard deviation of the noise.
+        If ``None``, the noise is sampled uniformly at random
+        in :math:`[\sigma_{\text{min}}, \sigma_{\text{max}}]`) during the forward pass. Default: ``None``.
+
+    """
+
+    def __init__(self, sigma_min=0.0, sigma_max=0.5, sigma=None):
+        super().__init__()
+        self.sigma_min = sigma_min
+        self.sigma_max = sigma_max
+        self.sigma = sigma
+
+    def forward(self, x):
+        r"""
+        Adds the noise to measurements x.
+
+        :param torch.Tensor x: measurements
+        :returns: noisy measurements.
+        """
+        if self.sigma is None:
+            sigma = (
+                torch.rand((x.shape[0], 1) + (1,) * (x.dim() - 2))
+                * (self.sigma_max - self.sigma_min)
+                + self.sigma_min
+            )
+            self.sigma = sigma.to(x.device)
+        noise = torch.randn_like(x) * self.sigma
+        return x + noise
+
+
+class PoissonNoise(torch.nn.Module):
+    r"""
+
+    Poisson noise :math:`y = \mathcal{P}(\frac{x}{\gamma})`
+    with gain :math:`\gamma>0`.
+
+    If ``normalize=True``, the output is divided by the gain, i.e., :math:`\tilde{y} = \gamma y`.
+
+    It can be added to a physics operator in its construction or by setting:
+    ::
+
+        >>> physics.noise_model = PoissonNoise()
+
+    :param float gain: gain of the noise.
+    :param bool normalize: normalize the output.
+
+    """
+
+    def __init__(self, gain=1.0, normalize=True):
+        super().__init__()
+        self.normalize = torch.nn.Parameter(
+            torch.tensor(normalize), requires_grad=False
+        )
+        self.gain = torch.nn.Parameter(torch.tensor(gain), requires_grad=False)
+
+    def forward(self, x):
+        r"""
+        Adds the noise to measurements x
+
+        :param torch.Tensor x: measurements
+        :returns: noisy measurements
+        """
+        y = torch.poisson(x / self.gain)
+        if self.normalize:
+            y *= self.gain
+        return y
+
+
+class PoissonGaussianNoise(torch.nn.Module):
+    r"""
+    Poisson-Gaussian noise :math:`y = \gamma z + \epsilon` where :math:`z\sim\mathcal{P}(\frac{x}{\gamma})`
+    and :math:`\epsilon\sim\mathcal{N}(0, I \sigma^2)`.
+
+    It can be added to a physics operator by setting
+
+    ::
+
+        >>> physics.noise_model = PoissonGaussianNoise()
+
+    :param float gain: gain of the noise.
+    :param float sigma: Standard deviation of the noise.
+
+    """
+
+    def __init__(self, gain=1.0, sigma=0.1):
+        super().__init__()
+        self.gain = torch.nn.Parameter(torch.tensor(gain), requires_grad=False)
+        self.sigma = torch.nn.Parameter(torch.tensor(sigma), requires_grad=False)
+
+    def forward(self, x):
+        r"""
+        Adds the noise to measurements x
+
+        :param torch.Tensor x: measurements
+        :returns: noisy measurements
+        """
+        y = torch.poisson(x / self.gain) * self.gain
+
+        y += torch.randn_like(x) * self.sigma
+        return y
+
+
+class UniformNoise(torch.nn.Module):
+    r"""
+    Uniform noise :math:`y = x + \epsilon` where :math:`\epsilon\sim\mathcal{U}(-a,a)`.
+
+    It can be added to a physics operator by setting
+    ::
+
+        >>> physics.noise_model = UniformNoise()
+
+    :param float a: amplitude of the noise.
+    """
+
+    def __init__(self, a=0.1):
+        super().__init__()
+        self.a = torch.nn.Parameter(torch.tensor(a), requires_grad=False)
+
+    def forward(self, x):
+        r"""
+        Adds the noise to measurements x
+
+        :param torch.Tensor x: measurements
+        :returns: noisy measurements
+        """
+        return x + (torch.rand_like(x) - 0.5) * 2 * self.a

deepinv/physics/range.py

@@ -0,0 +1,53 @@
+import torch
+from deepinv.physics.forward import DecomposablePhysics
+
+
+class Decolorize(DecomposablePhysics):
+    r"""
+    Converts RGB images to grayscale.
+
+    Follows the `rec601 <https://en.wikipedia.org/wiki/Rec._601>`_ convention.
+
+    Signals must be tensors with 3 colour (RGB) channels, i.e. [*,3,*,*]
+    The measurements are grayscale images.
+
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.mask = 1.0
+
+    def V_adjoint(self, x):
+        y = x[:, 0, :, :] * 0.2989 + x[:, 1, :, :] * 0.5870 + x[:, 2, :, :] * 0.1140
+        return y.unsqueeze(1)
+
+    def V(self, y):
+        return torch.cat([y * 0.2989, y * 0.5870, y * 0.1140], dim=1)
+
+
+# # test code
+# if __name__ == "__main__":
+#     device = "cuda:0"
+#
+#     import deepinv as dinv
+#     import matplotlib.pyplot as plt
+#     import torchvision
+#
+#     dinv.device = "cpu"
+#
+#     x = torchvision.io.read_image("../../datasets/celeba/img_align_celeba/085307.jpg")
+#     x = x.unsqueeze(0).float().to(dinv.device) / 255
+#     x = torchvision.transforms.Resize((128, 128))(x)
+#
+#     physics = Decolorize()
+#
+#     y = physics(x)
+#
+#     print(physics.adjointness_test(x))
+#     print(physics.compute_norm(x))
+#     xhat = physics.A_adjoint(y)
+#
+#     plot_results = False  # set to True to plot results
+#
+#     if plot_results:
+#         dinv.utils.plot([x, xhat, y])