1.在代码中 增加了s键开始追踪 e键结束追踪 显示移动距离(代码中可调标尺和像素的比值 以便接近实际距离)
2.绘制了监测区域 只在区域内的检测
3.规定了检测的类别 只有人类才绘制轨迹
import osimport cv2
from ultralytics import YOLO
from collections import defaultdict
import numpy as np
import json
import datetimedef drawTrajectory(boxes, track_ids, track_history, track_length, img, drawing, roi):# 绘制轨迹并计算轨迹长度for box, track_id in zip(boxes, track_ids):x, y, w, h = boxcenter = (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_lengthprint(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 = Falsewhile 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和类别IDboxes = 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' 类别的 IDperson_boxes.append(box)person_track_ids.append(track_id)# 检测开始信号和结束信号key = cv2.waitKey(1) & 0xFFif key == ord('s'): # 按下 's' 键表示开始信号drawing = Trueprint("开始绘制轨迹")# 清空轨迹历史和轨迹长度track_history.clear()track_length.clear()elif key == ord('e'): # 按下 'e' 键表示结束信号drawing = Falseprint("停止绘制轨迹")# 在保存截图前绘制轨迹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就改了