|
|
@@ -0,0 +1,196 @@
|
|
|
+import cv2
|
|
|
+import numpy as np
|
|
|
+from ultralytics import YOLO
|
|
|
+from collections import defaultdict, deque
|
|
|
+import datetime
|
|
|
+import time
|
|
|
+import math
|
|
|
+# Load the YOLO11 model
|
|
|
+model = YOLO("yolo11m.pt")
|
|
|
+
|
|
|
+# Open the video file
|
|
|
+video_path = r"E:\desktop_file\速度标定\run.mp4"
|
|
|
+# video_path = r"E:\wx_file\WeChat Files\wxid_1lcmt2w2jdwl22\FileStorage\File\2024-11\3.4-13时胶乳包装.mp4"
|
|
|
+cap = cv2.VideoCapture(video_path)
|
|
|
+
|
|
|
+# 存储最近的200帧用于回溯
|
|
|
+frame_buffer = deque(maxlen=200) # 新增帧缓冲区
|
|
|
+
|
|
|
+# Store the track history
|
|
|
+track_history = defaultdict(lambda: [])
|
|
|
+# 用于存储每个 track_id 最近的时间戳
|
|
|
+time_stamps = defaultdict(lambda: deque(maxlen=200)) # 固定长度为 50
|
|
|
+# 用于存储瞬时速度
|
|
|
+instantaneous_velocities = defaultdict(lambda: deque(maxlen=100))
|
|
|
+
|
|
|
+
|
|
|
+def apply_bias(position):
|
|
|
+ """
|
|
|
+ 偏置函数：使用 x/ln(1+x) 计算偏置
|
|
|
+ 已弃用
|
|
|
+ """
|
|
|
+ x, y = position
|
|
|
+ bias_x = np.log1p(x) if x > 0 else 0
|
|
|
+ bias_y = np.log1p(y) if y > 0 else 0
|
|
|
+ return np.array([bias_x, bias_y])
|
|
|
+
|
|
|
+
|
|
|
+def save_high_speed_video(buffer, trigger_time):
|
|
|
+ """将缓冲区中的帧保存为MP4文件"""
|
|
|
+ if len(buffer) < 1:
|
|
|
+ return
|
|
|
+
|
|
|
+ # 生成唯一文件名
|
|
|
+ timestamp = trigger_time.strftime("%Y%m%d%H%M%S%f")
|
|
|
+ output_path = f"high_speed_{timestamp}.mp4"
|
|
|
+
|
|
|
+ # 使用MP4编码（需确保OpenCV支持）
|
|
|
+ fourcc_mp4 = cv2.VideoWriter_fourcc(*'x264')
|
|
|
+ writer = cv2.VideoWriter(output_path, fourcc_mp4, fps, (frame_width, frame_height))
|
|
|
+
|
|
|
+ for frame in buffer:
|
|
|
+ writer.write(frame)
|
|
|
+ writer.release()
|
|
|
+
|
|
|
+
|
|
|
+def map_to_ellipse(position):
|
|
|
+ x, y = position
|
|
|
+ center_x = 640
|
|
|
+ center_y = 360
|
|
|
+ a = 580
|
|
|
+ b = 280
|
|
|
+
|
|
|
+ x_norm = x / 1280
|
|
|
+ y_norm = y / 720
|
|
|
+
|
|
|
+ d_norm = math.sqrt((x_norm - 0.5) ** 2 + (y_norm - 0.5) ** 2)
|
|
|
+ theta_norm = math.atan2(y_norm - 0.5, x_norm - 0.5)
|
|
|
+ f = d_norm
|
|
|
+ a_new = a * f
|
|
|
+ b_new = b * f
|
|
|
+
|
|
|
+ bias_x = center_x + a_new * math.cos(theta_norm)
|
|
|
+ bias_y = center_y + b_new * math.sin(theta_norm)
|
|
|
+
|
|
|
+ return np.array([bias_x, bias_y])
|
|
|
+
|
|
|
+# 创建 VideoWriter 对象以保存输出视频
|
|
|
+fourcc = cv2.VideoWriter_fourcc(*'XVID') # 视频编码格式
|
|
|
+output_file = "output_video.avi" # 输出文件名
|
|
|
+fps = 25 # 帧率
|
|
|
+frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
|
|
+frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
|
|
+out = cv2.VideoWriter(output_file, fourcc, fps, (frame_width, frame_height))
|
|
|
+
|
|
|
+speed_threshold = 30 # 速度阈值
|
|
|
+high_velocity_count_threshold = 20 # 高速度计数阈值
|
|
|
+
|
|
|
+# Loop through the video frames
|
|
|
+while cap.isOpened():
|
|
|
+ # 记录当前时间
|
|
|
+ current_time = time.time()
|
|
|
+
|
|
|
+ # Read a frame from the video
|
|
|
+ success, frame = cap.read()
|
|
|
+
|
|
|
+ if success:
|
|
|
+ # 将当前帧加入缓冲区（深拷贝避免覆盖）
|
|
|
+ frame_buffer.append(frame.copy()) # 新增
|
|
|
+
|
|
|
+ # Run YOLO11 tracking on the frame, persisting tracks between frames
|
|
|
+ results = model.track(frame, persist=True, classes=0, conf=0.6)
|
|
|
+
|
|
|
+ if results[0].boxes and results[0].boxes.id is not None:
|
|
|
+ # Get the boxes and track IDs
|
|
|
+ boxes = results[0].boxes.xywh.cpu()
|
|
|
+ track_ids = results[0].boxes.id.int().cpu().tolist()
|
|
|
+
|
|
|
+ for box, track_id in zip(boxes, track_ids):
|
|
|
+ x, y, w, h = box
|
|
|
+
|
|
|
+ # 绘制边界框
|
|
|
+ cv2.rectangle(frame, (int(x - w / 2), int(y - h / 2)), (int(x + w / 2), int(y + h / 2)), (0, 255, 0), 2)
|
|
|
+ # 计算左下角坐标
|
|
|
+ bottom_left_x = int(x - w / 2)
|
|
|
+ bottom_left_y = int(y + h / 2)
|
|
|
+
|
|
|
+ # 计算中心点
|
|
|
+ center_x = int(x)
|
|
|
+ center_y = int(y)
|
|
|
+
|
|
|
+ # 绘制中心点
|
|
|
+ cv2.circle(frame, (center_x, center_y), 5, (255, 0, 0), -1) # 红色中心点，半径为 5
|
|
|
+
|
|
|
+ # 记录位置
|
|
|
+ track_history[track_id].append((bottom_left_x, bottom_left_y))
|
|
|
+ if len(track_history[track_id]) > 100:
|
|
|
+ del track_history[track_id][:-50] # 维持历史长度
|
|
|
+
|
|
|
+ # 记录每一帧的时间
|
|
|
+ time_stamps[track_id].append(current_time)
|
|
|
+
|
|
|
+ # 计算时间间隔
|
|
|
+ if len(time_stamps[track_id]) > 1:
|
|
|
+ delta_time = time_stamps[track_id][-1] - time_stamps[track_id][-2] # 最近两帧的时间差
|
|
|
+ else:
|
|
|
+ delta_time = 0
|
|
|
+
|
|
|
+ instantaneous_velocity = 0
|
|
|
+ # 计算二维瞬时速度
|
|
|
+ if len(track_history[track_id]) >= 2:
|
|
|
+ pos1 = np.array(track_history[track_id][-1]) # 最新位置
|
|
|
+ pos2 = np.array(track_history[track_id][-2]) # 前一个位置
|
|
|
+
|
|
|
+ pos1 = map_to_ellipse(pos1)
|
|
|
+ pos2 = map_to_ellipse(pos2)
|
|
|
+ distance = np.linalg.norm(pos1 - pos2)
|
|
|
+
|
|
|
+ # 使用时间间隔进行速度计算
|
|
|
+ instantaneous_velocity = distance / delta_time if delta_time > 0 else np.zeros(2)
|
|
|
+
|
|
|
+ instantaneous_velocity_magnitude = round(np.linalg.norm(instantaneous_velocity), 1)
|
|
|
+ instantaneous_velocities[track_id].append(instantaneous_velocity_magnitude)
|
|
|
+ else:
|
|
|
+ instantaneous_velocity_magnitude = 0
|
|
|
+
|
|
|
+ # 判断是否有足够数量的高速度
|
|
|
+ high_velocity_count = sum(1 for velocity in instantaneous_velocities[track_id] if velocity > speed_threshold)
|
|
|
+
|
|
|
+ if high_velocity_count >= high_velocity_count_threshold:
|
|
|
+
|
|
|
+ # 原逻辑：截图，标红
|
|
|
+ # cv2.putText(frame, str(instantaneous_velocity_magnitude), (int(x), int(y)),
|
|
|
+ # cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
|
|
|
+ # data_time = str(datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S"))
|
|
|
+ # file_name = "high_speed_" + data_time + ".jpg"
|
|
|
+ # cv2.imwrite(file_name, frame)
|
|
|
+
|
|
|
+ # 新增逻辑：删除超过 speed_threshold 的瞬时速度
|
|
|
+ instantaneous_velocities[track_id] = deque(
|
|
|
+ [velocity for velocity in instantaneous_velocities[track_id] if velocity <= speed_threshold],
|
|
|
+ maxlen=100
|
|
|
+ )
|
|
|
+ # 新增保存视频逻辑
|
|
|
+ data_time = datetime.datetime.now()
|
|
|
+ save_high_speed_video(frame_buffer, data_time)
|
|
|
+ else:
|
|
|
+ cv2.putText(frame, str(instantaneous_velocity_magnitude), (int(x), int(y)),
|
|
|
+ cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 121, 23), 2)
|
|
|
+
|
|
|
+ # Save the annotated frame to the output video
|
|
|
+ out.write(frame) # 将处理后的视频帧写入文件
|
|
|
+
|
|
|
+ # Display the annotated frame
|
|
|
+ cv2.imshow("YOLO11 Tracking", frame)
|
|
|
+
|
|
|
+ # Break the loop if 'q' is pressed
|
|
|
+ if cv2.waitKey(1) & 0xFF == ord("q"):
|
|
|
+ break
|
|
|
+ else:
|
|
|
+ # Break the loop if the end of the video is reached
|
|
|
+ break
|
|
|
+
|
|
|
+# Release the video capture object and close the display window
|
|
|
+cap.release()
|
|
|
+out.release() # 释放 VideoWriter 对象
|
|
|
+cv2.destroyAllWindows()
|
tuantuan