Video inference using Detectron 2

I have gone through the notebook example from Detectron2 Object Detection Model. My question is how can I use my trained Detectron 2 model on video .mp4 file for inference?

Hi @Timothy_Malche :wave:

It is a very similar process to running that on a single image, you need to loop over frames and pass each of them separately to the model. Normally, if you would like to run the model for a single image, you would do it like that:

outputs = predictor(image)

This part of the code is identical for video prediction, but we need to read the source video and save the resulting video.

from typing import Generator
from dataclasses import dataclass

import numpy as np

import os
import cv2


# ↓ helper function to read frames from source video

def generate_frames(video_file: str) -> Generator[np.ndarray, None, None]:
    video = cv2.VideoCapture(video_file)

    while video.isOpened():
        success, frame = video.read()

        if not success:
            break

        yield frame

    video.release()

# ↓ stores information about output video file, width and height of the frame must be equal to input video

@dataclass(frozen=True)
class VideoConfig:
    fps: float
    width: int
    height: int

# ↓ reate cv2.VideoWriter object that we can use to save output video

def get_video_writer(target_video_path: str, video_config: VideoConfig) -> cv2.VideoWriter:
    video_target_dir = os.path.dirname(os.path.abspath(target_video_path))
    os.makedirs(video_target_dir, exist_ok=True)
    return cv2.VideoWriter(
        target_video_path, 
        fourcc=cv2.VideoWriter_fourcc(*"mp4v"), 
        fps=video_config.fps, 
        frameSize=(video_config.width, video_config.height), 
        isColor=True
    )

video_config = VideoConfig(
    fps=30, 
    width=1920, 
    height=1080)
video_writer = get_video_writer(
    target_video_path=TARGET_VIDEO_PATH, 
    video_config=video_config)

frame_iterator = iter(generate_frames(video_file=SOURCE_VIDEO_B_PATH))

for frame in frame_iterator:
    ...
    outputs = predictor(frame)
    
video_writer.release()

1 Like