动态绘制入侵检测区域矩形

import cv2 as cv
import numpy as np
import pygetwindow as gw
import mss
import torch

# 定义窗口标题
window_title = "Plants vs. Zombies"

# 定义置信度阈值
CONFIDENCE_THRESHOLD = 0.5

# 加载YOLOv5模型
model = torch.hub.load('..', 'custom', path='../runs/train/exp48/weights/best.pt', source='local')

# 全局变量，用于处理鼠标事件
start_point = None
end_point = None
drawing = False
monitoring_area = []

def get_window_rect(title):
    """
    获取指定标题窗口的矩形区域。
    """
    windows = gw.getWindowsWithTitle(title)
    if not windows:
        raise ValueError(f"Window with title '{title}' not found.")
    return windows[0]._rect

def capture_screen(rect):
    """
    捕获屏幕上指定矩形区域的图像。
    """
    with mss.mss() as sct:
        monitor = {"left": rect.left, "top": rect.top, "width": rect.width, "height": rect.height}
        img = sct.grab(monitor)
        img_np = np.array(img)
        img_np = cv.cvtColor(img_np, cv.COLOR_BGRA2BGR)  # 将BGRA图像转换为BGR图像
        return img_np

def detect_objects(img):
    """
    使用YOLOv5模型检测图像中的物体。
    """
    img_rgb = cv.cvtColor(img, cv.COLOR_BGR2RGB)  # 将图像从BGR转换为RGB
    results = model(img_rgb)  # 执行推断
    detections = results.pandas().xyxy[0]  # 解析检测结果
    detections = detections[detections['confidence'] >= CONFIDENCE_THRESHOLD]  # 过滤低置信度的检测结果
    return detections

def draw_boxes(img, detections):
    """
    在图像上绘制检测到的物体的边界框和标签。
    """
    for _, row in detections.iterrows():
        x1, y1, x2, y2 = int(row['xmin']), int(row['ymin']), int(row['xmax']), int(row['ymax'])
        conf = row['confidence']
        label = row['name']
        color = (0, 255, 0)  # 绿色，用于绘制边界框

        # 绘制边界框
        cv.rectangle(img, (x1, y1), (x2, y2), color, 2)
        # 绘制标签
        label_text = f"{label} {conf:.2f}"
        label_size, _ = cv.getTextSize(label_text, cv.FONT_HERSHEY_SIMPLEX, 0.5, 2)
        label_x1 = x1
        label_y1 = y1 - 10 if y1 - 10 > 10 else y1 + 10
        cv.rectangle(img, (label_x1, label_y1 - label_size[1]), (label_x1 + label_size[0], label_y1 + 2), color, -1)
        cv.putText(img, label_text, (label_x1, label_y1), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
    return img

def mouse_callback(event, x, y, flags, param):
    """
    处理鼠标事件，用于绘制矩形区域。
    """
    global start_point, end_point, drawing, monitoring_area

    if event == cv.EVENT_LBUTTONDOWN:
        # 鼠标左键按下时，记录起始点
        start_point = (x, y)
        drawing = True

    elif event == cv.EVENT_MOUSEMOVE:
        if drawing:
            # 鼠标移动时，更新矩形区域的终点
            end_point = (x, y)
            monitoring_area = [start_point, (end_point[0], start_point[1]),
                               end_point, (start_point[0], end_point[1])]

    elif event == cv.EVENT_LBUTTONUP:
        # 鼠标左键松开时，更新矩形区域
        end_point = (x, y)
        drawing = False
        monitoring_area = [start_point, (end_point[0], start_point[1]),
                           end_point, (start_point[0], end_point[1])]

def draw_rectangle(img, rect_points):
    """
    在图像上绘制矩形区域。
    """
    if rect_points:
        rect_points = np.array(rect_points, np.int32)
        rect_points = rect_points.reshape((-1, 1, 2))
        cv.polylines(img, [rect_points], isClosed=True, color=(255, 0, 0), thickness=2)  # 红色，用于绘制矩形区域
    return img

def main():
    """
    主函数，执行屏幕捕获、物体检测和显示操作。
    """
    global monitoring_area

    try:
        rect = get_window_rect(window_title)
        print("Monitoring Window Rect:", rect)

        # 创建窗口并设置鼠标回调
        cv.namedWindow("Plants vs. Zombies")
        cv.setMouseCallback("Plants vs. Zombies", mouse_callback)

        while True:
            img = capture_screen(rect)  # 捕获屏幕图像
            detections = detect_objects(img)  # 检测图像中的物体
            img = draw_boxes(img, detections)  # 绘制检测到的物体的边界框和标签
            # Print detections to console
            if not detections.empty:
                print("Detected objects:")
                for _, row in detections.iterrows():
                    print(f"Label: {row['name']}, Confidence: {row['confidence']:.2f}, "
                          f"Bounding Box: ({row['xmin']}, {row['ymin']}) to ({row['xmax']}, {row['ymax']})")
            # 绘制矩形区域
            img = draw_rectangle(img, monitoring_area)

            # 计算并显示入侵个数
            intrusion_count = 0
            for _, row in detections.iterrows():
                x1, y1, x2, y2 = int(row['xmin']), int(row['ymin']), int(row['xmax']), int(row['ymax'])
                if (start_point and end_point and
                    min(start_point[0], end_point[0]) < x1 < max(start_point[0], end_point[0]) and
                    min(start_point[1], end_point[1]) < y1 < max(start_point[1], end_point[1]) and
                    min(start_point[0], end_point[0]) < x2 < max(start_point[0], end_point[0]) and
                    min(start_point[1], end_point[1]) < y2 < max(start_point[1], end_point[1])):
                    intrusion_count += 1

            # 显示入侵个数
            cv.putText(img, f"Intrusion Count: {intrusion_count}", (50, 50), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)

            # 显示图像
            cv.imshow("Plants vs. Zombies", img)

            # 按 'q' 键退出循环
            if cv.waitKey(1) == ord('q'):
                break

    except Exception as e:
        print(f"Error: {e}")

    cv.destroyAllWindows()

if __name__ == "__main__":
    main()