Vous ne pouvez pas sélectionner plus de 25 sujets Les noms de sujets doivent commencer par une lettre ou un nombre, peuvent contenir des tirets ('-') et peuvent comporter jusqu'à 35 caractères.

183 lignes
9.1KB

  1. import argparse
  2. import time
  3. from pathlib import Path
  4. import cv2
  5. import torch
  6. import torch.backends.cudnn as cudnn
  7. from models.experimental import attempt_load
  8. from utils.datasets import LoadStreams, LoadImages
  9. from utils.general import check_img_size, check_requirements, check_imshow, non_max_suppression, apply_classifier, \
  10. scale_coords, xyxy2xywh, strip_optimizer, set_logging, increment_path, save_one_box
  11. from utils.plots import colors, plot_one_box
  12. from utils.torch_utils import select_device, load_classifier, time_synchronized
  13. def detect(opt):
  14. source, weights, view_img, save_txt, imgsz = opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
  15. save_img = not opt.nosave and not source.endswith('.txt') # save inference images
  16. webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith(
  17. ('rtsp://', 'rtmp://', 'http://', 'https://'))
  18. # Directories
  19. save_dir = increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok) # increment run
  20. (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
  21. # Initialize
  22. set_logging()
  23. device = select_device(opt.device)
  24. half = device.type != 'cpu' # half precision only supported on CUDA
  25. # Load model
  26. model = attempt_load(weights, map_location=device) # load FP32 model
  27. stride = int(model.stride.max()) # model stride
  28. imgsz = check_img_size(imgsz, s=stride) # check img_size
  29. names = model.module.names if hasattr(model, 'module') else model.names # get class names
  30. if half:
  31. model.half() # to FP16
  32. # Second-stage classifier
  33. classify = False
  34. if classify:
  35. modelc = load_classifier(name='resnet101', n=2) # initialize
  36. modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']).to(device).eval()
  37. # Set Dataloader
  38. vid_path, vid_writer = None, None
  39. if webcam:
  40. view_img = check_imshow()
  41. cudnn.benchmark = True # set True to speed up constant image size inference
  42. dataset = LoadStreams(source, img_size=imgsz, stride=stride)
  43. else:
  44. dataset = LoadImages(source, img_size=imgsz, stride=stride)
  45. # Run inference
  46. if device.type != 'cpu':
  47. model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters()))) # run once
  48. t0 = time.time()
  49. for path, img, im0s, vid_cap in dataset:
  50. img = torch.from_numpy(img).to(device)
  51. img = img.half() if half else img.float() # uint8 to fp16/32
  52. img /= 255.0 # 0 - 255 to 0.0 - 1.0
  53. if img.ndimension() == 3:
  54. img = img.unsqueeze(0)
  55. # Inference
  56. t1 = time_synchronized()
  57. pred = model(img, augment=opt.augment)[0]
  58. # Apply NMS
  59. pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
  60. t2 = time_synchronized()
  61. # Apply Classifier
  62. if classify:
  63. pred = apply_classifier(pred, modelc, img, im0s)
  64. # Process detections
  65. for i, det in enumerate(pred): # detections per image
  66. if webcam: # batch_size >= 1
  67. p, s, im0, frame = path[i], f'{i}: ', im0s[i].copy(), dataset.count
  68. else:
  69. p, s, im0, frame = path, '', im0s.copy(), getattr(dataset, 'frame', 0)
  70. p = Path(p) # to Path
  71. save_path = str(save_dir / p.name) # img.jpg
  72. txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}') # img.txt
  73. s += '%gx%g ' % img.shape[2:] # print string
  74. gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh
  75. imc = im0.copy() if opt.save_crop else im0 # for opt.save_crop
  76. if len(det):
  77. # Rescale boxes from img_size to im0 size
  78. det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
  79. # Print results
  80. for c in det[:, -1].unique():
  81. n = (det[:, -1] == c).sum() # detections per class
  82. s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string
  83. # Write results
  84. for *xyxy, conf, cls in reversed(det):
  85. if save_txt: # Write to file
  86. xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
  87. line = (cls, *xywh, conf) if opt.save_conf else (cls, *xywh) # label format
  88. with open(txt_path + '.txt', 'a') as f:
  89. f.write(('%g ' * len(line)).rstrip() % line + '\n')
  90. if save_img or opt.save_crop or view_img: # Add bbox to image
  91. c = int(cls) # integer class
  92. label = None if opt.hide_labels else (names[c] if opt.hide_conf else f'{names[c]} {conf:.2f}')
  93. plot_one_box(xyxy, im0, label=label, color=colors(c, True), line_thickness=opt.line_thickness)
  94. if opt.save_crop:
  95. save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)
  96. # Print time (inference + NMS)
  97. print(f'{s}Done. ({t2 - t1:.3f}s)')
  98. # Stream results
  99. if view_img:
  100. cv2.imshow(str(p), im0)
  101. cv2.waitKey(1) # 1 millisecond
  102. # Save results (image with detections)
  103. if save_img:
  104. if dataset.mode == 'image':
  105. cv2.imwrite(save_path, im0)
  106. else: # 'video' or 'stream'
  107. if vid_path != save_path: # new video
  108. vid_path = save_path
  109. if isinstance(vid_writer, cv2.VideoWriter):
  110. vid_writer.release() # release previous video writer
  111. if vid_cap: # video
  112. fps = vid_cap.get(cv2.CAP_PROP_FPS)
  113. w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
  114. h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
  115. else: # stream
  116. fps, w, h = 30, im0.shape[1], im0.shape[0]
  117. save_path += '.mp4'
  118. vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
  119. vid_writer.write(im0)
  120. if save_txt or save_img:
  121. s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
  122. print(f"Results saved to {save_dir}{s}")
  123. print(f'Done. ({time.time() - t0:.3f}s)')
  124. if __name__ == '__main__':
  125. parser = argparse.ArgumentParser()
  126. parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)')
  127. parser.add_argument('--source', type=str, default='data/images', help='source') # file/folder, 0 for webcam
  128. parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
  129. parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')
  130. parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')
  131. parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
  132. parser.add_argument('--view-img', action='store_true', help='display results')
  133. parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
  134. parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
  135. parser.add_argument('--save-crop', action='store_true', help='save cropped prediction boxes')
  136. parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
  137. parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
  138. parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
  139. parser.add_argument('--augment', action='store_true', help='augmented inference')
  140. parser.add_argument('--update', action='store_true', help='update all models')
  141. parser.add_argument('--project', default='runs/detect', help='save results to project/name')
  142. parser.add_argument('--name', default='exp', help='save results to project/name')
  143. parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
  144. parser.add_argument('--line-thickness', default=3, type=int, help='bounding box thickness (pixels)')
  145. parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels')
  146. parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences')
  147. opt = parser.parse_args()
  148. print(opt)
  149. check_requirements(exclude=('tensorboard', 'pycocotools', 'thop'))
  150. with torch.no_grad():
  151. if opt.update: # update all models (to fix SourceChangeWarning)
  152. for opt.weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']:
  153. detect(opt=opt)
  154. strip_optimizer(opt.weights)
  155. else:
  156. detect(opt=opt)