【MXNet】Segmentationで人を塗りつぶす （前回の続き）

前回は人を白く塗りつぶした。
touch-sp.hatenablog.com
今回は静止画と動画の両方でバックをつけてみた。

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

動画はこちらからダウンロードさせて頂いた。
（Video by Adesh Kumar Singh from Pixabay）
サイズは640×360を指定して「dance.mp4」に名前を変更した。

静止画をバックにした結果

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

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

背景画像のダウンロード

画像はこちらからダウンロードして「back.png」とした。

Pythonスクリプト

import numpy as np
import mxnet as mx
from mxnet.gluon.data.vision import transforms
import gluoncv

import decord
import cv2
from PIL import Image

transform_fn = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize([.485, .456, .406], [.229, .224, .225])
])

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(len(vr)))
fps = vr.get_avg_fps()
height, width = vr[0].shape[0:2]

back_img = Image.open('back.png')
back_img = back_img.resize((width, height))

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(len(vr)):

    front_img = Image.fromarray(vr[i].asnumpy())

    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 = back_img.copy()
    back.paste(front_img, b)
    img = np.array(back)
    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    out.write(img)

out.release()

動画をバックにした結果

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

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

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

背景動画のダウンロード

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

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()