논문스터디 : Towards Controllable and Photorealistic Region-wise Image Manipulation

(본 논문은 Swapping Autoencoder 논문의 직접적인 개선을 목표로 했기 때문에, 해당 정리를 읽으시기 전에 Swapping Autoencoder 리뷰 를 먼저 읽으시는 걸 추천드립니다.)

Intro

Swapping Autoencoder (이하 Swap.Ae) 는 딥러닝 기반 이미지 편집을, 이미지를 미시적으로 표현하는 2차원적 구조적 정보 (이하 content, Swap.Ae 에서는 structure로 표현)와, 거시적으로 표현하는 1차원적 스타일, 텍스쳐 정보 (이하 style, Swap.Ae 에서는 texture로 표현)으로 나눠 생각해서 해결하고자 했습니다.
(가령 성당 이미지의 경우, 성당의 위치와 구조물의 생김새, 창문 등의 디자인적 요소는 content에 해당하며, 이미지의 색조나 채도, 건물의 질감 등의 요소는 style에 해당합니다.)
(다만 이미지에 따라 해당 두 요소가 항상 말끔하게 떨어지는건 아니며, 그럼에도 최대한 말이 되는 방향으로 나누는 것이 이 모델의 목표입니다.)

이 두 요소가 나눠진 상태에서 content / style 에 해당하는 Latent code를 서로 다른 이미지들끼리 바꾸는 것으로 image translation 이 가능하고
(예 : 봄을 배경으로 한 들판에서 content를, 눈 덮인 설산을 style 로 삼는 것으로 눈이 덮인 들판 이미지를 생성 가능), 나아가서 2차원 content space 상에서 latent code를 조작하는 것으로 image editing이 가능합니다.
(예 : 산과 강이 있는 이미지에서 강에 해당하는 content 부분을 산이 있는 곳에 옮겨다 붙여넣는 걸로 강의 크기를 키우는 것이 가능).

이러한 editing이 자연스럽게 성립하기 위해서는 content와 style간의 분리가 제대로 된 것이 중요합니다; 그렇지 못한 경우, image translation 결과가 content 나 style 둘 중 하나에만 크게 영향을 받게 되어 좋은 결과가 나오기 어려울 것입니다. 저자들은 Swapping Autoencoder의 훈련이 대체로 잘되는 편이지만, 여러 인퍼런스 결과에서 content / style 간 분리가 제대로 되지 않아 구조적 일관성이 Image translation 문제에서 잘 지켜지지 않는 등, 분리를 위한 추가적인 조치들이 필요함을 지적합니다.

Swapping Autoencoder 실패 케이스의 한 예시. Content / Style 중에서 Style에 비중이 더 쏠려 있어 의미있는 편집이 어려운게 확인됩니다.

이를 위해 저자들은 self-supervised training에 기반한 2개의 loss를 추가적으로 제시합니다. 첫 번째로 code consistency loss는, 서로 다른 이미지들의 content / style code 로 생성된 이미지를 다시 encoder에 넣어서 content / style code를 획득했을 때 뽑힌 latent code 들이 입력에 사용한 code 들과 같은지를 계산합니다. 이 loss를 통해 content / style 의 일관성을 추가적으로 확보할 수 있었습니다.

두 번째로 content alignment loss는, 이미지의 마스크 단위 변환을 위해서 추가된 loss 입니다. (예 : 얼굴의 피부나 눈, 입술만 선택적으로 바꾸는 경우) 훈련과정에서 랜덤하게 생성한 흑백 마스크를 통해 두 이미지의 latent code들을 합치고, 이렇게 생성된 이미지가 (마스크 범위 안에서) 얼마나 이전 이미지들과 같은지를 계산해서, 구하기 힘든 semantic image dataset 없이도 비슷한 훈련을 수행할 수 있게끔 합니다.

이 두개의 loss term을 추가한 모델을 저자들은 TCP라고 부르며, FFHQ, LSUN-Church 같은 데이터셋에 대해 높은 성과를 기록한다고 주장합니다.

모델 구조

기본적인 모델 구조는 Swapping Autoencoder에서 크게 달라지지 않았기 때문에, 모델 구조에 대한 설명은 Swapping Autoencoder 리뷰 을 참고하시기 바라며,
해당 리뷰에서는 새롭게 제안된 2개의 loss에 대해서만 적도록 하겠습니다.

• 논문에서 StyleGAN2 Generator 구조를 따라갔다고 서술하는데, swap.ae 에서는 StyleGAN2에 여러 변형을 거친 (예 : output skip 제거, residual block 구조), StyleGAN2 Discriminator 구조와 유사한 모양을 가지고 있어 다를 수 있습니다.

Code consistency loss

기존의 Swap.Ae 는 Encoder 단계에서 style에서의 구조적 정보를 없애기 위해. average pooling을 통해 style을 1차원 벡터로 만들어 (비교적 단순하게) content / style간 분리를 진행했습니다. 그러나 앞선 예시에서 보인것처럼 여전히 style에 따라 구조적인 정보가 변하는 경우들이 존재하고, 저자들은 이는 content의 정보 상당수가 style로 섞여 들어가기 때문에 발생한다고 주장합니다.

Whitening and Coloring Batch Transform for GANs 에서 제시된 feature whitening을 content code에 사용해 content / style 간 분리를 강화할 경우, 중요한 구조적 정보가 사라진다는 문제가 발생하기 때문에 사용하지 않았습니다.

대신, 저자들은 비교적 단순한 self-supervised 방법론을 사용합니다; Encoder가 convergence를 이루었으면, 이미지 A에서 나온 content code와 B에서 나온 style code를 합쳤을 때 생성된 이미지 C를 다시 Encoder에 넣었을때, 추출된 content code 는 생성에 사용한 A의 content code와 같아야 하고, style code 역시 생성에 사용한 B의 style code와 같아야 하는 겁니다.

위의 가설을 충족하기 위해선 Encoder가 convergence를 이룬 후에 해당 loss를 적용해야 하지만, 실제로는 K iterations (K=16) 마다 적용하는 것으로 해당 문제를 우회합니다. (진짜로?)

BicycleGAN 에서 제시한 Conditional Latent Regression GAN 모델과 유사한 loss이지만, 정해진 Gaussian Distribution 이 아닌, 실제 입력값에서 나온 latent code를 기반으로 합니다.

Content alignment loss

Region-wise style transfer는, binary mask 가 입력 이미지들과 같이 주어진 상태에서, 마스크로 가려진 부분은 원래의 content / style을 유지하고, 그렇지 않은 부분은 content를 유지하면서 style만 바꿔지는걸 목표로 합니다.

이론적으론 별다른 훈련 없이 global style transfer를 마친 이미지를 마스크로 곱해서 더하면 되지만, 그러할 경우 마스크의 경계선 부분에 artifact가 생기는걸 막을 수 없기 때문에 저자들은 훈련중에 입/출력 이미지가 아닌, feature map 차원에서부터 마스크를 적용하여 훈련하는 것으로 이 문제를 해결하고자 합니다. (이 때, 마스크는 랜덤하게 생성합니다.)