spyrit.external.drunet.DRUNet

class spyrit.external.drunet.DRUNet(noise_level=5, n_channels=1, nc=[64, 128, 256, 512], nb=4, act_mode='R', downsample_mode='strideconv', upsample_mode='convtranspose', normalize=True)

Bases: UNetRes

Plug-and-Play Image Restoration with Deep Denoiser Prior. DRUNet is a pretrained plug-and-play denoising network that has been pretrained for a wide range of noise levels. It admits the noise level as an input, so it does not require training. DRUNet was proposed in the work

[ZhLZ21] K. Zhang et al., Plug-and-Play Image Restoration with Deep Denoiser Prior. In: IEEE Transactions on Pattern Analysis and Machine Intelligence, 44(10), 6360-6376, 2021.

Original Code: https://github.com/cszn/DPIR/blob/master/models/network_unet.py

Args:

noise_level (float): noise level value in the range [0, 255]. This is used to create a noise level map that is concatenated to the input images.

n_channels (int): number of image channels

nc (list of int): number of features

nb (int): number of residual blocks

act_mode (str): activation function mode

downsample_mode (str): downsample mode

upsample_mode (str): upsample mode

normalize (bool): Normalize to (-1,1). Defaults to True.

Input / Output:

x: input images with shape (\(B\), n_channels, \(H\), \(W\)).

output: denoised images with shape (\(B\), n_channels, \(H\), \(W\)).

normalize (bool): Normalize to (-1,1) if True.

Attributes:

noise_level (tensor): noise level tensor with shape \((1)\).

Note

DRUNet has been tested only with n_channels =1 but UNetRes can be used with n_channels >1.

Methods

concat_noise_map(x)	Concatenation of noise level map to reconstructed images
forward(x)	Define the computation performed at every call.
set_noise_level(noise_level)	Reset noise level value