You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

223 lines
11KB

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