Segmentationで人を塗りつぶす (前回、前々回の続き)

はじめに

touch-sp.hatenablog.com
touch-sp.hatenablog.com
今回はバックも動画にしてみた。

結果

いつものように「output.mp4」をGIFに変換したものをのせておく。

ffmpeg -i output.mp4 -vf scale=320:-1 output.gif

f:id:touch-sp:20201121000814g:plain
Music Videoに出てきそうないい感じのものができた。
今回ダンス動画を変更したのは人物が占める割合を大きくしたかったから。

背景動画のダウンロード

背景動画はこちらからからダウンロードした。
サイズは640×360を指定して「back.mp4」に名前を変更した。
セグメンテーションがうまくいっていない部分を隠すためPythonスクリプトの中でRGB→BGR変換をして赤っぽくしている。

ダンス動画ファイルのダウンロード

動画はこちらからダウンロードした。
サイズは1280×720を指定して「dance.mp4」に名前を変更した。Pythonスクリプトの中で640×360に変換している。

Pythonスクリプト

from mxnet.gluon.data.vision.transforms import Resize
import numpy as np
import mxnet as mx
from mxnet.gluon.data.vision import transforms
from numpy.core.fromnumeric import transpose
import gluoncv

import decord
import cv2
from PIL import Image

ctx = mx.gpu()
model = gluoncv.model_zoo.get_model('deeplab_resnet152_voc', pretrained=True, root='./models', ctx=ctx)

vr = decord.VideoReader('dance.mp4')
decord.bridge.set_bridge('mxnet')
all_frames = vr.get_batch(range(100,300))
fps = vr.get_avg_fps()

back_vr = decord.VideoReader('back.mp4')

height, width = (360, 640)
transform_fn = transforms.Compose([
    transforms.Resize((width, height)),
    transforms.ToTensor(),
    transforms.Normalize([.485, .456, .406], [.229, .224, .225])
])

imgs = transform_fn(all_frames).as_in_context(ctx)

fourcc = cv2.VideoWriter_fourcc('m','p','4','v')
out = cv2.VideoWriter('output.mp4',fourcc, fps, (width,height))

for i in range(all_frames.shape[0]):

    front_img = Image.fromarray(all_frames[i].asnumpy()).resize((width, height))
    back_img = Image.fromarray(back_vr[i].asnumpy()).resize((width, height))
    
    output = model.predict(imgs[i:i+1])
    
    predict = mx.nd.squeeze(mx.nd.argmax(output, 1)).asnumpy()
    person = np.where(predict ==15, 0, 255).astype('int32')

    if i == 0:
        before_person = person
    else:
        if ((np.sum(person==0) < np.sum(before_person==0)*0.5) or (np.sum(person==0) > np.sum(before_person==0)*1.5)):
            person = before_person
        else:
            before_person = person

    b = Image.fromarray(person).convert('L')
    
    back_img.paste(front_img, b)
    img = np.array(back_img)
    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    out.write(img)

out.release()


このエントリーをはてなブックマークに追加