Dataset objects (for training and testing)

The software below provides functionality (with PyTorch) to load datasets for training and testing.

Abstract base class

All dataset objects derive from this abstract base class, which lays out the notion of a dataset with pairs of images: one image is lensed (simulated or measured), and the other is lensless (simulated or measured).

class lensless.utils.dataset.DualDataset(indices=None, background=None, downsample=1, flip=False, flip_ud=False, flip_lr=False, transform_lensless=None, transform_lensed=None, input_snr=None, **kwargs)

Abstract class for defining a dataset of paired lensed and lensless images.

__init__(indices=None, background=None, downsample=1, flip=False, flip_ud=False, flip_lr=False, transform_lensless=None, transform_lensed=None, input_snr=None, **kwargs)

Dataset consisting of lensless and corresponding lensed image.

Parameters

• indices (range or int or None) – Indices of the images to use in the dataset (if integer, it should be interpreted as range(indices)), by default None.

• psf_path (str) – Path to the PSF of the imaging system, by default None.

• background (Tensor or None, optional) – If not None, background is removed from lensless images, by default None. If PSF is provided, background is estimated from the PSF.

• background_pix (tuple, optional) – Pixels to use for background estimation, by default (0, 15).

• downsample (int, optional) – Downsample factor of the lensless images, by default 1.

• flip (bool, optional) – If True, lensless images are flipped, by default False.

• transform_lensless (PyTorch Transform or None, optional) – Transform to apply to the lensless images, by default None. Note that this transform is applied on HWC images (different from torchvision).

• transform_lensed (PyTorch Transform or None, optional) – Transform to apply to the lensed images, by default None. Note that this transform is applied on HWC images (different from torchvision).

• input_snr (float, optional) – If not None, Poisson noise is added to the lensless images to match the given SNR.

abstract __len__()

Abstract method to get the length of the dataset. It should take into account the indices parameter.

abstract _get_images_pair(idx)

Abstract method to get the lensed and lensless images. Should return a pair (lensless, lensed) of numpy arrays with values in [0,1].

Parameters

idx (int) – images index

Measured dataset objects

class lensless.utils.dataset.HFDataset(huggingface_repo, split, n_files=None, psf=None, rotate=False, flipud=False, flip_lensed=False, downsample=1, downsample_lensed=1, input_snr=None, psf_snr=None, display_res=None, sensor='rpi_hq', slm='adafruit', alignment=None, return_mask_label=False, save_psf=False, simulation_config={}, simulate_lensless=False, force_rgb=False, cache_dir=None, single_channel_psf=False, random_flip=False, bg_snr_range=None, bg_fp=None, **kwargs)

__init__(huggingface_repo, split, n_files=None, psf=None, rotate=False, flipud=False, flip_lensed=False, downsample=1, downsample_lensed=1, input_snr=None, psf_snr=None, display_res=None, sensor='rpi_hq', slm='adafruit', alignment=None, return_mask_label=False, save_psf=False, simulation_config={}, simulate_lensless=False, force_rgb=False, cache_dir=None, single_channel_psf=False, random_flip=False, bg_snr_range=None, bg_fp=None, **kwargs)

Wrapper for lensless datasets on Hugging Face.

Parameters

• huggingface_repo (str) – Hugging Face repository ID.

• split (str or torch.utils.data.Dataset) – Split of the dataset to use: ‘train’, ‘test’, or ‘all’. If a Dataset object is given, it is used directly.

• n_files (int, optional) – Number of files to load from the dataset, by default None, namely all.

• psf (str, optional) – File name of the PSF at the repository. If None, it is assumed that there is a mask pattern from which the PSF is simulated, by default None.

• rotate (bool, optional) – If True, lensless images and PSF are rotated 180 degrees. Lensed/original image is not rotated! By default False.

• downsample (float, optional) – Downsample factor of the lensless images, by default 1.

• downsample_lensed (float, optional) – Downsample factor of the lensed images, by default 1.

• display_res (tuple, optional) – Resolution of images when displayed on screen during measurement.

• sensor (str, optional) – If psf not provided, the sensor to use for the PSF simulation, by default “rpi_hq”.

• slm (str, optional) – If psf not provided, the SLM to use for the PSF simulation, by default “adafruit”.

• alignment (dict, optional) – Alignment parameters between lensless and lensed data. If “top_left”, “height”, and “width” are provided, the region-of-interest from the reconstruction of lensless is extracted and lensed is reshaped to match. If “crop” is provided, the region-of-interest is extracted from the simulated lensed image, namely a simulation configuration should be provided within alignment.

• return_mask_label (bool, optional) – If multimask dataset, return the mask label (True) or the corresponding PSF (False).

• save_psf (bool, optional) – If multimask dataset, save the simulated PSFs.

• random_flip (bool, optional) – If True, randomly flip the lensless images vertically and horizonally with equal probability. By default, no flipping.

• simulation_config (dict, optional) – Simulation parameters for PSF if using a mask pattern.

• bg_snr_range (list, optional) – List [low, high] of range of possible SNRs for which to add the background. Used in conjunction with ‘bg’.

• bg_fp (string, optional) – File path of background to add to the data for simulating a measurement in ambient light.

extract_roi(reconstruction, lensed=None, axis=(1, 2), flip_lr=None, flip_ud=None, rotate_aug=False, shift_aug=None)

Parameters

• flip_lr (torch.Tensor, optional) – Tensor of booleans indicating whether to flip the reconstruction left-right, by default None.

• flip_ud (bool, optional) – Tensor of booleans indicating whether to flip the reconstruction up-down, by default None.

class lensless.utils.dataset.MeasuredDataset(root_dir, lensless_fn='diffuser', lensed_fn='lensed', image_ext='npy', **kwargs)

Dataset consisting of lensless and corresponding lensed image. It can be used with a PyTorch DataLoader to load a batch of lensless and corresponding lensed images. Unless the setup is perfectly calibrated, one should expect to have to use transform_lensed to adjust the alignment and rotation.

__init__(root_dir, lensless_fn='diffuser', lensed_fn='lensed', image_ext='npy', **kwargs)

Dataset consisting of lensless and corresponding lensed image. Default parameters are for the DiffuserCam Lensless Mirflickr Dataset (DLMD).

Parameters

• root_dir (str) – Path to the test dataset. It is expected to contain two folders: ones of lensless images and one of lensed images.

• lensless_fn (str, optional) – Name of the folder containing the lensless images, by default “diffuser”.

• lensed_fn (str, optional) – Name of the folder containing the lensed images, by default “lensed”.

• image_ext (str, optional) – Extension of the images, by default “npy”.

class lensless.utils.dataset.MeasuredDatasetSimulatedOriginal(measured_dir, original_dir, simulator, measurement_ext='png', original_ext='jpg', downsample=1, background=None, flip=False, **kwargs)

Abstract class for defining a dataset of paired lensed and lensless images.

Dataset consisting of lensless image captured from a screen and the corresponding image shown on the screen. Unlike lensless.utils.dataset.MeasuredDataset, the ground-truth lensed image is simulated using a lensless.utils.simulation.FarFieldSimulator object rather than measured with a lensed camera.

The class assumes that the measured_dir and original_dir have file names that match.

The method _get_images_pair must be defined.

__init__(measured_dir, original_dir, simulator, measurement_ext='png', original_ext='jpg', downsample=1, background=None, flip=False, **kwargs)

Dataset consisting of lensless image captured from a screen and the corresponding image shown on screen.

class lensless.utils.dataset.DiffuserCamTestDataset(data_dir=None, n_files=None, downsample=2)

Dataset consisting of lensless and corresponding lensed image. This is the standard dataset used for benchmarking.

__init__(data_dir=None, n_files=None, downsample=2)

Dataset consisting of lensless and corresponding lensed image. Default parameters are for the test set of DiffuserCam Lensless Mirflickr Dataset (DLMD).

Parameters

• data_dir (str, optional) – The path to DiffuserCam_Test dataset, by default looks inside the data folder.

• n_files (int, optional) – Number of image pairs to load in the dataset , by default use all.

• downsample (int, optional) – Downsample factor of the lensless images, by default 2. Note that the PSF has a resolution of 4x of the images.

Simulated dataset objects

These dataset objects can be used for training and testing with simulated data. The main assumption is that the imaging system is linear shift-invariant (LSI), and that the lensless image is the result of a convolution of the lensed image with a point-spread function (PSF). Check out this Medium post for more details on the simulation procedure.

With simulated data, we can avoid the hassle of collecting a large amount of data. However, it’s important to note that the LSI assumption can sometimes be too idealistic, in particular for large angles.

Nevertheless, simulating data is the only option of learning the mask / PSF.

class lensless.utils.dataset.SimulatedFarFieldDataset(dataset, simulator, pre_transform=None, dataset_is_CHW=False, flip=False, vertical_shift=None, horizontal_shift=None, crop=None, downsample=1, **kwargs)

Dataset of propagated images (through simulation) from a Torch Dataset. lensless.utils.simulation.FarFieldSimulator is used for simulation, assuming a far-field propagation and a shift-invariant system with a single point spread function (PSF).

__init__(dataset, simulator, pre_transform=None, dataset_is_CHW=False, flip=False, vertical_shift=None, horizontal_shift=None, crop=None, downsample=1, **kwargs)

Parameters

• dataset (torch.utils.data.Dataset) – Dataset to propagate. Should output images with shape [H, W, C] unless dataset_is_CHW is True (and therefore images have the dimension ordering of [C, H, W]).

• simulator (lensless.utils.simulation.FarFieldSimulator) – Simulator object used on images from dataset. Waveprop simulator to use for the simulation. It is expected to have is_torch = True.

• pre_transform (PyTorch Transform or None, optional) – Transform to apply to the images before simulation, by default None. Note that this transform is applied on HCW images (different from torchvision).

• dataset_is_CHW (bool, optional) – If True, the input dataset is expected to output images with shape [C, H, W], by default False.

• flip (bool, optional) – If True, images are flipped beffore the simulation, by default False.

class lensless.utils.dataset.SimulatedDatasetTrainableMask(mask, dataset, simulator, **kwargs)

Dataset of propagated images (through simulation) from a Torch Dataset with learnable mask. The waveprop package is used for the simulation, assuming a far-field propagation and a shift-invariant system with a single point spread function (PSF). To ensure autograd compatibility, the dataloader should have num_workers=0.

__init__(mask, dataset, simulator, **kwargs)

Parameters

• mask (lensless.hardware.trainable_mask.TrainableMask) – Mask to use for simulation. Should be a 4D tensor with shape [1, H, W, C]. Simulation of multi-depth data is not supported yet.

• dataset (torch.utils.data.Dataset) – Dataset to propagate. Should output images with shape [H, W, C] unless dataset_is_CHW is True (and therefore images have the dimension ordering of [C, H, W]).

• simulator (lensless.utils.simulation.FarFieldSimulator) – Simulator object used on images from dataset. Waveprop simulator to use for the simulation. It is expected to have is_torch = True.

set_psf(psf=None)

Set the PSF of the simulator.

Parameters

psf (torch.Tensor, optional) – PSF to use for the simulation. If None, the PSF of the mask is used.