SRGAN으로 애니메이션 업스케일링하기
Last updated: Dec 13, 2022
손건호, 소프트웨어학부 소프트웨어전공, [email protected]
박태완, 소프트웨어학부 컴퓨터전공, [email protected]
ㅤ
ㅤ
ㅤ
Proposal
GAN은 Generative Adversarial Networks(생성적 적대 신경망)의 줄임말로 실제 이미지 같은 가짜 이미지를 만들어내는 생성 모델이다. 이러한 GAN을 이용하면 해상도가 낮은 이미지(가짜 이미지)를 통해 해상도가 높은 이미지(실제 이미지)도 생성해낼 수 있다. 저해상도 이미지를 고해상도 이미지로 바꾸는 것을 Super Resolution이라고 하는데, 이러한 Super Resolution을 수행해주는 GAN을 SRGAN이라고 부른다. SRGAN은 현재 나와있는 Super Resolution(SR) 모델 중 MOS(Mean Opinion Score), 즉 사용자 만족도가 가장 높은 모델이다. 이를 통해 저해상도의 애니메이션을 고해상도로 바꾸어 QHD 해상도의 영상을 얻을 수 있도록 한다.
ㅤ
ㅤ
Datasets
이 프로젝트는 ‘애니메이션’을 Super Resolution하는데 최적화시키기 위해, 애니메이션으로부터 데이터를 추출하였다. 인터넷에서 구할 수 있는 애니메이션들로부터, 스크린샷들을 추출하는 방식으로 진행하였다. 하지만 그냥 모든 프레임의 스크린샷을 찍게되면 생기는 문제점이 있다. 바로 비슷한 장면들이 연속적으로 나타나거나 정지해 있는 화면에서는 중복으로 스크린샷이 여러장 캡쳐된다는 것이다.
이를 해결하기 위해 각 프레임별로 히스토그램을 생성하고, 해당 히스토그램을 상관관계로 비교해 상관관계가 이전 프레임의 일정 이하인 프레임만 추출하도록 코드를 작성하였다. 비록 상관관계로만 비교를 하는 것은 정확도가 높지 않으나 24분짜리 영상에는 24 * 60 * 24 = 34,560의 프레임이 있기에 유사도 비교 시간이 너무 길어서는 안된다.
def capture_frames(start_count: int, video_path: Path, image_path: Path, name_prefix: str, similarity: float) -> int:
video_absolute_path = str(video_path.absolute())
image_absolute_path = str(image_path.absolute())
# Capture Video
vidcap = cv2.VideoCapture(video_absolute_path)
success, image = vidcap.read()
if not success:
print('Done. Captured 0 frames.')
return start_count
cv2.imwrite(f'{image_absolute_path}/{name_prefix}{start_count}.jpg', image)
prev_img = image
relative_frame = 1
relative_capture = 1
while success:
sys.stdout.write(f'\rReading new frame: frame#{relative_frame} ')
success, image = vidcap.read()
if not success:
break
curr_hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
curr_hist = cv2.calcHist([curr_hsv], [0,1], None, [180,256], [0,180,0, 256])
cv2.normalize(curr_hist, curr_hist, 0, 1, cv2.NORM_MINMAX)
prev_hsv = cv2.cvtColor(prev_img, cv2.COLOR_BGR2HSV)
prev_hist = cv2.calcHist([prev_hsv], [0,1], None, [180,256], [0,180,0, 256])
cv2.normalize(prev_hist, prev_hist, 0, 1, cv2.NORM_MINMAX)
# Compare current frame and previous frame
ret = cv2.compareHist(prev_hist, curr_hist, cv2.HISTCMP_CORREL)
if ret <= similarity:
# Compared with the previous frame, current frame is not similar. Capturing frame.
cv2.imwrite(f'{image_absolute_path}/{name_prefix}{start_count}.jpg', image)
print(f'Captured a new frame: frame#{start_count}, similarity: {ret}')
start_count += 1
relative_capture += 1
prev_img = image
relative_frame += 1
ㅤ
Reading new frame: frame#18174 Captured a new frame: frame#13846, similarity: 0.496534625920312
Reading new frame: frame#18175 Captured a new frame: frame#13847, similarity: 0.5525413283322532
Reading new frame: frame#18176 Captured a new frame: frame#13848, similarity: 0.45070728123014864
Reading new frame: frame#18177 Captured a new frame: frame#13849, similarity: 0.6020107052682623
Reading new frame: frame#18178 Captured a new frame: frame#13850, similarity: 0.6506659616878512
Reading new frame: frame#18179 Captured a new frame: frame#13851, similarity: 0.5860715814537811
Reading new frame: frame#18181 Captured a new frame: frame#13852, similarity: 0.5452228452494848
Reading new frame: frame#18449 Captured a new frame: frame#13853, similarity: 0.5636435382384591
Reading new frame: frame#18580 Captured a new frame: frame#13854, similarity: 0.3280791418789342
Reading new frame: frame#18777 Captured a new frame: frame#13855, similarity: 0.3914499579764988
Reading new frame: frame#18964 Captured a new frame: frame#13856, similarity: 0.04331356312667525
Reading new frame: frame#18987 Captured a new frame: frame#13857, similarity: 0.2210866763214178
출력 결과를 보면 프레임을 비연속적으로 캡쳐하는 모습을 볼 수 있다. 이렇게 해서 데이터를 총 74221장을 확보했다. 샘플 데이터는 다음과 같다.
ㅤ
Methodology
Model
GAN은 Generator 모델과 Discriminator 모델로 이루어진다. Generator 모델은 아주 정교한 가짜 이미지를 만들기 위해 노력하고, Discriminator 모델은 가짜 이미지에게는 낮은 점수를, 진짜 이미지에게는 높은 점수를 부여하는 역할을 한다. Generator과 Discriminator 모델은 서로 경쟁하면서 학습하게 되는데 여기서 우리는 Discriminator보다는 Generator모델을 활용하여 정교한 이미지를 갖는 것을 목표로 한다.
GAN은 SR에서도 사용될 수 있고, 일명 SRGAN이라고 불리우며, 이 논문에서 사용된 모델은 아래 그림과 같다.
CNN을 이용한 SR 모델들은 PSNR과 SSIM을 기준으로 모델의 성능을 평가한다. 하지만 SRGAN 모델은 MOS값이 높은 것을 자랑한다. 애니메이션이라는 데이터는 컴퓨터가 봤을 때 좋아야 하는 데이터가 아니라 사람이 봤을 때 좋아야 하는 데이터이므로, MOS값이 높은 SRGAN을 선택하였다.
모델 코드는 아래와 같다.
import torch
import torch.nn as nn
import math
class Generator(nn.Module):
def __init__(self, scale_factor):
upsample_block_num = int(math.log(scale_factor, 2))
super(Generator, self).__init__()
self.block1 = nn.Sequential(
nn.Conv2d(3, 64, 9, 1, 4),
nn.PReLU()
)
self.block2 = ResidualBlock(64)
self.block3 = ResidualBlock(64)
self.block4 = ResidualBlock(64)
self.block5 = ResidualBlock(64)
self.block6 = ResidualBlock(64)
self.block7 = nn.Sequential(
nn.Conv2d(64, 64, 3, 1, 1),
nn.BatchNorm2d(64)
)
block8 = [UpsampleBlock(64, 2) for _ in range(upsample_block_num)]
block8.append(nn.Conv2d(64, 3, 9, 1, 4))
self.block8 = nn.Sequential(*block8)
def forward(self, x):
block1 = self.block1(x)
block2 = self.block2(block1)
block3 = self.block3(block2)
block4 = self.block4(block3)
block5 = self.block5(block4)
block6 = self.block6(block5)
block7 = self.block7(block6)
block8 = self.block8(block1 + block7)
# outputs range between 0~1
return (torch.tanh(block8) + 1) / 2
class Discriminator(nn.Module):
def __init__(self):
super(Discriminator, self).__init__()
self.net = nn.Sequential(
nn.Conv2d(3, 64, 3, 1, 1),
nn.LeakyReLU(0.2),
nn.Conv2d(64, 64, 3, 2, 1),
nn.BatchNorm2d(64),
nn.LeakyReLU(0.2),
nn.Conv2d(64, 128, 3, 2, 1),
nn.BatchNorm2d(128),
nn.LeakyReLU(0.2),
nn.Conv2d(128, 128, 3, 2, 1),
nn.BatchNorm2d(128),
nn.LeakyReLU(0.2),
nn.Conv2d(128, 256, 3, 1, 1),
nn.BatchNorm2d(256),
nn.LeakyReLU(0.2),
nn.Conv2d(256, 256, 3, 2, 1),
nn.BatchNorm2d(256),
nn.LeakyReLU(0.2),
nn.Conv2d(256, 512, 3, 1, 1),
nn.BatchNorm2d(512),
nn.LeakyReLU(0.2),
nn.Conv2d(512, 512, 3, 2 ,1),
nn.BatchNorm2d(512),
nn.LeakyReLU(0.2),
nn.AdaptiveAvgPool2d(1),
nn.Conv2d(512, 1024, 1, 1, 0),
nn.LeakyReLU(0.2),
nn.Conv2d(1024, 1, 1, 1, 0)
)
def forward(self, x):
batch_size = x.size(0)
return torch.sigmoid(self.net(x).view(batch_size))
class ResidualBlock(nn.Module):
def __init__(self, channels):
super(ResidualBlock, self).__init__()
self.resblock = nn.Sequential(
nn.Conv2d(channels, channels, 3, 1, 1),
nn.BatchNorm2d(channels),
nn.PReLU(),
nn.Conv2d(channels, channels, 3, 1, 1),
nn.BatchNorm2d(channels)
)
def forward(self, x):
residual = self.resblock(x)
return x + residual
class UpsampleBlock(nn.Module):
def __init__(self, in_channels, up_scale):
super(UpsampleBlock, self).__init__()
self.upblock = nn.Sequential(
nn.Conv2d(in_channels, in_channels*up_scale**2, 3, 1, 1),
nn.PixelShuffle(up_scale),
nn.PReLU()
)
def forward(self, x):
return self.upblock(x)
Loss Function
기존 GAN의 Loss Function은 MinMax Loss Function이다. Discriminator의 오차만을 Backpropagation 시키는 형태이다. 하지만, SRGAN은 기존 GAN의 Adversarial Loss에 추가해 Content Loss도 이용한다. Perceptual Loss는 Adversarial Loss, Content Loss와 MSE Loss로 구성된다.
MSE Loss는 HR(High Resolution)이미지와 SR(Super Resolution)이미지의 픽셀별 차이를 제곱하여 평균낸 값이고, Content Loss는 ‘VGG넷에 통과된 HR’과 ‘VGG넷에 통과된 SR’의 픽셀별 차이를 제곱하여 평균낸 값이다.
논문 저자에 의하면 MSE Loss는 텍스쳐를 흐릿하게 만들 수 있다고 하였는데, 저자는 이 때문에 Perceptual Loss를 제안하였다. 이미지를 VGG넷의 i번째 maxpooling 전, j번째 convolution까지만 통과시키면 이미지에 대한 Feature Map을 얻을 수 있고, 이를 통해 얻어낸 Content Loss와 MSE Loss를 잘 섞으면 기존 MSE Loss만 사용했을 때 style들의 평균을 내어 흐릿하게 Output하던 현상을 완화할 수 있다.
추가로, 우리는 TV Loss(antisotropic Total Variation Loss)를 사용하였다. 이 Loss는 모든 픽셀에 대하여 각각 아래 행, 오른쪽 열에 있는 픽셀과의 차이를 제곱하여 더한 후, (B/2) 제곱을 취한 값의 합이다. 여기서 B는 보통 1을 사용한다.
아래 코드는 Generator모델의 Loss function이다.
import torch
from torch import nn
from torchvision.models.vgg import vgg16
class GeneratorLoss(nn.Module):
def __init__(self):
super(GeneratorLoss, self).__init__()
vgg = vgg16(pretrained=True)
loss_network = nn.Sequential(*list(vgg.features)[:31]).eval()
for param in loss_network.parameters():
param.requires_grad = False
self.loss_network = loss_network
self.mse_loss = nn.MSELoss()
self.tv_loss = TVLoss()
def forward(self, out_labels, out_images, target_images):
adversarial_loss = torch.mean(1 - out_labels)
perception_loss = self.mse_loss(self.loss_network(out_images), self.loss_network(target_images))
image_loss = self.mse_loss(out_images, target_images)
tv_loss = self.tv_loss(out_images)
return 0.001*adversarial_loss + image_loss + 0.05*perception_loss + 2e-8*tv_loss
class TVLoss(nn.Module):
def __init__(self, tv_loss_weight = 1):
super(TVLoss, self).__init__()
self.tv_loss_weight = tv_loss_weight
def forward(self, x):
# 0 : batch_size
# 1 : channel_size
# 2 : img_height
# 3 : img_weight
batch_size = x.size()[0]
h_x = x.size()[2]
w_x = x.size()[3]
count_h = self.tensor_size(x[:, :, 1:, :])
count_w = self.tensor_size(x[:, :, :, 1:])
h_tv = torch.pow(x[:, :, 1:, :] - x[:, :, :h_x-1, :], 2).sum()
w_tv = torch.pow(x[:, :, :, 1:] - x[:, :, :, :w_x-1], 2).sum()
return self.tv_loss_weight * 2 * (h_tv / count_h + w_tv / count_w) / batch_size
@staticmethod
def tensor_size(t):
return t.size()[1] * t.size()[2] * t.size()[3]
if __name__ == "__main__":
g_loss = GeneratorLoss()
print(g_loss)
Discriminator 모델의 Loss function은 단순하다. 이미지가 Discriminator로부터 얻은 점수(0~1)를 통해 1 - (실제 이미지 점수) + (가짜 이미지 점수)를 최소화시키면 된다.
훈련코드에 구현되어있는 Descriminator의 Loss는 아래와 같다.
real_out = netD(real_img).mean()
fake_out = netD(fake_img).mean()
d_loss = 1 - real_out + fake_out
Evaluation & Analysis
영상의 대표적인 평가 방식으로는 PSNR, SSIM 두 가지가 있다.
PSNR
PSNR은 Peak Signal-to-Noise Ratio 의 약자로 최대 신호 대 잡음비라고 부른다. 신호가 가질 수 있는 최대값에 대한 잡음의 비율을 구하는 것인데 이때 최대값은 이미지가 가질 수 있는 최대값(01일때 1, 0255일때 255), 잡음은 MSE를 사용한다. PSNR의 공식은 아래와 같다.
공식을 보면 알 수 있듯이, MSE가 0이 되지 않도록 해야 하고, 인간보다는 기계가 봤을 때의 품질을 계산하므로, PSNR이 작게 나온 영상의 화질이 더 좋아보일수도 있다는 것을 알아야 한다.
PSNR 코드는 MSE를 조금 가공하는 수준으로 아래와 같이 간단히 구현할 수 있다.
valing_results['psnr'] = 10 * log10((1**2) / (valing_results['mse'] / valing_results['batch_sizes']))
SSIM
SSIM은 Structural Similarity Index Map의 약자로 두 이미지의 유사도를 luminance(휘도), contrast(대비), sturcture(구조) 3가지 요소로 확인하는 방법이다. 각각은 다음과 같이 구해진다.
이들을 각각 알파, 베타, 감마제곱해서 곱해주면, 최종 SSIM값은 다음과 같이 표현된다.
SSIM은 0~1 사이의 값을 가지며 1과 가까울수록 좋은 값이다. luminance, contrast, structure의 값을 이용하기 때문에 실제 인간의 시각 기관과 비슷하게 두 개의 이미지를 비교한다.
전체 SSIM 코드는 프로젝트 코드에서 pytorch_ssim/__init__.py에서 볼 수 있지만, 핵심적인 부분은 다음과 같다.
def _ssim(img1, img2, window, window_size, channel, size_average = True):
mu1 = F.conv2d(img1, window, padding = window_size // 2 , groups = channel)
mu2 = F.conv2d(img2, window, padding = window_size // 2 , groups = channel)
mu1_sq = mu1.pow(2)
mu2_sq = mu2.pow(2)
mu1_mu2 = mu1 * mu2
sigma1_sq = F.conv2d(img1*img1, window, padding = window_size // 2, groups = channel) - mu1_sq
sigma2_sq = F.conv2d(img2*img2, window, padding = window_size // 2, groups = channel) - mu2_sq
sigma12 = F.conv2d(img1*img2, window, padding = window_size // 2, groups = channel) - mu1_mu2
C1 = 0.01 ** 2
C2 = 0.03 ** 2
ssim_map = ((2 * mu1_mu2 + C1)*(2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2))
if size_average:
return ssim_map.mean()
else:
return ssim_map.mean(1).mean(1).mean(1)
Optimizer
Generator과 Discriminator의 Optimizer은 Adam을 사용하였다.
optimizerG = optim.Adam(netG.parameters())
optimizerD = optim.Adam(netD.parameters())
Train Result
훈련은 SSIM과 PSNR값 모두 높은 값을 찾아 목표로 하였고, Generator의 성능과 큰 연관이 있는 loss fuction의 loss 반영 비율을 수정해가면서 튜닝을 진행하였다.
튜닝과정에서 아래와 같은 비율을 사용하였다.
1st
0.001*adversarial_loss + image_loss + 0.05*perception_loss + 2e-8*tv_loss
ㅤ
2nd
0.001*adversarial_loss + image_loss + 0.06*perception_loss + 2e-7*tv_loss
ㅤ
Final
그리고 결국 아래와 같은 하이퍼 파라미터로, 100에폭 훈련 결과 91에폭에서 좋은 모델 파일을 얻을 수 있었다.
0.001*adversarial_loss + image_loss + 0.07*perception_loss + 2e-7*tv_loss
SSIM이 인간의 눈에 조금 더 적합한 평가지표로 알려져 있는데, 하이퍼파라미터 별로 SSIM을 비교해보면 마지막 튜닝에서 가장 안정적으로 높은 SSIM을 얻는 것을 확인할 수 있다.
Result Image Comparison
이렇게 해서 얻은 결과 이미지는 다음과 같다.
색감 차이가 조금 있지만, 자세히 보지 않으면 거의 차이가 나지 않는 것을 볼 수 있다.
Related Reference
Photo-Realistic Single Image Super-Resolution Using a Generative Adversarial Network (SRGAN Paper)
https://wikidocs.net/146367 (SRGAN Korean Explanation)
https://github.com/leftthomas/SRGAN (SRGAN Paper Implementation with TVLoss)
https://velog.io/@hyebbly/Deep-Learning-Loss-%EC%A0%95%EB%A6%AC-1-GAN-loss (GAN의 MinMax Loss Function)
https://m.blog.naver.com/mincheol9166/221771426327 (PSNR, SSIM) ㅤ
ㅤ
Conclusion
Convolution을 이용한 Generator 모델은 지역적인 특징을 충분히 반영하여 보다 나은 결과를 가져다 주었다. 특히, 특정 Task에 특화된 모델이다 보니 논문에서 사용된 데이터셋으로 학습된 모델보다 확실히 평가지표들의 점수가 높았다. 논문에서는 Set5 데이터셋에서의 PSNR과 SSIM이 가장 높았는데 각각 29.4, 0.8472였다. 하지만 우리의 모델은 35.027, 0.9579에 달하는 높은 점수를 얻을 수 있었고, 이는 SOTA 모델인 HAT-L의 Set5에 대한 점수보다 높은 값이다. 물론 데이터셋이 다르므로 직접 비교하는 것은 문제가 있겠지만, 애니메이션과 같은 단순한 이미지 데이터에서 주목할만한 성능을 보여준다는 것을 알 수 있다.
Train을 할 땐 이미지 중 일부를 잘라내서 Train하고 Validation을 할 땐 이미지의 큰 부분을 Resize해서 inference 하기 때문에 Train과 Validation 사이의 괴리감이 생긴다. 하지만 GPU의 VRAM 문제로 정상적인 학습을 위해서 이러한 문제를 포기해야 하기에 아쉬움이 남는다.
ㅤ
ㅤ
Member Roles
손건호 : 모델, 훈련, 테스트 코드 구현, 블로그 글 작성
박태완 : 데이터 수집, 전처리, 모델 파라미터 튜닝, 유튜브 녹화
ㅤ
ㅤ