+6

[Paper Explain] JoJoGAN for One Shot Face Stylization

# Intro

Style transfer là một bài toán thuộc lĩnh vực computer vision nhận được sự chú ý của nhiều nhà nghiên cứu bởi tính ứng dụng cao trong các ứng dụng chỉnh sửa ảnh áp dụng công nghệ AI. Việc huấn luyện một mạng nơ ron để thực hiện style transfer là rất khó khăn bởi vấn đề tìm kiếm dữ liệu. Paper JoJoGAN mà mình sẽ giới thiệu ở đây đã đề xuất một thủ tục để finetune mạng Generator của StyleGAN để thực hiện style transfer chỉ với duy nhất 1 ảnh reference.

Một số khái niệm cần biết

StyleGAN generator

Mạng generator của stylegan có 1 chút khác biệt với các mô hình GAN truyền thống. Thay vì nhận đầu vào trực tiếp là vector latent ziZz_i \in Z thì vector zz sẽ được đưa qua một mạng MLP 8 lớp để tạo ra vector wWw \in W với cùng số chiều. Vector ww được dùng để kiểm soát style của ảnh thông qua lớp Adaptive Instance Normalization. Một phép biến đối affline (học thông qua một lớp fully connected - ký hiệu A trong hình dưới) được áp dụng lên ww trước khi đưa vào mạng generator.

GAN inversion

Mạng generator của GAN sẽ tạo ra ảnh từ một vector latent zz. GAN inversion là quá trình ngược lại, từ ảnh đầu vào tìm vector latent tương ứng của nó.

image.png

Một số cách tiếp cận cho GAN inversion

  • Optimization based: sử dụng gradient descent để tối ưu vector zz sao cho ảnh khôi phục xrecx^{rec} giống với ảnh thật xx nhất.

z=argminz(x,G(z;θ))z ^ { * } = \underset { z } { \arg \min } \ell ( x , G ( z ; \theta ) )

Do quá trình tối ưu diễn ra lúc inference nên chi phí tính toán khá là tốn kém. Một số paper dựa trên cách tiếp cận này:

θE=argminθEnL(G(E(xn;θE)),xn)\theta _ { E } ^ { * } = \underset { \theta _ { E } } { \arg \min } \sum _ { n } L ( G ( E ( x _ { n } ; \theta _ { E } ) ) , x _ { n } )

Một số paper:

JoJoGAN

Workflow chính

image.png Workflow bao gồm 4 bước. Ký hiệu:

  • TT: ánh xạ GAN inversion
  • GG: generator của StyleGAN
  • ss: style parameter (output của biến đổi affline - các khối chữ A trong generator)
  • θ\theta: trọng số generator của StyleGAN.

Bước 1: GAN inversion. Một quan sát của tác giả trong quá trình inversion là với 1 ảnh style reference yy thì ảnh sau khi reconstruct lại trông khá là thật thay vì có nét stylized kiểu hoạt hình (step 1 trong hình trên). Nguyên nhân có thể do các mạng encoder trong GAN inversion được huấn luyện để tạo ra latent vector cho ảnh thật hơn là ảnh được stylized. Output của quá trình là vector w=T(y)w = T(y) trong không gian WW và một tập style parameter tương ứng s(w)s(w)

Bước 2: Tạo tập training. Dùng s(w)s(w) thu được ở bước 1 để tìm 1 tập style parameter SS khác gần giống với s(w)s(w). Việc tìm kiếm tập SS được thực hiện thông qua cơ chế style mixing của StyleGAN. Trong lúc huấn luyện StyleGAN, một phần ảnh trong tập training được tạo ra bằng cách sample 2 latent vector z1z_1z2z_2 thay vì một. Trong quá trình forward của mạng generator, ta lựa chọn ngẫu nhiên một điểm trong mạng để chuyển sử vector style w1w_1 sang w2w_2.
Style mixing trong JoJoGAN: Với kiến trúc StyleGAN2 gồm 26 modulation layer (Adaptive instance normalization), ta sẽ có style parameter kích thước 26×51226 \times 512. Ký hiệu: M{0,1}26M \in \{0, 1\}^{26} và vector mask; FC là lớp mạng mapping ZWZ \rightarrow W; ziN(0,1)z_i \sim N(0, 1). Tập style parameter mới được tính bằng công thức:

si=Ms+(1M)s(FC(zi))s _ { i } = M \cdot s + ( 1 - M ) \cdot s ( F C ( z _ { i } ) )

Lựa chọn vector M khác nhau sẽ dẫn đến các hiệu ứng khác nhau

image.png

Bước 3: Finetuning mạng generator để có G(si,θ^)yG(s_i, \hat \theta) \approx y

θ^=θargmin1NiNL(G(si;θ),y)\hat { \theta } = \stackrel { argmin } { \theta } \frac { 1 } { N } \sum _ { i } ^ { N } L \left ( G \left ( s _ { i } ; \theta \right ), y \right )

trong đó hàm mất mát L được sử dụng là perceptual loss. Thông thường perceptual loss được tính bằng backbone VGG được train trên ảnh kích thước 224×224224 \times 224, còn StyleGAN lại sinh ảnh kích thước 1024×10241024 \times 1024. Một hướng để xử lý là down sample ảnh sau khi sinh xuống để tính loss nhưng cách xử lý này sẽ dẫn tới mất mát đặc trưng ảnh. Hướng thứ 2 được sử dụng trong paper này là tính perceptual loss dựa trên activation của mạng discriminator. Quá trình training của StyleGAN giúp discriminator tính toán các feature map mà không bỏ qua các chỉ tiết nhỏ.

L(G(si;θ),y)=D(G(si;θ))D(y)1L ( G ( s _ { i } ; \theta ) , y ) = \| D ( G ( s _ { i } ; \theta ) ) - D ( y ) \| _ { 1 }

Bước 4: Inference. Với ảnh input uu, Ảnh được stylized được tính bằng ustylized=G(s(T(u));θ^)u_{stylized} = G ( s ( T ( u ) ) ; \hat { \theta } ). Ta cngx có thể tạo ra ảnh ngẫu nhiên với style chỉ định bằng cách sample từ phân phối N(0,1)N(0, 1) và đưa qua generator.

Một số biến thể

Controling identity

Một số ảnh style reference khiển ảnh input không giữ được identity (mất đi nhiều đặc trưng về structure). Trong trường hợp này, tác giả có sử dụng thêm hàm identity loss để bảo toàn thông tin về nhận dạng.

Lid=1sim(F(G(si;θ)),F(G(si;θ^)))L _ { i d } = 1 - s i m ( F ( G ( s _ { i } ; \theta ) ) , F ( G ( s _ { i } ; \hat { \theta } ) ) )

trong đó simsim là khoảng cách cosine và FF là mô hình arcface pretrained.

image.png

Kiểm soát style intensity bằng feature interpolation

Sử dụng feature interpolation cho phép ta kiểm soát cường độ của style được transfer sang input. Ký hiệu fiAf_i^A là feature map thứ ii của generator gốc và fiBf_i^B là feature map của generator sau khi fine tune. Feature map mới được tính bằng công thức f=(1α)fiA+αfiBf = ( 1 - \alpha ) f _ { i } ^ { A } + \alpha f _ { i } ^ { B }

image.png

Extreme Style References

Khi muốn transfer style với một ảnh reference out of distribution yy (ví dụ ảnh chó mèo trong khi StyleGAN train trên ảnh mặt người), thay vì sử dụng trực tiếp s(T(y))s(T(y)) để xây dựng tập style parameter, tác giả sử dụng mean style code s=110000s(FC(zN(0,I)))\overline { s } = \sum _ { 1 } ^ { 10000 } s ( F C ( z \sim N ( 0 , I ) ) ) để làm ước lượng cho s(T(y))s(T(y)) cho ảnh out of distribution.

image.png

Lựa chọn GAN inversion

Theo quan sát của tác giả, mếu thủ tục GAN inversion tạo ra ảnh realistic từ ảnh reference, JoJoGAN sẽ được train để ánh xạ sis_i sang ảnh được stylized với cường độ mạnh. Ngược lại nếu GAN inversion tạo ra ảnh được stylized JoJoGAN sẽ ánh xạ sis_i sang ảnh stylized nhưng với cường độ nhẹ hơn và giữ lại nhiều đặc trưng của input hơn. image.png

Như đã nhắc đến ở phần Extreme Style References, mean style code là ước lượng tốt nhất của s(T(y))s(T(y)) với ảnh out of distribution. Điều này cũng đúng khi áp dụng với một GAN inverter không tốt. Tác giả tạo ra một virtual inverter V(y)V(y) bằng cách cộng latent vector từ GAN inversion với mean style code bằng thủ tục ở bước 3 với các vector M khác nhau. Vector M được lựa chọn sao cho G(s(V(y));θ)G ( s ( V ( y ) ) ; \theta ) có các đặc trưng mong muốn (VD: output có mắt của ảnh reference thì G(s(V(y));θ)G ( s ( V ( y ) ) ; \theta ) phải có mắt trông thật. image.png

Qualitative result

Một số kết quả so sánh về mặt hình ảnh image.png

image.png

Reference


All Rights Reserved

Viblo
Let's register a Viblo Account to get more interesting posts.