写一个yolov5基于目标追踪的双向计数demo

当然，我可以为你提供一个基于目标追踪的Yolov5双向计数的demo。下面是一个简单的示例代码：

import cv2
import torch
from models.experimental import attempt_load
from utils.general import non_max_suppression, scale_coords
from utils.torch_utils import select_device

# 加载Yolov5模型
weights = 'yolov5s.pt'
device = select_device('')
model = attempt_load(weights, map_location=device)
stride = int(model.stride.max())

# 设置视频输入
video_path = 'input.mp4'
cap = cv2.VideoCapture(video_path)
cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
ret, frame = cap.read()

# 设置目标追踪器
tracker = cv2.legacy.TrackerKCF_create()

# 初始化计数器
total_count = 0
enter_count = 0
exit_count = 0

while ret:
    # 目标检测
    img = torch.from_numpy(frame).to(device)
    img = img.float() / 255.0
    if img.ndimension() == 3:
        img = img.unsqueeze(0)
    pred = model(img)[0]
    pred = non_max_suppression(pred, 0.4, 0.5)

    # 目标追踪和计数
    for det in pred[0]:
        if det is not None:
            det = scale_coords(img.shape[2:], det.unsqueeze(0), frame.shape).round()
            xmin, ymin, xmax, ymax = int(det[0][0]), int(det[0][1]), int(det[0][2]), int(det[0][3])
            bbox = (xmin, ymin, xmax - xmin, ymax - ymin)

            # 初始化目标追踪器
            tracker.init(frame, bbox)

            # 更新目标追踪
            track_success, bbox = tracker.update(frame)

            if track_success:
                # 绘制追踪框
                cv2.rectangle(frame, (int(bbox[0]), int(bbox[1])),
                              (int(bbox[0] + bbox[2]), int(bbox[1] + bbox[3])), (0, 255, 0), 2)

                # 判断目标进出
                center_x = int(bbox[0] + bbox[2] / 2)
                if center_x < frame.shape[1] / 2:
                    enter_count += 1
                else:
                    exit_count += 1
                total_count = enter_count - exit_count

    # 在图像上绘制计数信息
    cv2.putText(frame, f"Total Count: {total_count}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
    cv2.putText(frame, f"Enter Count: {enter_count}", (10, 70), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
    cv2.putText(frame, f"Exit Count: {exit_count}", (10, 110), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)

    # 显示图像
    cv2.imshow('Demo', frame)
    
    # 退出循环条件
    if cv2.waitKey(1) == ord('q'):
        break

    ret, frame = cap.read()

cap.release()
cv2.destroyAllWindows()

请确保已经安装了必要的依赖项，并且已经下载了Yolov5的权重文件（yolov5s.pt）。将视频文件命名为"input.mp4"并与代码放在同一个目录下。运行代码后，它将显示一个窗口，显示每个目标的追踪框以及计数信息。你可以根据自己的需求进行修改和定制。希望对你有帮助！