[Paper Review] StyleGAN3-Editing: Third Time’s the Charm? Image and Video Editing with StyleGAN3 논문 리뷰

업데이트:


StyleGAN3는 기존 모델(StyleGAN ver1, 2)의 문제점이었던 texture sticking의 문제를 해결하여 비디오를 효과적으로 생성할 수 있게 한 모델이다.

🤔 그렇다면, 기존 StyleGAN2 모델에서 파생되었던 다양한 연구들(ex. image editing, gan inversion)은 StyleGAN3에서도 적용이 될까?

본 논문에서는 StyleGAN3로 변화하면서 어떤 이점이 생겼는지, 그리고 기존의 연구들은 어떻게 변화할 것인지를 분석한다.

  • StyleGAN3의 latent space가 StyleGAN2에 비해 얼마나 disentanglement한지를 비교
  • 기존 image editing/GAN-inversion 연구들은 stylegan3에서 잘 동작하는지 분석
  • unaligned image에 대해서도 inversion을 할 수 있는 StyleGAN3 Encoder
  • 실험을 해보니 aligned G가 unaligned G보다 disentangle하게 inversion되고 editing이 잘된다고 함 → 이미지 생성에 aligned G 사용
  • StyleGAN3로 비디오를 inversion & recon 할 수 있는 방법론 제안


1. The StyleGAN3 Architecture

1.1 Latent Space

  • StyleGAN2
    • initial latent code $z \sim \mathcal{N}(0,1)^{512}$ → MLP → intermediate code $w$
  • StyleGAN3
    • StyleGAN2는 어떤 해상도의 이미지를 생성하냐에 따라 $w$의 dimension이 다름
      • 1024 - 18x512 , 256 - 14x512
    • StyleGAN3: output image resolution과 상관없이 $w$는 16개로 고정

      \[\left(w_{0}, \ldots, w_{15}\right)\]

1.2 Constant Input → Fourier Feature

  • StyleGAN2: 4x4x512 learned constant input

  • StyleGAN3: Fourier Feature

    • 왜 Fourier Feature를 사용했는지는 다음 링크를 참고
    • Fourier Feature를 input으로 사용하기 때문에 rotation과 translation이 용이
      • $w_{0}$ → learned affine layer → Fourier Feature
      • $w_{i}$ → learned affine layers에 넣어서 conv kernel weight를 조절
    • 이를 식으로 정리하면

      $y=G\left(w ;\left(r, t_{x}, t_{y}\right)\right)$

Image Editing

  • 기존 StyleGAN2는 intermediate latent code $w$를 조절함으로써 이미지를 editing 한다.
  • 이와 달리 StyleGAN3는 rotation이나 translation이 된 이미지를 얻기 위해 Fourier Feature를 조절

    ⇒ 😊 aligned image만으로 학습된 Generator라도, Fourier Feature을 조절하면 rotation이나 translation된 이미지 생성 가능!

$G_{unaligned}$ 로 aligned된 이미지 생성

Unaligned image로 학습된 Generator의 경우

  • Fourier Feature의 input으로 $w_{0}$를 사용하면 unaligned image가 생성되나
  • $(w_{1}, \ldots, w_{15})$ 의 평균, $\bar{w}$을 Fourier Feature의 input으로 사용하면 강압적으로 aligned이 된 이미지가 생성됨
    • Why? (1) training distribution의 input pose의 평균이 센터에 align 되어있을 수도 있고, (2) 주로 첫번째 layer에서 translation과 rotation 여부가 결정되기 때문에 이를 조절하면 아예 align이 된 것처럼 보이는 이미지가 생성됨


2. Analysis

2.1 Rotation Control

$w_{0}$는 이미지의 translation과 rotation을 결정하는 중요한 요소이다.

🤔 그렇다면 $w_{0}$ 가 어느 layer까지 이미지의 t&r를 변화시킬 수 있을까?

Experiment 01

  • $w_{1}$만을 변화시키고 나머지$(w_{0},w_{2},w_{3}…)$ 는 그대로 둠

    ⇒ (결과) 얼굴의 rotation이 조금씩 변화가 있었음 + face shape, eye와 같은 다른 특성들에도 변화가 있었음

Experiment 02

  • $w_{0}, w_{1}$은 고정하고 나머지 code $w^{*}$ 만 random으로 생성
\[G\left(\left(w_{0}, w_{1}, w^{*}, \ldots, w^{*}\right)\right)\]

⇒ (결과) translation이나 rotation은 그대로이고 identity만 바뀜

(결론) $w_{0}$는 이미지의 translation과 rotation을 결정하며, $w_{1}$도 이미지의 rotation에 일정부분 관여한다.

2.2 Disentanglement Analysis

위 논문에서 제안한 DCI metric을 통해 latent space의 disentanglement함을 측정한다.

  • D (Disentanglement): latent space가 얼마나 disentangle한지
    • 특정 방향으로의 editing이 용이해야함(특정 identity를 바꾸고자 했을 때 다른 특징도 변화하면 안됨)
  • C (Conpleteness): single code의 변화가 어떤 특성을 얼마나 변화시키는지
  • I (Informativeness): generative factor에 대해 code가 얼마나 유익한지(유의미한 변화를 끼치는지)

StyleGAN2, 3 모두 latent space가 StyleSpace에서 가장 disentanglement하며,

StyleGAN3의 W space가 StyleGAN2의 W space에 비해 별로..


3. Image Editing

3.1 Editing via Linear Latent Directions

InterfaceGAN 처럼 latent space에서 선형적 방향을 찾은 다음 이를 수정하여 이미지를 editing하는 방법

aligned image의 경우

  • stylegan2와 같은 방법으로 이미지 editing
\[G_{\text {aligned }}(w+\delta D ;(0,0,0))\]

unaligned image의 경우

  1. unaligned generator 사용

    • classifier를 사용하여 특정 attribute에 관여하는 방향 찾기
      • (문제점) 사용한 pre-trained classifier가 aligned image로 학습되어있어서 unaligned image에 대해서는 좋은 classification score를 구하지 못함 → latent space W 에서 이상한 direction을 학습
      • Fourier Feature의 input으로 $w_{0}$ 대신 $\bar{w}$ 를 써서 강제로 align된 듯 한 이미지를 생성한 다음에 pretrained classifier로 latent’s attribute score를 구함
  2. aligned generator 사용

    • aligned generator를 사용하는 대신 $(r, t_{x}, t_{y})$ 값을 조절

      ⇒ aligned & unaligned 이미지를 모두 생성할 수 있으며, 이들 모두 같은 latent direction 조절하여 이미지를 editing 할 수 있게 됨

      \[G_{\text {aligned }}(w+ \left.\delta D ;\left(r, t_{x}, t_{y}\right)\right)\]

Result

  • 실험 결과 $G_{unaligned}$의 latent space가 $G_{aligned}$보다 entangle함을 확인

    ⇒ $G_{aligned}$를 이용하여 차후 연구를 진행

3.2 Editing via Non-Linear Latent Paths

latent space에서 비선형적으로 latent code를 조절하면 더 disentangle하게 이미지 editing이 가능하다.

  • 관련논문
    • Styleflow: Attribute-conditioned exploration of stylegan-generated images using conditional continuous normalizing flows (2020)
    • Guidedstyle: Attribute knowledge guided style manipulation for semantic face editing (2020)

W+ space를 사용하는 StyleCLIP mapper를 봐도 이미지를 editing 했을 때 배경까지 같이 변화함 → 속성들끼리 entangle 한 것

⇒ 저자들은 W나 W+ 보다 disentangle한 S space를 사용하기로 결론

3.3 Editing via Latent Directions in S

StyleSpace에서의 code를 조절함으로써 이미지를 editing하면(Fig 7) W나 W+ space에서 이미지를 editing 했을 때(Fig 6)에 비해 disentangle하게 image editing이 가능하다.

  • Fig 6에 비해 Fig 7 는 원하는 속성만 딱 변화하고 배경은 거의 변하지 않음

4. StyleGAN3 Inversion

Optimize 방식을 사용하면 다음 loss를 통해 이미지를 inversion할 수 있다.

\[\hat{w}=\underset{w}{\arg \min } \mathcal{L}\left(x, G\left(w ;\left(r, t_{x}, t_{y}\right)\right)\right.\]

그러나 optimize를 해서 image를 inversion하면, 비디오로 봤을 때 consistency가 떨어진다.

⇒ Encoder로 사용!

encoder를 통해 여러 이미지들(비디오)을 inversion하면 이미지들 간에 consistency가 있기 때문에 비디오로 봤을 때 자연스럽다.

다음 loss function을 통해 여러 이미지들에 대해 inversion 할 수 있는 Encoder를 학습하였다.

\[\sum_{i=1}^{N} \mathcal{L}\left(x_{i}, G\left(E\left(x_{i}\right)\right)\right)\]
  • image inversion: $w=E\left(x\right)$
  • edited image: $y_{\text {edit }}=G\left(f(w) ;\left(r, t_{x}, t_{y}\right)\right)$

4.1 Designing the Encoder Network

  • Unaligned Generator로 Unaligned image에 대해 inversion 시킬 수 있는 Encoder를 학습시켜봄
    • 아래의 식: Encoder를 학습시키는 loss function
    • Generator: pretrained unaligned G
\[\sum_{i=1}^{N} \mathcal{L}\left(x_{\text {unaligned }}^{i}, G\left(E\left(x_{\text {unaligned }}^{i}\right)\right)\right)\]

(결론) unaligned image는 이미지의 분포가 너무 다양하기 때문에 inversion & recon이 어려움

그렇다면 aligned image에 대해서만 inversion하도록 Encoder를 단순화해보자!

→ Encoder는 unaligned image에 비해 비교적 단순한 분포를 가지는 이미지(=aligned image)를 인코딩하기 때문에 input의 identity를 좀 더 잘 포착할 수 있음

→ aligned image를 inversion하더라도 Generator를 이용하여 unaligned image처럼 만들 수 있음

🤔 여전히 StyleGAN3의 장점을 살려서 unaligned image에 대해서도 encoding을 하고 싶은데.. 이 경우에는 어떻게 할까?

4.1.1 Aligned Encoder로 Unaligned image inversion하기

  1. 우선 unaligned image $x_{\text {unaligned }}$에 대해 align을 맞춤 → $x_{\text {aligned }}$

  2. 1의 과정을 통해 align이 맞춰진 이미지에 대해 encoder로 inversion

    \[w_{\text {aligned }}=E\left(x_{\text {aligned }}\right)\]
  3. $x_{\text {unaligned }}$에서 얼마나 rotation & translation $(r, t_{x}, t_{y})$ 을 하면 $x_{\text {aligned }}$가 되는지를 계산

    How? (1) 두 이미지( $x_{\text {unaligned }}$와 $x_{\text {aligned }}$)에서 눈을 detect → 이를 비교하여 얼마나 rotation이 되었는지를 계산

    (2) (1)에서 계산한 값을 토대로 $x_{\text {unaligned }}$를 rotation

    (3) rotation이 된 $x_{\text {unaligned }}$와 $x_{\text {aligned }}$의 거리를 비교하여(vertical & horizontal distance) 얼마나 translation 되었는지를 계산

  4. 3에서 계산한 $(r, t_{x}, t_{y})$을 토대로 이미지 생성

    \[y_{\text {unaligned }}=G\left(w_{\text {aligned }} ;\left(r, t_{x}, t_{y}\right)\right)\]

4.1.2 Encoder는 어떤 구조를 사용할까?

저자들은 pSp, e4e, restyle Encoder를 사용했다고 한다.

🤔 저자들이 이 세 가지 encoder들의 코드 hierarchy가 똑같아서 얘들로만 비교 실험한 것 같다.(실험이 쉬우니까..) 다른 idinvert나 다른 stylegan encoder들도 실험해봤으면 좋았을 텐데.. 아숩..

aligned image에 대해서 inversion하는 Encoder를 학습시키므로 Encoder 구조는 StyleGAN2와 동일하다.

나머지는 다 그대로고 Encoder 학습에 필요한 Generator만 변화

  • 이전에는 StyleGAN2 였다면, 이번에는 pretrained aligned StyleGAN3 Generator를 사용
  • 모든 encoder들은 7만장의 FFHQ dataset으로 학습됨

pSp와 e4e 학습

  • 훈련 디테일: Due to the larger memory consumption required by StyleGAN3, our encoders are trained using a batch size of 2. To match the batch size used in the official implementations of the StyleGAN2 pSp and e4e encoders, we apply gradient accumulation to attain an effective batch size of 8 (i.e., an optimization step is performed every four batches). All encoders are trained using a single NVIDIA P40 GPU.
  • Loss function
\[\mathcal{L}(x)=\lambda_{l 2} \mathcal{L}_{2}(x)+\lambda_{\text {lpips }} \mathcal{L}_{L P I P S}(x)+\lambda_{i d} \mathcal{L}_{\mathrm{id}}(x)\]

restyle 학습

  • restyle encoder는 progressive하게 training하는데 저자들은 이를 안했다고 함
  • 대신 첫 training부터 16 latent codes를 한번에 예측! (이렇게 했더니 더 빠르게 수렴했다고 report)

4.2 Inverting Images into StyleGAN3

  • test set으로는 CelebA-HQ 를 사용

Qualitative Evalution

  • StyleGAN3가 2와 비교해도 괜찮게 이미지를 생성함
  • StyleGAN3: unaligned source image에 대해 align을 맞추고 encoding 했다가 계산된 r&t 값을 바탕으로 unaligned image를 생성하는데, 생성된 이미지가 꽤 괜찮음

Quantitative Evalution

  • StyleGAN3의 score가 StyleGAN2의 score보다 안좋음
    • 저자들은 이 원인이 (StyleGAN3에게) entangle한 W+ space으로 inversion 했기 때문이라고 주장.
    • 저자들은 S space로 inversion하면 StyleGAN3도 꽤 괜찮은 성능을 가질거라고 생각하나봄

Editability via Latent Space Manipulation

  • $Restyle_{e4e}$ encoder: 이미지의 identity도 보존되면서 image editing이 매우 잘됨
  • $Restyle_{pSp}$ encoder: 이미지 recon은 잘하지만, editing 성능이 떨어짐
  • 이런 성능차이가 StyleGAN3에서 더 심함

5. Inverting and Editing Videos

비디오 inversion은 매우 어렵다. Why?

  1. 비디오의 각 프레임들을 recon & editing 했을 때 consistency를 유지하기 어려움
  2. 눈을 감고 있거나 입을 벌리고 있다던가 하는 어려운 표정들이 존재

StyleGAN3는 다양한 위치의 이미지를 생성할 수 있고, texture sticking도 없기 때문에 비디오 생성에 적합하다. 다만, 비디오 자체에 제약은 있다. (unaligned image라고 하더라도 주어진 범위(fixed bounding box) 내에서 crop이 된 이미지여야함 - ex. vox1 dataset)

unaligned image에 대해 inversion & recon하는 과정은 다음 과정을 따르며, 이를 통해 비디오가 생성된다.

\[y_{i}=G\left(w_{i} ;\left(r_{i}, t_{x, i}, t_{y, i}\right)\right)\]

비디오 전체를 editing할 수도 있음

\[y_{i, edit}=G\left(f(w_{i}) ;\left(r_{i}, t_{x, i}, t_{y, i}\right)\right)\]

5.1 Latent Vector Smoothing

r&t 좌표로 이미지를 다시 원위치로 보낸다고 해도, 여러 이미지들간의 consistency를 확인하고 inversion하는게 아니라 이미지 각각에 대해서 inversion & recon을 하는 것이기 때문에 비디오를 recon했을 때 연속된 프레임 간에 inconsistency가 생길 수 있다. (shaking issue)

저자들은 이를 해결하기 위해 latent code들을 smoothing하는 과정을 진행한다.

\[\begin{aligned}w_{i, \text { smooth }} &=\sum_{j=i-2}^{i+2} \mu_{j} f\left(w_{j}\right) \\T_{i, \text { smooth }} &=\sum_{j=i-2}^{i+2} \mu_{j} T_{j}\end{aligned}\]

where

\[\left[\mu_{i-2}, \mu_{i-1}, \mu_{i}, \mu_{i+1}, \mu_{i+2}\right]=\frac{1}{3}[0.25,0.75,1,0.75,0.5]\]

5.2 PTI (Pivotal Tuning Inversion) method 적용

저자들은 unaligned image를 만들기 위해 PTI method를 채택했다고 한다.

PTI 논문 컨셉이 과녁을 맞추려고 노력하기 보다는 과녁의 위치를 옮겨주겠다 이런 컨셉인데..

저자들도 이 모델의 아이디어를 차용하여 강제로 unaligned image가 만들어지도록 Generator를 fine-tuning 해주었다.

\[y_{i, \mathrm{PTI}}=G_{\mathrm{PTI}}\left(w_{i} ;\left(r_{i}, t_{x, i}, t_{y, i}\right)\right)\]

결과 비디오를 잘 만드려고 engineering을 엄청 열심히 한 듯한….

Detail

For inverting and editing a given video we perform a pervideo fine-tuning of the StyleGAN3 generator network using the pivotal tuning technique from Roich et al. [56]. For each video, training is performed for a total of 8, 000 optimization steps with a batch size of 2 using the L2 pixel-wise loss and the LPIPS [56] loss, both with equal weight coefficients. For example, given a video consisting of 200 frames, each frame is observed an average of 40 times during training. During training, we do not alter the weights of the input Fourier features layer of the generator.

5.3 Bringing It All Together

마지막으로는 smoothed edited latent code, smoothed transformation, 그리고 fine-tuned generator를 다같이 사용!

\[y_{i, \text { final }}=G_{\mathrm{PTI}}\left(w_{i, s m o o t h} ; T_{i, s m o o t h}\right)\]

정리하자면..

  1. align & crop the frame: Cropped Frame($x_{\text {unaligned }}$)에서 align을 맞춰서 Aligned Frame($x_{\text {aligned }}$)을 얻음
  2. computing the Fourier feature transformations
    • Predict Transforms 모듈로 얼마나 rotation & translation 되었는지를 계산
    • 이때, 확장된 view에 대해서도 이미지를 생성하기 위해 추가로 y 방향으로 translation이 더 된 이미지에 대해서도 생성을 한다.
  3. align된 이미지($x_{\text {aligned }}$) Encoder로 inversion + editing 할거면 w 수정 (→ f(w))
  4. PTI-tuned StyleGAN3를 이용하여 이미지 생성
  5. 두 이미지를 합쳐서 확장된 edited result image 구하기

Result

5.4 Expanding the Field of View

논문에서 굉장히 흥미로운 파트이다. StyleGAN3의 이점을 톡톡히 활용한 것 같다.

저자들은 이미지를 shift해서 더 넓은 view에 대해 영상을 생성하였다.

\[\begin{aligned}y &=G_{\mathrm{PTI}}\left(w_{i, \text { smooth }} ; T_{i, s m o o t h}\right) \\y_{\text {shift }} &=G_{\mathrm{PTI}}\left(w_{i, s m o o t h} ; T_{\Delta} \cdot T_{i, s m o o t h}\right)\end{aligned}\]

확장 전

확장 후


6. Opinion

🤗 StyleGAN2를 바탕으로 진행되었던 다양한 연구들을 StyleGAN3에 적용해보고 싶었는데, 발빠르게 논문이 나와서 재미나게 읽었던 논문이다.

결과도 잘 나오는 것 같고, 여러 실험 결과도 흥미로웠던 논문! 아쉬운 점들도 일부 보이지만.. 그래도 idea develop에 많은 도움이 되었다.

얼른 코드 공개되었으면!

댓글남기기