ブルーバックとか使わずにクロマキー合成する

ハッピバースデートゥミ〜〜〜♪ハッピバースデートゥミ〜〜〜♪、ハッピバースデーディアわたし〜♪
ハッピバースデートゥミ〜〜〜♪

くら＠ちゅらデータ（43）です。

この投稿はちゅらデータアドベントカレンダー2023の15日目の記事です。

OBSで人と背景をクロマキー合成するときに背景を単色にする必要があるんですが大変なんでどうにかしたいなーって話です。

はじめに

オンラインミーティングだとmeetやzoomの機能で綺麗に人と背景を合成してくれて散らかった部屋を晒さなくて済むんですが、OBSで合成する時が困るんですよね。
OBSのインプットの時点で背景を単色にできればスムーズに合成できるなと

今回はmeetで使われているかもしれないMediaPipeを使って実現してみます
MediaPipeのデモで動作は確認できるんですが今回はローカルのPythonで実行させます

環境構築

MediaPipeのライブラリをインスールします

$ python -m pip install mediapipe

MediaPipe以外で必要なライブラリも入れておきます

$ python -m pip install opencv-python numpy

ついでに後で使うモデルもダウンロード

$ wget -O deeplabv3.tflite -q https://storage.googleapis.com/mediapipe-models/image_segmenter/deeplab_v3/float32/1/deeplab_v3.tflite

準備はこれでOK

早速動かしてみます

まず、OpneCVを使ってWebカメラから画像を読み込んで表示するプログラムを書きます

main1.py

import cv2
import time

def main():
    cap = cv2.VideoCapture(0)

    prev_time = time.time()
    while True:
        # カメラから画像を読み込む
        ret, frame = cap.read()
        if not ret:
            break

        # 表示
        cv2.imshow('Face Overlay', frame)

        # FPS算出
        current_time = time.time()
        print(f"FPS: {1 / (current_time - prev_time):.2f}", end='\r')
        prev_time = current_time

        # キーボードのqキーで終了
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()

ほぼカメラスペック通りの30FPSくらい出てる感じ

次にMediaPipeのラッパーを作成します。

MediaPipeSegmentWrapper1.py

import mediapipe as mp
import numpy as np

#背景と人物を分けるクラス
class MediaPipeSegmentWrapper:
    def __init__(self):
        # MediaPipe セルフセグメンテーションの初期化 
        self.segmenter = mp.solutions.selfie_segmentation.SelfieSegmentation(model_selection=1)

    def process_image(self, image):
        # イメージの処理
        results = self.segmenter.process(image)

        # セグメンテーションマスクの取得
        return np.stack((results.segmentation_mask,) * 3, axis=-1) > 0.1

main関数にラッパーを組み込んで、マスクの結果を使って背景を緑にする関数を通す処理に変更します

main1.py

import cv2
import numpy as np
import time
from MediaPipeSegmentWrapper3 import MediaPipeSegmentWrapper

#背景を緑にする関数
def apply_green_background(image, mask):
    green_bg = np.zeros_like(image)
    green_bg[:] = [0, 255, 0]  # 緑色
    return np.where(mask, image, green_bg)

def main():
    cap = cv2.VideoCapture(0)
    
　　　　　　　　segmenter = MediaPipeSegmentWrapper()#MediaPipeのラッパーを定義
    
    prev_time = time.time()
    while True:
        # カメラから画像を読み込む
        ret, frame = cap.read()
        if not ret:
            break

　　　　　　　　　　　　　　　　mask = segmenter.process_image(frame)#セグメンテーション
        frame = apply_green_background(frame, mask)

        # 表示
        cv2.imshow('Face Overlay', frame)

        # FPS算出
        current_time = time.time()
        print(f"FPS: {1 / (current_time - prev_time):.2f}", end='\r')
        prev_time = current_time

        # キーボードのqキーで終了
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()

結果がこちら、そこそこマスクできていますね。もうこれでいいんじゃないかという感じ
１５FPSくらいでした、サイトの情報だとCPUでのレイテンシーが３４msecくらいなので
処理分フレームレートが落ちたと思えばちょうどいい感じ

次はMediaPipeのサイトでサンプルとして上がっているDeepLabV3を使ってみましょう
モデルを指定するやり方は少し違っていて下記のようになります

MediaPipeSegmentWrapper2.py

import cv2
import numpy as np
import mediapipe as mp
from mediapipe.tasks import python
from mediapipe.tasks.python import vision

class MediaPipeSegmentWrapper:
    def __init__(self, model_path='./deeplabv3.tflite'):
        # MediaPipe セグメンテーションの初期化
        self.segmenter = vision.ImageSegmenter.create_from_options(
            vision.ImageSegmenterOptions(
                base_options=python.BaseOptions(model_asset_path=model_path),
                running_mode=mp.tasks.vision.RunningMode.IMAGE,
                output_category_mask=True))

    def process_image(self, image):
        # 画像フォーマットを指定してMediaPipe Imageオブジェクトを作成
        mp_image = mp.Image(image_format=mp.ImageFormat.SRGB,
                            data=cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        # セグメンテーションの実行
        segmentation_result = self.segmenter.segment(mp_image)

        # マスクの取得
        return np.stack((segmentation_result.category_mask.numpy_view(),) * 3, axis=-1) > 0.1

結果はちょっとよくなっている？顔の周りは改善されてますが体は微妙ですね
９FPSくらいなのですが、CPUのレイテンシーは124msecとのことなので早すぎますね。
何ででしょうか？

pythonの実行画面を OBSのソースとして取り込んでしまえば簡単に合成できます

MediaPipe自体には顔トラッキングやボディトラッキングなど他に色々遊べそうなものがあるので気が向いたら触ってみようと思います

おまけ

running_modeっていうパラメータあるんですがIMAGEを指定してます
これ他にVIDEOとLIVE_STREAMを指定できて、なんか動作早そうだったんで試したんですがサンプルコードだと動作しません。ネットで探してみたんですが見つけきれず😅
ドキュメントとコードを読んだ結果動かせるようになったので乗っけておきます、DeepLabV3でも18FPSくらいになりました

MediaPipeSegmentWrapper３.py

import cv2
import numpy as np
import mediapipe as mp
from mediapipe import Image
from mediapipe.tasks import python
from mediapipe.tasks.python import vision
from mediapipe.tasks.python.vision.image_segmenter import ImageSegmenterResult

class MediaPipeSegmentWrapper:
    def __init__(self, model_path='./deeplabv3.tflite', invert_mask=False):
        self.segmenter = vision.ImageSegmenter.create_from_options(
            vision.ImageSegmenterOptions(
                base_options=python.BaseOptions(model_asset_path=model_path),
                running_mode=mp.tasks.vision.RunningMode.LIVE_STREAM,
                output_category_mask=True,
                result_callback=self.callback))
        self.frame_timestamp_ms = 0
        self.processed_image = None 
        self.invert_mask = invert_mask

    def process_image(self, image):
        self.frame_timestamp_ms += 1
        mp_image = mp.Image(image_format=mp.ImageFormat.SRGB,
                            data=cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        self.segmenter.segment_async(mp_image, self.frame_timestamp_ms)
        return None

    def callback(self, result: ImageSegmenterResult, output_image: Image, timestamp_ms: int):
        self.processed_image = np.stack((result.category_mask.numpy_view(),) * 3, axis=-1) > 0.1
        if (self.invert_mask):
            self.processed_image = np.logical_not(self.processed_image)

    def get_processed_image(self):
        return self.processed_image