如何在视频上测试基于Detectron2和PyTorch的自定义Faster RCNN模型

核心逻辑：视频本质是连续帧序列，你只需要把原来逐张读取图片的逻辑替换为逐帧读取视频、处理后逐帧写入输出视频即可，原有模型初始化代码不需要做任何改动。

完整可运行代码 #

第一步：复用你原有模型初始化代码 #

这部分和你现在的图片推理初始化逻辑完全一致，直接保留即可：

cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
cfg.DATASETS.TEST = ("my_dataset_test", )
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7   # set the testing threshold for this model
predictor = DefaultPredictor(cfg)
test_metadata = MetadataCatalog.get("my_dataset_test")

第二步：替换原有图片遍历逻辑为视频处理逻辑 #

import cv2
from detectron2.utils.visualizer import Visualizer
from google.colab.patches import cv2_imshow

# 替换为你自己的输入输出路径
input_video_path = "/content/input.mp4" 
output_video_path = "/content/output_detected.mp4" 

# 读取输入视频并获取基础参数
cap = cv2.VideoCapture(input_video_path)
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = int(cap.get(cv2.CAP_PROP_FPS))

# 初始化视频写入器，Colab环境推荐使用mp4v编码
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))

# 逐帧处理视频
while cap.isOpened():
    ret, frame = cap.read()
    # 读到视频末尾自动退出
    if not ret:
        break
    
    # 推理和可视化逻辑和你原有单图处理逻辑完全一致
    outputs = predictor(frame)
    v = Visualizer(frame[:, :, ::-1],
                  metadata=test_metadata, 
                  scale=1.0 # 保持1即可保证输出帧和原视频尺寸匹配，避免写入失败
                   )
    out_frame = v.draw_instance_predictions(outputs["instances"].to("cpu"))
    # 转换回OpenCV兼容的BGR格式写入视频
    result_frame = out_frame.get_image()[:, :, ::-1]
    out.write(result_frame)

# 释放所有资源
cap.release()
out.release()
cv2.destroyAllWindows()

print(f"检测完成，结果已保存至：{output_video_path}")

注意事项 #

• 运行代码前需要先把待检测的视频上传到Colab文件目录，再把input_video_path替换为你的实际上传路径
• 如果需要压缩输出视频体积，可以把Visualizer的scale参数调小，同时修改VideoWriter初始化时的宽高为(int(width*scale), int(height*scale))，保证帧尺寸匹配
• 长视频建议提前切分为短片段分批处理，避免Colab运行中断丢失结果
• 处理完成后可以在Colab左侧文件栏找到输出视频，右键即可下载到本地

内容的提问来源于stack exchange，提问作者Nidhi