Source-linked AI summary
SUNet: Swin Transformer UNet for Image Denoising
Chi-Mao Fan, Tsung-Jung Liu, Kuan-Hsien Liu
TL;DR
Image restoration is an ill-posed problem, and CNN-based methods have limitations in content-independent filtering and long-range dependency modeling. SUNet applies Swin Transformer layers within a UNet for image denoising, using dual up-sampling to avoid checkerboard artifacts and achieving competitive denoising results.
Problem
Image restoration is an ill-posed problem, while CNN convolution kernels can be content-independent and capture local rather than global information.
Method
SUNet integrates Swin Transformer Blocks into a UNet architecture for image denoising and adds a dual up-sample block combining subpixel and bilinear methods.
Results
SUNet achieves competitive denoising results on two common datasets and maintains good PSNR and SSIM scores with fewer parameters than UNet-based comparators.
Takeaways & Limitations
The paper reports that Swin Transformer remains promising for image restoration, while stating that it is too early to conclude it can replace convolution.
Abstract
from arXiv · showhide
Image restoration is a challenging ill-posed problem which also has been a long-standing issue. In the past few years, the convolution neural networks (CNNs) almost dominated the computer vision and had achieved considerable success in different levels of vision tasks including image restoration. However, recently the Swin Transformer-based model also shows impressive performance, even surpasses the CNN-based methods to become the state-of-the-art on high-level vision tasks. In this paper, we proposed a restoration model called SUNet which uses the Swin Transformer layer as our basic block and then is applied to UNet architecture for image denoising. The source code and pre-trained models are available at https://github.com/FanChiMao/SUNet.
I. INTRODUCTION
Image restoration addresses corrupted images that can affect downstream vision tasks, while CNN-based methods face limitations in content-independent local filtering and long-range dependency modeling. SUNet integrates Swin Transformer with UNet for denoising, introducing dual up-sampling and reporting competitive results on two common datasets.
- Image restoration can improve performance in object detection, image segmentation, and image classification.
- CNN-based restoration methods use content-independent convolution kernels and may lose global information during long-range dependency modeling.
- SUNet integrates Swin Transformer with UNet for image denoising, following the paper’s claimed first incorporation of these components in denoising.
- The proposed dual up-sample block combines subpixel and bilinear up-sampling to prevent checkerboard artifacts and outperforms transpose-convolution up-sampling in experiments.
- SUNet achieves competitive results on two common image-denoising datasets.
II. RELATED WORK
Learning-based methods, particularly CNNs, became mainstream in image restoration because of their reported performance advantages over conventional methods.
- Learning-based methods defeat conventional model-based methods in execution speed and performance as hardware develops.
A. Image Restoration
Traditional restoration methods rely on priors or algorithms but can be slow, computationally expensive, and difficult to apply to complex textures. Learning-based methods, especially CNNs, became mainstream because of their performance.
- Traditional image restoration methods use image priors or algorithms, including self-similarity, sparse coding, and total variation.
- Conventional restoration methods are time-consuming, computationally expensive, and difficult to use for restoring complex image textures.
- Learning-based methods, especially CNNs, became mainstream in image restoration because of their impressive performance.
B. UNet
UNet provides hierarchical multi-scale features and encoder–decoder skip connections for image reconstruction, while Swin Transformer addresses transformer limitations in vision through shifted windows and supports pixel-wise tasks.
- UNet uses hierarchical feature maps to obtain rich multi-scale contextual features and skip connections to enhance image reconstruction.
- UNet is widely applied to segmentation and restoration and can incorporate different extractive blocks through its adaptive backbone.
- Directly applying transformers to vision faces scale and dense-prediction challenges, while Swin Transformer reduces parameters with shifted windows and achieves state-of-the-art performance in pixel-wise tasks.
A. SUNet
SUNet combines shallow feature extraction, Swin-Transformer-based UNet feature extraction, and convolutional reconstruction to denoise noisy images. It takes a noisy image as input and produces a noise-free image.
- Architecture: SUNet contains shallow feature extraction, UNet feature extraction, and reconstruction modules.The shallow module extracts low-frequency color or texture information, while the UNet module extracts high-level multi-scale deep features.
- Shallow feature extraction: A 3 × 3 convolution extracts shallow features from the noisy input image.The resulting shallow feature has spatial resolution H × W and C channels, with C set to 96 in the experiments.
- UNet feature extraction: The UNet feature extractor uses Swin Transformer Blocks to obtain high-level and multi-scale deep features.Each Swin Transformer Block contains 8 Swin Transformer Layers and replaces convolutions in the UNet architecture.
- Reconstruction: A final 3 × 3 convolution reconstructs the noise-free image from the deep features.The network output is the estimated clean image, while the ground-truth clean image corresponds to the noisy input’s reference image.
B. Swin Transformer Block
The Swin Transformer Block replaces traditional convolution layers with paired window-based and shifted-window self-attention layers. Its layers also use normalization, multilayer perceptrons, residual connections, and cyclic shifting.
- Attention structure: Each Swin Transformer Layer pair uses window multi-head self-attention and shifted-window multi-head self-attention.The number of layers is always a multiple of two, assigning one layer to each attention type.
- Block design: Swin Transformer Blocks substitute traditional convolution layers in the UNet feature-extraction module.The block contains 8 Swin Transformer Layers in the paper’s experiments.
- Shifted windows: Cyclic shifting reduces computation time while retaining characteristics of convolution, including translation invariance.
- Layer operations: Each layer applies Layer Normalization, a two-layer GELU-activated MLP, and residual additions around its transformations.The prose defines LN as Layer Normalization and MLP as a multilayer perceptron with two fully connected layers and GELU activation.
C. Resizing module
SUNet uses patch merging for down-sampling and a dual up-sample module for up-sampling. The dual module combines bilinear and sub-pixel methods to address checkerboard artifacts associated with transpose convolution.
- Resizing modules: SUNet uses patch merging for down-sampling and dual up-sampling for up-sampling between UNet feature-map scales.These resizing modules are necessary because UNet processes feature maps at different scales.
- Down-sampling: Patch merging concatenates each group of 2 × 2 neighboring patches and applies a linear layer to produce the desired output channels.The operation can also be viewed as the unfolding step of a convolution.
- Up-sampling: The proposed dual up-sample module combines bilinear up-sampling and PixelShuffle up-sampling.The figure identifies the second method as sub-pixel up-sampling.
- Up-sampling: The dual up-sample design is intended to prevent checkerboard artifacts caused by transpose-convolution-like patch expansion.
IV. EXPERIMENTS
SUNet is trained end-to-end without pretrained networks and evaluated quantitatively with PSNR and SSIM, where higher values indicate better performance.
- Implementation Details: SUNet is an end-to-end trainable model implemented in PyTorch 1.8.0 without pretrained networks.The implementation uses a single NVIDIA GTX 1080Ti GPU.
- Evaluation Metrics: The quantitative evaluation uses Peak Signal-to-Noise Ratio and Structure Similarity Index metrics.PSNR is measured in decibels, and both metrics are higher-is-better.
B. Experiment Datasets
SUNet is trained on DIV2K image patches with additive Gaussian noise and evaluated on the CBSD68 and Kodak24 denoising datasets. Visual comparisons include CBSD68 image 126007 corrupted with AWGN at σ = 50.
- Training Set: DIV2K provides 800 training images and 100 testing images for the denoising experiments.The training images have an average resolution of about 1920 × 1080.
- Evaluation: Table I reports dataset-average denoising scores, with FLOPs measured on 256 × 256 color images.Best and second-best scores are highlighted and underlined, respectively.
- Visual Comparison: Figure 4 compares denoising outputs for CBSD68 image 126007 corrupted by AWGN with σ = 50.PSNR and SSIM values below the subfigures are calculated by patches.
- Training Set: Training uses 256 × 256 patches with randomly added AWGN at noise levels from σ = 5 to σ = 50.Validation uses σ = 10, σ = 30, and σ = 50.
- Testing Set: Evaluation uses CBSD68 with 68 color images and Kodak24 with 24 images.CBSD68 images have resolution 768 × 512, while Kodak24 images have size 321 × 481.
C. Image Denoising Performance
SUNet combines Swin Transformer layers with a UNet architecture for image denoising and achieves competitive results across the reported comparisons. It also reduces computational costs relative to UNet-based models, while real-world noise and blur remain future targets.
- Comparison with CNN-based Methods: SUNet achieves the best PSNR and SSIM results among the compared CNN-based methods with almost the same FLOPs.The comparison includes DnCNN, IrCNN, and FFDNet.
- Comparison with UNet-based Methods: SUNet uses 60% fewer parameters and 3% fewer FLOPs than the compared UNet-based methods while retaining good PSNR and SSIM scores.The comparison is among SUNet, DHDN, and RDUNet.
- Model Complexity: SUNet has 99M parameters because its self-attention operation does not share kernel weights across layers.The paper presents different kernel values for features in different layers as more reasonable.
- Architecture: SUNet combines a Swin Transformer backbone with UNet architecture for image denoising.The architecture is presented as a restoration model using Swin Transformer layers as its basic blocks.
- Scope and Future Work: The authors identify real-world noise and real-world blur as future restoration tasks while retaining Swin Transformer layers.The conclusion states that it is too early to say Swin Transformer can replace convolution.