SRGANでDVDの映像を高画質化したかった(tensorflow)

SRGANを使ってDVDの高画質化を目指す。

といっても、自分でやったのは学習用の画像を用意することとtensorflowが吐き出すエラーを眺めて変顔をしていたことくらいである。（僕はよくわからんエラーと直面すると顔が変になってしまう）

SRGANとは

SRGAN(Super-Resolution Usinga Generative Adversarial Network)とは、GAN：敵対的生成ネットワーク（Generative Adversarial Network）の構造を用いて低解像度画像から高解像度画像に生成する機械学習モデルである。
低解像度の画像に処理を行い高解像度にする技術を超解像というらしい。

以下がその論文
[1609.04802v5] Photo-Realistic Single Image Super-Resolution Using a Generative Adversarial Network

以下の記事ではGANについて詳しく書かれてる。SRGANについての記述もある。日本語。
elix-tech.github.io

今回は、SRGANにDVDのキャプチャを入力し超解像していく。超解像後の画像サイズは縦横がそれぞれ4倍になる。

モデルはGitHubの公開リポジトリのものを利用した。
github.com
こちらのファイルをすべてダウンロードした。

学習用画像の準備

今回、高画質化したいDVDはライブ映像が収録されたものであるため、学習用画像の作成にもライブ映像を使うことにした。

そこで、Blu-rayのライブ映像をキャプチャしたものを学習用として用意した。
コードはこんなかんじ

import cv2

input_video = 'video.m2ts'
cap = cv2.VideoCapture(input_video)

frame_count = int(cap.get(7))
for num, i in enumerate(range(0, frame_count, 1000)):
    cap.set(1, i)
    _, frame = cap.read()
    cv2.imwrite('%d.png' %(num), frame)

参考
動画を扱う — OpenCV-Python Tutorials 1 documentation

あらかじめBlu-ray映像をvideo.m2tsとして保存してある。
ちなみに今回学習用に用意したBlu-rayディスクはこれ

NMB48 渡辺美優紀卒業コンサート in ワールド記念ホール ~最後までわるきーでゴメンなさい~ [Blu-ray]

出版社/メーカー: laugh out loud records
発売日: 2017/10/11
メディア: Blu-ray
この商品を含むブログを見る

何を隠そう、僕はNMB48のファンである。

これで1920×1080のライブ画像が1400枚くらい取得できた。
ここから適当にtrain用に1000枚、valid用に400枚選びそれぞれのフォルダに保存。

どうやら、これらの画像を1/4倍した画像も必要らしいので適当にresizeして保存しておいた。（たぶんなくても学習はできるっぽい）

というわけで以下のフォルダに学習用画像の準備ができた。

1080_1920_train ... 1080×1920の画像1000枚
1080_1920_valid ... 1080×1920の画像400枚
270_480_train ... 270×480の画像1000枚
270_480_valid ... 270×480の画像400枚

学習

モデルなどはGitHubのものをほぼそのまま使った。config.pyを少しいじって準備した画像を利用できるようにした。

config.TRAIN.hr_img_path = "1080_1920_train"
config.TRAIN.lr_img_path = "270_480_train"
config.VALID.hr_img_path = "1080_1920_valid"
config.VALID.lr_img_path = "270_480_valid"

これだけ。

python main.py

で学習開始したが、このままではOOMエラーがでてしまった。
batch_sizeはデフォルトで16だが、4まで減らしたらOOMがでなくなった。
僕のPCはgtx1060 6GBを使っているので、もっと高性能なGPUを積んでいるPCなら16のままいけるかも。

学習は全部終わるのに5日くらいかかった。

DVDの高画質化

ようやく今回の目的であるDVDの高画質化をする。

高画質化するDVDはこれ

NMB48 GRADUATION CONCERT ~KEI JONISHI/SHU YABUSHITA/REINA FUJIE~ [DVD]

出版社/メーカー: laugh out loud records
発売日: 2017/09/20
メディア: DVD
この商品を含むブログを見る

dvd.VOBとしてあらかじめ保存しておく。

DVDのサイズは480×720なので、SRGANに通すと1920×2880まで大きくなるというわけである。これで、超高画質で楽しめるとわくわくしながら以下のコードを書いた

import time
import cv2
import numpy as np
import tensorflow as tf
import tensorlayer as tl
from model import SRGAN_g

checkpoint_dir = "checkpoint"

t_image = tf.placeholder('float32', [1, None, None, 3], name='input_image')
net_g = SRGAN_g(t_image, is_train=False, reuse=False)

sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True, log_device_placement=False))
tl.layers.initialize_global_variables(sess)
tl.files.load_and_assign_npz(sess=sess, name=checkpoint_dir + '/g_srgan.npz', network=net_g)

input_video = 'dvd.VOB'
output_video = 'dvd_hr.avi'

cap = cv2.VideoCapture(input_video)

frame_count = int(cap.get(7))
fps = 30
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) * 4
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) * 4

fourcc = cv2.VideoWriter_fourcc(*'XVID')
vw = cv2.VideoWriter(output_video, fourcc, fps, (width, height))

while(True):
    ret, img = cap.read()
    if ret == True:
        img = (img / 127.5) - 1
        img = sess.run(net_g.outputs, {t_image: [img]})
        img = (img + 1) * 127.5
        img = img.astype(np.uint8)

        vw.write(img[0])

    else :    
        break
        
cap.release()
cv2.destroyAllWindows()