import torch
import cv2
import numpy as np
from shapely.geometry import Point, Polygon
import os

# # 自定义模型文件的路径
# model_path = 'yolov5s.pt'  # 假设模型文件名为 yolov5s.pt 并与执行文件在同一个目录
# # 本地YOLOv5仓库路径
# repo_path = '../base_model/yolov5'
# 自定义模型文件的路径
print(f"Current working directory (yolov5.py): {os.getcwd()}")
model_path = 'D:/Project/FristProject/model/mode_test/yolov5s.pt'  # 假设模型文件名为 yolov5s.pt 并与执行文件在同一个目录
# 本地YOLOv5仓库路径
repo_path = 'D:/Project/FristProject/model/base_model/yolov5'

# 加载自定义模型
model = torch.hub.load(repo_path, 'custom', path=model_path, source='local')

# 定义监控区域（多边形顶点）
region_points = [(100, 100), (500, 100), (500, 400), (100, 400)]
polygon = Polygon(region_points)

# 打开摄像头
cap = cv2.VideoCapture(0)

def is_point_in_polygon(point, polygon):
    return polygon.contains(Point(point))

while True:
    ret, frame = cap.read()
    if not ret:
        break

    # 进行推理
    results = model(frame)
    detections = results.pandas().xyxy[0]

    # 绘制监控区域
    cv2.polylines(frame, [np.array(region_points, np.int32)], isClosed=True, color=(0, 255, 0), thickness=2)

    for _, row in detections.iterrows():
        if row['name'] == 'person':
            # 获取人员检测框的中心点
            x_center = (row['xmin'] + row['xmax']) / 2
            y_center = (row['ymin'] + row['ymax']) / 2

            if is_point_in_polygon((x_center, y_center), polygon):
                # 触发报警
                cv2.putText(frame, 'Alert: Intrusion Detected', (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
                print('Alert: Intrusion Detected')

            # 绘制检测框
            cv2.rectangle(frame, (int(row['xmin']), int(row['ymin'])), (int(row['xmax']), int(row['ymax'])), (255, 0, 0), 2)
            cv2.circle(frame, (int(x_center), int(y_center)), 5, (0, 0, 255), -1)

    # 显示结果
    cv2.imshow('Frame', frame)

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()