deepinv 0.1.0.dev0__py3-none-any.whl

deepinv/__about__.py

@@ -0,0 +1,17 @@
+__all__ = [
+    "__title__",
+    "__summary__",
+    "__url__",
+    "__version__",
+    "__author__",
+    "__email__",
+    "__license__",
+]
+
+__title__ = "deepinv"
+__summary__ = "Deep Learning for Inverse Problems Library for PyTorch"
+__url__ = "https://github.com/edongdongchen/deepinv"
+__version__ = "0.0.1"
+__author__ = "Dongdong Chen, Julian Tachella"
+__email__ = "echendongdong@gmail.com"
+__license__ = "BSD 3-Clause Clear"

deepinv/__init__.py

@@ -0,0 +1,71 @@
+from .__about__ import *
+import torch
+
+__all__ = [
+    "__title__",
+    "__summary__",
+    "__url__",
+    "__version__",
+    "__author__",
+    "__email__",
+    "__license__",
+]
+
+from deepinv import models
+
+__all__ += ["models"]
+
+from deepinv import optim
+
+__all__ += ["optim"]
+
+from deepinv import loss
+
+__all__ += ["loss"]
+
+from deepinv import utils
+
+__all__ += ["utils"]
+
+from deepinv import models
+
+__all__ += ["iterative"]
+
+from deepinv import physics
+
+__all__ += ["physics"]
+
+from deepinv import datasets
+
+__all__ += ["datasets"]
+
+from deepinv import transform
+
+__all__ += ["transform"]
+
+from deepinv import sampling
+
+__all__ += ["sampling"]
+
+from deepinv.loss import metric
+
+__all__ += ["metric"]
+
+from deepinv import unfolded
+
+__all__ += ["unfolded"]
+
+from deepinv.training_utils import train, test
+
+# GLOBAL PROPERTY
+dtype = torch.float
+
+# if torch.cuda.is_available():
+#    try:
+#        free_gpu_id = get_freer_gpu()
+#        device = torch.device(f"cuda:{free_gpu_id}")
+#    except:
+#        device = torch.device("cuda")
+#        print("unable to get GPU info")
+# else:
+#    device = "cpu"

deepinv/datasets/__init__.py

@@ -0,0 +1 @@
+from .datagenerator import generate_dataset, HDF5Dataset

deepinv/datasets/datagenerator.py

@@ -0,0 +1,238 @@
+from tqdm import tqdm
+import os
+import h5py
+import torch
+from torch.utils.data import DataLoader, Subset
+from torch.utils import data
+
+
+class HDF5Dataset(data.Dataset):
+    r"""
+    DeepInverse HDF5 dataset with signal/measurement pairs.
+
+    :param str path: Path to the folder containing the dataset (one or multiple HDF5 files).
+    :param bool train: Set to ``True`` for training and ``False`` for testing.
+    """
+
+    def __init__(self, path, train=True):
+        super().__init__()
+        self.data_info = []
+        self.data_cache = {}
+        self.unsupervised = False
+
+        hd5 = h5py.File(path, "r")
+        if train:
+            if "x_train" in hd5:
+                self.x = hd5["x_train"]
+            else:
+                self.unsupervised = True
+            self.y = hd5["y_train"]
+        else:
+            self.x = hd5["x_test"]
+            self.y = hd5["y_test"]
+
+    def __getitem__(self, index):
+        y = torch.from_numpy(self.y[index]).type(torch.float)
+
+        x = y
+        if not self.unsupervised:
+            x = torch.from_numpy(self.x[index]).type(torch.float)
+
+        return x, y
+
+    def __len__(self):
+        return len(self.y)
+
+
+def generate_dataset(
+    train_dataset,
+    physics,
+    save_dir,
+    test_dataset=None,
+    device="cpu",
+    train_datapoints=None,
+    test_datapoints=None,
+    dataset_filename="dinv_dataset",
+    batch_size=4,
+    num_workers=0,
+    supervised=True,
+):
+    r"""
+    Generates dataset of signal/measurement pairs from base dataset.
+
+    It generates the measurement data using the forward operator provided by the user.
+    The dataset is saved in HD5 format and can be easily loaded using the HD5Dataset class.
+    The generated dataset contains a train and test splits.
+
+    :param torch.data.Dataset train_dataset: base dataset (e.g., MNIST, CelebA, etc.)
+        with images used for generating associated measurements
+        via the chosen forward operator. The generated dataset is saved in HD5 format and can be easily loaded using the
+        HD5Dataset class.
+    :param deepinv.physics.Physics physics: Forward operator used to generate the measurement data.
+        It can be either a single operator or a list of forward operators. In the latter case, the dataset will be
+        assigned evenly across operators.
+    :param str save_dir: folder where the dataset and forward operator will be saved.
+    :param torch.data.Dataset test_dataset: if included, the function will also generate measurements associated to the
+        test dataset.
+    :param torch.device device: which indicates cpu or gpu.
+    :param int, None train_datapoints: Desired number of datapoints in the training dataset. If set to ``None``, it will use the
+        number of datapoints in the base dataset. This is useful for generating a larger train dataset via data
+        augmentation (which should be chosen in the train_dataset).
+    :param int, None test_datapoints: Desired number of datapoints in the test dataset. If set to ``None``, it will use the
+        number of datapoints in the base test dataset.
+    :param str dataset_filename: desired filename of the dataset.
+    :param int batch_size: batch size for generating the measurement data
+        (it only affects the speed of the generating process)
+    :param int num_workers: number of workers for generating the measurement data
+        (it only affects the speed of the generating process)
+    :param bool supervised: Generates supervised pairs (x,y) of measurements and signals.
+        If set to ``False``, it will generate a training dataset with measurements only (y)
+        and a test dataset with pairs (x,y)
+
+    """
+    if os.path.exists(os.path.join(save_dir, dataset_filename)):
+        print(
+            "WARNING: Dataset already exists, this will overwrite the previous dataset."
+        )
+
+    if test_dataset is None and train_dataset is None:
+        raise ValueError("No train or test datasets provided.")
+
+    if not os.path.exists(save_dir):
+        os.makedirs(save_dir)
+
+    if not (type(physics) in [list, tuple]):
+        physics = [physics]
+        G = 1
+    else:
+        G = len(physics)
+
+    if train_dataset is not None:
+        if train_datapoints is not None:
+            datapoints = int(train_datapoints)
+        else:
+            datapoints = len(train_dataset)
+
+        n_train = datapoints  # min(len(train_dataset), datapoints)
+        n_train_g = int(n_train / G)
+        n_dataset_g = int(min(len(train_dataset), datapoints) / G)
+
+    if test_dataset is not None:
+        test_datapoints = (
+            test_datapoints if test_datapoints is not None else len(test_dataset)
+        )
+        n_test = min(len(test_dataset), test_datapoints)
+        n_test_g = int(n_test / G)
+
+    hf_paths = []
+
+    for g in range(G):
+        hf_path = f"{save_dir}/{dataset_filename}{g}.h5"
+        hf_paths.append(hf_path)
+        hf = h5py.File(hf_path, "w")
+
+        hf.attrs["operator"] = physics[g].__class__.__name__
+
+        if train_dataset is not None:
+            x = train_dataset[0]
+        elif test_dataset is not None:
+            x = test_dataset[0]
+
+        x = x[0] if isinstance(x, list) or isinstance(x, tuple) else x
+        x = x.to(device).unsqueeze(0)
+
+        # choose operator and generate measurement
+        y = physics[g](x)
+
+        torch.save(physics[g].state_dict(), f"{save_dir}/physics{g}.pt")
+
+        if train_dataset is not None:
+
+            hf.create_dataset("y_train", (n_train_g,) + y.shape[1:], dtype="float")
+            if supervised:
+                hf.create_dataset("x_train", (n_train_g,) + x.shape[1:], dtype="float")
+
+            if G > 1:
+                print(
+                    f"Computing train measurement vectors from base dataset of operator {g + 1} out of {G}..."
+                )
+            else:
+                print("Computing train measurement vectors from base dataset...")
+
+            index = 0
+
+            epochs = int(n_train_g / len(train_dataset)) + 1
+            for e in tqdm(range(epochs)):
+                train_dataloader = DataLoader(
+                    Subset(
+                        train_dataset,
+                        indices=list(range(g * n_dataset_g, (g + 1) * n_dataset_g)),
+                    ),
+                    batch_size=batch_size,
+                    num_workers=num_workers,
+                    pin_memory=False if device == "cpu" else True,
+                )
+
+                for i, x in enumerate(train_dataloader):
+                    x = x[0] if isinstance(x, list) or isinstance(x, tuple) else x
+                    x = x.to(device)
+
+                    # choose operator and generate measurement
+                    y = physics[g](x)
+
+                    # Add new data to it
+                    bsize = x.size()[0]
+
+                    if bsize + index > n_train_g:
+                        bsize = n_train_g - index
+
+                    hf["y_train"][index : index + bsize] = y[:bsize, :].to("cpu").numpy()
+                    if supervised:
+                        hf["x_train"][index : index + bsize] = (
+                            x[:bsize, :, :, :].to("cpu").numpy()
+                        )
+                    index = index + bsize
+
+        if test_dataset is not None:
+            index = 0
+            test_dataloader = DataLoader(
+                Subset(
+                    test_dataset, indices=list(range(g * n_test_g, (g + 1) * n_test_g))
+                ),
+                batch_size=batch_size,
+                num_workers=num_workers,
+                pin_memory=True,
+            )
+
+            if G > 1:
+                print(
+                    f"Computing test measurement vectors from base dataset of operator {g + 1} out of {G}..."
+                )
+            else:
+                print("Computing test measurement vectors from base dataset...")
+
+            for i, x in enumerate(tqdm(test_dataloader)):
+                x = x[0] if isinstance(x, list) or isinstance(x, tuple) else x
+                x = x.to(device)
+
+                # choose operator
+                y = physics[g](x)
+
+                if i == 0:  # create dict
+                    hf.create_dataset(
+                        "x_test", (n_test_g,) + x.shape[1:], dtype="float"
+                    )
+                    hf.create_dataset(
+                        "y_test", (n_test_g,) + y.shape[1:], dtype="float"
+                    )
+
+                # Add new data to it
+                bsize = x.size()[0]
+                hf["x_test"][index : index + bsize] = x.to("cpu").numpy()
+                hf["y_test"][index : index + bsize] = y.to("cpu").numpy()
+                index = index + bsize
+        hf.close()
+
+    print("Dataset has been saved in " + str(save_dir))
+
+    return hf_paths[0] if G == 1 else hf_paths

deepinv/loss/__init__.py

@@ -0,0 +1,10 @@
+from deepinv.loss.mc import MCLoss
+from deepinv.loss.ei import EILoss
+from deepinv.loss.moi import MOILoss
+from deepinv.loss.sup import SupLoss
+from deepinv.loss.score import ScoreLoss
+from deepinv.loss.tv import TVLoss
+from deepinv.loss.sure import SureGaussianLoss, SurePoissonLoss, SurePGLoss
+from deepinv.loss.regularisers import JacobianSpectralNorm, FNEJacobianSpectralNorm
+from deepinv.loss.measplit import SplittingLoss, Neighbor2Neighbor
+from deepinv.loss.metric import LpNorm, CharbonnierLoss

deepinv/loss/ei.py

@@ -0,0 +1,76 @@
+import torch
+import torch.nn as nn
+
+
+class EILoss(nn.Module):
+    r"""
+    Equivariant imaging self-supervised loss.
+
+    Assumes that the set of signals is invariant to a group of transformations (rotations, translations, etc.)
+    in order to learn from incomplete measurement data alone https://https://arxiv.org/pdf/2103.14756.pdf.
+
+    The EI loss is defined as
+
+    .. math::
+
+        \| T_g \hat{x} - \inverse{\forw{T_g \hat{x}}}\|^2
+
+
+    where :math:`\hat{x}=\inverse{y}` is a reconstructed signal and
+    :math:`T_g` is a transformation sampled at random from a group :math:`g\sim\group`.
+
+    By default, the error is computed using the MSE metric, however any other metric (e.g., :math:`\ell_1`)
+    can be used as well.
+
+    :param deepinv.Transform, torchvision.transforms transform: Transform to generate the virtually
+        augmented measurement. It can be any torch-differentiable function (e.g., a ``torch.nn.Module``).
+    :param torch.nn.Module metric: Metric used to compute the error between the reconstructed augmented measurement and the reference
+        image.
+    :param bool apply_noise: if ``True``, the augmented measurement is computed with the full sensing model
+        :math:`\sensor{\noise{\forw{\hat{x}}}}` (i.e., noise and sensor model),
+        otherwise is generated as :math:`\forw{\hat{x}}`.
+    :param float weight: Weight of the loss.
+    :param bool no_grad: if ``True``, the gradient does not propagate through :math:`T_g`. Default: ``True``.
+    """
+
+    def __init__(
+        self,
+        transform,
+        metric=torch.nn.MSELoss(),
+        apply_noise=True,
+        weight=1.0,
+        no_grad=True,
+    ):
+        super(EILoss, self).__init__()
+        self.name = "ei"
+        self.metric = metric
+        self.weight = weight
+        self.T = transform
+        self.noise = apply_noise
+        self.no_grad = no_grad
+
+    def forward(self, x_net, physics, model, **kwargs):
+        r"""
+        Computes the EI loss
+
+        :param torch.Tensor x_net: Reconstructed image :math:`\inverse{y}`.
+        :param deepinv.physics.Physics physics: Forward operator associated with the measurements.
+        :param torch.nn.Module model: Reconstruction function.
+        :return: (torch.Tensor) loss.
+        """
+
+        if self.no_grad:
+            with torch.no_grad():
+                x2 = self.T(x_net)
+        else:
+            x2 = self.T(x_net)
+
+        if self.noise:
+            y = physics(x2)
+        else:
+            y = physics.A(x2)
+
+        x3 = model(y, physics)
+
+        loss_ei = self.weight * self.metric(x3, x2)
+        return loss_ei

deepinv/loss/mc.py

@@ -0,0 +1,39 @@
+import torch
+import torch.nn as nn
+
+
+class MCLoss(nn.Module):
+    r"""
+    Measurement consistency loss
+
+    This loss enforces that the reconstructions are measurement-consistent, i.e., :math:`y=\forw{\inverse{y}}`.
+
+    The measurement consistency loss is defined as
+
+    .. math::
+
+        \|y-\forw{\inverse{y}}\|^2
+
+    where :math:`\inverse{y}` is the reconstructed signal and :math:`A` is a forward operator.
+
+    By default, the error is computed using the MSE metric, however any other metric (e.g., :math:`\ell_1`)
+    can be used as well.
+
+    :param torch.nn.Module metric: metric used for computing data consistency, which is set as the mean squared error by default.
+    """
+
+    def __init__(self, metric=torch.nn.MSELoss()):
+        super(MCLoss, self).__init__()
+        self.name = "mc"
+        self.metric = metric
+
+    def forward(self, y, x_net, physics, **kwargs):
+        r"""
+        Computes the measurement splitting loss
+
+        :param torch.Tensor y: measurements.
+        :param torch.Tensor x_net: reconstructed image :math:`\inverse{y}`.
+        :param deepinv.physics.Physics physics: forward operator associated with the measurements.
+        :return: (torch.Tensor) loss.
+        """
+        return self.metric(physics.A(x_net), y)

deepinv/loss/measplit.py

@@ -0,0 +1,219 @@
+import torch
+from deepinv.physics import Inpainting
+import numpy as np
+
+
+class SplittingLoss(torch.nn.Module):
+    r"""
+    Measurement splitting loss.
+
+    Splits the measurement and forward operator (of size :math:`m`)
+    into two smaller pairs  :math:`(y_1,A_1)` (of size :math:`m_1`) and  :math:`(y_2,A_2)` (of size :math:`m_2`) ,
+    to compute the self-supervised loss:
+
+    .. math::
+
+        \frac{m}{m_2}\| y_2 - A_2 \inversef{y_1,A_1}\|^2
+
+    where :math:`R` is the trainable network. See https://pubmed.ncbi.nlm.nih.gov/32614100/.
+
+    By default, the error is computed using the MSE metric, however any other metric (e.g., :math:`\ell_1`)
+    can be used as well.
+
+    :param torch.nn.Module metric: metric used for computing data consistency,
+        which is set as the mean squared error by default.
+    :param float split_ratio: splitting ratio, should be between 0 and 1. The size of :math:`y_1` increases
+        with the splitting ratio.
+    :param bool regular_mask: If ``True``, it will use a regular mask, otherwise it uses a random mask.
+    """
+
+    def __init__(self, metric=torch.nn.MSELoss(), split_ratio=0.9, regular_mask=False):
+        super().__init__()
+        self.name = "ms"
+        self.metric = metric
+        self.regular_mask = regular_mask
+        self.split_ratio = split_ratio
+
+    def forward(self, y, physics, model, **kwargs):
+        r"""
+        Computes the measurement splitting loss
+
+        :param torch.Tensor y: Measurements.
+        :param deepinv.physics.Physics physics: Forward operator associated with the measurements.
+        :param torch.nn.Module model: Reconstruction function.
+        :return: (torch.Tensor) loss.
+        """
+        tsize = y.size()[1:]
+
+        # sample a splitting
+        mask = torch.ones(tsize).to(y.device)
+        if not self.regular_mask:
+            mask[torch.rand_like(mask) > self.split_ratio] = 0
+        else:
+            stride = int(1 / (1 - self.split_ratio))
+            start = np.random.randint(stride)
+            mask[..., start::stride, start::stride] = 0.0
+
+        # create inpainting masks
+        inp = Inpainting(tsize, mask)
+        inp2 = Inpainting(tsize, 1 - mask)
+
+        # concatenate operators
+        physics1 = inp * physics  # A_1 = P*A
+        physics2 = inp2 * physics  # A_2 = (I-P)*A
+
+        # divide measurements
+        y1 = inp.A(y)
+        y2 = inp2.A(y)
+
+        loss_ms = self.metric(physics2.A(model(y1, physics1)), y2)
+        loss_ms /= 1 - self.split_ratio  # normalize loss
+
+        return loss_ms
+
+
+class Neighbor2Neighbor(torch.nn.Module):
+    r"""
+    Neighbor2Neighbor loss.
+
+    Splits the noisy measurements using two masks :math:`A_1` and :math:`A_2`, each choosing a different neighboring
+    map (see details in `"Neighbor2Neighbor: Self-Supervised Denoising from Single Noisy Images"
+    <https://openaccess.thecvf.com/content/CVPR2021/papers/Huang_Neighbor2Neighbor_Self-Supervised_Denoising_From_Single_Noisy_Images_CVPR_2021_paper.pdf>`_).
+
+    The self-supervised loss is computed as:
+
+    .. math::
+
+        \| A_2 y - R(A_1 y)\|^2 + \gamma \| A_2 y - R(A_1 y) - (A_2 R(y) - A_1 R(y))\|^2
+
+    where :math:`R` is the trainable denoiser network, :math:`\gamma>0` is a regularization parameter
+    and no gradient is propagated when computing :math:`R(y)`.
+
+    By default, the error is computed using the MSE metric, however any other metric (e.g., :math:`\ell_1`)
+    can be used as well.
+
+    The code has been adapted from the repository https://github.com/TaoHuang2018/Neighbor2Neighbor.
+
+    :param torch.nn.Module metric: metric used for computing data consistency,
+        which is set as the mean squared error by default.
+    :param float gamma: regularization parameter :math:`\gamma`.
+    """
+
+    def __init__(self, metric=torch.nn.MSELoss(), gamma=2.0):
+        super().__init__()
+        self.name = "neigh2neigh"
+        self.metric = metric
+        self.gamma = gamma
+
+    def space_to_depth(self, x, block_size):
+        n, c, h, w = x.size()
+        unfolded_x = torch.nn.functional.unfold(x, block_size, stride=block_size)
+        return unfolded_x.view(n, c * block_size**2, h // block_size, w // block_size)
+
+    def generate_mask_pair(self, img):
+        # prepare masks (N x C x H/2 x W/2)
+        n, c, h, w = img.shape
+        mask1 = torch.zeros(
+            size=(n * h // 2 * w // 2 * 4,), dtype=torch.bool, device=img.device
+        )
+        mask2 = torch.zeros(
+            size=(n * h // 2 * w // 2 * 4,), dtype=torch.bool, device=img.device
+        )
+        # prepare random mask pairs
+        idx_pair = torch.tensor(
+            [[0, 1], [0, 2], [1, 3], [2, 3], [1, 0], [2, 0], [3, 1], [3, 2]],
+            dtype=torch.int64,
+            device=img.device,
+        )
+        rd_idx = torch.zeros(
+            size=(n * h // 2 * w // 2,), dtype=torch.int64, device=img.device
+        )
+        torch.randint(low=0, high=8, size=(n * h // 2 * w // 2,), out=rd_idx)
+        rd_pair_idx = idx_pair[rd_idx]
+        rd_pair_idx += torch.arange(
+            start=0,
+            end=n * h // 2 * w // 2 * 4,
+            step=4,
+            dtype=torch.int64,
+            device=img.device,
+        ).reshape(-1, 1)
+        # get masks
+        mask1[rd_pair_idx[:, 0]] = 1
+        mask2[rd_pair_idx[:, 1]] = 1
+        return mask1, mask2
+
+    def generate_subimages(self, img, mask):
+        n, c, h, w = img.shape
+        subimage = torch.zeros(
+            n, c, h // 2, w // 2, dtype=img.dtype, layout=img.layout, device=img.device
+        )
+        # per channel
+        for i in range(c):
+            img_per_channel = self.space_to_depth(img[:, i : i + 1, :, :], block_size=2)
+            img_per_channel = img_per_channel.permute(0, 2, 3, 1).reshape(-1)
+            subimage[:, i : i + 1, :, :] = (
+                img_per_channel[mask].reshape(n, h // 2, w // 2, 1).permute(0, 3, 1, 2)
+            )
+        return subimage
+
+    def forward(self, y, physics, model, **kwargs):
+        r"""
+        Computes the neighbor2neighbor loss.
+
+
+        :param torch.Tensor y: Measurements.
+        :param deepinv.physics.Physics physics: Forward operator associated with the measurements.
+        :param torch.nn.Module model: Reconstruction function.
+        :return: (torch.Tensor) loss.
+        """
+
+        assert len(y.shape) == 4, "Input measurements should be images"
+        assert (
+            y.shape[2] % 2 == 0 and y.shape[3] % 2 == 0
+        ), "Image dimensions should be even"
+
+        mask1, mask2 = self.generate_mask_pair(y)
+
+        y1 = self.generate_subimages(y, mask1)
+        xhat1 = model(y1, physics)
+        y2 = self.generate_subimages(y, mask2)
+
+        xhat = model(y, physics).detach()
+        y1_hat = self.generate_subimages(xhat, mask1)
+        y2_hat = self.generate_subimages(xhat, mask2)
+
+        loss_n2n = self.metric(xhat1, y2) + self.gamma * self.metric(
+            xhat1 - y1_hat, y2 - y2_hat
+        )
+
+        return loss_n2n
+
+
+# if __name__ == "__main__":
+#     import deepinv as dinv
+#
+#     sigma = 0.1
+#     physics = dinv.physics.Denoising()
+#     physics.noise_model = dinv.physics.GaussianNoise(sigma)
+#
+#     # choose a reconstruction architecture
+#     backbone = dinv.models.MedianFilter()
+#     f = dinv.models.ArtifactRemoval(backbone)
+#     batch_size = 1
+#     imsize = (3, 128, 128)
+#
+#     for split_ratio in np.linspace(0.7, 0.99, 10):
+#         x = torch.ones((batch_size,) + imsize, device=dinv.device)
+#         y = physics(x)
+#
+#         # choose training losses
+#         loss = SplittingLoss(split_ratio=split_ratio, regular_mask=True)
+#         x_net = f(y, physics)
+#         mse = dinv.metric.mse()(physics.A(x), physics.A(x_net))
+#         split_loss = loss(y, physics, f)
+#
+#         print(
+#             f"split_ratio:{split_ratio:.2f}  mse: {mse:.2e}, split-loss: {split_loss:.2e}"
+#         )
+#         rel_error = (split_loss - mse).abs() / mse
+#         print(f"rel_error: {rel_error:.2f}")