105 lines
3.6 KiB
Python
105 lines
3.6 KiB
Python
|
import numpy as np
|
||
|
|
import cv2
|
||
|
|
from ultralytics import YOLO
|
||
|
|
|
||
|
|
def point_to_line_distance(point, line_start, line_end):
|
||
|
|
# 线的方向向量
|
||
|
|
line_dir = np.array(line_end) - np.array(line_start)
|
||
|
|
# 点到线起点的向量
|
||
|
|
point_to_start = np.array(point) - np.array(line_start)
|
||
|
|
|
||
|
|
# 计算叉积（只适用于二维空间）
|
||
|
|
cross_product = line_dir[0] * point_to_start[1] - line_dir[1] * point_to_start[0]
|
||
|
|
|
||
|
|
# 叉积的符号会表明点在直线的哪一侧
|
||
|
|
return cross_product
|
||
|
|
|
||
|
|
def check_crossing_detection_box(box, start_point, end_point):
|
||
|
|
x1, y1, x2, y2 = box
|
||
|
|
points = [(x1, y1), (x2, y1), (x1, y2), (x2, y2)]
|
||
|
|
distances = [point_to_line_distance(p, start_point, end_point) for p in points]
|
||
|
|
|
||
|
|
# 检查是否有正负距离，即检测框的点分别位于斜线两侧
|
||
|
|
crossing = any(dist < 0 for dist in distances) and any(dist > 0 for dist in distances)
|
||
|
|
return crossing
|
||
|
|
|
||
|
|
def process_frame(frame):
|
||
|
|
height, width = frame.shape[:2]
|
||
|
|
desired_height = 500
|
||
|
|
scale = desired_height / height
|
||
|
|
resized_frame = cv2.resize(frame, (int(width * scale), desired_height))
|
||
|
|
frame = resized_frame
|
||
|
|
|
||
|
|
height, width = frame.shape[:2]
|
||
|
|
start_point = (0, height // 2)
|
||
|
|
end_point = (width, height // 3)
|
||
|
|
|
||
|
|
cv2.line(frame, start_point, end_point, (255, 0, 0), 3)
|
||
|
|
|
||
|
|
results = model(frame, verbose=False)
|
||
|
|
for result in results:
|
||
|
|
boxes = result.boxes
|
||
|
|
classes = model.names
|
||
|
|
for box in boxes:
|
||
|
|
x1, y1, x2, y2 = box.xyxy[0].tolist()
|
||
|
|
class_id = int(box.cls.item())
|
||
|
|
if classes[class_id] == 'person':
|
||
|
|
cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
|
||
|
|
if check_crossing_detection_box((x1, y1, x2, y2), start_point, end_point):
|
||
|
|
cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 0, 255), 2)
|
||
|
|
|
||
|
|
return frame
|
||
|
|
|
||
|
|
def main():
|
||
|
|
global model
|
||
|
|
try:
|
||
|
|
model = YOLO('yolov8s.pt')
|
||
|
|
print("模型加载成功。")
|
||
|
|
except Exception as e:
|
||
|
|
print(f"加载模型时出现错误: {e}")
|
||
|
|
model = None
|
||
|
|
|
||
|
|
if model is not None:
|
||
|
|
video_path = "./image/a.mp4"
|
||
|
|
print("开始视频捕获...")
|
||
|
|
cap = cv2.VideoCapture(video_path)
|
||
|
|
if not cap.isOpened():
|
||
|
|
print("打开视频流或文件时出错")
|
||
|
|
else:
|
||
|
|
print("视频流已成功打开。")
|
||
|
|
|
||
|
|
# 获取视频的帧率和尺寸
|
||
|
|
fps = cap.get(cv2.CAP_PROP_FPS)
|
||
|
|
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||
|
|
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||
|
|
|
||
|
|
# 创建VideoWriter对象
|
||
|
|
fourcc = cv2.VideoWriter_fourcc(*'mp4v') # 使用mp4v编码器
|
||
|
|
out = cv2.VideoWriter('output.mp4', fourcc, fps, (width, height)) # 输出文件名为output.mp4
|
||
|
|
|
||
|
|
while cap.isOpened():
|
||
|
|
ret, frame = cap.read()
|
||
|
|
if not ret:
|
||
|
|
break
|
||
|
|
|
||
|
|
processed_frame = process_frame(frame)
|
||
|
|
|
||
|
|
# 将处理后的帧写入视频文件
|
||
|
|
out.write(processed_frame)
|
||
|
|
|
||
|
|
cv2.imshow('Processed Video', processed_frame)
|
||
|
|
key = cv2.waitKey(1)
|
||
|
|
if key & 0xFF == ord('q'):
|
||
|
|
break
|
||
|
|
|
||
|
|
# 释放VideoWriter资源
|
||
|
|
out.release()
|
||
|
|
|
||
|
|
cap.release()
|
||
|
|
cv2.destroyAllWindows()
|
||
|
|
print("视频处理完成。")
|
||
|
|
else:
|
||
|
|
print("模型未加载，无法继续处理。")
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
main()
|