SwinIRを使って高解像度の動画を作成するswin_video

Last updated at 2024-02-03Posted at 2024-01-28

はじめに

AIにより低解像度の静止画を綺麗に拡大することができるSwinIRを使って動画の解像度を2倍または4倍にするPython用のツールを作成しました。
古いSD動画などをドラッグアンドドロップすることで 2倍または4倍の動画を作成します。

処理結果

ピクサベイの動画を使っていくつか試してみました。
それぞれバイリニアとSwinIRで2倍にリサイズして比較します。

街の空撮

olhao_-_14088(360p)_it_x2.mp4

海岸の桟橋

1198860324_it_x2.mp4

ドイツの森

1114708233_sbs_x2.mp4

動作条件

Windows10 / 11
4GB以上のnVIDIA製ビデオカード
Python 3.11.4

ダウンロード

次のものをダウンロードしてください

SwinIR一式
https://github.com/JingyunLiang/SwinIR
・'SwinIR-main'フォルダ
SwinIR用学習データ 2つ
https://github.com/JingyunLiang/SwinIR/releases
・003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x2_GAN.pth
・003_realSR_BSRGAN_DFOWMFC_s64w8_SwinIR-L_x4_GAN.pth
ffmpeg
https://github.com/GyanD/codexffmpeg/releases/tag/2020-12-20-git-ab6a56773f
ffmpeg-2020-12-20-git-ab6a56773f-full_build.zip
・ffmpeg.exe
・ffprobe.exe
ffstream
ffstream1_1_0.zip
・'ffstream'フォルダ
swin_video
swin_video1_0_0.zip

swin_video.zipを解凍してできる'swin_video'フォルダへそれぞれコピーしてください。
全てコピーすると以下のようになります。

Pythonのインストール

このツールはPython用のスクリプトになっていますので、Pythonを公式サイトからダウンロードしてインストールしてください。

https://www.python.org/downloads/release/python-3114/
・python-3.11.4-amd64.exe

インストール画面ではカスタマイズを選択し Pythonへパスを通すオプションをONにしてください。

インストールが終わったらPCを再起動します（再起動しないとパス設定が有効にならない）。

Windows PowerShellを起動してPythonが使えることを確認します。Pythonから抜けるには[Ctrl] + [Z]。

AppStoreのPythonのページが表示されるようならパスの設定ができていません。
インストーラーで修復か再インストールしてください。

Pythonで必要なモジュール

requirements.txt

ffmpeg-python==0.2.0
keyboard==0.13.5
numpy==1.24.4
opencv-contrib-python==4.8.0.76
requests==2.28.1
timm==0.9.12
tqdm==4.66.1
#torch==2.2.0.dev20230915+cu121
#torchvision==0.17.0.dev20230915+cu121

同梱のmodule_install.batで一括インストールできます。

また、PyTorchのCUDA 12版が必要です。CPU版や CUDA 11版では動かない場合があります。
動かない場合は次の様に上書きインストールしてください。

pip install --upgrade --force-reinstall torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

使い方

以下のバッチファイルへ動画をD&Dしてください。
カレントフォルダ（swin_videoのフォルダ）へ2倍または4倍に拡大された動画が作成されます。

swin_video_x2.bat
2倍に拡大
swin_video_x4.bat
4倍に拡大
swin_video_sbs.bat
バイリニアで拡大した画像とSwinIRを左右に並べた動画を作成
swin_video_it.bat
バイリニアとSwinIRを1秒毎に切り替える動画を作成

注意
SwinIRは静止画用で 1枚の画像を処理するのに数秒かかります。
動画のように数万枚処理するにはとんでもなく時間がかかりますので途中でESCキーで終了してください。
RTX4090で640x480 30fps 1分の動画を処理（2倍に拡大）するのに 50分程度かかります。

サンプルプログラム

SwinIRを使って動画を2倍にするには次のようにします。

超解像モデル SwinIRの準備

コマンドライン引数を管理するargsへSwinIRに関する設定を追加してモデルを取得します。

parser.add_argument('-s', '--scale', type=int, default=2, choices=[2, 4], help='scale factor: 2, 4')
args = parser.parse_args()
args.task = 'real_sr'
args.tile = None
if args.scale==2:
    # 2倍用のモデル
    args.model_path = '003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x2_GAN.pth'
    args.large_model = False
elif args.scale==4:
    # 4倍用のモデル
    args.model_path = '003_realSR_BSRGAN_DFOWMFC_s64w8_SwinIR-L_x4_GAN.pth'
    args.large_model = True
    
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')   # 基本CUDA専用
model = main_test_swinir.define_model(args)
model.eval()
model = model.to(device)

倍率毎に使う学習データが異なります。

超解像処理 SwinIR

def sr(model, args, frame_rgb:Tensor)->Tensor:
    with torch.no_grad():
        img_lq = frame_rgb.unsqueeze(0)  # CHW-RGB to NCHW-RGB
        window_size = 8
        # pad input image to be a multiple of window_size
        _, _, h_old, w_old = img_lq.size()
        h_pad = (h_old // window_size + 1) * window_size - h_old
        w_pad = (w_old // window_size + 1) * window_size - w_old
        img_lq = torch.cat([img_lq, torch.flip(img_lq, [2])], 2)[:, :, :h_old + h_pad, :]
        img_lq = torch.cat([img_lq, torch.flip(img_lq, [3])], 3)[:, :, :, :w_old + w_pad]
        output = main_test_swinir.test(img_lq, model, args, window_size)
        output = output[..., :h_old * args.scale, :w_old * args.scale]
        return output.squeeze()

# super resolution
output = sr(model, args, frame_rgb)

関数sr()へモデルとargsとRGB画像(0.0～1.0の実数)を渡せば 2倍または4倍にされたRGB画像が作成されます。
作成されるRGB画像は0.0～1.0の実数ですが範囲外になる場合がありますので、必要に応じて制限してください。

output = output.clamp(0, 1)    # 0.0-1.0に制限する

FFStreamの場合は出力処理内で制限されるので不要です。

動画の入出力

ffmpeg-pythonを利用したFFStreamを使います。
最終的には別プロセスで実行しているffmpeg.exeで処理されます。

# 入力パス 出力パス 画質20 ピクセルフォーマットRGB 音声がある場合はコピーする
# 出力画像サイズを2倍,4倍にする フレーム処理に使うデバイスを指定する
ff = FFStream(input_path, output_path, crf = 20, pix_fmt='rgb24',
            copy_audio_stream=True, image_scale=args.scale, device=device)
    
bar = tqdm(total=ff.frames, dynamic_ncols=True)
for i in range(ff.frames):
    if keyboard.is_pressed('escape'):
        break
        
    # FFmpegからRGBフレームを受け取る
    frame_rgb = ff.get_rgb()
    if (frame_rgb == None) :
        break                                       # 動画の最後に達した

    # super resolution
    output = sr(model, args, frame_rgb)

    # FFmpegの出力ストリームへRGBフレームを送る
    ff.put_rgb(output)

    bar.update(1)
bar.close()
ff.close()

FFStreamについて記事にしました。

全体

swin_video_s.py

version=(1,0,0)
# 2024.1.28	Ver1.0.0 初版
import os
import sys
import argparse
import torch
from torch import Tensor
import keyboard
from tqdm import tqdm

# SwinIRのフォルダを参照するようにする
sys.path.append(os.path.join(os.path.dirname(__file__), 'SwinIR-main'))
import main_test_swinir

from ffstream.ffstream import FFStream      #import ffmpeg-python

# 超解像処理
def sr(model, args, frame_rgb:Tensor)->Tensor:
    with torch.no_grad():
        img_lq = frame_rgb.unsqueeze(0)  # CHW-RGB to NCHW-RGB
        window_size = 8
        # pad input image to be a multiple of window_size
        _, _, h_old, w_old = img_lq.size()
        h_pad = (h_old // window_size + 1) * window_size - h_old
        w_pad = (w_old // window_size + 1) * window_size - w_old
        img_lq = torch.cat([img_lq, torch.flip(img_lq, [2])], 2)[:, :, :h_old + h_pad, :]
        img_lq = torch.cat([img_lq, torch.flip(img_lq, [3])], 3)[:, :, :, :w_old + w_pad]
        output = main_test_swinir.test(img_lq, model, args, window_size)
        output = output[..., :h_old * args.scale, :w_old * args.scale]
        return output.squeeze()

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    # 画像を2倍または4倍にします
    parser.add_argument('-s', '--scale', type=int, default=2, choices=[2, 4], help='scale factor: 2, 4')
    args = parser.parse_args()
    args.task  = 'real_sr'
    args.tile  = None

    if args.scale==2:
        # 2倍用のモデル
        args.model_path = '003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x2_GAN.pth'
        args.large_model = False
    elif args.scale==4:
        # 4倍用のモデル
        args.model_path = '003_realSR_BSRGAN_DFOWMFC_s64w8_SwinIR-L_x4_GAN.pth'
        args.large_model = True

    # 超解像のモデル
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')   # 基本CUDA専用
    model = main_test_swinir.define_model(args)
    model.eval()
    model = model.to(device)

    # 動画の入出力
    input_path  = 'input.mp4'
    output_path = 'output.mp4'
    print(f'------------------------------------------------------')
    print(f'version :swin_video={version},  ffstream={FFStream.version},  ycbcr={FFStream.ycbcr_discription}{FFStream.ycbcr_version}')
    print(f'device  :{device}')
    print(f'input   :{input_path}')
    print(f'output  :{output_path}')
    print(f'press [escape] key to cancel')
    # 入力パス 出力パス 画質20 ピクセルフォーマットRGB 音声がある場合はコピーする 出力画像サイズを2倍,4倍にする フレーム処理に使うデバイスを指定する
    ff = FFStream(input_path, output_path, crf = 20, pix_fmt='rgb24', copy_audio_stream=True, image_scale=args.scale, device=device)
    
    bar = tqdm(total=ff.frames, dynamic_ncols=True)
    for i in range(ff.frames):
        if keyboard.is_pressed('escape'):
            break
        
        # FFmpegからRGBフレームを受け取る
        frame_rgb = ff.get_rgb()
        if (frame_rgb == None) :
            break                                       # 動画の最後に達した

        # super resolution
        output = sr(model, args, frame_rgb)

        # FFmpegの出力ストリームへRGBフレームを送る
        ff.put_rgb(output)

        bar.update(1)
    bar.close()
    ff.close()

さいごに

処理時間がとんでもないことになっていますが、SwinIRで得られる画像はとても綺麗で「ここまで出来るんだ」と驚きました。

10年くらい前高速道路の違反車を特定するのに監視カメラの映像からナンバープレートを読み取る技術がテレビ番組で紹介されていました。一枚一枚の画像は解像度が低く文字の判別は出来ないけど、複数の画像から文字が視認できるまで情報を復元する技術に「これが自力でできたらいいのに」と思っていた夢の技術がもしかしたら・・・今手元にあるもので可能なのかもしれません。

You get articles that match your needs
You can efficiently read back useful information
You can use dark theme

What you can do with signing up