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()