[논문 리뷰] SFT-GAN : Spatial Feature Transform

[Recovering Realistic Texture in Image Super-resolution by Deep Spatial Feature Transform]

https://openaccess.thecvf.com/content_cvpr_2018/papers/Wang_Recovering_Realistic_Texture_CVPR_2018_paper.pdf

 

Super-resolution 분야에 해당하는 논문이다. Transform이란 단어 때문에 Transformer model을 응용한 논문인 줄 알았으나 아니었다. 아무튼 super-resolution의 이전 모델들까지 공부해야 제대로 이해할 수 있어서, 새로운 분야를 파고들 수 있던 좋은 기회였다.

 

---

<1.> Introduction

 

기존의 super-resolution에 딥러닝 모델을 도입한 연구들로 인해 많은 발전이 있었다. 그러나 pixel-wise MSE loss를 사용하다 보니, blurry 하거나 overly-smooth 한 결과들이 나타나는 것은 여전한 한계점으로 남아있었다. 이에 대한 해결책으로 perceptual loss가 도입되었고, (SRGAN에서 사용함) 이 한계점도 극복할 수 있었고 visual quality가 크게 향상되었다.

 

그러나 texture recovery는 Super Resolution에서 중요한 문제로 남아있었고, 이 논문에서는 어떻게 하면 texture recovery를 더 자연스럽게 할 수 있을지에 대한 해결책을 제시한다.

 

이 논문에서 제시한 해결책은 바로 categorical prior. 논문에서는 이미지의 각 영역에 대한 semantic class에 대한 정보를 담고 있는 것이라 했는데, 쉽게 말하면 픽셀이 어떤 class에 속하는지를 미리 아는 것이다. 이렇게 prior를 제공하는 것은 다르게 말해서 texture에 대한 constraint 주고 더 만족스러운 결과가 나오게끔 하는 것으로 생각해볼 수 있다. 

 

그렇다면 이 semantic한 정보를 어떻게 줄 것인가부터 시작한다. CNN으로 이를 해결하기에는, 한 이미지 내에서는 많고 다양한 클래스, 크기의 객체들이 있기 때문에 굉장히 까다로웠다. 따라서, semantic segmentation map을 categorical map으로 사용하기로 한다. semantic segmentation은 pixel 단위로 class에 대한 정보를 충분히 담고 있기 때문이다.

 

그다음은 이 prior를 어떻게 네트워크 안에 잘 녹여내서 texture를 더 잘 생성할 것인가? 에 대한 문제로 넘어간다. 논문에서 spatial feature transform(이하 SFT)이라는 방법을 해결책으로서 사용한다. 이후에 다시 살펴보겠지만, 이 방법으로 네트워크의 feature map에 affine transformation을 적용할 수 있다는 것이 효과라고 한다.

 

SFT의 장점으로 3가지를 소개하고 있다. 1) 파라미터면에서 효율적이다. 2) 기존의 SR 네트워크에 도입될 수 있다. 3) depth map과 같은 prior로도 확장시킬 수 있다.

 

---

<2.> Methodology

 

모델에 대한 수식적인 설명인데 풀어서 써보겠다.

 

Super resolution에 쓰이는 CNN 기반 모델은 low resolution image x가 주어졌을 때, ground truth에 가까운 high resolution image로 mapping 하는 함수라고 볼 수 있다. 학습을 통해 이 함수의 parameter(weight)들을 update 하는 것이다. (1번 수식)

 

논문에서는 결국 하고자 하는 것은 categorical prior을 사용하자는 것인데, 그래서 조건부 확률처럼 categorical prior를 알고 있고, x가 주어지는 것으로 표현한 것 같다. (4번 수식) 

 

쉽게 이해하자면 semantic segmentation 정보와 low resolution image 를 갖고 high resolution image를 결과로 만들 수 있도록 모델을 학습시키는 것이다.

 

이제 구체적인 방법으로 넘어가기 전에, 대부분의 SR model에서는 low resolution에서 대부분의 연산들이 수행되고, 거의 마지막쯤 upsampling을 통해 resolution을 키운다.

 

또한 모델 전체적으로, SR-GAN과 매우 비슷한 구조를 하고 있다.

---

<3.> Spatial Feature Transform

SFT layer

위 그림처럼 레이어의 configuration까지 구체적으로 서술되어 있는 것은 supplementary에서 확인할 수 있다. SFT layer의 input, 즉 condition은 모든 SFT layer에 동일하게 들어간다. efficiency를 위해 이렇게 적용했다고 한다.

 

segmentation map을 1x1 크기의 convolution layer들에 통과시킨 결과들을 사용한다. 왜 1x1을 사용하냐면, CNN 특성상 이웃하는 픽셀들을 고려하기 때문에 서로 다른 category들이 영향을 줄 수 있는 것을 막기 위함이다. (예를 들어, 두 영역이 맞닿는 경계선을 처리하는 경우를 생각해볼 수 있겠다.) 

 

SFT layer에서는 우선 condition을 2개의 branch를 통과시킨다. 하나는 r(감마)를 만들기 위한 branch고 다른 하나는 b(베타)를 만들기 위한 branch다. 각 branch 또한 위에서 언급한 이유와 동일한 이유로, 1x1 convolution을 사용한다. 감마와 베타 둘 다 convolution layer를 통과한 feature map과 크기가 동일하다.

 

이 feature map과 감마에 대해 element-wise multiplication(Hadamard multiplication)을 수행하고, 그 결과에 베타를 더해준다. 

Affine transformation by SFT layer

이 과정은 feature map을 scaling하고 shift 하는 과정이라고 소개하고 있다.

 

논문을 읽으면서 가장 많이 고민한 부분은 이 SFT layer가 결국 어떤 역할을 할 것으로 생각했길래 이렇게 설계를 했을까... 였다. 

 

나름대로 고민해본 결과를 생각해보면, 기하학적으로 Affine transformation이라 하면 scaling이나 rotation, shift, 등을 해도 선행성과 평행이 보장되는 변환을 의미한다. 여기서는 feature map에 감마와 베타를 각각 곱하고 더해도 feature map의 크기와 차원 수가 그대로 유지되기 때문에 이런 면에서도 affine transfomation이라 할 수 있겠다. 공간적인 정보들을 처리하게 되니까 spatial affine transformation인가...

 

categorical prior관점에서는 image feature에 semantic segmentation 정보를 pixel-wise 하게 곱하고 더하는 방식으로 추가해준다고 생각해보았다. 즉, SFT-layer와 conv layer로 이뤄진 residual block을 통과할수록 feature map에는 image의 feature에 대한 정보뿐 아니라 어떤 class에 속하는지에 대한 정보도 있게 된다. 이러한 면에서 class-conditional SR model이 될 수 있다고 보았다.

 

---

 

<4.> GAN-based SR model

model의 마지막에 upsampling하고 다시 convolution layer들을 통과시켜서 최종 결과물, 즉 high resolution image를 만든다. 이 결과가 잘 만들어졌는지를 어떻게 평가할 것인가? 가 중요한 대목이다.

 

이전의 연구들은 앞서 언급했듯이 ground truth와 비교해서 pixel-wise MSE loss를 사용했다. 여기서는 SR-GAN에서 사용했던 perceptual loss와 Discriminator를 사용한 GAN loss를 사용한다.

 

perceptual loss는 간략히 설명하면 '사람이 보기에 자연스러운 이미지인가?'를 loss로 어떻게 나타낼까? 에 대한 해결로 등장한 loss다. SR model의 feature map(마지막쯤 있는 feature map을 사용함)과 Ground truth를 다른 pre-trained 된 feature extractor model에 통과시켰을 때 나오는 feature map(역시 마지막쯤 있는 것을 사용) 사이의 차이를 MSE loss로 나타낸다.

 

GAN loss(adversarial loss)는 생성한 high-resolution image를 discriminator에 통과시켰을 때, real인지 fake인지 구분하는지에 대한 loss다. Discriminator는 SR로 만든 image를 fake로, Ground truth image를 real로 구분하도록 학습된다.

 

이미지 해상도가 크면 real-fake에 대한 구분이 좀 더 힘들어지지 않을까?하고 생각해서 봤더니 Dataset을 좀 더 crop 해서 low-resolution으로는 24x24, high-resolution으로는 96x96을 사용한다.

 

---

<5.> Result & Discussion

Result

결과를 보면 확실히 훨씬 자연스럽다. SRGAN과는 크게 차이가 없어보이지만, 왼쪽의 결과와는 확연한 차이가 보인다. MSE loss를 사용하지 않고 perceptual loss, adversarial loss 뿐 아니라 SFT layer를 도입했기 때문이다. Ablation study에선 이 SFT layer의 성능에 대해서도 연구를 했는데 확실히 성능 향상이 있었다.

 

다만 읽으면서 들었던 모델의 아쉬운 점은 segmentation map이라는 다른 input 또한 준비를 해야 한다는 점과, image resolution이 96x96 보다 큰 것에 대해서는 언급이 없었다는 점이다.