denoise.model3d

3D U-Net for Noise2Inverse volumetric denoising.

Architecture source: SSD_3D (Laugros et al., bioRxiv 2025)

“Self-supervised image restoration in coherent X-ray neuronal microscopy” https://doi.org/10.1101/2025.02.10.633538

Original U-Net implementation: ELEKTRONN3 (Martin Drawitsch, MPG)

https://github.com/ELEKTRONN/elektronn3 Based on https://github.com/jaxony/unet-pytorch (Jackson Huang, MIT License)

Modifications in this file:

• Removed test utilities

• Added unet3d() factory function for N2I-compatible instantiation

class denoise.model3d.UNet(*args: Any, **kwargs: Any)

Bases: Module

3D U-Net with skip connections for volumetric image restoration.

Input: [B, in_channels, D, H, W] Output: [B, out_channels, D, H, W]

Spatial dimensions must each be divisible by 2**n_blocks.

forward(x)

forward_gradcp(x)

Forward pass with gradient checkpointing (saves ~20-50% memory).

static weight_init(m)

denoise.model3d.unet3d(in_channels: int = 1, out_channels: int = 1, n_blocks: int = 3, start_filts: int = 32, normalization: str = 'layer', activation: str = 'relu') → UNet

Factory function matching the configuration used in Laugros et al. 2025.

Input/output shape: [B, 1, D, H, W] D, H, W must each be divisible by 2**n_blocks (e.g. 64 for n_blocks=3).

Parameters:

• in_channels (int) – Number of input channels (1 for single-channel volumes).

• out_channels (int) – Number of output channels (1 for denoising).

• n_blocks (int) – Encoder depth. Controls receptive field and GPU memory use. Default 3 is a good balance; use 4 for deeper context (needs more memory).

• start_filts (int) – Filter count of the first encoder block. Subsequent blocks double this.

• normalization (str) – Normalisation type: ‘layer’, ‘group’, ‘instance’, ‘batch’, or ‘none’.

• activation (str) – Activation function: ‘relu’, ‘leaky’, ‘silu’, etc.