本記事では、A Video Restoration Transformer(VRT)と呼ばれる機械学習手法を用いて、動画の超解像を行う方法を紹介します。

出典: jingyunliang/vrt

VRT

概要

Video Restoration Transformer(VRT)とは、Transformerベースのビデオの復元(ビデオの超解像)手法です。

Video Restorationは、低解像度フレームから高解像度フレームを復元することを目的としています。 単一画像の超解像とは異なり、Video Restorationでは複数のビデオフレームからの時間情報を利用する必要があります。

既存手法は、一般に、sliding windowまたはRNNを利用することでこれに対処します。これは、フレームごとの復元によって制限されるか、長距離モデリング機能が不足しています。

VRTでは、並列フレーム予測と長距離時間依存性モデリング能力を備えています。
より具体的には、VRTは複数のスケールで構成され、各スケールは、時間的相互セルフアテンション（TMSA）と並列ワーピングの2種類のモジュールで構成されています。

TMSAは、ビデオを小さなクリップに分割します。このクリップでは、関節の動きの推定、特徴の位置合わせ、特徴の融合に相互注意が適用され、特徴の抽出にはセルフアテンションが使用されます。
クロスクリップインタラクションを有効にするために、ビデオシーケンスは1つおきのレイヤーにシフトされます。さらに、並列ワーピングは、並列フィーチャワーピングによって隣接するフレームからの情報をさらに融合するために使用されます。

ビデオの超解像、ビデオのブレ除去、ビデオのノイズ除去を含む3つのタスクの実験結果は、VRTが9つのベンチマークデータセットで最先端の方法を大幅に上回っていることを示しています。

フレームワーク

VRTのフレームワークは下記のとおりです。

VRTは、特徴抽出(feature extraction)部分と、再構成(reconstruction)部分で構成されています。

特徴抽出(feature extraction)

最初にconv2Dにて浅い特徴を抽出し、続いて異なる画像解像度でフレームを整列させるマルチスケールネットワークを構成しています。
ここでは、(合計スケール数-1)回の特徴をダウンサンプリングした後に、特徴を元のサイズに戻すことにより、特徴を徐々にアップサンプリングします。

このようにして時間的相互セフルアテンション（TMSA）と並列ワーピング（PA）の2種類のモジュールによって、特徴を抽出し、さまざまなスケールでオブジェクトまたはカメラの動きを処理できます。

最後に、マルチスケールの特徴抽出、位置合わせ、融合の後、さらに特徴を改良するためにいくつかのTMSAモジュールを追加し、深い特徴を抽出しています。

再構成(reconstruction)

浅い特徴と深い特徴を追加して高解像度フレームを再構築します。

異なるフレームは、対応する機能に基づいて独立して再構築されます。
さらに、特徴学習の負担を軽減するために、global residual learningを採用し、双線形にアップサンプリングされた低解像度シーケンスとground-truth高解像度シーケンスの間の残余のみを予測します。

実際には、さまざまな復元モジュールがさまざまな復元タスクに使用されています。

デモ(Colaboratory)

イメージが付きづらいと思いますので、実際に動かしながらVRTができることを見ていきます。
ソースコードは本記事にも記載していますが、下記のGitHubでも取得可能です。
GitHub - Colaboratory demo

また、下記から直接Google Colaboratoryで開くこともできます。

なお、このデモはPythonで実装しています。
Pythonの実装に不安がある方、Pythonを使った機械学習について詳しく勉強したい方は、以下の書籍やオンライン講座などがおすすめです。

環境セットアップ

それではセットアップしていきます。 Colaboratoryを開いたら下記を設定しGPUを使用するようにしてください。

次にライブラリのインストール及び、VRTのソースコードをGitHubから取得します。

%cd /content/
!rm -r VRT

%cd /content/
!rm -r VRT

# GitHubからVRTのcodeをclone
!git clone https://github.com/JingyunLiang/VRT.git

%cd /content/VRT

# Commits on Oct 5, 2022
!git checkout 94a5f504eb84aedf1314de5389f45f4ba1c2d022

%cd /content/VRT
!pip install -r requirements.txt
!pip install moviepy==0.2.3.5 imageio==2.4.1

最後に使用するライブラリをインポートします。

import os
import glob
from google.colab import files
import shutil
import cv2
import matplotlib.pyplot as plt
from pathlib import Path

from moviepy.editor import *
from moviepy.video.fx.resize import resize

以上で環境セットアップは完了です。

テスト動画のセットアップ

次にテスト動画をセットアップします。
本記事では、下記の動画を事前にPCにダウンロードしておき、以下のセルを実行委手Google Colaboratory上にアップロードして使用しています。

超解像、ノイズ除去に使用した動画(動画を再生すると音が流れます。ご注意ください。)
blur除去に使用した動画

%cd /content/VRT

# upload images
!rm -r /content/VRT/testsets/*
!rm -r /content/VRT/results/*

os.makedirs('/content/VRT/testsets/uploaded/frames/000', exist_ok=True)
input_folder = '/content/VRT/testsets/uploaded/'
frame_folder = '/content/VRT/testsets/uploaded/frames/000'
frame_path = '/content/VRT/testsets/uploaded/frames/000/frame%08d.png'

uploaded = files.upload()
for filename in uploaded.keys():
  # 動画のリサイズ
  clip = VideoFileClip(filename)
  print('video size width:', clip.w, 'height:', clip.h)
  #@markdown 動画サイズ変更  両方0で変更なし
  resize_w =  0#@param {type:"integer"}
  resize_h =  0#@param {type:"integer"}
  if resize_h == 0 and resize_w == 0:
    pass
  else:
    if resize_w == 0:
      clip = resize(clip, height=resize_h)
    elif resize_h == 0:
      clip = resize(clip, width=resize_w)
    else:
      clip = resize(clip, width=resize_w, height=resize_h)
      
  #@markdown 動画の長さ変更(秒)  end=0は変更なし
  start = 0 #@param {type:"integer"}
  end =  5#@param {type:"integer"}
  if end != 0:
    if start < end and end <= clip.duration:
      clip = clip.subclip(start, end)

  print('move', filename, 'to', input_folder)
  input_path = os.path.join(input_folder, filename)
  clip.write_videofile(input_path)
   
  if not filename.split('.')[-1].lower() in ['png','jpeg','jpg','tiff','bmp','tif',]: # image or video
    print('動画のフレーム画像を抽出 dst:', frame_folder)
    !ffmpeg -i {input_path} -qscale:v 1 -qmin 1 -qmax 1 -vsync 0 {frame_path}
  else:
    !cp {input_path} {frame_folder}

  !rm {filename}

!ls testsets/uploaded

ここで、動画を指定のサイズに変更し、指定の秒数だけ切り出しています。
動画によって上下しますが、5秒ほどの動画を超解像するだけでおよそ30分ほどかかるため、
まずは短い時間で試されることをおすすめします。

動画の超解像

はじめに、動画の超解像を試してみます。
320x240のサイズの動画を4倍の1280x960に超解像します。

#@markdown uploadかベンチマークか選択
type = "upload" #@param ["upload", "BT001", "BT002", "BT003", "BT004", "BT005", "BT006", "BT007", "BT008"] {allow-input: false}
#@markdown uploadを選択した場合に超解像、blur除去、ノイズ除去か選択
process_type = "blur" #@param ["sr", "blur", "noise"] {allow-input: false}

# アップロードしたビデオで超解像
if type == "upload":
  if process_type == "sr":
    !python main_test_vrt.py \
      --task 001_VRT_videosr_bi_REDS_6frames \
      --folder_lq "/content/VRT/testsets/uploaded/frames/" \
      --tile 6 128 128 --tile_overlap 2 20 20
  if process_type == "blur":
    !python main_test_vrt.py \
      --task 005_VRT_videodeblurring_DVD \
      --folder_lq "/content/VRT/testsets/uploaded/frames/" \
      --folder_gt "/content/VRT/testsets/uploaded/frames/" \
      --tile 12 128 128 --tile_overlap 2 20 20
  if process_type == "noise":
    !python main_test_vrt.py \
      --task 008_VRT_videodenoising_DAVIS \
      --sigma 10 \
      --folder_lq "/content/VRT/testsets/uploaded/frames/" \
      --folder_gt "/content/VRT/testsets/uploaded/frames/" \
      --tile 12 128 128 --tile_overlap 2 20 20


# 以下ベンチマークテスト
# 001, video sr trained on REDS (6 frames), tested on REDS4
elif type == "BT001":
  !python main_test_vrt.py \
    --task 001_VRT_videosr_bi_REDS_6frames \
    --folder_lq testsets/REDS4/sharp_bicubic \
    --folder_gt testsets/REDS4/GT \
    --tile 40 64 64 --tile_overlap 2 20 20
# 002, video sr trained on REDS (16 frames), tested on REDS4
elif type == "BT002":
  !python main_test_vrt.py -\
    -task 002_VRT_videosr_bi_REDS_16frames \
    --folder_lq testsets/REDS4/sharp_bicubic \
    --folder_gt testsets/REDS4/GT \
    --tile 40 64 64 --tile_overlap 2 20 20
# 003, video sr trained on Vimeo (bicubic), tested on Vid4 and Vimeo
elif type == "BT003":
  !python main_test_vrt.py \
    --task 003_VRT_videosr_bi_Vimeo_7frames \
    --folder_lq testsets/Vid4/BIx4 \
    --folder_gt testsets/Vid4/GT \
    --tile 32 64 64 --tile_overlap 2 20 20
  !python main_test_vrt.py \
    --task 003_VRT_videosr_bi_Vimeo_7frames \
    --folder_lq testsets/vimeo90k/vimeo_septuplet_matlabLRx4/sequences \
    --folder_gt testsets/vimeo90k/vimeo_septuplet/sequences \
    --tile 8 0 0 --tile_overlap 0 20 20
# 004, video sr trained on Vimeo (blur-downsampling), tested on Vid4, UDM10 and Vimeo
elif type == "BT004":
  !python main_test_vrt.py \
    --task 004_VRT_videosr_bd_Vimeo_7frames \
    --folder_lq testsets/Vid4/BDx4 \
    --folder_gt testsets/Vid4/GT \
    --tile 32 64 64 --tile_overlap 2 20 20
  !python main_test_vrt.py \
    --task 004_VRT_videosr_bd_Vimeo_7frames \
    --folder_lq testsets/UDM10/BDx4 \
    --folder_gt testsets/UDM10/GT \
    --tile 32 64 64 --tile_overlap 2 20 20
  !python main_test_vrt.py \
    --task 004_VRT_videosr_bd_Vimeo_7frames \
    --folder_lq testsets/vimeo90k/vimeo_septuplet_BDLRx4/sequences \
    --folder_gt testsets/vimeo90k/vimeo_septuplet/sequences \
    --tile 8 0 0 --tile_overlap 0 20 20
# 005, video deblurring trained and tested on DVD
elif type == "BT005":
  !python main_test_vrt.py \
    --task 005_VRT_videodeblurring_DVD \
    --folder_lq testsets/DVD10/test_GT_blurred \
    --folder_gt testsets/DVD10/test_GT \
    --tile 12 128 128 --tile_overlap 2 20 20
# 006, video deblurring trained and tested on GoPro
elif type == "BT006":
  !python main_test_vrt.py \
    --task 006_VRT_videodeblurring_GoPro \
    --folder_lq testsets/GoPro11/test_GT_blurred \
    --folder_gt testsets/GoPro11/test_GT \
    --tile 12 128 128 --tile_overlap 2 20 20
# 007, video deblurring trained on REDS, tested on REDS4
elif type == "BT007":
  !python main_test_vrt.py \
    --task 007_VRT_videodeblurring_REDS \
    --folder_lq testsets/REDS4/blur \
    --folder_gt testsets/REDS4/GT \
    --tile 12 128 128 --tile_overlap 2 20 20
# 008, video denoising trained on DAVIS (noise level 0-50) and tested on Set8 and DAVIS
elif type == "BT008":
  !python main_test_vrt.py \
    --task 008_VRT_videodenoising_DAVIS \
    --sigma 10 --folder_lq testsets/Set8 \
    --folder_gt testsets/Set8 \
    --tile 12 128 128 --tile_overlap 2 20 20
  !python main_test_vrt.py \
    --task 008_VRT_videodenoising_DAVIS \
    --sigma 10  \
    --folder_lq testsets/DAVIS-test \
    --folder_gt testsets/DAVIS-test \
    --tile 12 128 128 --tile_overlap 2 20 20

超解像の結果は以下の通りです。

元動画と超解像動画のフレーム画像同士を比較してみましょう。

非常に分かり辛いですが、顎のあたりのジャギーがやや解消されています。

その他、Blur除去、ノイズ除去も試してみましたが、モデルが学習するノイズ等のパターンと異なっているためか、ほぼ改善は見られませんでした。

この点、求める結果に応じてモデルの追加学習が必要だと考えられます。

まとめ

本記事では、VRTを使って動画の超解像の方法を紹介しました。
Google Colaboratoryではメモリやストレージの制限から、やや実行に難ありですが、視覚的に効果が確認しやすいため知的好奇心をくすぐられます。

また本記事では、機械学習を動かすことにフォーカスしてご紹介しました。
もう少し学術的に体系立てて学びたいという方には以下の書籍などがお勧めです。ぜひご一読下さい。

また動かせるだけから理解して応用できるエンジニアの足掛かりに下記のUdemyなどもお勧めです。

参考文献

1.  論文 - VRT: A Video Restoration Transformer

2. GitHub - jingyunliang/vrt