LION.models.CNNs.drunet

Classes

DRUNet([model_parameters])	DRUNet (Zhang et al., 2021): the denoiser of Deep Plug-and-Play Image Restoration, a U-Net of residual blocks over four scales, without biases by default, optionally taking the noise level as an extra input channel.
ResBlock([n_channels, kernel_size, ...])	A residual block of DRUNet: two convolutions with act between them, added to the input.

Functions

downsample_strideconv([in_channels, ...])	A strided convolution that halves the image size (DRUNet's downsampling).
upsample_convtranspose([in_channels, ...])	A transposed convolution that doubles the image size (DRUNet's upsampling).

---

class LION.models.CNNs.drunet.DRUNet(model_parameters=None)

Bases: LIONmodel

DRUNet (Zhang et al., 2021): the denoiser of Deep Plug-and-Play Image Restoration, a U-Net of residual blocks over four scales, without biases by default, optionally taking the noise level as an extra input channel. Takes images.

model_parameters (default_parameters()):

• in_channels, out_channels: 1, 1;

• int_channels: the first scale’s channels, doubled at each scale (64);

• kernel_size: 3x3;

• n_blocks: residual blocks per scale (4);

• use_noise_level: the noise level as an input channel, as in the paper (False);

• bias_free: no biases (True);

• act: a torch.nn.functional activation (“leaky_relu”);

• enforce_positivity: a ReLU on the output (False).

Initialize internal Module state, shared by both nn.Module and ScriptModule.

static default_parameters()

The paper’s architecture (64 channels, 4 residual blocks per scale, no biases), without the noise level input by default (see the class).

classmethod pretrained(verbose=1)

The authors’ trained grayscale DRUNet (drunet_gray.pth, from their KAIR repository): trained on natural images in [0, 1] with Gaussian noise levels up to 50/255, given as the noise level input. Their architecture is this one with act="relu" and use_noise_level=True.

The weights (130 MB) are downloaded once into torch’s hub cache (torch.hub.get_dir()) and checked against their hash; verbose (1) shows the download’s progress. The model is on LION’s device. Give it images in [0, 1] (e.g. set_normalisation) and noise levels in those units.

static cite(cite_format='bib')

Where DRUNet comes from: its citation, BibTeX (“bib”) or “MLA”.

forward(x0, noise_level=None)

noise_level (with use_noise_level): one number, or one per image of the batch.

class LION.models.CNNs.drunet.ResBlock(n_channels=64, kernel_size=(3, 3), bias_free=True, act=<function leaky_relu>)

Bases: Module

A residual block of DRUNet: two convolutions with act between them, added to the input.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

x plus the block’s two convolutions of it.

LION.models.CNNs.drunet.downsample_strideconv(in_channels=1, out_channels=64, kernel_size=2, stride=2, bias_free=True)

A strided convolution that halves the image size (DRUNet’s downsampling).

LION.models.CNNs.drunet.upsample_convtranspose(in_channels=64, out_channels=1, kernel_size=2, stride=2, bias_free=True)

A transposed convolution that doubles the image size (DRUNet’s upsampling).