roboflow的trackers追踪

susu10830189112026-04-02 14:29

git国内地址： https://gitcode.com/GitHub_Trending/track/trackers

github地址：https://github.com/roboflow/trackers

官网文档地址：https://trackers.roboflow.com/develop/trackers/bytetrack/

环境可以在之前的一篇rf-detr环境基础上安装即可。

pip install trackers

额外的库：

pip install "trackers[detection]"

从源代码安装

pip install git+https://github.com/roboflow/trackers.git

运行检测追踪代码：

复制代码

import cv2
import supervision as sv
# from rfdetr import RFDETRLarge
from rfdetr.platform.models import RFDETR2XLarge

from trackers import ByteTrackTracker
tracker = ByteTrackTracker()

# from trackers import SORTTracker
# tracker = SORTTracker()


model = RFDETR2XLarge()

box_annotator = sv.BoxAnnotator()
label_annotator = sv.LabelAnnotator()

video_capture = cv2.VideoCapture("input.mp4")
if not video_capture.isOpened():
    raise RuntimeError("Failed to open video source")

while True:
    success, frame_bgr = video_capture.read()
    if not success:
        break

    frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
    detections = model.predict(frame_rgb)
    detections = tracker.update(detections)

    annotated_frame = box_annotator.annotate(frame_bgr, detections)
    annotated_frame = label_annotator.annotate(
        annotated_frame,
        detections,
        labels=detections.tracker_id,
    )

    cv2.imshow("RF-DETR + ByteTrack", annotated_frame)
    if cv2.waitKey(1) & 0xFF == ord("q"):
        break

video_capture.release()
cv2.destroyAllWindows()

注意如果想使用RFDETR2XLarge模型：需要安装扩展：

pip install rfdetr[plus]

调用：

from rfdetr.platform.models import RFDETR2XLarge