概要

前提
コード

前提

configファイルがある
学習済みモデル（チェックポイントファイル）がある
動画ファイルがある
検出するクラスが1つのみ（多クラスはこのコードを改造すればいけるかも．．．）

コード

以下のコードを参考にした．

createMaskVideo.py

from argparse import ArgumentParser
import cv2
import numpy as np
from mmengine.model.utils import revert_sync_batchnorm
from mmengine.visualization.utils import tensor2ndarray
from mmseg.apis import inference_model, init_model
from mmseg.apis.inference import show_result_pyplot

def main():
    parser = ArgumentParser()
    parser.add_argument('video', help='Video file or webcam id')
    parser.add_argument('config', help='Config file')
    parser.add_argument('checkpoint', help='Checkpoint file')
    parser.add_argument(
        '--device', default='cuda:0', help='Device used for inference')
    parser.add_argument(
        '--output-file', default=None, type=str, help='Output video file path')
    args = parser.parse_args()

    assert args.output_file, \
        'At least one output should be enabled.'

    # build the model from a config file and a checkpoint file
    model = init_model(args.config, args.checkpoint, device=args.device)
    if args.device == 'cpu':
        model = revert_sync_batchnorm(model)

    # build input video
    # 動画を読み込む
    if args.video.isdigit():
        args.video = int(args.video)
    cap = cv2.VideoCapture(args.video)
    assert (cap.isOpened())
    input_height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
    input_width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
    input_fps = cap.get(cv2.CAP_PROP_FPS)

    # init output video
    # 出力動画の設定をする
    writer = None
    output_height = None
    output_width = None
    if args.output_file is not None:
        fourcc = cv2.VideoWriter_fourcc('m','p','4','v')
        output_fps = input_fps
        output_height = int(input_height)
        output_width = int(input_width)
        writer = cv2.VideoWriter(args.output_file, fourcc, output_fps,
                                 (output_width, output_height), False)

    # start looping
    try:
        while True:
            # frameを取得する
            flag, frame = cap.read()
            if not flag:
                break

            # test a single image
            # セマンティックセグメンテーションを行う
            result = inference_model(model, frame)

            # クラスIDが0のマスクを取り出す
            ids = np.unique(result.pred_sem_seg.cpu().data)[::-1]
            labels = np.array(ids, dtype=np.int64)
            #labels[1]やlabels[2]に変えると，別のクラスのマスク動画を作ることができるようになる
            binary_masks = result.pred_sem_seg.cpu().data == labels[0]
            binary_masks = tensor2ndarray(binary_masks)
            binary_masks = binary_masks.astype('uint8') * 255
            mask = np.zeros_like(frame)
            for b in binary_masks:
                m = np.zeros_like(frame)
                m[...] = (1, 1, 1)
                m = cv2.bitwise_and(m, m, mask = b)
                mask = mask + m
            _, mask = cv2.threshold(mask, 0, 255, cv2.THRESH_BINARY)
            mask = mask[:, :, 0]

            # 動画を出力する
            # この出力動画をffmpegに通して動画サイズを小さくしても良い
            if writer:
                writer.write(np.array(mask, dtype = 'uint8'))
    finally:
        if writer:
            writer.release()
        cap.release()


if __name__ == '__main__':
    main()

ターミナルで以下のコマンドにて使用できる．

$ python createMaskVideo.py [動画ファイルパス] [モデルのconfigファイルパス] [モデルのチェックポイントファイルパス] --output-file [出力動画ファイルパス]

mmsegmentation v3で動画からあるクラスのマスク動画を作る

概要

前提

コード