Detectron2 实时摄像头目标检测指南

2023-12-03 08:01:05

探索 Detectron2：使用摄像头进行实时目标检测

在计算机视觉领域，物体检测是一项必不可少的技术，它为图像中的物体识别和定位提供了依据。Detectron2 是 Facebook AI Research 开发的一个灵活且高效的物体检测框架，它使得在实际场景中轻松实现物体检测成为可能。本文将逐步指导您使用 Detectron2 在实时视频流中进行目标检测，让您亲身体验这一强大技术的魅力。

Prerequisites

使用摄像头进行目标检测，你需要以下配置：

安装好 Python 3.6 或更高版本。
安装 Detectron2，请参考官方文档。
配置一个摄像头，确保它已连接到您的计算机。

初始化

在开始之前，让我们导入必要的库并初始化 Detectron2：

import cv2
import detectron2
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg

加载配置和模型

Detectron2 提供了预先训练好的模型，存储在 model zoo 中。在此示例中，我们将使用在 COCO 数据集上训练的 R50-FPN 3x 模型：

cfg = get_cfg()
cfg.merge_from_file("configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
cfg.MODEL.WEIGHTS = "detectron2://COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x/137849600/model_final_f10217.pkl"
predictor = DefaultPredictor(cfg)

打开摄像头

要访问摄像头，我们将使用 OpenCV 库：

cap = cv2.VideoCapture(0)

实时目标检测

现在，我们准备开始实时目标检测：

while True:
    ret, frame = cap.read()
    outputs = predictor(frame)
    visualized_output = visualize_output(frame, outputs)
    cv2.imshow("Detection", visualized_output)
    if cv2.waitKey(1) & 0xFF == ord("q"):
        break
cap.release()
cv2.destroyAllWindows()

输出可视化

import numpy as np

def visualize_output(frame, outputs):
    predictions = outputs["instances"].to("cpu")
    boxes = predictions.pred_boxes.tensor.numpy()
    classes = predictions.pred_classes.numpy()
    scores = predictions.scores.numpy()

    for box, cls, score in zip(boxes, classes, scores):
        if score < 0.5:
            continue
        cls_name = cfg.MODEL.ROI_HEADS.NAME_CLASS_MAP[cls]
        color = np.random.rand(3) * 255
        cv2.rectangle(frame, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), color, 2)
        cv2.putText(frame, cls_name, (int(box[0]), int(box[1] - 5)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)

    return frame