1.在代码中 增加了s键开始追踪 e键结束追踪 显示移动距离(代码中可调标尺和像素的比值 以便接近实际距离)
2.绘制了监测区域 只在区域内的检测
3.规定了检测的类别 只有人类才绘制轨迹
import os
import cv2
from ultralytics import YOLO
from collections import defaultdict
import numpy as np
import json
import datetime
def drawTrajectory(boxes, track_ids, track_history, track_length, img, drawing, roi):
# 绘制轨迹并计算轨迹长度
for box, track_id in zip(boxes, track_ids):
x, y, w, h = box
center = (int(x), int(y)) # 检测框的中心点
# 检查中心点是否在 ROI 内
if roi[0] < center[0] < roi[2] and roi[1] < center[1] < roi[3]:
if drawing:
track = track_history[track_id]
track.append((float(x), float(y))) # 添加中心点到轨迹历史
# 计算轨迹长度
if len(track) > 1:
for i in range(1, len(track)):
track_length[track_id] += np.linalg.norm(np.array(track[i]) - np.array(track[i - 1]))
# 绘制轨迹(无论是否正在更新轨迹历史)
if track_id in track_history:
track = track_history[track_id]
if len(track) > 1:
points = np.hstack(track).astype(np.int32).reshape((-1, 1, 2))
cv2.polylines(img, [points], isClosed=False, color=(230, 230, 230), thickness=2)
# 在图像上显示轨迹长度
actual_length = 0.5 # 实际长度(单位:米)
pixel_length = 1000 # 标尺在图像中的像素长度
pixel_to_meter_ratio = actual_length / pixel_length
print(f"ID:{track_id},移动了轨迹长度{track_length[track_id] * pixel_to_meter_ratio:.2f}")
cv2.putText(img, f"ID: {track_id}: length={track_length[track_id] * pixel_to_meter_ratio:.2f} m",
(int(x), int(y)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
if __name__ == "__main__":
# 加载配置文件
with open("config.json", "r", encoding="utf-8") as f:
config = json.load(f)
# 从配置文件中读取参数
video_path = config["video_path"]
roi = config["roi"]
model_path = config["model_path"]
# 加载 YOLO 模型
model = YOLO(model=model_path)
# 打开视频文件
cap = cv2.VideoCapture(video_path)
# 用于存储轨迹历史
track_history = defaultdict(lambda: [])
# 用于存储轨迹长度
track_length = defaultdict(lambda: 0)
# 状态标志,表示是否正在绘制轨迹
drawing = False
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
# 运行目标追踪(禁用默认的边界框绘制)
result = model.track(source=frame, persist=True, show=False, show_boxes=False)
# img = frame.copy() # 使用原始帧,而不是 YOLO 绘制的帧
img = result[0].plot()
# 获取边界框、轨迹ID和类别ID
boxes = result[0].boxes.xywh.cpu()
track_ids = result[0].boxes.id.int().cpu().tolist()
class_ids = result[0].boxes.cls.int().cpu().tolist()
# 过滤出类别为 'person' 的检测结果
person_boxes = []
person_track_ids = []
for box, track_id, class_id in zip(boxes, track_ids, class_ids):
if class_id == 0: # 0 是 'person' 类别的 ID
person_boxes.append(box)
person_track_ids.append(track_id)
# 检测开始信号和结束信号
key = cv2.waitKey(1) & 0xFF
if key == ord('s'): # 按下 's' 键表示开始信号
drawing = True
print("开始绘制轨迹")
# 清空轨迹历史和轨迹长度
track_history.clear()
track_length.clear()
elif key == ord('e'): # 按下 'e' 键表示结束信号
drawing = False
print("停止绘制轨迹")
# 在保存截图前绘制轨迹
drawTrajectory(person_boxes, person_track_ids, track_history, track_length, img, drawing, roi)
# 定义文件夹名称
output_folder = "output_images"
# 如果文件夹不存在,则创建文件夹
if not os.path.exists(output_folder):
os.makedirs(output_folder)
# 获取当前时间戳并格式化为字符串
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
# 将时间戳拼接到文件名中,并保存到指定文件夹
output_image_path = os.path.join(output_folder, f"output_frame_{timestamp}.png")
cv2.imwrite(output_image_path, img)
print(f"当前帧已保存为: {output_image_path}")
elif key == 27: # 按下 ESC 键退出
break
# 绘制 ROI 矩形
cv2.rectangle(img, (roi[0], roi[1]), (roi[2], roi[3]), (0, 255, 0), 2)
# 绘制轨迹并计算轨迹长度(仅对 ROI 内的 persons)
drawTrajectory(person_boxes, person_track_ids, track_history, track_length, img, drawing, roi)
# 显示图像
cv2.imshow("demo", img)
cap.release()
cv2.destroyAllWindows()
源码如上 现在AI遍地都是 想改写复制源码交给AI就改了