yolov5识别屏幕指定区域

import cv2
import numpy as np
from numpy import random
from PIL import ImageGrab
import cv2 as CV2
import time
import win32api
import torch

from yolov5_camera_detect.models.experimental import attempt_load
from yolov5_camera_detect.utils.general import (check_img_size, non_max_suppression, scale_coords, plot_one_box)
from yolov5_camera_detect.utils.torch_utils import select_device

print('Setup complete. Using torch %s %s' % (
torch.__version__, torch.cuda.get_device_properties(0) if torch.cuda.is_available() else 'CPU'))

# Initialize
device = select_device()
frame_h = 480
frame_w = 800
obj_count = 0  # 警戒区目标
obj_count_old = 0  # 警戒区旧目标
take_photo_num = 0;  # 拍照次数
# 每个监测不一定都检测得到，所以做个缓冲区用于取平均值,因为要避免某帧的目标丢失，会造成目标数量的跳变，引发拍照记录
obj_count_buf = np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0])  # 10个值
# Load model
model = attempt_load('weights/yolov5s.pt', map_location=device)  # load FP32 model cuda
# Get names and colors
names = model.module.names if hasattr(model, 'module') else model.names
colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(names))] #people框的颜色
# imgsz = check_img_size(486, s=model.stride.max())   # check img_size


frame_mask = np.zeros((frame_h, frame_w, 3), dtype=np.uint8)  # 做一个相同尺寸格式的图片mask   frame_mask隐藏的那个
postion = [(200, 100), (275, 391), (632, 381), (600, 70)]  # 警戒区位置点,逆时针从左上角开始 # 识别区坐标
CV2.fillPoly(frame_mask, [np.array(postion)], (0, 0, 255))  # 警戒区内数字填充255，0，0成为mask   在frame_mask上绘制多边形并对其填充，位置、颜色


def process_img(original_image):  # 原图处理函数
    processed_img = CV2.cvtColor(original_image, CV2.COLOR_BGR2RGB)  # BGR格式转换RGB
    processed_img = CV2.resize(processed_img, (frame_w, frame_h))  # 改变输入尺寸
    return processed_img


def MouseEvent(a, b, c, d, e):  # 鼠标处理事件响应函数
    if (a == 1):  # 获取左键点击坐标点
        print(b, c)


CV2.namedWindow('frame')
CV2.setMouseCallback('frame', MouseEvent)  # 窗口与回调函数绑定 # 点击窗口（图片执行mouseevent函数）

while (1):
    # get a frame
    # frame = cv2.imread('4.jpg')
    # frame = process_img(frame)
    # cv2.imshow('nihao',frame)

    # frame = np.array(cv2.VideoCapture(0))
    # print(frame)
    frame = np.array(ImageGrab.grab(bbox=(0, 100, 800, 600)))  # frame为一个数组  ，np.array,将抓的图转换为数组，frame是数组
    # bbox：（x_左，y_上，x_右，y_下），数字是线的坐标

    if np.shape(frame):  # frame有数据才能往下执行
        # processing
        frame = process_img(frame)  # BGR格式转换RGB
        img = frame.copy()  # img为gpu格式，常规方法不能读取，im0为img的copy版可直接读取
        # print("img:",np.shape(img))  #img: (480, 800, 3)，480个[]，每个800个数据，3维(高度、宽度、通道数)
        # print(img),数组
        img = np.transpose(img, (2, 0, 1))  # torch.Size([480, 800， 3])转torch.Size([3, 480, 800])
        # print("img:",np.shape(img))
        img = torch.from_numpy(img).to(device)
        img = img.float()  # uint8 to fp32
        img /= 255.0  # 0 - 255 to 0.0 - 1.0
        # print(np.shape(img))#>>>torch.Size([3, 416, 352])
        if img.ndimension() == 3:
            img = img.unsqueeze(0)  # 这个函数主要是对数据维度进行扩充,在0的位置加了一维
        # print(np.shape(img))#>>>torch.Size([1, 3, 416, 352])
        pred = model(img)[0]
        # Apply NMS 非极大值抑制
        pred = non_max_suppression(pred, 0.5, 0.5)  # 大于0.4阈值的输出，只显示classes:>= 1，不能显示0？

        # # 绘图  包含区域红色，标记的两种圆圈
        if pred != [None]:
            for i, det in enumerate(pred):
                # Rescale boxes from img_size to im0 size
                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], frame.shape).round()
                # Write results
                for *xyxy, conf, cls in reversed(det):
                    if cls == 0:  # 只显示0（person）的标签，因为non_max_suppression(只显示classes:>= 1）的标签
                        label = '%s %.2f' % (names[int(cls)], conf)
                        plot_one_box(xyxy, frame, label=label, color=colors[int(cls)],
                                     line_thickness=1)  # utils.general专用画框标注函数
                        xy = torch.tensor(xyxy).tolist()  # 张量转换成列表形式
                        x, y, x1, y1 = int(xy[0]), int(xy[1]), int(xy[2]), int(xy[3])  # 获取左顶右底坐标
                        center_xy = (int(np.average([x, x1])), int(np.average([y, y1])))  # 计算中心点
                        if (frame_mask[(center_xy[1], center_xy[0])] == [0, 0, 255]).all():  # 中心点在警戒区
                            obj_color = (0, 0, 0)  # 改变中心点颜色，警戒区域外颜色
                            obj_count += 1
                        else:
                            obj_color = (0, 0, 255)  # 改变中心点颜色   警戒区域内颜色
                        CV2.circle(frame, center_xy, 10, obj_color, 4)  # 开始画点,画圆圈
        obj_count_buf = np.append(obj_count_buf[1:], obj_count)  # 保持更新10个缓冲区
        cbr = int(np.around(np.average(obj_count_buf)))#cbr是当前区域内人数
        CV2.putText(frame, 'obj_count :%s obj take_photo: %s' % (cbr, take_photo_num), (100, 20),
                    CV2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 255), 2)  # 文字信息显示
        frame = CV2.addWeighted(frame, 1.0, frame_mask, 0.1, 0.0)  # 叠加掩码图片进实时图，将frame_mask叠加到frame中
        if (obj_count_old != cbr):#警戒区人数有变化
            take_photo_num += 1
            CV2.imwrite("./photo/%s.jpg" % take_photo_num, frame, [int(CV2.IMWRITE_JPEG_QUALITY), 50])  # 保存图片
            print('take photo number :%s' % take_photo_num)  # 显示记录的照片张数
            CV2.putText(frame, 'take photo', (100, 300), CV2.FONT_HERSHEY_SIMPLEX, 3, (0, 0, 255), 3)  # 文字信息显示

        obj_count_old = cbr  # 保存上个数据
        obj_count = 0  # 目标显示清零，等待下次探测

        # show a frame
        # CV2.imshow("capture", frame[:,:,::-1])
        CV2.imshow("frame", frame)
        CV2.imshow("frame_mask", frame_mask[:,:,::-1])

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

CV2.destroyAllWindows()