Datasets#

The datasets module lets you use datasets with DeepInverse, for testing and training.

Datasets can be either:

Loaded in from your own data, see base datasets for desired format;
Paired dataset generated synthetically from ground-truth, see generating datasets;
A public dataset, such as DIV2K or FastMRI, see predefined datasets.

Base Datasets#

Datasets can return optionally ground-truth images x, measurements y, or physics parameters params, or any combination of these, in one of the following ways:

x i.e a dataset that returns only ground truth;
(x, y) i.e. a dataset that returns pairs of ground truth and measurement. x can be equal to torch.nan if your dataset is ground-truth-free.
(x, params) i.e. a dataset of ground truth and dict of physics parameters. Useful for training with online measurements.
(x, y, params) i.e. a dataset that returns ground truth, measurements and dict of params.

Tip

If you have a dataset of measurements only (y) or (y, params) you should modify it such that it returns (torch.nan, y) or (torch.nan, y, params)

If you have your own dataset (e.g. a PyTorch Dataset), check that it is compatible using the function deepinv.datasets.check_dataset() (e.g. to be used with deepinv.Trainer or deepinv.test).

Generating Datasets#

You can generate a dataset associated with a certain forward operator using deepinv.datasets.generate_dataset() using a base dataset. Your base dataset can be any dataset that returns ground truth, i.e. either one of our predefined datasets, your own data in the format provided by one of our base datasets, or other external datasets.

For example, here we generate a dataset of inpainting measurements from the deepinv.datasets.Set14HR dataset:

Note

We support all data types supported by h5py, including complex numbers.

>>> import deepinv as dinv
>>> from torchvision.transforms import ToTensor, Compose, CenterCrop
>>> save_dir = dinv.utils.demo.get_data_home() / 'set14'
>>>
>>> # Define base train dataset
>>> dataset = dinv.datasets.Set14HR(save_dir, download=True, transform=Compose([CenterCrop(128), ToTensor()]))
...
>>>
>>> # Define forward operator
>>> physics = dinv.physics.Inpainting(img_size=(3, 128, 128), mask=0.8, noise_model=dinv.physics.GaussianNoise(sigma=.05))
>>>
>>> # Generate paired dataset
>>> path = dinv.datasets.generate_dataset(dataset, physics, save_dir=save_dir, verbose=False)

The datasets are saved in .h5 (HDF5) format, and can be easily loaded to PyTorch’s standard torch.utils.data.DataLoader:

>>> from torch.utils.data import DataLoader
>>>
>>> train_dataset = dinv.datasets.HDF5Dataset(path)
>>> dataloader = DataLoader(train_dataset, batch_size=4, shuffle=True)
>>> x, y = next(iter(dataloader))
>>> x.shape, y.shape
(torch.Size([4, 3, 128, 128]), torch.Size([4, 3, 128, 128]))
>>> train_dataset.close()

We can also use physics generators to randomly generate physics params for data, and save and load the physics params into the dataset:

>>> physics_generator = dinv.physics.generator.SigmaGenerator()
>>> path = dinv.datasets.generate_dataset(dataset, physics, physics_generator=physics_generator, save_dir=save_dir, verbose=False)
>>> train_dataset = dinv.datasets.HDF5Dataset(path, load_physics_generator_params=True)
>>> dataloader = DataLoader(train_dataset, batch_size=4, shuffle=True)
>>> x, y, params = next(iter(dataloader))
>>> print(params['sigma'].shape)
torch.Size([4])

Predefined Datasets#

Multiple popular easy-to-download datasets are available, which all either return ground-truth only x, paired data (x, y) or with params (e.g. mask) (x, y, params).

All these datasets inherit from deepinv.datasets.ImageDataset.

Table 26 Datasets Overview#
Dataset	Returns	Dataset Size	Tensor Sizes	Description
`DIV2K`	`x`	800 (train) + 100 (val) images	RGB, up to 2040x2040 pixels (variable)	A widely-used dataset for natural image restoration.
`Urban100HR`	`x`	100 images	up to 1200x1280 pixels (variable)	Contains diverse high-resolution urban scenes, typically used for testing super-resolution algorithms.
`Set14HR`	`x`	14 high-resolution images	RGB, 248×248 to 512×768 pixels.	A small benchmark dataset for super-resolution tasks, containing a variety of natural images.
`CBSD68`	`x`	68 images	RGB, 481x321 pixels	A subset of the color Berkeley Segmentation Dataset.
`FastMRISliceDataset`	`(x, y)` or `(x, y, {'mask': mask, 'coil_maps': coil_maps})`	Over 100,000 MRI slices	Complex, varying shape approx. 640x320	Raw MRI knee and brain fully-sampled or undersampled k-space data and optional RSS targets from the FastMRI dataset.
`SimpleFastMRISliceDataset`	`x`	973 (knee) and 455 (brain) images	320x320 fully-sampled reconstructed slices	Easy-to-use in-memory prepared subset of 2D slices from the full FastMRI slice dataset for knees and brains, padded to standard size.
`CMRxReconSliceDataset`	`(x, y)` or `(x, y, {'mask': mask})`	300 patients, each with 8-13 slices	Padded to 512x256x12 time steps	Dynamic MRI sequences of cardiac cine from short axis (5-10 slices) and long axis (3 views) split by patient, from the CMRxRecon challenge.
`SKMTEASliceDataset`	`(x, y, {'mask': mask, 'coil_maps': coil_maps})`	25,000 slices from 155 patients	Complex double-echo with 8 coils of shape 512x160.	Raw MRI knee multicoil undersampled k-space data and fully-sampled ground truth from the Stanford SKM-TEA dataset, with precomputed Poisson disc masks from 4x to 16x acceleration, and pre-estimated coil maps.
`LidcIdriSliceDataset`	`x`	Over 200,000 CT scan slices	Slices 512x512 voxels	A comprehensive dataset of lung CT scans with annotations, used for medical image processing and lung cancer detection research.
`Flickr2kHR`	`x`	2,650 images	RGB, up to 2000x2000 pixels (variable)	A dataset from Flickr containing high-resolution images for tasks like super-resolution and image restoration.
`LsdirHR`	`x`	84499 (train) + 1000 (val) images	RGB, up to 2160x2160 pixels (variable)	A dataset with high-resolution images, often used for training large reconstruction models.
`FMD`	`x`	12000 images	512x512 pixels	The Fluorescence Microscopy Dataset (FMD) is a dataset of real fluorescence microscopy images.
`Kohler`	`(x, y)`	48 blurry + 9547 sharp images	800x800 RGB	A blind-deblurring dataset consists of blurry shots and sharp frames, each blurry shot being associated with about 200 sharp frames.
`NBUDataset`	`x` Tensor or TensorList	510 images across 6 satellites	Cx256x256 multispectral (C=4 or 8) and 1x1024x1024 panchromatic	Multispectral satellite images of urban scenes from 6 different satellites.

Data Transforms#

We provide some torchvision-style transforms for use when loading data:

Table 27 Data Transforms Overview#
Transform	Description
`deepinv.datasets.utils.Rescale`	Min-max or clip value rescaling.
`deepinv.datasets.utils.ToComplex`	Add empty imaginary dimension to image.
`deepinv.datasets.utils.Crop`	Crop image in corner or with arbitrary crop position and/or size.
`deepinv.datasets.MRISliceTransform`	Transform raw FastMRI data by simulating masks and estimating coil maps.

Dataset	Description
`deepinv.datasets.ImageDataset`	Base abstract dataset class
`deepinv.datasets.ImageFolder`	Dataset that loads images (ground-truth, measurements or both) from a folder
`deepinv.datasets.TensorDataset`	Dataset that returns tensor(s) passed in at input: either tensor(s) for a single observation or a whole dataset of them
`deepinv.datasets.HDF5Dataset`	Dataset of measurements generated using `deepinv.datasets.generate_dataset()`, see below for how to use.