Detectron 2 | Video object Detection | Copy and Run

detectron2_video_od.py

# install detectron2
# git clone https://github.com/facebookresearch/detectron2.git
# cd detectron2
# pip install -e .
# cd ..


import uuid
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.data import MetadataCatalog
from detectron2.utils. visualizer import ColorMode, Visualizer
from detectron2.utils.logger import setup_logger
from detectron2 import model_zoo
from detectron2.utils.video_visualizer import VideoVisualizer
from PIL import Image

import cv2
import time,tqdm,os,ast,uuid,urllib
import numpy as np
tempathVideo  = "" #video save path
cfgDetectronVideo = get_cfg()
cfgDetectronVideo.MODEL.DEVICE = "cpu"
cfgDetectronVideo.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfgDetectronVideo.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7  # set threshold for this model
cfgDetectronVideo.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
predictor = DefaultPredictor(cfgDetectronVideo)
v = VideoVisualizer(MetadataCatalog.get(cfgDetectronVideo.DATASETS.TRAIN[0]), ColorMode.IMAGE)

print("\nVideo Object Detection module loaded\n")

# Initialize visualizer
def runOnVideo(video, maxFrames):
        """ Runs the predictor on every frame in the video (unless maxFrames is given),
        and returns the frame with the predictions drawn.
        """

        readFrames = 0
        while True:
            hasFrame, frame = video.read()
            if not hasFrame:
                break

            # Get prediction results for this frame
            outputs = predictor(frame)

            # Make sure the frame is colored
            frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)

            # Draw a visualization of the predictions using the video visualizer
            visualization = v.draw_instance_predictions(frame, outputs["instances"].to("cpu"))

            # Convert Matplotlib RGB format to OpenCV BGR format
            visualization = cv2.cvtColor(visualization.get_image(), cv2.COLOR_RGB2BGR)

            yield visualization

            readFrames += 1
            if readFrames > maxFrames:
                break                                               
            
def showPerson(file_path): #run showpath with path of mp4 video
    
    video = cv2.VideoCapture(file_path)
    width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
    frames_per_second = video.get(cv2.CAP_PROP_FPS)
    num_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))

    # Initialize video writer
    path = os.path.join(tempathVideo, f'out_{str(uuid.uuid4())}_.mp4') 
    video_writer = cv2.VideoWriter(path, fourcc=cv2.VideoWriter_fourcc(*"mp4v"), fps=float(frames_per_second), frameSize=(width, height), isColor=True)

    

    # Create a cut-off for debugging
    num_frames = 120

    # Enumerate the frames of the video
    for visualization in tqdm.tqdm(runOnVideo(video, num_frames), total=num_frames):
        # Write to video file
        video_writer.write(visualization)

    # Release resources
    video.release()
    video_writer.release()
    return path
  
showPerson("path of video") #return path of saved video with OD