SwinIR

class deepinv.models.SwinIR(img_size=128, patch_size=1, in_chans=3, embed_dim=180, depths=[6, 6, 6, 6, 6, 6], num_heads=[6, 6, 6, 6, 6, 6], window_size=8, mlp_ratio=2, qkv_bias=True, qk_scale=None, drop_rate=0.0, attn_drop_rate=0.0, drop_path_rate=0.1, norm_layer=<class 'torch.nn.modules.normalization.LayerNorm'>, ape=False, patch_norm=True, use_checkpoint=False, upscale=1, img_range=1.0, upsampler='', resi_connection='1conv', pretrained='download', pretrained_noise_level=15, **kwargs)[source]

Bases: Module

SwinIR denoising network.

The Swin Image Restoration (SwinIR) denoising network was introduced in SwinIR: Image Restoration Using Swin Transformer. This code is adapted from the official implementation by the authors.

Parameters:

  • img_size (int, tuple) – Input image size. Default 128.

  • patch_size (int, tuple) – Patch size. Default: 1.

  • in_chans (int) – Number of input image channels. Default: 3.

  • embed_dim (int) – Patch embedding dimension. Default: 180.

  • depths (tuple) – Depth of each Swin Transformer layer.

  • num_heads (tuple) – Number of attention heads in different layers.

  • window_size (int) – Window size. Default: 8.

  • mlp_ratio (float) – Ratio of mlp hidden dim to embedding dim. Default: 2.

  • qkv_bias (bool) – If True, add a learnable bias to query, key, value. Default: True.

  • qk_scale (float) – Override default qk scale of head_dim ** -0.5 if set. Default: None.

  • drop_rate (float) – Dropout rate. Default: 0.

  • attn_drop_rate (float) – Attention dropout rate. Default: 0.

  • drop_path_rate (float) – Stochastic depth rate. Default: 0.1.

  • norm_layer (torch.nn.Module) – Normalization layer. Default: nn.LayerNorm.

  • ape (bool) – If True, add absolute position embedding to the patch embedding. Default: False.

  • patch_norm (bool) – If True, add normalization after patch embedding. Default: True.

  • use_checkpoint (bool) – Whether to use checkpointing to save memory. Default: False.

  • upscale (int) – Upscale factor. 2/3/4/8 for image SR, 1 for denoising and compress artifact reduction

  • img_range (float) – Image range. 1. or 255. Default: 1.

  • upsampler (str, None) – The reconstruction module. ‘’/’pixelshuffle’/’pixelshuffledirect’/’nearest+conv’/None. Default: ‘’.

  • resi_connection (str) – The convolutional block before residual connection. Should be either ‘1conv’ or ‘3conv’. Default: ‘1conv’.

  • pretrained (str, None) – Use a pretrained network. If pretrained=None, the weights will be initialized at random using PyTorch’s default initialization. If pretrained='download', the weights will be downloaded from the authors’ online repository https://github.com/JingyunLiang/SwinIR/releases/tag/v0.0 (only available for the default architecture). Finally, pretrained can also be set as a path to the user’s own pretrained weights. Default: ‘download’. See pretrained-weights for more details.

  • pretrained_noise_level (int) – The noise level of the pretrained model to be downloaded (in 0-255 scale). This value is directly concatenated to the download url; should be chosen in the set {15, 25, 50}. Default: 15.

forward(x, sigma=None)[source]

Run the denoiser on noisy image. The noise level is not used in this denoiser.

Parameters:

  • x (torch.Tensor) – noisy image, of shape B, C, W, H.

  • sigma (float) – noise level (not used).