"""PaDIS patch-based denoising loss and training utilities."""
from __future__ import annotations
from typing import Sequence
import torch
import torch.nn as nn
import torch.nn.functional as F
[docs]
def validate_patch_schedule(
patch_sizes: Sequence[int], patch_probabilities: Sequence[float]
) -> None:
"""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.
"""
if len(patch_sizes) != len(patch_probabilities):
raise ValueError("patch_sizes and patch_probabilities must have same length.")
if not patch_sizes:
raise ValueError("At least one patch size is required.")
if any(size % 8 != 0 for size in patch_sizes):
raise ValueError("All PaDIS patch sizes must be divisible by 8.")
prob_sum = float(sum(patch_probabilities))
if abs(prob_sum - 1.0) > 1e-6:
raise ValueError("patch_probabilities must sum to 1.")
[docs]
def build_position_grid(
batch_size: int,
height: int,
width: int,
*,
device: torch.device,
dtype: torch.dtype,
) -> torch.Tensor:
"""Create normalised ``(x, y)`` position channels for an image batch.
Parameters
----------
batch_size : int
Number of grids to return.
height, 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
-------
torch.Tensor
Tensor of shape ``(batch_size, 2, height, width)`` spanning ``[-1, 1]``
in each spatial direction.
"""
y = torch.linspace(-1.0, 1.0, height, device=device, dtype=dtype)
x = torch.linspace(-1.0, 1.0, width, device=device, dtype=dtype)
yy, xx = torch.meshgrid(y, x, indexing="ij")
grid = torch.stack((xx, yy), dim=0).unsqueeze(0)
return grid.expand(batch_size, -1, -1, -1)
[docs]
def zero_pad_images(images: torch.Tensor, pad_width: int) -> torch.Tensor:
"""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
-------
torch.Tensor
The original tensor when ``pad_width`` is zero, otherwise a padded
tensor preserving batch and channel dimensions.
"""
if pad_width == 0:
return images
return F.pad(images, (pad_width, pad_width, pad_width, pad_width), mode="constant")
[docs]
def sample_patch_size(
patch_sizes: Sequence[int],
patch_probabilities: Sequence[float],
*,
device: torch.device,
) -> int:
"""Draw one patch size from a validated categorical schedule."""
probs = torch.as_tensor(patch_probabilities, device=device, dtype=torch.float32)
index = torch.multinomial(probs, num_samples=1).item()
return int(patch_sizes[index])
[docs]
def sample_patch_pair(
images: torch.Tensor,
positions: torch.Tensor,
patch_size: int,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Sample aligned image and position patches independently per batch item.
Parameters
----------
images, 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 : tuple of torch.Tensor
Batched crops sharing the same random top-left coordinate for each
input image.
"""
batch_size, _, height, width = images.shape
if patch_size > height or patch_size > width:
raise ValueError("patch_size cannot exceed padded image dimensions.")
top = torch.randint(0, height - patch_size + 1, (batch_size,), device=images.device)
left = torch.randint(0, width - patch_size + 1, (batch_size,), device=images.device)
rows = top[:, None] + torch.arange(patch_size, device=images.device)[None, :]
cols = left[:, None] + torch.arange(patch_size, device=images.device)[None, :]
batch = torch.arange(batch_size, device=images.device)[:, None, None]
image_patch = images.permute(1, 0, 2, 3)[
:, batch, rows[:, :, None], cols[:, None, :]
].permute(1, 0, 2, 3)
position_patch = positions.permute(1, 0, 2, 3)[
:, batch, rows[:, :, None], cols[:, None, :]
].permute(1, 0, 2, 3)
return image_patch, position_patch
[docs]
def sample_image_patch_with_position_channels(
images: torch.Tensor,
patch_size: int,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Sample image patches and construct their absolute position channels.
Unlike :func:`sample_patch_pair`, this routine avoids materialising a full
position grid and derives coordinates only for sampled pixels.
"""
batch_size, _, height, width = images.shape
if patch_size > height or patch_size > width:
raise ValueError("patch_size cannot exceed padded image dimensions.")
top = torch.randint(0, height - patch_size + 1, (batch_size,), device=images.device)
left = torch.randint(0, width - patch_size + 1, (batch_size,), device=images.device)
rows = top[:, None] + torch.arange(patch_size, device=images.device)[None, :]
cols = left[:, None] + torch.arange(patch_size, device=images.device)[None, :]
batch = torch.arange(batch_size, device=images.device)[:, None, None]
image_patch = images.permute(1, 0, 2, 3)[
:, batch, rows[:, :, None], cols[:, None, :]
].permute(1, 0, 2, 3)
x_pos = (cols.to(images.dtype) / (width - 1) - 0.5) * 2.0
y_pos = (rows.to(images.dtype) / (height - 1) - 0.5) * 2.0
x_pos = x_pos[:, None, None, :].expand(-1, 1, patch_size, -1)
y_pos = y_pos[:, None, :, None].expand(-1, 1, -1, patch_size)
position_patch = torch.cat((x_pos, y_pos), dim=1)
return image_patch, position_patch
[docs]
def sample_image_patch(
images: torch.Tensor,
patch_size: int,
) -> torch.Tensor:
"""Sample one independent square crop from every image in a batch."""
batch_size, _, height, width = images.shape
if patch_size > height or patch_size > width:
raise ValueError("patch_size cannot exceed padded image dimensions.")
top = torch.randint(0, height - patch_size + 1, (batch_size,), device=images.device)
left = torch.randint(0, width - patch_size + 1, (batch_size,), device=images.device)
rows = top[:, None] + torch.arange(patch_size, device=images.device)[None, :]
cols = left[:, None] + torch.arange(patch_size, device=images.device)[None, :]
batch = torch.arange(batch_size, device=images.device)[:, None, None]
return images.permute(1, 0, 2, 3)[
:, batch, rows[:, :, None], cols[:, None, :]
].permute(1, 0, 2, 3)
[docs]
def score_from_denoiser(
noisy_image_patch: torch.Tensor, denoised_patch: torch.Tensor, sigma: torch.Tensor
) -> torch.Tensor:
"""Convert denoiser predictions into a variance-exploding score estimate.
The returned score is ``(denoised - noisy) / sigma**2`` with ``sigma``
broadcast across non-batch dimensions.
"""
while sigma.ndim < noisy_image_patch.ndim:
sigma = sigma.unsqueeze(-1)
return (denoised_patch - noisy_image_patch) / sigma.square()
[docs]
class PaDISDenoisingLoss(nn.Module):
"""EDM-style denoising objective used to train PaDIS priors.
Parameters
----------
sigma_min, 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, 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.
"""
def __init__(
self,
sigma_min: float = 0.002,
sigma_max: float = 40.0,
sigma_distribution: str = "edm_lognormal_truncated",
P_mean: float = -1.2,
P_std: float = 1.2,
sigma_data: float = 0.5,
reduction: str = "batch_mean_sum",
augment_pipe: nn.Module | None = None,
) -> None:
super().__init__()
if sigma_min <= 0 or sigma_max <= sigma_min:
raise ValueError("Require 0 < sigma_min < sigma_max.")
if sigma_distribution == "bounded_edm_lognormal":
sigma_distribution = "edm_lognormal_truncated"
if sigma_distribution not in (
"edm_lognormal",
"edm_lognormal_truncated",
"log_uniform",
):
raise ValueError(
"sigma_distribution must be edm_lognormal, "
"edm_lognormal_truncated, or log_uniform."
)
if reduction not in ("batch_mean_sum", "mean"):
raise ValueError("reduction must be batch_mean_sum or mean.")
self.sigma_min = float(sigma_min)
self.sigma_max = float(sigma_max)
self.sigma_distribution = sigma_distribution
self.P_mean = float(P_mean)
self.P_std = float(P_std)
self.sigma_data = float(sigma_data)
self.reduction = reduction
self.augment_pipe = augment_pipe
[docs]
def sample_sigma(self, batch_size: int, device: torch.device) -> torch.Tensor:
"""Sample a column vector of training noise levels."""
if self.sigma_distribution in ("edm_lognormal", "edm_lognormal_truncated"):
sigma = self._sample_edm_lognormal(batch_size, device)
if self.sigma_distribution == "edm_lognormal":
return sigma
return self._truncate_sigma(sigma, device)
log_min = torch.log(torch.tensor(self.sigma_min, device=device))
log_max = torch.log(torch.tensor(self.sigma_max, device=device))
log_sigma = log_min + torch.rand(batch_size, device=device) * (
log_max - log_min
)
return torch.exp(log_sigma)
def _sample_edm_lognormal(
self, batch_size: int, device: torch.device
) -> torch.Tensor:
rnd_normal = torch.randn(batch_size, device=device)
return torch.exp(rnd_normal * self.P_std + self.P_mean)
def _truncate_sigma(
self, sigma: torch.Tensor, device: torch.device
) -> torch.Tensor:
invalid = (sigma < self.sigma_min) | (sigma > self.sigma_max)
attempts = 0
while bool(invalid.any().item()) and attempts < 16:
sigma[invalid] = self._sample_edm_lognormal(
int(invalid.sum().item()), device
)
invalid = (sigma < self.sigma_min) | (sigma > self.sigma_max)
attempts += 1
return sigma.clamp(self.sigma_min, self.sigma_max)
[docs]
def forward(
self,
model: nn.Module,
clean_patch: torch.Tensor,
position_patch: torch.Tensor | None = None,
augment_pipe: nn.Module | None = None,
) -> torch.Tensor:
"""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``.
Returns
-------
torch.Tensor
Scalar reduced training loss.
"""
pipe = self.augment_pipe if augment_pipe is None else augment_pipe
target_patch, augment_labels = (
pipe(clean_patch) if pipe is not None else (clean_patch, None)
)
sigma = self.sample_sigma(target_patch.shape[0], target_patch.device)
sigma_view = sigma.reshape(
target_patch.shape[0], *([1] * (target_patch.ndim - 1))
)
noisy_patch = target_patch + sigma_view * torch.randn_like(target_patch)
sigma_data = torch.as_tensor(
self.sigma_data, device=target_patch.device, dtype=target_patch.dtype
)
c_skip = sigma_data.square() / (sigma_view.square() + sigma_data.square())
c_out = (
sigma_view * sigma_data / (sigma_view.square() + sigma_data.square()).sqrt()
)
c_in = 1 / (sigma_data.square() + sigma_view.square()).sqrt()
c_noise = sigma.log() / 4
if position_patch is not None:
model_input = torch.cat((c_in * noisy_patch, position_patch), dim=1)
else:
model_input = c_in * noisy_patch
if augment_labels is None:
model_output = model(model_input, c_noise)
else:
model_output = model(model_input, c_noise, augment_labels=augment_labels)
denoised = c_skip * noisy_patch + c_out * model_output
weight = (sigma_view.square() + sigma_data.square()) / (
sigma_view * sigma_data
).square()
loss = weight * (denoised - target_patch).square()
if self.reduction == "batch_mean_sum":
return loss.flatten(1).sum(dim=1).mean()
return torch.mean(loss)