Splatoon2 の粗い動画の ESPCN を用いた高解像度化

目的

Splatoon2 の配信動画を観ててたまに「画質が粗い」と感じることがないだろうか。

画質がひどく粗い例：

原因としては、配信環境が貧弱でオリジナル動画データが 720p（Switch自体は 1080p/60fps まで対応） だったり、配信プラットフォーム側のエンコード処理が上手くいっていないなど色々考えられる。

それでもやはり Splatoon2 の画質の粗い配信動画は観ててストレスフルなので、何とか高解像度で観られるようにならないものだろうか。

また、将来的にどの動画コンテンツも 2K, 4K, 8K といった高画質が当たり前になり、現在の Switch の 1080p（Full HD）画質も粗いと感じ始めたときどうすればよいのか。

今回やったことを3行で

• Splatoon2 の動画データから学習データと評価データを作った
• 深層学習を用いた画像の高解像度化手法の一つである ESPCN を使って Splatoon2 の粗い解像度を画像を高解像度化するモデルを作った
• 手元に作った高解像度化モデルを使って粗い動画を実際に高解像度化できるか確かめた

結果

高解像度化した動画はオリジナルに近い動画になっているような気がする（低解像度のものよりははっきりした動画になっている、と思う）

高解像度化したもの（output 動画）

3倍解像度を落としたもの（low resolution 動画）

オリジナル（original 動画）

手順

基本的に ESPCN の こちら の実装を使います。

1. 学習データと評価データの準備

まず学習データとなる動画データを用意します。

今回自分の過去の Splatoon2 のプレイ動画（1080p/30fps）を学習データとして用いました（ネットワークサービスにおける任天堂の著作物の利用に関するガイドライン 遵守）。

また Youtube などからダウンロードすることは公式チャンネル等除いて禁止されているらしいので注意が必要です（参考）。

今回動画データを 30 フレームに1回ごとに画像としてキャプチャしました。フレームの切り出し方法は こちら を参考。

今回切り出した画像を こちら を参考に学習データと訓練用データに 8:2 の割合で分けます。学習データ 5591枚, 評価データ 1398枚用意しました。

こちら を使って画像ファイルを HDF5形式のファイルに変換します。下記コマンドで実行します※
※ HDF5ファイルを生成する際にメモリが 130 GB 程度必要だったので、SSD上の仮想メモリをそれに合わせて増やして実行しました。

python prepare.py --images-dir "<学習データ or 評価データのフォルダパス>" --output-path "<HDF5ファイルの出力先>" --scale 3 

以上を問題がなく実施できれば学習データ train.h5, 評価データ test.h5 という2つの HDF5形式のファイルが作られます。

2. ESPCN を用いた高解像度化モデルの学習

こちら を用いて ESPCN を用いた高解像度化モデルの学習を行います。実行コマンドは下記です。

--train-file に手順1.で作成した学習データ, --eval-file に評価データを指定します。--num-workers はマシンの論理コア数に合わせて 16 としました。

python train.py --train-file "data/train.h5" \
                --eval-file "data/test.h5" \
                --outputs-dir "weight" \
                --scale 3 \
                --lr 1e-3 \
                --batch-size 4096 \
                --num-epochs 200 \
                --num-workers 16 \
                --seed 123           

学習が始まると下記のようなログと一緒に 1 epoch ごとに weightフォルダの中に学習済みのモデルデータが出力されます。

epoch: 0/199: : 7245936it [17:13, 7014.19it/s, loss=0.001938]                   
eval psnr: 35.59
epoch: 1/199: : 7245936it [17:17, 6981.19it/s, loss=0.000400]                   
eval psnr: 36.16
epoch: 2/199: : 7245936it [17:13, 7011.23it/s, loss=0.000365]                   
eval psnr: 36.02
epoch: 3/199: : 7245936it [17:16, 6990.12it/s, loss=0.000346]                   
eval psnr: 36.65
epoch: 4/199: : 7245936it [17:21, 6957.72it/s, loss=0.000414]                   
eval psnr: 36.65
epoch: 5/199: : 7245936it [17:17, 6985.19it/s, loss=0.000315]                   
eval psnr: 36.83
epoch: 6/199: : 7245936it [17:20, 6965.47it/s, loss=0.000315]                   
eval psnr: 36.44
epoch: 7/199: : 7245936it [17:18, 6979.01it/s, loss=0.000311]                   
eval psnr: 36.70
epoch: 8/199: : 7245936it [17:18, 6980.38it/s, loss=0.000462]                   
eval psnr: 36.55
epoch: 9/199: : 7245936it [17:17, 6983.96it/s, loss=0.000293]                   
eval psnr: 37.22
epoch: 10/199: : 7245936it [17:17, 6981.65it/s, loss=0.000282]                  
eval psnr: 37.37
epoch: 11/199: : 7245936it [17:18, 6975.79it/s, loss=0.000274]                  
eval psnr: 37.49
epoch: 12/199: : 7245936it [17:18, 6974.52it/s, loss=0.000348]                  
eval psnr: 37.51
epoch: 13/199: : 7245936it [17:39, 6839.09it/s, loss=0.000264]                  
eval psnr: 37.60
epoch: 14/199: : 7245936it [18:15, 6615.75it/s, loss=0.000262]                  
eval psnr: 37.47
epoch: 15/199: : 7245936it [17:20, 6962.83it/s, loss=0.000288]                  
eval psnr: 37.70

loss, psnr といったパラメータの具合を見ながら良さそうな値に落ち着いてきたら終了し、結果の良さそうな学習モデルデータ（"epoch_xx.pth"）を以下では利用します。

3. 高解像度化モデルを使って粗い動画と実際に高解像度化した動画を出力させる

2.で作成した高解像度化モデルデータを使って、1.で用意したSplatoon2 のプレイ動画（1080p/30fps）を用いて粗い動画と実際に高解像度化した動画の2種類を出力させてみます。

こちら のコードを下記のように動画の入出力に合わせて書き換えて Google Colab 上で実行しました。画像データの Pillow と OpenCV の間の変換は こちら をそのまま使わせていただきました。

途中の学習データと入力or出力動画データのパスについては適宜ご自身の環境に読み替えが必要です。

import torch
import torch.backends.cudnn as cudnn
import numpy as np
import cv2
import PIL.Image as pil_image
from PIL import Image

from models import ESPCN
from utils import convert_ycbcr_to_rgb, preprocess, calc_psnr

def pil2cv(image):
    ''' PIL型 -> OpenCV型 '''
    new_image = np.array(image, dtype=np.uint8)
    if new_image.ndim == 2:  # モノクロ
        pass
    elif new_image.shape[2] == 3:  # カラー
        new_image = cv2.cvtColor(new_image, cv2.COLOR_RGB2BGR)
    elif new_image.shape[2] == 4:  # 透過
        new_image = cv2.cvtColor(new_image, cv2.COLOR_RGBA2BGRA)
    return new_image

def cv2pil(image):
    ''' OpenCV型 -> PIL型 '''
    new_image = image.copy()
    if new_image.ndim == 2:  # モノクロ
        pass
    elif new_image.shape[2] == 3:  # カラー
        new_image = cv2.cvtColor(new_image, cv2.COLOR_BGR2RGB)
    elif new_image.shape[2] == 4:  # 透過
        new_image = cv2.cvtColor(new_image, cv2.COLOR_BGRA2RGBA)
    new_image = Image.fromarray(new_image)
    return new_image

cudnn.benchmark = True
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

scale = 3

model = ESPCN(scale_factor=scale).to(device)

state_dict = model.state_dict()
for n, p in torch.load('weight/x3/epoch_xx.pth', map_location=lambda storage, loc: storage).items():
  if n in state_dict.keys():
    state_dict[n].copy_(p)
  else:
    raise KeyError(n)

model.eval()

frame_rate = 30.0 # frame rate
size_lr = (640, 360)
size_output = (1920, 1080)

cap = cv2.VideoCapture('data/original.mp4')
fmt = cv2.VideoWriter_fourcc('m', 'p', '4', 'v') # mp4
writer_lr = cv2.VideoWriter('data/low_resolution.mp4', fmt, frame_rate, size_lr)
writer_output = cv2.VideoWriter('data/output.mp4', fmt, frame_rate, size_output)

while(cap.isOpened()):
  flag, image = cap.read()  # Capture frame-by-frame
  if flag == False:  # Is a frame left?
    break
  image = cv2pil(image)

  image_width = (image.width // scale) * scale
  image_height = (image.height // scale) * scale

  hr = image.resize((image_width, image_height), resample=pil_image.BICUBIC)
  lr = hr.resize((hr.width // scale, hr.height // scale), resample=pil_image.BICUBIC)
  lr_cv2 = pil2cv(lr)
  writer_lr.write(lr_cv2)
  bicubic = lr.resize((lr.width * scale, lr.height * scale), resample=pil_image.BICUBIC)

  lr, _ = preprocess(lr, device)
  hr, _ = preprocess(hr, device)
  _, ycbcr = preprocess(bicubic, device)

  with torch.no_grad():
    preds = model(lr).clamp(0.0, 1.0)

  psnr = calc_psnr(hr, preds)
  print('PSNR: {:.2f}'.format(psnr))

  preds = preds.mul(255.0).cpu().numpy().squeeze(0).squeeze(0)

  output = np.array([preds, ycbcr[..., 1], ycbcr[..., 2]]).transpose([1, 2, 0])
  output = np.clip(convert_ycbcr_to_rgb(output), 0.0, 255.0).astype(np.uint8)
  output = pil_image.fromarray(output)
  output_cv2 = pil2cv(output)
  writer_output.write(output_cv2)

writer_lr.release()
writer_output.release()

それぞれ出力された動画データは高解像度化したもの（output 動画）や3倍解像度を落としたもの（low resolution 動画）という形で Youtube にアップロードしてありますので、適宜そちらのリンクをご確認いただけます※

※ Youtube の解像度はデフォルトだと回線速度に合わせて自動で決められてしまうので、動画右下の歯車マークから選べる最高の解像度を選択した上でご確認いただく必要があります（そうしないと高解像度化されたかどうか正しく評価できません）

参考

この記事は以下の情報を参考にして執筆しました。

• Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network
• yjn870/ESPCN-pytorch
• ネットワークサービスにおける任天堂の著作物の利用に関するガイドライン
• YouTube動画はダウンロード・保存できる？違法性まで徹底解説
• Python, OpenCVで動画ファイルからフレームを切り出して保存
• 画像ディレクトリをtrain_test_splitする関数
• 【Python】Pillow ↔ OpenCV 変換