我想用 Detectron2 创建身体姿势估计器。我一直在尝试在本地运行以下代码来读取视频文件,逐帧进行预测并使用处理后的帧录制视频。但它不能正常工作。如果我可以这样说,它会创建空文件并显示此类错误(下图)错误。所以我寻求帮助,因为我不明白出了什么问题。
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('D:/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('D:/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()