HistoGAN은 StyleGAN2를 기반으로 색상을 Control 할 수 있게 수정한 모델이다.
GAN을 통해 이미지를 생성할 때 색상을 조절하는 것은 상당히 까다로운 문제인데, 이 문제를 Color Histogram을 입력으로 받아 Histogram에 맞는 이미지를 생성하는 모델을 제안한다.
log-chroma 공간의 2D Histogram을 사용하였으며 Recoloring을 위한 novelty 있는 손실 함수를 제안하여 Encoder-Decoder 모델인 ReHistoGAN도 설계하였다.
Paper: https://arxiv.org/abs/2011.11731
Github: https://github.com/mahmoudnafifi/HistoGAN
Introduction
색상을 조절하기 위해서는 입력으로 사용할 만한 색상을 제대로 표현할 수 있는 어떤 특징이 필요하다.
Color Histogram은 도메인과 무관하게 이미지의 색상을 설명하는 직관적인 방법이므로 생성을 위한 입력으로 사용하기에 적절하다.
이 입력을 생성하려는 이미지에 적용하기 위한 방법은 StyleGAN이 대표적으로 잘 설명하고 있다.
StyleGAN은 입력을 Style Code로 Mapping 한 후 이를 AdaIN 또는 ModConv와 같은 방법으로 통계량을 조정하여 생성 중인 이미지에 반영한다.
HistoGAN은 이를 적용하여 Histogram을 Color Code로 Mapping 한 후 ModConv를 통해 반영한다.
또 이 접근법을 통해 실제 사진을 다시 채색하기 위한 방법을 제안한다.
재 채색된 GT이미지 없이 학습하는 novelty 있는 GAN loss를 적용한 Encoder-Decoder 구조를 만들어 입력한 사진을 손상 없이 재 채색한 결과물을 보여주며 이 구조를 ReHistoGAN으로 명명하였다.
Figure 1의 맨 아래 이미지는 목표 색상을 특정하지 않고 자동으로 재 채색한 결과물이다.
Method
Histogram feature
Color Histogram을 통해서 학습을 하기 위해서는 미분 가능한(Backprop이 가능한) histogram 생성 함수가 필요하다.
이 histogram 생성함수는 관련 논문인 SIIE(Github, matlab), Correcting improperly white-balanced images(Github, python), Convolutional color constancy(Github, matlab) 에서 조명 변화에 영향을 적게 받는 log-chroma 공간의 histogram을 사용하는 방법을 응용한 것이다.
3개 모두 코드를 살펴보았는데 각 논문에서 histogram을 만드는 데 큰 차이점은 없었으며, Convolutional color constancy의 method를 그대로 사용하고 있었다.
코드를 보면 RGB 색공간에서 각 채널을 다른 채널로 나누고 log를 취한 후 각 픽셀의 세기를 곱하여 histogram을 그린다. R 채널의 u, v를 수식으로 나타내면 다음과 같다.
$$I_{uR}(x)=\log{\frac{I_R(x)+\epsilon}{I_G(x)+\epsilon}}, I_{vR}(x)=\log{\frac{I_R(x)+\epsilon}{I_B(x)+\epsilon}}$$
RGB to YUV 수식과도 다르고 다른 색공간에서도 찾을 수 없어서 log를 취한 결과물을 왜 u, v라고 부르는지는 모르겠으나, xyz 좌표를 한 번 변환한 것을 uvw로 부르는 것처럼 2D 좌표 xy를 한번 변환했으니까 uv라고 부르는 것 같다.
# Correcting improperly white-balanced images
# https://github.com/mahmoudnafifi/WB_sRGB/blob/master/WB_sRGB_Python/classes/WBsRGB.py#L59
import numpy as np
def rgb_uv_hist(I, h=60):
""" Computes an RGB-uv histogram tensor. """
# Keep only pixels whose values in each channel are all greater than zero
# and flatten them along the channel axis.
I_reshaped = I[(I > 0).all(axis=-1)]
eps = 6.4 / h
hist = np.zeros((h, h, 3)) # histogram will be stored here
Iy = np.linalg.norm(I_reshaped, axis=-1) # intensity vector
for i in range(3): # for each histogram layer, do
r = [] # excluded channels will be stored here
for j in range(3): # for each color channel do
if j != i:
r.append(j)
Iu = np.log(I_reshaped[:, i] / I_reshaped[:, r[1]])
Iv = np.log(I_reshaped[:, i] / I_reshaped[:, r[0]])
hist[:, :, i], _, _ = np.histogram2d(
Iu, Iv, bins=h,
range=((-3.2 - eps / 2, 3.2 - eps / 2),) * 2, # thresholding
weights=Iy)
norm_ = hist[:, :, i].sum()
hist[:, :, i] = np.sqrt(hist[:, :, i] / norm_) # (hist/norm)^(1/2)
return hist
SIIE가 위 코드와 다른 부분은 thresholding 하지 않고 exp(-x)를 취한 것(gaussian kernel을 곱해주는 것)이었다.
HistoGAN은 이를 응용하여 inverse quadratic kernel을 사용해서 histogram을 그렸으며 kernel을 바꾼 것이 학습 안정성을 크게 높여줬다고 한다.
inverse quadratic 커널 $k$, histogram bin을 조절하는 $\tau$에 대해서 수식으로 나타내면 다음과 같다.
$$k(I_{uc}, I_{vc}, u, v)=\frac{1}{1+(|I_{uc}-u|/\tau)^2} \times \frac{1}{1+(|I_{vc}-v|/\tau)^2}$$
이후 총합계가 1이 되도록 normalize 하여 최종 histogram을 생성한다.
Color Controlled Image Generation
앞서 설명했듯이 생성한 Histogram을 StyleGAN에서 style code를 주입하는 것과 동일한 방식으로 사용한다.
마지막 2개의 블록에서만 style code 대신 histogram code를 사용한다고 하는데 github 코드를 보면 첫 블록부터 두 latent code를 concat 하여 사용하고 있다. 왜 논문과 코드를 다르게 서술했는지는 모르겠다.
출력 이미지에서 generated histogram $H_g$를 계산해서 입력으로 줬던 target histogram $H_t$와의 loss를 계산한다.
두 histogram 간의 손실 함수 $C$는 Hellinger 거리를 사용하였으며 수식은 다음과 같다.
$$C(H_g, H_t)=\frac{1}{\sqrt{2}}\lVert H_g^{1/2} - H_t^{1/2} \rVert_2$$
여기서 $H^{1/2}$는 element-wise square root다.
논문에서 왜 Hellinger 거리를 사용하였는지에 대한 설명이 전혀 없다.
Hellinger 거리는 두 histogram이 얼마나 겹치는지를 비교하는 것이기에 loss로 사용하기 적합한데 아래처럼 수식을 전개하면 Bhattacharyya 거리의 형태로 나타낼 수 있다.
$$\begin{align} C(H_g, H_t)^2 &= \left( \frac{1}{\sqrt{2}}\lVert H_g^{1/2} - H_t^{1/2} \rVert_2 \right)^2 \\ &= \frac{1}{2} \sum_{i=1}^k \left( \sqrt{H_{gi}} - \sqrt{H_{ti}} \right) \\ &= \frac{1}{2} \int \left( \sqrt{H_{gi}} - \sqrt{H_{ti}} \right)^2 di \\ &= \frac{1}{2} \int \left( H_{gi} + H_{ti} - 2\sqrt{H_{gi}H_{ti}} \right) di \\ &= \frac{1}{2} \int \left( 2 - 2\sqrt{H_{gi}H_{ti}} \right) di \qquad \left( \because \int H_g = \int H_t = 1 \right) \\ &= 1 - \int \sqrt{H_{gi} H_{ti}} di \end{align}$$
또 논외지만 KL Divergence와 유사한 형태로 전개해 볼 수도 있다. (참조)
Table 1에서도 FID와 함께 KL Divergence와 Hellinger Distance를 정량지표로 평가하였기에 이를 정리해보았다.
$$\begin{align} d_{KL}(H_g, H_t) &= \int \log \frac{H_g}{H_t} H_g \\ &= 2 \int \log \frac{\sqrt{H_g}}{\sqrt{H_t}} H_g \\ &= 2 \int - \log \frac{\sqrt{H_t}}{\sqrt{H_g}} H_g \\ &\ge 2 \int \left( 1 -\frac{\sqrt{H_t}}{\sqrt{H_g}} \right) H_g \\ &= 2 \left( \int H_g - \int \sqrt{H_g}\sqrt{H_t} \right) \\ &= 2 - 2 \int \sqrt{H_g}\sqrt{H_t} \qquad \left( \because \int H_g = 1 \right) \\ &= \int \left( \sqrt{H_g} - \sqrt{H_t} \right)^2 = 2 C(H_g, H_t)^2 \end{align}$$
이 histogram match loss $C$에 GAN loss $D(I_g)$를 더하여 generator의 loss로 사용한다.
$$L_g=D(I_g)+\alpha C(H_g, H_t)$$
그런데 target histogram을 어디서 가져올 것인가는 또 다른 문제다.
입력 이미지의 histogram을 그대로 사용한다면 histogram의 정보는 완전히 쓸모가 없어진다. 때문에 다른 이미지 2장으로부터 histogram을 연산한 후 interpolation 하여 target으로 사용한다.
이렇게 interpolation을 통해 augmentation 하는 것이 좀 더 강건한 결과물을 얻게 해준다고 한다.
entropy가 아주 높은(극단적으로 어질러져 있는 상태 = 다양성이 높은) 이미지의 histogram을 예로 들어보자.
극단적으로 가정하여 gaussian noise와 같이 rgb 색상 값이 0~255 범위의 임의 값이라면 다양성은 최대로 볼 수 있고, 이때 histogram은 평평할 수밖에 없다. 그래서 2D log-chroma histogram을 그렸을 때 각 채널이 모두 유사하여 유의미한 색상 정보를 얻을 수 없게 된다.
이런 다양성이 아주 많은 landscape와 같은 도메인을 학습하는 경우에도 augmentation을 쓰는 것이 유효했다고 하는데, 아무리 다양하다고 해도 위 예시의 noise처럼 극단적이진 않을 테니 augmentation이 유효한 것은 당연해 보인다.
Image Recoloring
이전까지는 생성 모델인 HistoGAN에 대한 설명이었으며 이후는 HistoGAN의 접근 방식을 재 채색이라는 목적으로 확장한 모델인 ReHistoGAN에 대한 설명이다. 참고로 ReHistoGAN의 코드는 현재(2021.08.01) 공개되어있지 않다.
Optimize 방식이 아닌 Encoder-Decoder 방식으로 입력한 실제 사진을 pretrained HistoGAN의 Head(출력에 가장 가까운 부분)의 입력으로 사용되는 latent로 mapping 한다.
이 Encoder-Decoder는 U-Net 구조를 사용하였는데, 이는 이미지의 가장 디테일한 정보를 갖고 있는 첫 2블록의 latent feature를 skip connection을 통해 HistoGAN head의 입력으로 사용하기 위한 것이다.
Figure 6을 보면 skip connection을 통해서 색상 정보뿐만 아니라 눈가 주름, 치아 형태 등 훨씬 세세한 결과물을 얻음을 볼 수 있다.
또 입력 이미지 외에 target histogram은 어디로 입력하나 싶은데 Figure 4에서 맨 위의 $H(h \times h \times 3)$가 target histogram이며 이로부터 U-Net 출력과는 별도의 latent를 만들어 HistoGAN head의 convolution weight를 mod-demod 한다. 이 부분은 HistoGAN에서 설명한 것과 동일하게 반영되는 것이다.
ReHistoGAN의 loss는 기존 loss에 Reconstruction loss $R(I_i, I_r)$와 Variance loss $V(I_i, I_r)$를 추가한 것이다.
여기서 $I_i$는 입력 이미지, $I_r$은 reconstructed 이미지다.
Reconstruction loss는 일반적으로 $|I_i - I_r|$ 또는 $(I_i-I_r)^2$의 l1 또는 l2 loss를 사용한다.
그런데 이렇게 사용할 경우 각 픽셀 값이 그대로 전달되면서 색상 정보를 강제하게 된다.
때문에 HistoGAN은 각 이미지에 Laplacian 연산 $*L$을 취한 후 l1 loss를 구하는 방식으로 사용한다.
$$R(I_i, I_r)=\lVert I_i * L - I_r * L \rVert_1$$
이미지에서 w, h 축에 대한 Laplacian은 $L(f)=\frac{\partial^2 f}{\partial w^2} + \frac{\partial^2 f}{\partial h^2}$ 이며, 수식에서 보듯이 변곡점을 알 수 있어 흔히 이미지에서 edge를 검출하기 위해 사용한다.
색상 정보 없이 edge 간에 l1 loss를 적용하여 원래 형태를 잘 유지하는지를 나타내는 loss로 이해하면 되겠다.
이로써 얼마나 실제 데이터 같은가 $D$, 색상을 얼마나 잘 따라가는가 $C$, content를 잘 보존하는가 $R$로 학습하는 데 필요한 조건인 loss는 모두 설정되었다.
그런데 몇몇 경우 $C$를 최소화하기 위해서 히스토그램을 알맞게 조정해 색상이 각 물체에 알맞게 칠해지는 것이 아니라 히스토그램이 이동만 하는 경우가 있었고 이를 수정하기 위해 색상을 고르게 분산시키는 Variance loss $V(I_i, I_r)$를 추가하였다.
$$V(I_i, I_r) = - w \sum_{c \in \{ R, G, B \} } \left| \sigma (I_{ic} * G) - \sigma(I_{rc} * G) \right| $$
여기서 $w=\lVert H_t - H_i \rVert_1$이며 $G$는 gaussian blur kernel이다.
픽셀 차이를 계산하되 $w$를 곱함으로써 histogram이 다른 부분에 집중할 수 있게 되고 blur kernel을 적용함으로써 색상을 평활화하여 그 차이를 더 증가시키게 된다.
Figure 5를 보면 확실히 더 자연스러운 재 채색 결과를 볼 수 있다.
Results
HistoGAN
ReHistoGAN
ReHistoGAN의 재 채색 결과물은 정말 인상적이다.
무엇보다 다른 방법들은 현실에 없을 법한 결과물만 내놓지만 ReHistoGAN은 물체에 대한 지각 정보(Perceptual Information)를 학습해서 현실에 있을 법한 색상을 갖도록 조절한다.
Figure 10에서 땅은 붉은빛이 돌지 않고 하늘만 석양이 지는 것으로 재 채색한 것과 Figure 11의 현실적인 결과물이 지각 정보를 아주 잘 학습하고 있다는 것을 보여준다.
Limitation
이외에도 아주 많은 실험 결과물이 논문에 첨부되어있는데 Color Constancy에 대해서 깊게 탐구해 본 논문으로 어떻게 만들었는지, 어떻게 실험했는지, 결과물은 어떤지 모두가 흥미롭다.
결론
- Histogram을 입력하여 원하는 색상의 결과물만 생성하는 모델 제시
- Histogram에 맞게 입력 사진을 재 채색하는 모델 제시
- Histogram 생성 시 inverse-quadratic kernel을 취하여 좀 더 안정적인 학습을 함
개인적인 생각
- GAN 모델을 다루다 보면 원치 않는 색상의 영향 때문에 고통을 받는 경우가 많다.
특히 I2I는 질감만 변환하고 싶어도 색상 변환을 피할 수가 없는데 이를 완화할 아이디어를 얻을 수 있었다. - ReHistoGAN에서 input image의 지각 정보만 사용하는 것이 아니라 target image의 지각 정보도 반영한 재 채색 모델을 설계하는 것이 Future work로 보이는데 이 또한 정말 흥미롭다.
댓글