PaDIS.py

Source: LION/losses/PaDIS.py

PaDIS patch-based denoising loss and training utilities.

LION.losses.PaDIS.validate_patch_schedule(patch_sizes, patch_probabilities)[source]

Validate a discrete PaDIS patch-size distribution.

Parameters:

  • patch_sizes (sequence of int) – Candidate square patch widths. Every size must be divisible by eight.

  • patch_probabilities (sequence of float) – Sampling probability corresponding to each patch size. Values must sum to one within numerical tolerance.

Raises:

ValueError – If the schedule is empty, lengths differ, a size is incompatible with NCSN++ downsampling, or probabilities do not sum to one.

Return type:

None

LION.losses.PaDIS.build_position_grid(batch_size, height, width, *, device, dtype)[source]

Create normalised (x, y) position channels for an image batch.

Parameters:

  • batch_size (int) – Number of grids to return.

  • height (int) – Spatial grid dimensions.

  • width (int) – Spatial grid dimensions.

  • device (torch.device) – Device on which to allocate the grid.

  • dtype (torch.dtype) – Floating-point type of the returned coordinates.

Returns:

Tensor of shape (batch_size, 2, height, width) spanning [-1, 1] in each spatial direction.

Return type:

torch.Tensor

LION.losses.PaDIS.zero_pad_images(images, pad_width)[source]

Pad images with zeros on every spatial boundary.

Parameters:

  • images (torch.Tensor) – Batched image tensor in NCHW layout.

  • pad_width (int) – Number of pixels added to each boundary.

Returns:

The original tensor when pad_width is zero, otherwise a padded tensor preserving batch and channel dimensions.

Return type:

torch.Tensor

LION.losses.PaDIS.sample_patch_size(patch_sizes, patch_probabilities, *, device)[source]

Draw one patch size from a validated categorical schedule.

Parameters:

  • patch_sizes (Sequence[int])

  • patch_probabilities (Sequence[float])

  • device (device)

Return type:

int

LION.losses.PaDIS.sample_patch_pair(images, positions, patch_size)[source]

Sample aligned image and position patches independently per batch item.

Parameters:

  • images (torch.Tensor) – Aligned NCHW tensors with identical spatial dimensions.

  • positions (torch.Tensor) – Aligned NCHW tensors with identical spatial dimensions.

  • patch_size (int) – Width and height of each square crop.

Returns:

image_patch, position_patch – Batched crops sharing the same random top-left coordinate for each input image.

Return type:

tuple of torch.Tensor

LION.losses.PaDIS.sample_image_patch_with_position_channels(images, patch_size)[source]

Sample image patches and construct their absolute position channels.

Unlike sample_patch_pair(), this routine avoids materialising a full position grid and derives coordinates only for sampled pixels.

Parameters:

  • images (Tensor)

  • patch_size (int)

Return type:

tuple[Tensor, Tensor]

LION.losses.PaDIS.sample_image_patch(images, patch_size)[source]

Sample one independent square crop from every image in a batch.

Parameters:

  • images (Tensor)

  • patch_size (int)

Return type:

Tensor

LION.losses.PaDIS.score_from_denoiser(noisy_image_patch, denoised_patch, sigma)[source]

Convert denoiser predictions into a variance-exploding score estimate.

The returned score is (denoised - noisy) / sigma**2 with sigma broadcast across non-batch dimensions.

Parameters:

  • noisy_image_patch (Tensor)

  • denoised_patch (Tensor)

  • sigma (Tensor)

Return type:

Tensor

class LION.losses.PaDIS.PaDISDenoisingLoss(sigma_min=0.002, sigma_max=40.0, sigma_distribution='edm_lognormal_truncated', P_mean=-1.2, P_std=1.2, sigma_data=0.5, reduction='batch_mean_sum', augment_pipe=None)[source]

Bases: Module

EDM-style denoising objective used to train PaDIS priors.

Parameters:

  • sigma_min (float) – Inclusive training-noise bounds.

  • sigma_max (float) – Inclusive training-noise bounds.

  • sigma_distribution ({"edm_lognormal", "edm_lognormal_truncated", "log_uniform"}) – Distribution used to sample per-image noise levels.

  • P_mean (float) – Location and scale of the EDM log-normal distribution.

  • P_std (float) – Location and scale of the EDM log-normal distribution.

  • sigma_data (float) – Assumed standard deviation of clean training data.

  • reduction ({"batch_mean_sum", "mean"}) – Reduction compatible with the released PaDIS objective or a global elementwise mean.

  • augment_pipe (torch.nn.Module, optional) – Optional augmentation callable applied before corrupting images.

sample_sigma(batch_size, device)[source]

Sample a column vector of training noise levels.

Parameters:

  • batch_size (int)

  • device (device)

Return type:

Tensor

forward(model, clean_patch, position_patch=None, augment_pipe=None)[source]

Evaluate the weighted PaDIS denoising loss.

Parameters:

  • model (torch.nn.Module) – Denoiser accepting noisy images, noise levels, and optional position channels.

  • images (torch.Tensor) – Clean image or patch batch in normalised-intensity units.

  • position_channels (torch.Tensor, optional) – Position channels aligned with images.

  • clean_patch (Tensor)

  • position_patch (Tensor | None)

  • augment_pipe (Module | None)

Returns:

Scalar reduced training loss.

Return type:

torch.Tensor