peal.data.transformations.DiffusionAugmentation

class peal.data.transformations.DiffusionAugmentation

Bases: object

Augment images by a stochastic encode/decode round trip through a generator.

The generator (built with get_generator and moved to CUDA if available) encodes the image up to sampling_time_fraction of the diffusion process with stochastic="fully" and decodes it again stochastically, yielding a plausible variation of the input. Value ranges are converted via the dataset’s project_to_pytorch_default / project_from_pytorch_default on both sides.

Parameters:

• generator (str or dict or GeneratorInterface) – Generator specification accepted by get_generator.

• sampling_time_fraction (float) – Fraction of the diffusion time the image is noised to.

• num_discretization_steps (int) – Number of sampler steps for encoding and decoding.

• dataset (PealDataset, optional) – Dataset whose value range the input lives in. If None the input is assumed to already be in the pytorch default range.

Methods

__init__(generator, sampling_time_fraction, num_discretization_steps, dataset=None)

Build the generator on the best available device and store the knobs.

__call__(img)

Return a diffusion-resampled variant of img.

Parameters:

img (torch.Tensor) – Image [C, H, W] or batch [B, C, H, W] in the value range of dataset (or the pytorch default range if no dataset is set).

Returns:

Reconstructed image(s) with the same shape and value range. Note that the result stays on the generator’s device.

Return type:

torch.Tensor