如何在视频上正确运行Detectron2? [英] How can I properly run Detectron2 on Videos?
问题描述
我目前正在与Detectron2合作,用于视频中的人侦听器,我一直在尝试运行以下代码来读取视频文件,逐帧进行预测并以处理后的帧录制视频,但是我得到一个空的视频文件.我为此创建的环境位于Colaboratory中,并具有以下版本(python 3.6,opencv 4.2.30). 我是新手,但如果有人可以给我一个主意,我将不胜感激
I'm currently working with Detectron2 for people detectios in videos, I've been trying to run the following code to read a video file, make the prediction frame by frame and record a video with the processed frames, but I am getting an empty video file. The environment that I've created for this is located in Colaboratory and has the following versions (python 3.6 , opencv 4.2.30). I'm new on this but I would truly appreciate if someone could give me an idea
这是代码
#!/usr/bin/env python3
# -- coding: utf-8 --
import detectron2
from detectron2.utils.logger import setup_logger
setup_logger()
# import some common libraries
import numpy as np
import cv2
# import some common detectron2 utilities
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog
import time
cap = cv2.VideoCapture('piloto legger 1.mp4')
hasFrame, frame = cap.read()
FPS = cap.get(cv2.CAP_PROP_FPS)
frame_width = frame.shape[1]
frame_height = frame.shape[0]
video_writer = cv2.VideoWriter('out.mp4', cv2.VideoWriter_fourcc('m', 'p', '4', 'v'), FPS, (frame_width, frame_height))
while cv2.waitKey(1) < 0:
hasFrame, frame = cap.read()
if not hasFrame:
cv2.waitKey()
break
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-Keypoints/keypoint_rcnn_R_50_FPN_3x.yaml"))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7 # set threshold for this model
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Keypoints/keypoint_rcnn_R_50_FPN_3x.yaml")
predictor = DefaultPredictor(cfg)
outputs = predictor(frame)
v = Visualizer(frame[:,:,::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
imagen = (v.get_image()[:, :, ::-1])
cv2.imwrite('POSE detectron2.png', imagen)
video_writer.write(imagen)
cap.release()
video_writer.release()
cv2.destroyAllWindows()
推荐答案
我以您的代码为起点,并从Detectron2示例中获取了一些想法,以使其能够正常工作.
I used your code as a starting point, and took some ideas from the Detectron2 examples in order to make it work.
问题似乎与VideoWriter
的fourcc参数有关,但也可能与使用Visualizer
而不是VideoVisualizer
的代码有关(并且标度为1.2,使得图像的VideoWriter
尺寸错误).
The problem seems to have been something with the fourcc-argument of the VideoWriter
, but may also have been related to your code using Visualizer
instead of VideoVisualizer
(and with a scale of 1.2, which made the image the wrong size for the VideoWriter
).
下面的代码对我有用(而且预测器和可视化器是在循环外部定义的,而且速度也快得多):
The code below works for me (and is also a lot faster, as the predictor and visualizer are defined outside of the loop):
#!/usr/bin/env python3
# -- coding: utf-8 --
import detectron2
from detectron2.utils.logger import setup_logger
setup_logger()
# import some common libraries
import numpy as np
import tqdm
import cv2
# import some common detectron2 utilities
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.video_visualizer import VideoVisualizer
from detectron2.utils.visualizer import ColorMode, Visualizer
from detectron2.data import MetadataCatalog
import time
# Extract video properties
video = cv2.VideoCapture('video-input.mp4')
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
video_writer = cv2.VideoWriter('out.mp4', fourcc=cv2.VideoWriter_fourcc(*"mp4v"), fps=float(frames_per_second), frameSize=(width, height), isColor=True)
# Initialize predictor
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7 # set threshold for this model
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
predictor = DefaultPredictor(cfg)
# Initialize visualizer
v = VideoVisualizer(MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), ColorMode.IMAGE)
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
# 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 test image
cv2.imwrite('POSE detectron2.png', visualization)
# Write to video file
video_writer.write(visualization)
# Release resources
video.release()
video_writer.release()
cv2.destroyAllWindows()
这篇关于如何在视频上正确运行Detectron2?的文章就介绍到这了,希望我们推荐的答案对大家有所帮助,也希望大家多多支持IT屋!