本記事では、Towards An End-to-End Framework for Flow-Guided Video Inpaintingと呼ばれる機械学習手法を用いて、ビデオから特定の物体を除去する(Video Inpainting)する方法をご紹介します。

出典: MCG-NKU/E2FGVI

Towards An End-to-End Framework for Flow-Guided Video Inpainting

概要

Towards An End-to-End Framework for Flow-Guided Video Inpainting(E2FGVI)は、Flow-Guided Video InpatingをEnd to Endに実現するフレームワークです。

簡単な理解としては、動画をフレームごとに分割したフレーム画像と、フレーム画像ごとに修復したい領域を示したマスク画像を入力に、修復したフレーム画像を出力する技術です。
この処理を全フレーム画像に施すことでVideo Inpaintingを実現します。

flow completion, feature propagation, content hallucination modulesと呼ばれる3つのモジュールを共同で最適化可能としたことで、End to End(モデル一つ)でFlow-Guided Video Inpatingタスクを実現しています。

この構成は、 YouTube-VOS、DAVISのデータセットで最先端のパフォーマンス(SOTA)であると論文で示されています。

出典: Towards An End-to-End Framework for Flow-Guided Video Inpainting

詳細はこちらの論文をご参照ください。

本記事では、上記手法を用いて、Flow-Guided Video Inpaintingする方法をご紹介します。

デモ(Colaboratory)

それでは、実際に動かしながらFlow-Guided Video Inpaintingを行っていきます。
ソースコードは本記事にも記載していますが、下記のGitHubでも取得可能です。
GitHub - Colaboratory demo

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

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

環境セットアップ

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

初めに、論文発表元のGithubからソースコードを取得します

%cd /content

!git clone https://github.com/MCG-NKU/E2FGVI.git

次に各種ライブラリをインストールします。

%cd /content

# Install Pytorch
!pip install torch==1.5.1+cu101 torchvision==0.6.1+cu101 -f https://download.pytorch.org/whl/torch_stable.html 
# Install MMCV
!pip install mmcv-full -f https://download.openmmlab.com/mmcv/dist/cu101/torch1.5/index.html
# Install gdown
!pip install --upgrade gdown

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

学習済みモデルのセットアップ

続いて学習済みモデルをセットアップしていきます。
gdownを用いてGoogle Driveからモデルをダウンロードします。

%cd /content/E2FGVI

!gdown 'https://drive.google.com/uc?id=1tNJMTJ2gmWdIXJoHVi5-H504uImUiJW9'
!unzip E2FGVI_CVPR22_models.zip

/content/E2FGVI直下に学習済みモデルがダウンロードされます。

その他のセットアップ

ライブラリをインストールします。

import matplotlib.pyplot as plt
from matplotlib import animation

import cv2
from PIL import Image
import numpy as np
import importlib
import os
import argparse
from tqdm import tqdm
import torch

from core.utils import to_tensors

mask画像を読み込むための関数などを定義します。

# global variables
w, h = 432, 240
ref_length = 10  # ref_step
num_ref = -1
neighbor_stride = 5


# sample reference frames from the whole video 
def get_ref_index(f, neighbor_ids, length):
    ref_index = []
    if num_ref == -1:
        for i in range(0, length, ref_length):
            if i not in neighbor_ids:
                ref_index.append(i)
    else:
        start_idx = max(0, f - ref_length * (num_ref//2))
        end_idx = min(length, f + ref_length * (num_ref//2))
        for i in range(start_idx, end_idx+1, ref_length):
            if i not in neighbor_ids:
                if len(ref_index) > num_ref:
                    break
                ref_index.append(i)
    return ref_index


# read frame-wise masks
def read_mask(mpath):
    masks = []
    mnames = os.listdir(mpath)
    mnames.sort()
    for mp in mnames:
        m = Image.open(os.path.join(mpath, mp))
        m = m.resize((w, h), Image.NEAREST)
        m = np.array(m.convert('L'))
        m = np.array(m > 0).astype(np.uint8)
        m = cv2.dilate(m, cv2.getStructuringElement(
            cv2.MORPH_CROSS, (3, 3)), iterations=4)
        masks.append(Image.fromarray(m*255))
    return masks


#  read frames from video
def read_frame_from_videos(video_path):
    vname = video_path
    frames = []
    lst = os.listdir(vname)
    lst.sort()
    fr_lst = [vname+'/'+name for name in lst]
    for fr in fr_lst:
        image = cv2.imread(fr)
        image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        frames.append(image.resize((w, h)))
    return frames

先ほどダウンロードした学習済みモデルをロードします。

# set up models
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net = importlib.import_module('model.e2fgvi')
model = net.InpaintGenerator().to(device)
ckpt_path = 'E2FGVI-CVPR22.pth'
data = torch.load(ckpt_path, map_location=device)
model.load_state_dict(data)
print(f'Loading model from: {ckpt_path}')
model.eval()

続いて、入力データとなるフレーム画像を動画から生成します。

%cd /content/E2FGVI/
!mkdir -p /content/E2FGVI/examples/schoolgirls

!ffmpeg -i examples/schoolgirls.mp4 examples/schoolgirls/%05d.png

最後に、入力データのパスを設定します。

# prepare dataset
video_path = 'examples/schoolgirls'
mask_path = 'examples/schoolgirls_mask'
print(f'Loading videos and masks from: {video_path}')
frames = read_frame_from_videos(video_path)
video_length = len(frames)
imgs = to_tensors()(frames).unsqueeze(0) * 2 - 1
frames = [np.array(f).astype(np.uint8) for f in frames]

masks = read_mask(mask_path)
binary_masks = [np.expand_dims((np.array(m) != 0).astype(np.uint8), 2)
                for m in masks]
masks = to_tensors()(masks).unsqueeze(0)
imgs, masks = imgs.to(device), masks.to(device)
comp_frames = [None] * video_length

Flow-Guided Video Inpainting

それでは、いよいよFlow-Guided Video Inpaintingを実行します。

# completing holes by e2fgvi
print(f'Start test...')
for f in tqdm(range(0, video_length, neighbor_stride)):
    neighbor_ids = [i for i in range(max(0, f-neighbor_stride), min(video_length, f+neighbor_stride+1))]
    ref_ids = get_ref_index(f, neighbor_ids, video_length)
    selected_imgs = imgs[:1, neighbor_ids+ref_ids, :, :, :]
    selected_masks = masks[:1, neighbor_ids+ref_ids, :, :, :]
    with torch.no_grad():
        masked_imgs = selected_imgs*(1-selected_masks)
        pred_img, _ = model(masked_imgs, len(neighbor_ids))

        pred_img = (pred_img + 1) / 2
        pred_img = pred_img.cpu().permute(0, 2, 3, 1).numpy() * 255
        for i in range(len(neighbor_ids)):
            idx = neighbor_ids[i]
            img = np.array(pred_img[i]).astype(
                np.uint8)*binary_masks[idx] + frames[idx] * (1-binary_masks[idx])
            if comp_frames[idx] is None:
                comp_frames[idx] = img
            else:
                comp_frames[idx] = comp_frames[idx].astype(
                    np.float32)*0.5 + img.astype(np.float32)*0.5

推論結果がcomp_framesにnumpy配列として出力されます。

推論結果をフレーム画像として保存します。

%cd /content/E2FGVI
!mkdir results

import matplotlib.pyplot as plt

# 推論結果出力
for i, frame in enumerate(comp_frames):
  plt.imsave('results/frames_%06d.png'%(i), frame.astype(np.uint8))

最後に出力結果を動画に変換しMoviepyで表示させます。

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

frames_path = "results/frames_%06d.png"
result_video = "results/result.mp4"

!ffmpeg -i {frames_path} -c:v libx264 -vf "fps=25,format=yuv420p" {result_video}

clip = VideoFileClip(result_video)
resize_clip = resize(clip, height=400)
resize_clip.ipython_display()

出力結果は以下の通りです。

入力動画

推論結果の動画

動画からすべての子供たちがいなくなり物悲しい動画になってしまいました。
子どもたちを除いた箇所は自然に修復され違和感がありません。

まとめ

本記事では、E2FGVIを用いたFlow-Guided Video Inpaintingを行いました。
撮影中に映りこんでしまった通行人の除去などに役立ちそうです。

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

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

参考文献

1.  論文 - Towards An End-to-End Framework for Flow-Guided Video Inpainting

2. GitHub - MCG-NKU/E2FGVI