[논문 리뷰] SUNet: Swin Transformer UNet for Image Denoising

Chi-Mao Fan, Tsung-Jung Liu외 1명|arXiv (Cornell University)|2022. 02. 28.

Image and Signal Denoising Methods인용 수 4

한 줄 요약

이 논문은 장거리 상관관계와 전반적 맥락을 포착하기 위해 기존의 컨volutional 레이어를 Swin Transformer 블록으로 대체한 Swin Transformer 기반 U-Net 아키텍처인 SUNet을 제안한다. 이 모델은 기준 데이터셋에서 최고 성능을 기록하며 PSNR와 SSIM이 향상되었고, 체스기타 무늬 아티팩트를 줄이고 복원 품질을 향상시키기 위해 새로운 이중 업샘플링 블록을 도입하였다.

ABSTRACT

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.

연구 동기 및 목표

• CNN 기반 이미지 디노이징 모델의 한계, 즉 국소적 특징 종속성과 고정 커널 적용을 해결하기 위해.
• 특히 전반적 맥락 모델링이 핵심적인 이미지 디노이징과 같은 저수준 시각 작업에서 Swin Transformer의 잠재력을 탐색하기 위해.
• 특징 표현과 복원 정밀도를 향상시키기 위해 Swin Transformer 블록으로 강화된 U-Net 아키텍처를 설계하기 위해.
• 서브픽셀 및 이중선형 업샘플링 방법을 조합한 이중 업샘플링 블록을 도입하여 특징 업샘플링 시 발생하는 체스기타 무늬 아티팩트를 완화하기 위해.
• 표준 벤치마크에서 검증된 U-Net 프레임워크 내에서 Swin Transformer 백본을 사용하여 이미지 디노이징 분야의 새로운 최고 성능을 확립하기 위해.

제안 방법

• SUNet은 U-Net 인코더와 디코더의 표준 컨볼루션 레이어를 Swin Transformer 블록으로 대체하여 장거리 상관관계를 모델링하고 전반적 맥락을 포착한다.
• 주요 Swin Transformer 기반 U-Net 인코더-디코더 경로로 전달하기 전에 초기 특징을 추출하기 위해 얕은 3×3 컨볼루션 레이어를 사용한다.
• 체스기타 무늬 아티팩트를 줄이고 특징 복원 품질을 향상시키기 위해 서브픽셀 및 이중선형 업샘플링을 조합한 이중 업샘플링 블록을 도입한다.
• Swin Transformer 블록은 이동 윈도우를 사용하여 계산 복잡도를 감소시키면서도 국소 및 전반적 수신 영역에서 자기주의 주목을 유지한다.
• DIV2K 데이터셋에서 추가된 가우시안 백색 잡음(AWGN)을 다양한 수준으로 적용한 패치를 사용하여 평균 제곱오차(MSE) 손실을 기반으로 엔드 투 엔드로 학습한다.
• 최종 출력은 3×3 컨볼루션 레이어를 통해 디노이징된 이미지를 재구성한다.

실험 결과

연구 질문

• RQ1 Swin Transformer 기반 U-Net 아키텍처가 기존의 CNN 기반 및 U-Net 기반 모델보다 이미지 디노이징 작업에서 승리할 수 있는가?
• RQ2 Swin Transformer를 U-Net 아키텍처에 통합할 경우 특징 표현과 복원 품질에 어떤 영향을 미치는가?
• RQ3 제안된 이중 업샘플링 블록이 전치 컨볼루션 업샘플링에 비해 체스기타 무늬 아티팩트를 효과적으로 줄일 수 있는가?
• RQ4 기존의 U-Net 변종에 비해 더 적은 파라미터와 FLOPs로도 경쟁 가능한 성능을 달성할 수 있는가?
• RQ5 모델은 다양한 잡음 수준과 실제 이미지 데이터셋에 대해 얼마나 잘 일반화되는가?

주요 결과

• CBSD68 데이터셋에서 σ=50 잡음 조건에서 SUNet은 PSNR 36.79 dB, SSIM 0.953을 기록하여 비교된 모든 CNN 기반 및 U-Net 기반 모델을 능가한다.
• Kodak24 데이터셋에서 σ=50 조건에서 SUNet은 PSNR 29.54 dB, SSIM 0.810을 달성하여 다양한 이미지 세트에 걸쳐 강력한 일반화 능력을 보였다.
• 대부분의 기준 모델보다 높은 9900만 개의 파라미터를 가짐에도 불구하고, DHDN 및 RDUNet에 비해 FLOPs는 3% 감소하고 파라미터 수는 60% 감소하여 효율성 향상을 보였다.
• 이중 업샘플링 블록은 시각적 비교를 통해 더 매끄러운 질감과 더 선명한 에지로 체스기타 무늬 아티팩트를 크게 줄임을 입증하였다.
• CBSD68 데이터셋에서 σ=50 조건에서 SUNet은 최고의 SSIM 0.958을 기록하여 더 뛰어난 정서적 품질과 구조적 정밀도를 보였다.
• 모델은 모든 잡음 수준(σ=10, 30, 50)에서 일관된 개선을 보이며 기존 방법에 비해 PSNR 및 SSIM에서 뚜렷한 향상을 보였다.