prior.py¶

Source: LION/reconstructors/diffusion/padis/prior.py

Patch layout and prior-assembly internals for the PaDIS reconstructor.

class LION.reconstructors.diffusion.padis.prior.PatchLayout(indices, image_height, image_width)[source]¶

Bases: object

Describe a set of image patches and its unpadded image extent.

Parameters:
  • indices (list[tuple[int, int, int, int]])

  • image_height (int)

  • image_width (int)

class LION.reconstructors.diffusion.padis.prior.PaDISPrior[source]¶

Bases: object

Assemble position-aware patch denoising into a whole-image prior.

patch_layout(image_shape, params, device, generator=None, *, fixed_offset=False)[source]¶

Build the randomly offset PaDIS patch partition for one image.

Parameters:
  • image_shape (tuple[int, int])

  • device (device)

  • generator (Generator | None)

  • fixed_offset (bool)

Return type:

PatchLayout

denoise_patches(x, sigma, layout, params, generator=None)[source]¶

Denoise and reassemble one randomly offset PaDIS patch layout.

Parameters:
  • x (Tensor)

  • sigma (Tensor)

  • layout (PatchLayout)

  • generator (Generator | None)

Return type:

Tensor

fixed_overlap_patch_layout(image_shape, params)[source]¶

Build the deterministic fixed-overlap comparison layout.

Parameters:

image_shape (tuple[int, int])

Return type:

PatchLayout

denoise_fixed_overlap_patches(x, sigma, layout, params, *, assembly, generator=None)[source]¶

Denoise patches and combine them by averaging or stitching.

Parameters:
  • x (Tensor)

  • sigma (Tensor)

  • layout (PatchLayout)

  • assembly (Literal['fixed_average', 'fixed_stitch'])

  • generator (Generator | None)

Return type:

Tensor

denoise_whole_image(x, sigma, params)[source]¶

Denoise a complete image with a whole-image prior.

Parameters:
  • x (Tensor)

  • sigma (Tensor)

Return type:

Tensor

denoise_prior(x, sigma, params, image_shape, generator=None, *, layout_override=None)[source]¶

Dispatch denoising to the configured patch or whole-image prior.

Parameters:
  • x (Tensor)

  • sigma (Tensor)

  • image_shape (tuple[int, int])

  • generator (Generator | None)

  • layout_override (PatchLayout | None)

Return type:

Tensor