[Review] Multi-Domain Image Completion for Random Missing Input Data (하루논문 P12)
하루에 논문 한편씩, 핵심만 읽고 요약하기 : Paper 12
Paper Link : Link (CVPR 2020 Workshop)
1. Abstract
Multi-domain 데이터는 다양한 modalities으로부터 보완적인 정보의 장점을 이용하기 위해 컴퓨터 비전에서 널리 사용되고 있습니다. 하지만, 데이터의 손상이나 각 이미지의 프로토콜 간의 차이로 인하여 각 domain의 이미지들을 이용할 수 있는지 여부는 데이터 종류에 따라 크게 차이날 수 있으며 이러한 한계는 범용적인 모델을 만드는 것을 더욱 어렵게 합니다. 이러한 문제를 해결하기 위하여, 실제 applications에서 random missing-domain을 완성할 수 있는 GAN을 이용한 새로운 multi-domain 이미지 생성 방법을 제안합니다. 해당 방법은 다양한 domain이 모두 공유하는 "skeleton"과 각 domain끼리 구분되는 "flesh"를 추출하기 위한 representational disentanglement 개념을 사용합니다. 더 나아가 shared content encoder를 이용하여 image completion과 segmentation이 가능한 unified framework를 통해 해당 방법이 segmentation과 같은 high-level task에도 활용될 수 있는 것을 보여줍니다. 해당 방법을 이용하여 총 3가지 데이터셋(brain tumor segmentation, prostate segmentation, 그리고 facial expression image completion)에 실험을 진행하였습니다.
2. Introduction
Multi-domain이미지는 서로 보완적인 정보를 제공해줄 수 있기 때문에 다양한 컴퓨터 비전 task에서 사용되고 있습니다. 예를 들어, 4종류의 MRI modality (T1, T1-weighted, T2-weighted, FLAIR)는 brain tumor segmentation에서 tumor 영역을 정확하게 segment하기 위한 standard protocol로 사용되고 있습니다.
하지만, 이 중 몇몇 domain이 누락될 수 있는 것이 현실입니다. 특히 대규모 다기관 연구의 경우, 모든 기관에서 누락되는 domain 없이 데이터를 수집하는 것은 일반적으로 매우 어려운 일입니다. 이러한 경우를 고려하여 누락된 domain의 이미지를 생성하고 완성하는 알고리즘을 만드는 것은 매우 중요한 일입니다. 직관적으로 생각할 수 있는 방법 중 하나는 누락된 이미지에 대해서 동일 domain에 속하는 다른 이미지들을 바탕으로 nearest neighbor과 같은 단순한 방법으로 새로운 이미지를 생성해내는 것입니다. 하지만 이러한 방법은 단순히 pixel 수준의 유사성만 고려하기 때문에 semantic consistency가 부족하다는 한계가 있습니다.
위의 그림은 nearest neighbors로 생성한 이미지와 논문에서 제시하는 방법으로 생성한 이미지
각각을 이용하여 segmentation prediction을 했을 때의 결과입니다.
위에서 언급했던 바와 같이 nearest neighbor를 이용했을 때에는 semantic한 정보가 부족하기
때문에 ground truth와 매우 동떨어진 결과를 보이고 있습니다.
해당 논문에서는 Representational disentanglement scheme for Multi-domain Image Completion (ReMIC)에 기반한 general n-to-n image completion framework를 제안합니다.
3. Related Work
- Image-to-Image Translation
1. CycleGAN : Cycle-consistency loss를 이용하여 unpaired dataset에 대해 image-to-image translation을 가능하게 한 논문입니다. 하지만 해당 논문은 1-to-1 mapping만 가능하다는 한계가 있습니다.
2. StarGAN & RadialGAN : 단순히 두개의 domain 사이의 mapping만 가능한 CycleGAN과 달리 multi-domain image generation이 가능한 모델입니다. 하지만 두 모델 모두 1-to-n mapping만 가능하기 때문에 생성된 결과가 오직 하나의 input image을 source domain으로 사용하여 제한적이라는 한계가 있습니다.
3. CollaGAN : StarGAN과 RadialGAN과 마찬가지로 multi-domain mapping이 가능하지만 두 모델과 달리 CollaGAN은 n-to-1 mapping이 가능하다는 장점이 있습니다. 즉, 새로운 이미지를 생성할 때 단순히 하나의 이미지가 아니라 해당 이미지를 제외한 나머지 domain의 이미지를 모두 종합(collaborate)하여 사용한다는 특성을 갖고 있습니다.
- Learning Disentangled Representations
1. DRIT & MUNIT : 두 모델은 image translation을 수행할 때 두 개의 domain이 공유하는 content와 서로 구별되는 attribute features를 disentangle하는 방법을 사용합니다. 하지만 CycleGAN과 마찬가지로 1-to-1 mapping만 가능하다는 한계가 있습니다.
2. Liu et al : 해당 논문 또한 multi-domain image translation을 다루긴 하였으나, 각 domain에 특화된 스타일에 대한 논의는 거의 없었으며 domain adaptation을 위한 cross-domain latent code에 대한 학습에 집중하였습니다.
- Medical Image Synthesis
Missing modalities가 존재할 때 segmentation을 하기 위해 multi-domain에서 어떻게 representations을 얻어낼 것인가에 대한 연구(Nguyen et al, Havaei et al, Chartasias et al)는 이전에도 있었습니다. 하지만 이러한 연구들은 representation disentanglement의 관점이 아니라 단순히 multiple modalities에서 얻어낸 features를 종합(fuse)해주는 식으로 구성했다는 점에서 차이가 있습니다.
위의 그림은 해당 논문에서 제시하는 방법(ReMIC)과 다른 방법들을 비교하여 나타낸 것입니다.
다른 기존의 방법들과 달리 ReMIC은 1-to-1, n-to-1, 1-to-n 심지어 n-to-n이 가능한 것을 확인할 수 있습니다.
4. Method
3.1 Unified Image Completion and Segmentation
Fig. 3을 참고하면 ReMIC은 unified content encoder 
위의 그림은 왼쪽 상단 4개의 이미지를 input으로 사용하였을 때 content encoder을 통해 얻어낸 content codes를 시각화한 결과입니다. Content code의 output channels 256개 중에서 임의로 8개를 골라서 시각화하였다고 합니다. 각 channel이 종양, 뇌, 두개골과 같은 서로 다른 해부학적인 구조를 각기 반영하고 있는 것을 확인할 수 있습니다.
이렇게 얻어낸 style code와 content code를 바탕으로 generator 
추가로 저자들은 missing data 문제가 있는 경우에도 segmentation을 하는 것과 같이 image completion framework를 더욱 실용적인 상황에도 연장하기 위하여 segmentation generator 
3.2 Training Loss
5. Experiments & Results
실험에 사용한 데이터셋은 BraTs, ProstateX, 그리고 RaFD 총 3가지 데이터셋입니다.
5.1 Results of Multi-Domain Image Completion
해당 분석은 크게 두가지로 구분되어 진행되었는데 'multi domains 중 한 개의 domain만 missing된 경우 (n-to-1)'와 '1개 이상의 domain이 missing된 경우 (n-to-n)'로 나뉘어집니다. 그리고 사용된 evaluation metrics는 NRMSE, SSIM, PSNR 총 3가지입니다. (여기서 IS나 FID와 같이 최근에 generative model을 평가할 때 사용하는 metric을 사용하지 않은 점이 개인적으로 의문입니다.)
1~5번째 methods는 오직 하나의 domain만 missing된 경우의 결과이고,
나머지 3개의 methods는 1개 이상의 domain이 missing된 경우입니다. (k=* : *개 만큼의 domain이 input image로 존재하는 경우)
여기서 ReMIC w/o Reconstruction은 ReMIC에 reconstruction loss를 사용하지 않은 모델입니다.
BraTs와 ProstateX 데이터셋에 대해서는 ReMIC이 모든 경우에서 기존 모델보다 성능이 좋은 것을 확인할 수 있으며 심지어 1개 이상의 domain이 missing된 경우에도 매우 좋은 지표를 보이고 있습니다.
다음은 정성적으로 결과를 확인해보도록 하겠습니다.
왼쪽 (a)에서 첫번째 열 Target에 있는 4개의 이미지는 원본 이미지이며 나머지 4개의 열에 있는 결과는 각 model을 사용하였을 때의 결과입니다.
기존 모델인 MUNIT, StarGAN, CollaGAN에 비해서 논문에서 제시하는 ReMIC이 정성적으로 확인하기에도 ground truth인 원본 이미지와 가장 비슷한 결과를 보여주고 있으며 특히 tumor와 그 주변부가 잘 반영되어 있습니다.
오른쪽 그림은 ReMIC을 바탕으로 1개 이상의 domain이 missing 되었을 때를 예측한 결과입니다.
위의 table과 마찬가지로 k=*일 때, *개의 domain이 visible한 경우인데
1개의 domain만 사용가능한 k=1 경우에도 (a)에 있는 ground truth 이미지와 흡사한 것을 확인할 수 있습니다.
아래는 ProstateX와 RaFD에 대한 결과입니다.
5.2 Results of Missing-Domain Segmentation
Missing-domain image가 있을 때의 segmentation 정확도를 비교한 결과입니다. 이때 사용한 평가 지표는 Dice score입니다.
Oracle은 missing 이미지 없이 fully supervised로 학습된 2D U-Net style의 network입니다.
그리고 Oracle + *는 missing domain image를 * 방법을 이용하여 생성해내고 해당 이미지를 바탕으로 pre-trained Oracle 모델로 예측하였을 때를 나타낸 것입니다.
ReMIC+Seg는 image generation에 사용한 ReMIC의 content encoder와 별개로 segmentation을 위한 encoder를 두어서 학습한 결과이고,
이와 달리 ReMIC+Joint는 3.1에서 언급했던 것과 같이 image completion에 사용한 content encoder를 sharing하여 generation과 segmentation을 joint training한 경우입니다.
모든 경우에 대해서 ReMIC이 기존 방법보다 좋은 성능을 보이며 특히 BraTS 데이터셋의 경우에는 joint training한 것이 가장 좋은 성능을 보였습니다.
6. Conclusion
해당 논문에서는 1개 이상의 input domain이 missing되었을 때에도 사용가능한 multi-domain image completion framework을 제안하였습니다. 그리고 representational disentanglement를 이용하여 모든 domain이 공유하는 content code와 각 domain에 specific한 style code를 구분하여 학습하였습니다. 기존의 연구가 1-n, n-1, 혹은 1-1이었던 것에 비해 해당 논문은 실제로 인공지능을 사용할 때 벌어질 수 있는 상황인 n-n case에 대해서도 사용가능한 모델을 제시하였다는 것이 인상깊었습니다.