[논문 리뷰] Multi-Stage Progressive Image Restoration, Syed Waqas Zamir et al., 2021 | MPRNet 아키텍쳐 분석

https://arxiv.org/pdf/2102.02808.pdf

Multi-Stage Progressive Image Restoration, Syed Waqas Zamir et al., 2021을 읽고 MPRNet 아키텍쳐를 분석한 글입니다.

---

외부 환경에서 촬영된 영상은 비나 눈, 안개 등과 같은 날씨/환경에 의해 영상 내의 객체가 변형되거나 흐려지게 된다. 이러한 degradation 된 이미지를 복원하기 위해 기존에 제시되었던 CNN(Convolutional neural network)은 single stage가 적용되었다. 하지만 이미지를 복원하는 작업은 공간의 세부 정보와 높은 수준의 contextualized 정보 간의 균형이 매우 중요시되기 때문에 single stage가 아닌 Multi stage 구조를 적용한 MPRNet(Multi-Stage Progressive Image Restoration)이 제시되었다.

MPRNet의 구조 , https://arxiv.org/pdf/2102.02808.pdf

Multi-stage image restoration은 아래와 같은 architectural bottleneck들이 존재한다.

첫째, Encoder-decoder는 높은 수준의 contextualized 정보를 보존하기 어려우며 Single-scale pipeline은 공간의 세부 정보는 정확하지만 안정적이지 못한 결과물을 도출하게 된다. 때문에 효과적인 이미지 복원을 위해 Multi-stage 구조에서는 Encoder-decoder와 Single-scale pipeline을 선택적으로 조합하는 것이 필요하다.

둘째, 단순하게 One stage의 output을 다음 stage로 전달하는 것은 suboptimal result를 낳게 된다. 점진적인 복원을 위해선 각 stage에서의 supervision이 중요하다.

셋째, Multi-stage processing동안 Encoder-decoder branch에서 contextualized 특징을 보존하기 위해 intermediate features를 이전 stage에서 이후 stage로 propagate하는 매커니즘이 필요하다.

 

MPRNet은 이러한 architectural bottleneck을 극복하기 위해 아래와 같은 구조를 사용한다.

(1) Encoder-decoder Subnetwork

초기 단계에는 U-Net을 기반으로 한 Encoder-decoder Subnetwork를 사용하여 Multi-scale contextual 정보를 학습한다.

Encoder-decoder Subnetwork, https://arxiv.org/pdf/2102.02808.pdf

 

각 scale에서 feature를 추출하기 위해 CABs(channel attention blocks)를 추가하였고 U-Net skip connections에서 feature maps은 CAB와 함께 진행하였다. Decoder 단계에서는 feature의 공간 해상도를 향상시키기 위해 한 개의 convolution에서 bilinear upsampling 방법을 사용하였다.

 

(2) Original Resolution Subnetwork (ORSNet)

마지막 단계에는 ORSNet을 도입하여 input 이미지에서 output 이미지까지 세부 정보들을 보존하였다. ORSNet은 downsampling operation을 사용하지 않고 고해상도 feature를 만들어냈다. 이는 ORBs(original resolution blocks)로 구성되어 있으며 각 block에는 CABs가 추가로 포함되어 있다.

Original Resolution block, https://arxiv.org/pdf/2102.02808.pdf

(3) Supervised Attention Module

모든 두 stage 사이에 ground-truth에 의해 다음 stage로 전달되기 전, 이전 stage의 feature를 개선하기 위해서 Supervised attention map(SAM)을 도입했다. SAM은 각 stage에서 점진적인 이미지 복원에 유용한 ground truth supervisory 신호를 제공하며 현재 stage에서 덜 중요한 feature들을 억제하고 유용한 feature들만 다음 stage로 전달하도록 attention map을 생성한다.

Supervised Attention Module, https://arxiv.org/pdf/2102.02808.pdf

 

Supervised Attention Module을 사용한 feature 전달 과정, https://arxiv.org/pdf/2102.02808.pdf

 

(4) Cross-stage Feature Fusion (CSFF)

stage1과 stage2 사이의 CSFF (왼쪽), stage2와 마지막 stage 사이의 CSFF(오른쪽)

 

Cross-stage Feature Fusion(CSFF) 모듈은 Encoder-decoder에서 Up-sampling과 Down-sample을 반복적으로 사용해 네트워크의 정보 손실을 줄인다. 또 한 stage의 Multi-scale feature들은 다음 stage의 feature들을 풍부하게 하는데 도움이 되며 공간 정보의 흐름을 안정적으로 만들어 전체 구조에 여러 stage를 추가할 수 있게 된다. 이러한 장점들이 있는 CSFF 모듈을 위와 같이 두 Encoder-Decoer 사이와 Encoder-decoder와 ORSNet 사이에 도입했다.

 

---

 

MPRNEt Github : https://github.com/swz30/MPRNet

GitHub - swz30/MPRNet: [CVPR 2021] Multi-Stage Progressive Image Restoration. SOTA results for Image deblurring, deraining, and

해당 논문을 읽고 깃허브 클론을 하여 모델을 사용해 봤다.

 

1) MPRNet Github에 명시된 내용에 따라 환경을 세팅한다.

 

2) 데이터셋을 준비한다.

MPRNet을 적용하여 해당 모델을 테스트하기 위해 All-Weather 데이터셋을 사용하였다. 이 데이터셋은 Outdoor-Rain, RainDrop, Snow100K가 포함된 총 18,069장의 이미지로 구성되어 있으며 원본 이미지(ground truth)와 원본 이미지에 인위적으로 노이즈를 형성하거나 카메라 렌즈에 물방울을 뿌리는 등의 효과를 준 input 이미지가 함께 포함되어 있다. MPRNet에서의 deraining 기법을 사용하기 위해 All-Weather 데이터셋에서 Outdoor-Rain과 RainDrop 데이터만 테스트 이미지로 사용하였다.

 

3) MPRNet을 적용 결과

All-Weather 데이터셋의 RainDrop 데이터 중 빗줄기가 있는 이미지에 MPRNet의 deraining 기법을 적용한 결과는 아래와 같다. deraining 기법이 적용되면서 강한 빗줄기가 대부분 제거되었음을 볼 수 있다.

All-Weather 데이터셋 RainDrop 샘플 이미지

 

deraining 기법을 적용한 RainDrop 샘플 이미지

 

4) 느낀 점

MPRNet의 deraining 기법을 적용하기 위해 해당 모델을 학습하기 위해서는 input 이미지와 함께 ground truth가 필요하다. 때문에 ground truth를 구하기 힘든 실생활 영상에서는 이미지가 제대로 복원이 되지 않아 화질이 저하되거나 빗방울이 제대로 제거가 되지 않는다는 단점이 있다.

 

 

---

논문을 읽은 후 주관을 가지고 요약, 정리한 글이므로 해당 포스팅에 대한 이의 제기, 다른 의견 제시 등 다양한 지적, 의견은 언제나 대환영입니다!

댓글로 남겨주시면 늦더라도 확인하겠습니다!!!