YOLOV5实时检测屏幕

赵小心

YOLOV5实时检测屏幕


目录

• YOLOV5实时检测屏幕
  
  
  
  • 思考部分
    
    
    
    • 先把原本的detect.py的代码贴在这里
      
    • 分析代码并删减不用的部分
      
    
    
  • 把屏幕的截图通过OpenCV进行显示
    
    
    
    • 写一个屏幕截图的文件
      
    • 用OpenCV绘制窗口并显示
      
    
    
  • 最终代码
    
  
  

注：此为笔记目的：保留模型加载和推理部分，完成实时屏幕检测
实现思路：
1. 写一个实时截取屏幕的函数
2. 将截取的屏幕在窗口显示出来
3. 用OpenCV绘制一个窗口用来显示截取的屏幕
4. 在detect找出推理的代码，推理完成后得到中心点的xy坐标，宽高组成box
5. 在创建的OpenCV窗口用得到的推理结果绘制方框
实现效果：

思考部分

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

1. import argparse
2. import os
3. import platform
4. import sys
5. from pathlib import Path
7. import torch
9. FILE = Path(__file__).resolve()
10. ROOT = FILE.parents[0]  # YOLOv5 root directory
11. if str(ROOT) not in sys.path:
12.     sys.path.append(str(ROOT))  # add ROOT to PATH
13. ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
15. from models.common import DetectMultiBackend
16. from utils.dataloaders import IMG_FORMATS, VID_FORMATS, LoadImages, LoadScreenshots, LoadStreams
17. from utils.general import (LOGGER, Profile, check_file, check_img_size, check_imshow, check_requirements, colorstr, cv2,
18.                            increment_path, non_max_suppression, print_args, scale_boxes, strip_optimizer, xyxy2xywh)
19. from utils.plots import Annotator, colors, save_one_box
20. from utils.torch_utils import select_device, smart_inference_mode
23. @smart_inference_mode()
24. def run(
25.         weights=ROOT / 'yolov5s.pt',  # model path or triton URL
26.         source=ROOT / 'data/video/',
27.         data=ROOT / 'data/coco128.yaml',  # dataset.yaml path
28.         imgsz=(640, 640),  # inference size (height, width)
29.         conf_thres=0.25,  # confidence threshold
30.         iou_thres=0.45,  # NMS IOU threshold
31.         max_det=1000,  # maximum detections per image
32.         device='',  # cuda device, i.e. 0 or 0,1,2,3 or cpu
33.         view_img=False,  # show results
34.         save_txt=False,  # save results to *.txt
35.         save_conf=False,  # save confidences in --save-txt labels
36.         save_crop=False,  # save cropped prediction boxes
37.         nosave=False,  # do not save images/videos
38.         classes=None,  # filter by class: --class 0, or --class 0 2 3
39.         agnostic_nms=False,  # class-agnostic NMS
40.         augment=False,  # augmented inference
41.         visualize=False,  # visualize features
42.         update=False,  # update all models
43.         project=ROOT / 'runs/detect',  # save results to project/name
44.         name='exp',  # save results to project/name
45.         exist_ok=False,  # existing project/name ok, do not increment
46.         line_thickness=3,  # bounding box thickness (pixels)
47.         hide_labels=False,  # hide labels
48.         hide_conf=False,  # hide confidences
49.         half=False,  # use FP16 half-precision inference
50.         dnn=False,  # use OpenCV DNN for ONNX inference
51.         vid_stride=1,  # video frame-rate stride
52. ):
53.     source = str(source)
54.     save_img = not nosave and not source.endswith('.txt')  # save inference images
55.     is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)
56.     is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://'))
57.     webcam = source.isnumeric() or source.endswith('.streams') or (is_url and not is_file)
58.     screenshot = source.lower().startswith('screen')
59.     if is_url and is_file:
60.         source = check_file(source)  # download
62.     # Directories
63.     save_dir = increment_path(Path(project) / name, exist_ok=exist_ok)  # increment run
64.     (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir
66.     # Load model
67.     device = select_device(device)
68.     model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half)
69.     stride, names, pt = model.stride, model.names, model.pt
70.     imgsz = check_img_size(imgsz, s=stride)  # check image size
72.     # Dataloader
73.     bs = 1  # batch_size
74.     if webcam:
75.         view_img = check_imshow(warn=True)
76.         dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
77.         bs = len(dataset)
78.     elif screenshot:
79.         dataset = LoadScreenshots(source, img_size=imgsz, stride=stride, auto=pt)
80.     else:
81.         dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
82.     vid_path, vid_writer = [None] * bs, [None] * bs
84.     # Run inference
85.     model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *imgsz))  # warmup
86.     seen, windows, dt = 0, [], (Profile(), Profile(), Profile())
87.     for path, im, im0s, vid_cap, s in dataset:
88.         with dt[0]:
89.             im = torch.from_numpy(im).to(model.device)
90.             im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32
91.             im /= 255  # 0 - 255 to 0.0 - 1.0
92.             if len(im.shape) == 3:
93.                 im = im[None]  # expand for batch dim
95.         # Inference
96.         with dt[1]:
97.             visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False
98.             pred = model(im, augment=augment, visualize=visualize)
100.         # NMS
101.         with dt[2]:
102.             pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
104.         # Second-stage classifier (optional)
105.         # pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)
107.         # Process predictions
108.         for i, det in enumerate(pred):  # per image
109.             seen += 1
110.             if webcam:  # batch_size >= 1
111.                 p, im0, frame = path[i], im0s[i].copy(), dataset.count
112.                 s += f'{i}: '
113.             else:
114.                 p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)
116.             p = Path(p)  # to Path
117.             save_path = str(save_dir / p.name)  # im.jpg
118.             txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}')  # im.txt
119.             s += '%gx%g ' % im.shape[2:]  # print string
120.             gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
121.             imc = im0.copy() if save_crop else im0  # for save_crop
122.             annotator = Annotator(im0, line_width=line_thickness, example=str(names))
123.             if len(det):
124.                 # Rescale boxes from img_size to im0 size
125.                 det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()
127.                 # Print results
128.                 for c in det[:, 5].unique():
129.                     n = (det[:, 5] == c).sum()  # detections per class
130.                     s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string
132.                 # Write results
133.                 for *xyxy, conf, cls in reversed(det):
134.                     if save_txt:  # Write to file
135.                         xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
136.                         line = (cls, *xywh, conf) if save_conf else (cls, *xywh)  # label format
137.                         with open(f'{txt_path}.txt', 'a') as f:
138.                             f.write(('%g ' * len(line)).rstrip() % line + '\n')
140.                     if save_img or save_crop or view_img:  # Add bbox to image
141.                         c = int(cls)  # integer class
142.                         label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')
143.                         annotator.box_label(xyxy, label, color=colors(c, True))
144.                     if save_crop:
145.                         save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)
147.             # Stream results
148.             im0 = annotator.result()
149.             if view_img:
150.                 if platform.system() == 'Linux' and p not in windows:
151.                     windows.append(p)
152.                     cv2.namedWindow(str(p), cv2.WINDOW_NORMAL | cv2.WINDOW_KEEPRATIO)  # allow window resize (Linux)
153.                     cv2.resizeWindow(str(p), im0.shape[1], im0.shape[0])
154.                 cv2.imshow(str(p), im0)
155.                 cv2.waitKey(1)  # 1 millisecond
157.             # Save results (image with detections)
158.             if save_img:
159.                 if dataset.mode == 'image':
160.                     cv2.imwrite(save_path, im0)
161.                 else:  # 'video' or 'stream'
162.                     if vid_path[i] != save_path:  # new video
163.                         vid_path[i] = save_path
164.                         if isinstance(vid_writer[i], cv2.VideoWriter):
165.                             vid_writer[i].release()  # release previous video writer
166.                         if vid_cap:  # video
167.                             fps = vid_cap.get(cv2.CAP_PROP_FPS)
168.                             w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
169.                             h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
170.                         else:  # stream
171.                             fps, w, h = 30, im0.shape[1], im0.shape[0]
172.                         save_path = str(Path(save_path).with_suffix('.mp4'))  # force *.mp4 suffix on results videos
173.                         vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
174.                     vid_writer[i].write(im0)
176.         # Print time (inference-only)
177.         LOGGER.info(f"{s}{'' if len(det) else '(no detections), '}{dt[1].dt * 1E3:.1f}ms")
179.     # Print results
180.     t = tuple(x.t / seen * 1E3 for x in dt)  # speeds per image
181.     LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}' % t)
182.     if save_txt or save_img:
183.         s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
184.         LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
185.     if update:
186.         strip_optimizer(weights[0])  # update model (to fix SourceChangeWarning)
189. def parse_opt():
190.     parser = argparse.ArgumentParser()
191.     parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model path or triton URL')
192.     parser.add_argument('--source', type=str, default=ROOT / '0', help='file/dir/URL/glob/screen/1(webcam)')
193.     parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='(optional) dataset.yaml path')
194.     parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='inference size h,w')
195.     parser.add_argument('--conf-thres', type=float, default=0.45, help='confidence threshold')
196.     parser.add_argument('--iou-thres', type=float, default=0.2, help='NMS IoU threshold')
197.     parser.add_argument('--max-det', type=int, default=1000, help='maximum detections per image')
198.     parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
199.     parser.add_argument('--view-img', action='store_true', help='show results')
200.     parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
201.     parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
202.     parser.add_argument('--save-crop', action='store_true', help='save cropped prediction boxes')
203.     parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
204.     parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --classes 0, or --classes 0 2 3')
205.     parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
206.     parser.add_argument('--augment', action='store_true', help='augmented inference')
207.     parser.add_argument('--visualize', action='store_true', help='visualize features')
208.     parser.add_argument('--update', action='store_true', help='update all models')
209.     parser.add_argument('--project', default=ROOT / 'runs/detect', help='save results to project/name')
210.     parser.add_argument('--name', default='exp', help='save results to project/name')
211.     parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
212.     parser.add_argument('--line-thickness', default=3, type=int, help='bounding box thickness (pixels)')
213.     parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels')
214.     parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences')
215.     parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
216.     parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')
217.     parser.add_argument('--vid-stride', type=int, default=1, help='video frame-rate stride')
218.     opt = parser.parse_args()
219.     opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1  # expand
220.     print_args(vars(opt))
221.     return opt
224. def main(opt):
225.     check_requirements(exclude=('tensorboard', 'thop'))
226.     run(**vars(opt))
229. if __name__ == '__main__':
230.     opt = parse_opt()
231.     main(opt)

复制代码

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

1. import argparse
2. import os
3. import platform
4. import sys
5. from pathlib import Path
7. import torch
9. FILE = Path(__file__).resolve()
10. ROOT = FILE.parents[0]  # YOLOv5 root directory
11. if str(ROOT) not in sys.path:
12.     sys.path.append(str(ROOT))  # add ROOT to PATH
13. ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
15. from models.common import DetectMultiBackend
16. from utils.dataloaders import IMG_FORMATS, VID_FORMATS, LoadImages, LoadScreenshots, LoadStreams
17. from utils.general import (LOGGER, Profile, check_file, check_img_size, check_imshow, check_requirements, colorstr, cv2,
18.                            increment_path, non_max_suppression, print_args, scale_boxes, strip_optimizer, xyxy2xywh)
19. from utils.plots import Annotator, colors, save_one_box
20. from utils.torch_utils import select_device, smart_inference_mode

复制代码

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

1. if __name__ == '__main__':
2.     opt = parse_opt()
3.     main(opt)

复制代码

从if __name__ == '__main__开始
opt = parse_opt 就是一个获取命令行参数的函数，我们并不需要,可以删
进入main函数

1. def main(opt):
2.     check_requirements(exclude=('tensorboard', 'thop'))
3.     run(**vars(opt))

复制代码

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

1.     source = str(source)
2.     save_img = not nosave and not source.endswith('.txt')  # save inference images
3.     is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)
4.     is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://'))
5.     webcam = source.isnumeric() or source.endswith('.streams') or (is_url and not is_file)
6.     screenshot = source.lower().startswith('screen')
7.     if is_url and is_file:
8.         source = check_file(source)  # download
10.     # Directories
11.     save_dir = increment_path(Path(project) / name, exist_ok=exist_ok)  # increment run
12.     (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir

复制代码

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

1. # Load model
2. device = select_device(device)
3. 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
  

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

1. def LoadModule():
2.     device = select_device('')
3.     weights = 'yolov5s.pt'
4.     model = DetectMultiBackend(weights, device=device, dnn=False, fp16=False)
5.     return model

复制代码

继续往下读

1. bs = 1  # batch_size
2.     if webcam:
3.         view_img = check_imshow(warn=True)
4.         dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
5.         bs = len(dataset)
6.     elif screenshot:
7.         dataset = LoadScreenshots(source, img_size=imgsz, stride=stride, auto=pt)
8.     else:
9.         dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
10.     vid_path, vid_writer = [None] * bs, [None] * bs

复制代码

这里如果是使用网络摄像头作为输入，会通过LoadStreams类加载视频流，根据图像大小和步长采样，如果使用截图作为输入，则通过LoadScreenshots加载截图，都不是则通过LoadImages类加载图片文件
这是YOLOV5提供的加载dataset的部分，我们可以添加自己的dataset，所以删掉
继续往下

1. # Run inference
2.     model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *imgsz))  # warmup
3.     seen, windows, dt = 0, [], (Profile(), Profile(), Profile())
4.     for path, im, im0s, vid_cap, s in dataset:
5.         with dt[0]:
6.             im = torch.from_numpy(im).to(model.device)
7.             im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32
8.             im /= 255  # 0 - 255 to 0.0 - 1.0
9.             if len(im.shape) == 3:
10.                 im = im[None]  # expand for batch dim
12.         # Inference
13.         with dt[1]:
14.             visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False
15.             pred = model(im, augment=augment, visualize=visualize)
17.         # NMS
18.         with dt[2]:
19.             pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
21.         # Second-stage classifier (optional)
22.         # pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)
24.         # 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())
未知作用，删了

1. for path, im, im0s, vid_cap, s in dataset:
2.         with dt[0]:
3.             im = torch.from_numpy(im).to(model.device)
4.             im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32
5.             im /= 255  # 0 - 255 to 0.0 - 1.0
6.             if len(im.shape) == 3:
7.                 im = im[None]  # expand for batch dim
9.         # Inference
10.         with dt[1]:
11.             visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False
12.             pred = model(im, augment=augment, visualize=visualize)
14.         # NMS
15.         with dt[2]:
16.             pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)

复制代码

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

1. im = torch.from_numpy(im).to(model.device)
2. im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32
3. im /= 255  # 0 - 255 to 0.0 - 1.0
4. if len(im.shape) == 3:
5.     im = im[None]  # expand for batch dim
6. visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False
7. pred = model(im, augment=augment, visualize=visualize)
8. 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中可以看到

1. if self.transforms:
2.     im = self.transforms(im0)  # transforms
3. else:
4.     im = letterbox(im0, self.img_size, stride=self.stride, auto=self.auto)[0]  # padded resize
5.     im = im.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
6.     im = np.ascontiguousarray(im)  # contiguous
8. 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
则得到改写代码为

1. im = letterbox(img0, 640, stride=32, auto=True)[0]  # padded resize
2.     im = im.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
3.     im = np.ascontiguousarray(im)  # contiguous
4.     im = torch.from_numpy(im).to(model.device)
5.     im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32
6.     im /= 255  # 0 - 255 to 0.0 - 1.0
7.     if len(im.shape) == 3:
8.         im = im[None]  # expand for batch dim
9.     pred = model(im, augment=False, visualize=False)
10.     pred = non_max_suppression(pred, conf_thres=conf_thres, iou_thres=iou_thres, classes=None, agnostic=False,
11.                                max_det=1000)

复制代码

继续向下

1. for i, det in enumerate(pred):  # per image
2.             seen += 1
3.             if webcam:  # batch_size >= 1
4.                 p, im0, frame = path[i], im0s[i].copy(), dataset.count
5.                 s += f'{i}: '
6.             else:
7.                 p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)
9.             p = Path(p)  # to Path
10.             save_path = str(save_dir / p.name)  # im.jpg
11.             txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}')  # im.txt
12.             s += '%gx%g ' % im.shape[2:]  # print string
13.             gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
14.             imc = im0.copy() if save_crop else im0  # for save_crop
15.             annotator = Annotator(im0, line_width=line_thickness, example=str(names))
16.             if len(det):
17.                 # Rescale boxes from img_size to im0 size
18.                 det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()
20.                 # Print results
21.                 for c in det[:, 5].unique():
22.                     n = (det[:, 5] == c).sum()  # detections per class
23.                     s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string
25.                 # Write results
26.                 for *xyxy, conf, cls in reversed(det):
27.                     if save_txt:  # Write to file
28.                         xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
29.                         line = (cls, *xywh, conf) if save_conf else (cls, *xywh)  # label format
30.                         with open(f'{txt_path}.txt', 'a') as f:
31.                             f.write(('%g ' * len(line)).rstrip() % line + '\n')
33.                     if save_img or save_crop or view_img:  # Add bbox to image
34.                         c = int(cls)  # integer class
35.                         label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')
36.                         annotator.box_label(xyxy, label, color=colors(c, True))
37.                     if save_crop:
38.                         save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)

复制代码

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

1. boxs=[]
2.     for i, det in enumerate(pred):  # per image
3.         im0 = img0.copy()
4.         s = ' '
5.         s += '%gx%g ' % im.shape[2:]  # print string
6.         gn = torch.tensor(img0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
7.         imc = img0  # for save_crop
9.         if len(det):
10.             # Rescale boxes from img_size to im0 size
11.             det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()
12.             # Print results
13.             for c in det[:, 5].unique():
14.                 n = (det[:, 5] == c).sum()  # detections per class
15.                 s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string
17.         # Write results
18.         for *xyxy, conf, cls in reversed(det):
19.             xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
20.             line = (cls, *xywh)  # label format
21.             box = ('%g ' * len(line)).rstrip() % line
22.             box = box.split(' ')
23.             boxs.append(box)

复制代码

就此完成了推理部分的删减和重写
把屏幕的截图通过OpenCV进行显示

写一个屏幕截图的文件

写成 grabscreen.py

1. # 文件名:grabscreen.py
2. import cv2
3. import numpy as np
4. import win32gui
5. import win32print
6. import win32ui
7. import win32con
8. import win32api
9. import mss
12. def grab_screen_win32(region):
13.     hwin = win32gui.GetDesktopWindow()
14.     left, top, x2, y2 = region
15.     width = x2 - left + 1
16.     height = y2 - top + 1
18.     hwindc = win32gui.GetWindowDC(hwin)
19.     srcdc = win32ui.CreateDCFromHandle(hwindc)
20.     memdc = srcdc.CreateCompatibleDC()
21.     bmp = win32ui.CreateBitmap()
22.     bmp.CreateCompatibleBitmap(srcdc, width, height)
23.     memdc.SelectObject(bmp)
24.     memdc.BitBlt((0, 0), (width, height), srcdc, (left, top), win32con.SRCCOPY)
26.     signedIntsArray = bmp.GetBitmapBits(True)
27.     img = np.fromstring(signedIntsArray, dtype='uint8')
28.     img.shape = (height, width, 4)
30.     srcdc.DeleteDC()
31.     memdc.DeleteDC()
32.     win32gui.ReleaseDC(hwin, hwindc)
33.     win32gui.DeleteObject(bmp.GetHandle())
35.     return cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)

复制代码

通过img0 = grab_screen_win32(region=(0, 0, 1920, 1080))来作为im的参数传入，即可让屏幕截图作为推理图片
用OpenCV绘制窗口并显示

1. if len(boxs):
2.     for i, det in enumerate(boxs):
3.         _, x_center, y_center, width, height = det
4.         x_center, width = re_x * float(x_center), re_x * float(width)
5.         y_center, height = re_y * float(y_center), re_y * float(height)
6.         top_left = (int(x_center - width / 2.), int(y_center - height / 2.))
7.         bottom_right = (int(x_center + width / 2.), int(y_center + height / 2.))
8.         color = (0, 0, 255)  # RGB
9.         cv2.rectangle(img0, top_left, bottom_right, color, thickness=thickness)
11. 和
13. cv2.namedWindow('windows', cv2.WINDOW_NORMAL)
14. cv2.resizeWindow('windows', re_x // 2, re_y // 2)
15. cv2.imshow('windows', img0)
16. HWND = win32gui.FindWindow(None, "windows")
17. win32gui.SetWindowPos(HWND, win32con.HWND_TOPMOST, 0, 0, 0, 0, win32con.SWP_NOMOVE | win32con.SWP_NOSIZE)

复制代码

结合在一起
最终代码

1. import torch, pynputimport numpy as npimport win32gui, win32con, cv2from grabscreen import grab_screen_win32 # 本地文件from utils.augmentations import letterboxfrom models.common import DetectMultiBackendfrom utils.torch_utils import select_devicefrom utils.general import non_max_suppression, scale_boxes, xyxy2xywh# 可调参数conf_thres = 0.25iou_thres = 0.05thickness = 2x, y = (1920, 1080)re_x, re_y = (1920, 1080)def LoadModule():
2.     device = select_device('')
3.     weights = 'yolov5s.pt'
4.     model = DetectMultiBackend(weights, device=device, dnn=False, fp16=False)
5.     return modelmodel = 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
6.     im = im.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
7.     im = np.ascontiguousarray(im)  # contiguous
8.     im = torch.from_numpy(im).to(model.device)
9.     im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32
10.     im /= 255  # 0 - 255 to 0.0 - 1.0
11.     if len(im.shape) == 3:
12.         im = im[None]  # expand for batch dim
13.     pred = model(im, augment=False, visualize=False)
14.     pred = non_max_suppression(pred, conf_thres=conf_thres, iou_thres=iou_thres, classes=None, agnostic=False,
15.                                max_det=1000)    boxs=[]
16.     for i, det in enumerate(pred):  # per image
17.         im0 = img0.copy()
18.         s = ' '
19.         s += '%gx%g ' % im.shape[2:]  # print string
20.         gn = torch.tensor(img0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
21.         imc = img0  # for save_crop
23.         if len(det):
24.             # Rescale boxes from img_size to im0 size
25.             det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()
26.             # Print results
27.             for c in det[:, 5].unique():
28.                 n = (det[:, 5] == c).sum()  # detections per class
29.                 s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string
31.         # Write results
32.         for *xyxy, conf, cls in reversed(det):
33.             xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
34.             line = (cls, *xywh)  # label format
35.             box = ('%g ' * len(line)).rstrip() % line
36.             box = box.split(' ')
37.             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.

来源:https://www.cnblogs.com/water-wells/p/17448591.html
免责声明：由于采集信息均来自互联网，如果侵犯了您的权益，请联系我们【E-Mail:cb@itdo.tech】 我们会及时删除侵权内容，谢谢合作！