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就改了?