2、YOLOv5社交距离项目
yolov5检测要检测的视频流中的所有人，然后再计算所有检测到的人之间的相互“距离”，和现实生活中用“m”这样的单位衡量距离不一样的是，在计算机中，简单的方法是用检测到的两个人的质心，也就是检测到的目标框的中心之间相隔的像素值作为计算机中的“距离”来衡量视频中的人之间的距离是否超过安全距离。
构建步骤：
- 使用目标检测算法检测视频流中的所有人，得到位置信息和质心位置；
- 计算所有检测到的人质心之间的相互距离；
- 设置安全距离，计算每个人之间的距离对，检测两个人之间的距离是否小于N个像素，小于则处于安全距离，反之则不处于。
项目架构：
detect.py代码注释如下：
import argparse from utils.datasets import * from utils.utils import * def detect(save_img=False): out, source, weights, view_img, save_txt, imgsz = \ opt.output, opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size webcam = source == '0' or source.startswith('rtsp') or source.startswith('http') or source.endswith('.txt') # Initialize device = torch_utils.select_device(opt.device) if os.path.exists(out): shutil.rmtree(out) # delete output folder os.makedirs(out) # make new output folder half = device.type != 'cpu' # half precision only supported on CUDA # 下载模型 google_utils.attempt_download(weights) # 加载权重 model = torch.load(weights, map_location=device)['model'].float() # torch.save(torch.load(weights, map_location=device), weights) # update model if SourceChangeWarning # model.fuse() # 设置模型为推理模式 model.to(device).eval() if half: model.half() # to FP16 # Second-stage classifier classify = False if classify: modelc = torch_utils.load_classifier(name='resnet101', n=2) # initialize modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']) # load weights modelc.to(device).eval() # 设置 Dataloader vid_path, vid_writer = None, None if webcam: view_img = True torch.backends.cudnn.benchmark = True # set True to speed up constant image size inference dataset = LoadStreams(source, img_size=imgsz) else: save_img = True dataset = LoadImages(source, img_size=imgsz) # 获取检测类别的标签名称 names = model.names if hasattr(model, 'names') else model.modules.names # 定义颜色 colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(names))] # 开始推理 t0 = time.time() # 初始化一张全为0的图片 img = torch.zeros((1, 3, imgsz, imgsz), device=device) _ = model(img.half() if half else img) if device.type != 'cpu' else None for path, img, im0s, vid_cap in dataset: img = torch.from_numpy(img).to(device) img = img.half() if half else img.float() # uint8 to fp16/32 img /= 255.0 # 0 - 255 to 0.0 - 1.0 if img.ndimension() == 3: img = img.unsqueeze(0) # 预测结果 t1 = torch_utils.time_synchronized() pred = model(img, augment=opt.augment)[0] # 使用NMS pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, fast=True, classes=opt.classes, agnostic=opt.agnostic_nms) t2 = torch_utils.time_synchronized() # 进行分类 if classify: pred = apply_classifier(pred, modelc, img, im0s) people_coords = [] # 处理预测得到的检测目标 for i, det in enumerate(pred): if webcam: p, s, im0 = path[i], '%g: ' % i, im0s[i].copy() else: p, s, im0 = path, '', im0s save_path = str(Path(out) / Path(p).name) s += '%gx%g ' % img.shape[2:] # print string gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh if det is not None and len(det): # 把boxes resize到im0的size det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round() # 打印结果 for c in det[:, -1].unique(): n = (det[:, -1] == c).sum() # detections per class s += '%g %ss, ' % (n, names[int(c)]) # add to string # 书写结果 for *xyxy, conf, cls in det: if save_txt: # xyxy2xywh ==> 把预测得到的坐标结果[x1, y1, x2, y2]转换为[x, y, w, h]其中 xy1=top-left, xy2=bottom-right xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh with open(save_path[:save_path.rfind('.')] + '.txt', 'a') as file: file.write(('%g ' * 5 + '\n') % (cls, *xywh)) # label format if save_img or view_img: # Add bbox to image label = '%s %.2f' % (names[int(cls)], conf) if label is not None: if (label.split())[0] == 'person': # print(xyxy) people_coords.append(xyxy) # plot_one_box(xyxy, im0, line_thickness=3) plot_dots_on_people(xyxy, im0) # 通过people_coords绘制people之间的连接线 # 这里主要分为"Low Risk "和"High Risk" distancing(people_coords, im0, dist_thres_lim=(200, 250)) # Print time (inference + NMS) print('%sDone. (%.3fs)' % (s, t2 - t1)) # Stream results if view_img: cv2.imshow(p, im0) if cv2.waitKey(1) == ord('q'): # q to quit raise StopIteration # Save results (image with detections) if save_img: if dataset.mode == 'images': cv2.imwrite(save_path, im0) else: if vid_path != save_path: # new video vid_path = save_path if isinstance(vid_writer, cv2.VideoWriter): vid_writer.release() # release previous video writer 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)) vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*opt.fourcc), fps, (w, h)) vid_writer.write(im0) if save_txt or save_img: print('Results saved to %s' % os.getcwd() + os.sep + out) if platform == 'darwin': # MacOS os.system('open ' + save_path) print('Done. (%.3fs)' % (time.time() - t0)) if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--weights', type=str, default='./weights/yolov5s.pt', help='model.pt path') parser.add_argument('--source', type=str, default='./inference/videos/', help='source') # file/folder, 0 for webcam parser.add_argument('--output', type=str, default='./inference/output', help='output folder') # output folder parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)') parser.add_argument('--conf-thres', type=float, default=0.4, help='object confidence threshold') parser.add_argument('--iou-thres', type=float, default=0.5, help='IOU threshold for NMS') parser.add_argument('--fourcc', type=str, default='mp4v', help='output video codec (verify ffmpeg support)') parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') parser.add_argument('--view-img', action='store_true', help='display results') parser.add_argument('--save-txt', action='store_true', help='save results to *.txt') parser.add_argument('--classes', nargs='+', type=int, help='filter by class') parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS') parser.add_argument('--augment', action='store_true', help='augmented inference') opt = parser.parse_args() opt.img_size = check_img_size(opt.img_size) print(opt) with torch.no_grad(): detect() ，时长05:00 ，时长00:13
参考
[1].https://zhuanlan.zhihu.com/p/172121380
[2].https://blog.csdn.net/weixin_45192980/article/details/108354169
[3].https://github.com/ultralytics/yoloV5
[4].https://github.com/Akbonline/Social-Distancing-using-YOLOv5
