【opencv-python】opencv-pythonとPyTorchでカメラ入力を操作する（背景をぼかす）

はじめに

前回までになんちゃってズームとなんちゃって首振りを実装しました。

なんちゃってズーム

touch-sp.hatenablog.com

なんちゃって首振り

touch-sp.hatenablog.com



今回はPyTorchを使って人物以外の背景をぼかすことに挑戦しました。


以前このような記事を書いています。そちらのスクリプトを動画用に書き換えました。
touch-sp.hatenablog.com

Pythonスクリプト

今まで実装したすべてを詰め込みました。
1280×720の入力を720×720で出力するようにしています。
カメラがその入力に対応していなければ解像度を変更する必要があります。
こちらを参照して下さい。

import numpy as np
import torch
from torchvision import transforms
import cv2

device = 'cuda' if torch.cuda.is_available() else 'cpu'

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

model = torch.hub.load('pytorch/vision:v0.10.0', 'deeplabv3_resnet50', pretrained=True)
model.eval().to(device)

leftside  = 280
cropscale = 0
bokasi_strength = 0

w = 1280
h = 720
cap=cv2.VideoCapture(0)
cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'))
cap.set(cv2.CAP_PROP_FRAME_WIDTH, w)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, h)

while True:
    ret, frame = cap.read()

    frame = frame[:, leftside:leftside+720, :]

    if cropscale != 0:
        crop_h = 30 * cropscale
        crop_w = 30 * cropscale
        frame = cv2.resize(frame[crop_h:-crop_h, crop_w:-crop_w, :], dsize = (720, 720))

    if bokasi_strength != 0:
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        input_tensor = transform_fn(frame_rgb)
        input_batch = input_tensor.unsqueeze(0)
    
        with torch.no_grad():
            output = model(input_batch.to(device))['out'][0]
        
        predict = output.argmax(0).to('cpu').numpy()

        mask_1 = np.where(predict == 15, 1, 0)[...,np.newaxis]
        mask_2 = np.where(predict == 15, 0, 1)[...,np.newaxis]
        blurred_img = cv2.blur(frame, (bokasi_strength, bokasi_strength))
        frame = (frame * mask_1 + blurred_img * mask_2).astype('uint8')

    cv2.imshow('demo', frame)

    k = cv2.waitKey(1) & 0xFF
    if k == 27:
        break
    elif k == ord('a'):
        if leftside < 560:
            leftside += 70
    elif k == ord('d'):
        if leftside > 0:
            leftside -= 70
    elif k == ord('w'):
        if cropscale < 4:
            cropscale += 1
    elif k == ord('s'):
        if cropscale > 0:
            cropscale -= 1
    elif k == ord('q'):
        if bokasi_strength < 15:
            bokasi_strength += 5
    elif k == ord('e'):
        if bokasi_strength > 0:
            bokasi_strength -= 5

cap.release()
cv2.destroyAllWindows()

使い方

「a」「d」入力で首振り。
「w」「s」入力でズーム。
「q」「e」入力でぼかし強度の変更。