YOLOV5实时检测屏幕

YOLOV5实时检测屏幕

注：此为笔记

目的：保留模型加载和推理部分，完成实时屏幕检测

实现思路：

1. 写一个实时截取屏幕的函数

2. 将截取的屏幕在窗口显示出来

3. 用OpenCV绘制一个窗口用来显示截取的屏幕

4. 在detect找出推理的代码，推理完成后得到中心点的xy坐标，宽高组成box

5. 在创建的OpenCV窗口用得到的推理结果绘制方框

实现效果：

思考部分

先把原本的detect.py的代码贴在这里

import argparse

import os

import platform

import sys

from pathlib import Path

import torch

FILE = Path(__file__).resolve()

ROOT = FILE.parents[0]  # YOLOv5 root directory

if str(ROOT) not in sys.path:

    sys.path.append(str(ROOT))  # add ROOT to PATH

ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative

from models.common import DetectMultiBackend

from utils.dataloaders import IMG_FORMATS, VID_FORMATS, LoadImages, LoadScreenshots, LoadStreams

from utils.general import (LOGGER, Profile, check_file, check_img_size, check_imshow, check_requirements, colorstr, cv2,

                           increment_path, non_max_suppression, print_args, scale_boxes, strip_optimizer, xyxy2xywh)

from utils.plots import Annotator, colors, save_one_box

from utils.torch_utils import select_device, smart_inference_mode

@smart_inference_mode()

def run(

        weights=ROOT / 'yolov5s.pt',  # model path or triton URL

        source=ROOT / 'data/video/',

        data=ROOT / 'data/coco128.yaml',  # dataset.yaml path

        imgsz=(640, 640),  # inference size (height, width)

        conf_thres=0.25,  # confidence threshold

        iou_thres=0.45,  # NMS IOU threshold

        max_det=1000,  # maximum detections per image

        device='',  # cuda device, i.e. 0 or 0,1,2,3 or cpu

        view_img=False,  # show results

        save_txt=False,  # save results to *.txt

        save_conf=False,  # save confidences in --save-txt labels

        save_crop=False,  # save cropped prediction boxes

        nosave=False,  # do not save images/videos

        classes=None,  # filter by class: --class 0, or --class 0 2 3

        agnostic_nms=False,  # class-agnostic NMS

        augment=False,  # augmented inference

        visualize=False,  # visualize features

        update=False,  # update all models

        project=ROOT / 'runs/detect',  # save results to project/name

        name='exp',  # save results to project/name

        exist_ok=False,  # existing project/name ok, do not increment

        line_thickness=3,  # bounding box thickness (pixels)

        hide_labels=False,  # hide labels

        hide_conf=False,  # hide confidences

        half=False,  # use FP16 half-precision inference

        dnn=False,  # use OpenCV DNN for ONNX inference

        vid_stride=1,  # video frame-rate stride

):

    source = str(source)

    save_img = not nosave and not source.endswith('.txt')  # save inference images

    is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)

    is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://'))

    webcam = source.isnumeric() or source.endswith('.streams') or (is_url and not is_file)

    screenshot = source.lower().startswith('screen')

    if is_url and is_file:

        source = check_file(source)  # download

    # Directories

    save_dir = increment_path(Path(project) / name, exist_ok=exist_ok)  # increment run

    (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir

    # Load model

    device = select_device(device)

    model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half)

    stride, names, pt = model.stride, model.names, model.pt

    imgsz = check_img_size(imgsz, s=stride)  # check image size

    # Dataloader

    bs = 1  # batch_size

    if webcam:

        view_img = check_imshow(warn=True)

        dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)

        bs = len(dataset)

    elif screenshot:

        dataset = LoadScreenshots(source, img_size=imgsz, stride=stride, auto=pt)

    else:

        dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)

    vid_path, vid_writer = [None] * bs, [None] * bs

    # Run inference

    model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *imgsz))  # warmup

    seen, windows, dt = 0, [], (Profile(), Profile(), Profile())

    for path, im, im0s, vid_cap, s in dataset:

        with dt[0]:

            im = torch.from_numpy(im).to(model.device)

            im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32

            im /= 255  # 0 - 255 to 0.0 - 1.0

            if len(im.shape) == 3:

                im = im[None]  # expand for batch dim

        # Inference

        with dt[1]:

            visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False

            pred = model(im, augment=augment, visualize=visualize)

        # NMS

        with dt[2]:

            pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)

        # Second-stage classifier (optional)

        # pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)

        # Process predictions

        for i, det in enumerate(pred):  # per image

            seen += 1

            if webcam:  # batch_size >= 1

                p, im0, frame = path[i], im0s[i].copy(), dataset.count

                s += f'{i}: '

            else:

                p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)

            p = Path(p)  # to Path

            save_path = str(save_dir / p.name)  # im.jpg

            txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}')  # im.txt

            s += '%gx%g ' % im.shape[2:]  # print string

            gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh

            imc = im0.copy() if save_crop else im0  # for save_crop

            annotator = Annotator(im0, line_width=line_thickness, example=str(names))

            if len(det):

                # Rescale boxes from img_size to im0 size

                det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()

                # Print results

                for c in det[:, 5].unique():

                    n = (det[:, 5] == c).sum()  # detections per class

                    s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string

                # Write results

                for *xyxy, conf, cls in reversed(det):

                    if save_txt:  # Write to file

                        xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh

                        line = (cls, *xywh, conf) if save_conf else (cls, *xywh)  # label format

                        with open(f'{txt_path}.txt', 'a') as f:

                            f.write(('%g ' * len(line)).rstrip() % line + '\n')

                    if save_img or save_crop or view_img:  # Add bbox to image

                        c = int(cls)  # integer class

                        label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')

                        annotator.box_label(xyxy, label, color=colors(c, True))

                    if save_crop:

                        save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)

            # Stream results

            im0 = annotator.result()

            if view_img:

                if platform.system() == 'Linux' and p not in windows:

                    windows.append(p)

                    cv2.namedWindow(str(p), cv2.WINDOW_NORMAL | cv2.WINDOW_KEEPRATIO)  # allow window resize (Linux)

                    cv2.resizeWindow(str(p), im0.shape[1], im0.shape[0])

                cv2.imshow(str(p), im0)

                cv2.waitKey(1)  # 1 millisecond

            # Save results (image with detections)

            if save_img:

                if dataset.mode == 'image':

                    cv2.imwrite(save_path, im0)

                else:  # 'video' or 'stream'

                    if vid_path[i] != save_path:  # new video

                        vid_path[i] = save_path

                        if isinstance(vid_writer[i], cv2.VideoWriter):

                            vid_writer[i].release()  # release previous video writer

                        if vid_cap:  # video

                            fps = vid_cap.get(cv2.CAP_PROP_FPS)

                            w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))

                            h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

                        else:  # stream

                            fps, w, h = 30, im0.shape[1], im0.shape[0]

                        save_path = str(Path(save_path).with_suffix('.mp4'))  # force *.mp4 suffix on results videos

                        vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))

                    vid_writer[i].write(im0)

        # Print time (inference-only)

        LOGGER.info(f"{s}{'' if len(det) else '(no detections), '}{dt[1].dt * 1E3:.1f}ms")

    # Print results

    t = tuple(x.t / seen * 1E3 for x in dt)  # speeds per image

    LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}' % t)

    if save_txt or save_img:

        s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''

        LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")

    if update:

        strip_optimizer(weights[0])  # update model (to fix SourceChangeWarning)

def parse_opt():

    parser = argparse.ArgumentParser()

    parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model path or triton URL')

    parser.add_argument('--source', type=str, default=ROOT / '0', help='file/dir/URL/glob/screen/1(webcam)')

    parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='(optional) dataset.yaml path')

    parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='inference size h,w')

    parser.add_argument('--conf-thres', type=float, default=0.45, help='confidence threshold')

    parser.add_argument('--iou-thres', type=float, default=0.2, help='NMS IoU threshold')

    parser.add_argument('--max-det', type=int, default=1000, help='maximum detections per image')

    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')

    parser.add_argument('--view-img', action='store_true', help='show results')

    parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')

    parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')

    parser.add_argument('--save-crop', action='store_true', help='save cropped prediction boxes')

    parser.add_argument('--nosave', action='store_true', help='do not save images/videos')

    parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --classes 0, or --classes 0 2 3')

    parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')

    parser.add_argument('--augment', action='store_true', help='augmented inference')

    parser.add_argument('--visualize', action='store_true', help='visualize features')

    parser.add_argument('--update', action='store_true', help='update all models')

    parser.add_argument('--project', default=ROOT / 'runs/detect', help='save results to project/name')

    parser.add_argument('--name', default='exp', help='save results to project/name')

    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')

    parser.add_argument('--line-thickness', default=3, type=int, help='bounding box thickness (pixels)')

    parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels')

    parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences')

    parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')

    parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')

    parser.add_argument('--vid-stride', type=int, default=1, help='video frame-rate stride')

    opt = parser.parse_args()

    opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1  # expand

    print_args(vars(opt))

    return opt

def main(opt):

    check_requirements(exclude=('tensorboard', 'thop'))

    run(**vars(opt))

if __name__ == '__main__':

    opt = parse_opt()

    main(opt)

分析代码并删减不用的部分

import argparse

import os

import platform

import sys

from pathlib import Path

import torch

FILE = Path(__file__).resolve()

ROOT = FILE.parents[0]  # YOLOv5 root directory

if str(ROOT) not in sys.path:

    sys.path.append(str(ROOT))  # add ROOT to PATH

ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative

from models.common import DetectMultiBackend

from utils.dataloaders import IMG_FORMATS, VID_FORMATS, LoadImages, LoadScreenshots, LoadStreams

from utils.general import (LOGGER, Profile, check_file, check_img_size, check_imshow, check_requirements, colorstr, cv2,

                           increment_path, non_max_suppression, print_args, scale_boxes, strip_optimizer, xyxy2xywh)

from utils.plots import Annotator, colors, save_one_box

from utils.torch_utils import select_device, smart_inference_mode

做了一些包的导入，定义了一些全局变量，先保留下来，没用的最后删

向下

if __name__ == '__main__':

    opt = parse_opt()

    main(opt)

从if __name__ == '__main__开始

opt = parse_opt 就是一个获取命令行参数的函数，我们并不需要,可以删

进入main函数

def main(opt):

    check_requirements(exclude=('tensorboard', 'thop'))

    run(**vars(opt))

check_requirements函数检查requirements是否全都安装好了，无用，删了

进入run函数

    source = str(source)

    save_img = not nosave and not source.endswith('.txt')  # save inference images

    is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)

    is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://'))

    webcam = source.isnumeric() or source.endswith('.streams') or (is_url and not is_file)

    screenshot = source.lower().startswith('screen')

    if is_url and is_file:

        source = check_file(source)  # download

    # Directories

    save_dir = increment_path(Path(project) / name, exist_ok=exist_ok)  # increment run

    (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir

判断source的类型，即要要推理的源是什么，判断源是文件还是url还是webcam或者screenshot ，定义保存文件夹，我不需要保存，只需要实时检测屏幕，删除

继续向下，是加载模型的代码

# Load model

device = select_device(device)

model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half)

得知加载模型需要几个参数，分别是weights, device=device, dnn=dnn, data=data, fp16=half

通过开始的形参可知：

weights=ROOT / 'yolov5s.pt' 也就是模型的名称
device通过select_device函数得到
dnn和fp16在run函数里的参数都是FALSE

故加载模型的代码可以改写成

def LoadModule():

    device = select_device('')

    weights = 'yolov5s.pt'

    model = DetectMultiBackend(weights, device=device, dnn=False, fp16=False)

    return model

继续往下读

 bs = 1  # batch_size

    if webcam:

        view_img = check_imshow(warn=True)

        dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)

        bs = len(dataset)

    elif screenshot:

        dataset = LoadScreenshots(source, img_size=imgsz, stride=stride, auto=pt)

    else:

        dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)

    vid_path, vid_writer = [None] * bs, [None] * bs

这里如果是使用网络摄像头作为输入，会通过LoadStreams类加载视频流，根据图像大小和步长采样，如果使用截图作为输入，则通过LoadScreenshots加载截图，都不是则通过LoadImages类加载图片文件

这是YOLOV5提供的加载dataset的部分，我们可以添加自己的dataset，所以删掉

继续往下

# Run inference

    model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *imgsz))  # warmup

    seen, windows, dt = 0, [], (Profile(), Profile(), Profile())

    for path, im, im0s, vid_cap, s in dataset:

        with dt[0]:

            im = torch.from_numpy(im).to(model.device)

            im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32

            im /= 255  # 0 - 255 to 0.0 - 1.0

            if len(im.shape) == 3:

                im = im[None]  # expand for batch dim

        # Inference

        with dt[1]:

            visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False

            pred = model(im, augment=augment, visualize=visualize)

        # NMS

        with dt[2]:

            pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)

        # Second-stage classifier (optional)

        # pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)

        # Process predictions

model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *imgsz))

用于模型预热，传入形状为(1, 3, *imgsz)的图像进行预热操作，没用删了

seen, windows, dt = 0, [], (Profile(), Profile(), Profile())

未知作用，删了

for path, im, im0s, vid_cap, s in dataset:

        with dt[0]:

            im = torch.from_numpy(im).to(model.device)

            im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32

            im /= 255  # 0 - 255 to 0.0 - 1.0

            if len(im.shape) == 3:

                im = im[None]  # expand for batch dim

        # Inference

        with dt[1]:

            visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False

            pred = model(im, augment=augment, visualize=visualize)

        # NMS

        with dt[2]:

            pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)

上面这段for循环用于遍历数据集中的每个图像或视频帧进行推理，在循环的开头，将路径、图像、原始图像、视频捕获对象和步长传递给path, im, im0s, vid_cap, s。推理实时屏幕只需要传一张图片，所以不存在将遍历推理，所以要进行改写，改写成

im = torch.from_numpy(im).to(model.device)

im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32

im /= 255  # 0 - 255 to 0.0 - 1.0

if len(im.shape) == 3:

    im = im[None]  # expand for batch dim

visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False

pred = model(im, augment=augment, visualize=visualize)

pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)

这里是对 im 进行转换和推理，而改写的代码中没有im变量，则寻找im的来源

for path, im, im0s, vid_cap, s in dataset:

im来源于dataset

dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)

dataset来源于LoadImages的返回值

查看LoadImages的函数返回值和返回值的来源

在dataloaders.py中可以看到

if self.transforms:

    im = self.transforms(im0)  # transforms

else:

    im = letterbox(im0, self.img_size, stride=self.stride, auto=self.auto)[0]  # padded resize

    im = im.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB

    im = np.ascontiguousarray(im)  # contiguous

return path, im, im0, self.cap, s

如果transforms存在，则转换，如果transforms不存在，则调用letterbox函数对图像im0进行缩放和填充，使其符合模型要求的图像大小，将图像的通道顺序由HWC转换为CHW，将图像的通道顺序由BGR转换为RGB，将图像转换为连续的内存布局

其中需要的参数是im0, self.img_size, stride=self.stride, auto=self.auto

im0则是未经处理的图片，img_size填640（因为模型的图片大小训练的是640），stride填64（默认参数为64），auto填True

则得到改写代码为

im = letterbox(img0, 640, stride=32, auto=True)[0]  # padded resize

    im = im.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB

    im = np.ascontiguousarray(im)  # contiguous

    im = torch.from_numpy(im).to(model.device)

    im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32

    im /= 255  # 0 - 255 to 0.0 - 1.0

    if len(im.shape) == 3:

        im = im[None]  # expand for batch dim

    pred = model(im, augment=False, visualize=False)

    pred = non_max_suppression(pred, conf_thres=conf_thres, iou_thres=iou_thres, classes=None, agnostic=False,

                               max_det=1000)

继续向下

for i, det in enumerate(pred):  # per image

            seen += 1

            if webcam:  # batch_size >= 1

                p, im0, frame = path[i], im0s[i].copy(), dataset.count

                s += f'{i}: '

            else:

                p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)

            p = Path(p)  # to Path

            save_path = str(save_dir / p.name)  # im.jpg

            txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}')  # im.txt

            s += '%gx%g ' % im.shape[2:]  # print string

            gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh

            imc = im0.copy() if save_crop else im0  # for save_crop

            annotator = Annotator(im0, line_width=line_thickness, example=str(names))

            if len(det):

                # Rescale boxes from img_size to im0 size

                det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()

                # Print results

                for c in det[:, 5].unique():

                    n = (det[:, 5] == c).sum()  # detections per class

                    s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string

                # Write results

                for *xyxy, conf, cls in reversed(det):

                    if save_txt:  # Write to file

                        xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh

                        line = (cls, *xywh, conf) if save_conf else (cls, *xywh)  # label format

                        with open(f'{txt_path}.txt', 'a') as f:

                            f.write(('%g ' * len(line)).rstrip() % line + '\n')

                    if save_img or save_crop or view_img:  # Add bbox to image

                        c = int(cls)  # integer class

                        label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')

                        annotator.box_label(xyxy, label, color=colors(c, True))

                    if save_crop:

                        save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)

这段代码将推理后的结果进行转换，转换为label format，成为人能看懂的格式，删去输出结果，留下写入结果中的，格式转换，删掉保存为txt文件，得到需要的box，然后自己写一个boxs=[]，将结果append进去，方便在OpenCV中绘画识别方框，改写结果为

boxs=[]

    for i, det in enumerate(pred):  # per image

        im0 = img0.copy()

        s = ' '

        s += '%gx%g ' % im.shape[2:]  # print string

        gn = torch.tensor(img0.shape)[[1, 0, 1, 0]]  # normalization gain whwh

        imc = img0  # for save_crop

        if len(det):

            # Rescale boxes from img_size to im0 size

            det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()

            # Print results

            for c in det[:, 5].unique():

                n = (det[:, 5] == c).sum()  # detections per class

                s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string

        # Write results

        for *xyxy, conf, cls in reversed(det):

            xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh

            line = (cls, *xywh)  # label format

            box = ('%g ' * len(line)).rstrip() % line

            box = box.split(' ')

            boxs.append(box)

就此完成了推理部分的删减和重写

把屏幕的截图通过OpenCV进行显示

写一个屏幕截图的文件

写成 grabscreen.py

# 文件名:grabscreen.py

import cv2

import numpy as np

import win32gui

import win32print

import win32ui

import win32con

import win32api

import mss

def grab_screen_win32(region):

    hwin = win32gui.GetDesktopWindow()

    left, top, x2, y2 = region

    width = x2 - left + 1

    height = y2 - top + 1

    hwindc = win32gui.GetWindowDC(hwin)

    srcdc = win32ui.CreateDCFromHandle(hwindc)

    memdc = srcdc.CreateCompatibleDC()

    bmp = win32ui.CreateBitmap()

    bmp.CreateCompatibleBitmap(srcdc, width, height)

    memdc.SelectObject(bmp)

    memdc.BitBlt((0, 0), (width, height), srcdc, (left, top), win32con.SRCCOPY)

    signedIntsArray = bmp.GetBitmapBits(True)

    img = np.fromstring(signedIntsArray, dtype='uint8')

    img.shape = (height, width, 4)

    srcdc.DeleteDC()

    memdc.DeleteDC()

    win32gui.ReleaseDC(hwin, hwindc)

    win32gui.DeleteObject(bmp.GetHandle())

    return cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)

通过img0 = grab_screen_win32(region=(0, 0, 1920, 1080))来作为im的参数传入，即可让屏幕截图作为推理图片

用OpenCV绘制窗口并显示

if len(boxs):

    for i, det in enumerate(boxs):

        _, x_center, y_center, width, height = det

        x_center, width = re_x * float(x_center), re_x * float(width)

        y_center, height = re_y * float(y_center), re_y * float(height)

        top_left = (int(x_center - width / 2.), int(y_center - height / 2.))

        bottom_right = (int(x_center + width / 2.), int(y_center + height / 2.))

        color = (0, 0, 255)  # RGB

        cv2.rectangle(img0, top_left, bottom_right, color, thickness=thickness)

和

cv2.namedWindow('windows', cv2.WINDOW_NORMAL)

cv2.resizeWindow('windows', re_x // 2, re_y // 2)

cv2.imshow('windows', img0)

HWND = win32gui.FindWindow(None, "windows")

win32gui.SetWindowPos(HWND, win32con.HWND_TOPMOST, 0, 0, 0, 0, win32con.SWP_NOMOVE | win32con.SWP_NOSIZE)

结合在一起

最终代码

import torch, pynput

import numpy as np

import win32gui, win32con, cv2

from grabscreen import grab_screen_win32 # 本地文件

from utils.augmentations import letterbox

from models.common import DetectMultiBackend

from utils.torch_utils import select_device

from utils.general import non_max_suppression, scale_boxes, xyxy2xywh

# 可调参数

conf_thres = 0.25

iou_thres = 0.05

thickness = 2

x, y = (1920, 1080)

re_x, re_y = (1920, 1080)

def LoadModule():

    device = select_device('')

    weights = 'yolov5s.pt'

    model = DetectMultiBackend(weights, device=device, dnn=False, fp16=False)

    return model

model = LoadModule()

while True:

    names = model.names

    img0 = grab_screen_win32(region=(0, 0, 1920, 1080))

    im = letterbox(img0, 640, stride=32, auto=True)[0]  # padded resize

    im = im.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB

    im = np.ascontiguousarray(im)  # contiguous

    im = torch.from_numpy(im).to(model.device)

    im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32

    im /= 255  # 0 - 255 to 0.0 - 1.0

    if len(im.shape) == 3:

        im = im[None]  # expand for batch dim

    pred = model(im, augment=False, visualize=False)

    pred = non_max_suppression(pred, conf_thres=conf_thres, iou_thres=iou_thres, classes=None, agnostic=False,

                               max_det=1000)

    boxs=[]

    for i, det in enumerate(pred):  # per image

        im0 = img0.copy()

        s = ' '

        s += '%gx%g ' % im.shape[2:]  # print string

        gn = torch.tensor(img0.shape)[[1, 0, 1, 0]]  # normalization gain whwh

        imc = img0  # for save_crop

        if len(det):

            # Rescale boxes from img_size to im0 size

            det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()

            # Print results

            for c in det[:, 5].unique():

                n = (det[:, 5] == c).sum()  # detections per class

                s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string

        # Write results

        for *xyxy, conf, cls in reversed(det):

            xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh

            line = (cls, *xywh)  # label format

            box = ('%g ' * len(line)).rstrip() % line

            box = box.split(' ')

            boxs.append(box)

        if len(boxs):

            for i, det in enumerate(boxs):

                _, x_center, y_center, width, height = det

                x_center, width = re_x * float(x_center), re_x * float(width)

                y_center, height = re_y * float(y_center), re_y * float(height)

                top_left = (int(x_center - width / 2.), int(y_center - height / 2.))

                bottom_right = (int(x_center + width / 2.), int(y_center + height / 2.))

                color = (0, 0, 255)  # RGB

                cv2.rectangle(img0, top_left, bottom_right, color, thickness=thickness)

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

        cv2.destroyWindow()

        break

    cv2.namedWindow('windows', cv2.WINDOW_NORMAL)

    cv2.resizeWindow('windows', re_x // 2, re_y // 2)

    cv2.imshow('windows', img0)

    HWND = win32gui.FindWindow(None, "windows")

    win32gui.SetWindowPos(HWND, win32con.HWND_TOPMOST, 0, 0, 0, 0, win32con.SWP_NOMOVE | win32con.SWP_NOSIZE)

End.