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

Python



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のファンである。

これで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"

これだけ。で、学習開始したが、このままでは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]

NMB48 GRADUATION CONCERT ~KEI JONISHI/SHU YABUSHITA/REINA FUJIE~ [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()

結果

たしかにサイズは大きくなっているはずだが、動画を再生しても高画質化されてるという実感は得られなかった。

原因として考えられること(他に原因がわかる方がいらしたらおしえてください!)

  • DVDを4倍にしたら1920×2880になるが、学習用の画像サイズが1080×1920なのでそれ以上大きいサイズに超解像する場合にはうまくいかないのではないか
  • 今回もちいたDVDの映像自体がすでに綺麗なので、高画質化した効果が実感しづらい
  • ディスプレイが小さいため効果が実感しづらい

ていうか、そもそも学習がうまくいっていない可能性があるので別の画像で試してみた。

f:id:dululuttu:20180829030948p:plain

左がもともと150×150の画像、右がSRGANにかけて600×600に超解像したもの。
SRGANめちゃめちゃすごいじゃん。車の光沢感がしっかりとでてる。

このように150×150などの小さい画像を超解像することはうまくいったので、学習用の画像のサイズを大きくしたらDVD映像も超解像できるかもしれないと思った。



コメント

タイトルとURLをコピーしました