基于yolov5+byte_track实现计数的代码

以下是基于 YOLOv5 和 ByteTrack 实现计数的代码示例：

import cv2
import torch
from yolov5.models.experimental import attempt_load
from yolov5.utils.general import non_max_suppression
from bytetrack.core.evaluator import create_evaluator

# 加载 YOLOv5 模型
model = attempt_load('models/yolov5s.pt', map_location='cpu')
model.eval()

# 创建 ByteTrack 模型评估器
evaluator = create_evaluator('bytetrack/configs/byte_track.yaml')

# 打开视频文件
cap = cv2.VideoCapture('path/to/video/file')

# 定义计数器
counter = 0

while True:
    # 读取视频帧
    ret, frame = cap.read()

    if not ret:
        break

    # 使用 YOLOv5 模型检测目标
    results = model(frame)

    # 过滤掉低置信度的目标
    results = non_max_suppression(results, conf_thres=0.5, iou_thres=0.5)

    # 根据目标框位置，使用 ByteTrack 进行目标跟踪
    tracked_results = evaluator.update(results, frame)

    # 统计目标数量
    num_targets = len(tracked_results)

    # 在视频帧上绘制目标框和计数器
    cv2.putText(frame, f'Targets: {num_targets}', (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2, cv2.LINE_AA)
    for result in tracked_results:
        x1, y1, x2, y2, track_id = result
        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cv2.putText(frame, str(track_id), (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)

    # 显示视频帧
    cv2.imshow('Frame', frame)
    if cv2.waitKey(1) == ord('q'):
        break

    # 更新计数器
    counter += num_targets

# 显示最终的计数结果
print(f'Total Targets: {counter}')

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

在上述代码中，我们首先加载 YOLOv5 模型和 ByteTrack 模型评估器，然后打开视频文件并定义计数器。接着，我们开始读取视频帧，并使用 YOLOv5 模型检测目标。我们过滤掉低置信度的目标，并使用 ByteTrack 进行目标跟踪。然后，我们统计目标数量并在视频帧上绘制目标框和计数器。最终，我们更新计数器并在程序结束时显示最终的计数结果。