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.

297 lines
15KB

  1. # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
  2. """
  3. Run inference on images, videos, directories, streams, etc.
  4. Usage:
  5. $ python path/to/detect.py --source path/to/img.jpg --weights yolov5s.pt --img 640
  6. """
  7. import argparse
  8. import sys
  9. from pathlib import Path
  10. import cv2
  11. import numpy as np
  12. import torch
  13. import torch.backends.cudnn as cudnn
  14. FILE = Path(__file__).resolve()
  15. ROOT = FILE.parents[0] # YOLOv5 root directory
  16. if str(ROOT) not in sys.path:
  17. sys.path.append(str(ROOT)) # add ROOT to PATH
  18. ROOT = ROOT.relative_to(Path.cwd()) # relative
  19. from models.experimental import attempt_load
  20. from utils.datasets import LoadImages, LoadStreams
  21. from utils.general import apply_classifier, check_img_size, check_imshow, check_requirements, check_suffix, colorstr, \
  22. increment_path, non_max_suppression, print_args, save_one_box, scale_coords, set_logging, \
  23. strip_optimizer, xyxy2xywh
  24. from utils.plots import Annotator, colors
  25. from utils.torch_utils import load_classifier, select_device, time_sync
  26. @torch.no_grad()
  27. def run(weights=ROOT / 'yolov5s.pt', # model.pt path(s)
  28. source=ROOT / 'data/images', # file/dir/URL/glob, 0 for webcam
  29. imgsz=640, # inference size (pixels)
  30. conf_thres=0.25, # confidence threshold
  31. iou_thres=0.45, # NMS IOU threshold
  32. max_det=1000, # maximum detections per image
  33. device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
  34. view_img=False, # show results
  35. save_txt=False, # save results to *.txt
  36. save_conf=False, # save confidences in --save-txt labels
  37. save_crop=False, # save cropped prediction boxes
  38. nosave=False, # do not save images/videos
  39. classes=None, # filter by class: --class 0, or --class 0 2 3
  40. agnostic_nms=False, # class-agnostic NMS
  41. augment=False, # augmented inference
  42. visualize=False, # visualize features
  43. update=False, # update all models
  44. project=ROOT / 'runs/detect', # save results to project/name
  45. name='exp', # save results to project/name
  46. exist_ok=False, # existing project/name ok, do not increment
  47. line_thickness=3, # bounding box thickness (pixels)
  48. hide_labels=False, # hide labels
  49. hide_conf=False, # hide confidences
  50. half=False, # use FP16 half-precision inference
  51. ):
  52. source = str(source)
  53. save_img = not nosave and not source.endswith('.txt') # save inference images
  54. webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith(
  55. ('rtsp://', 'rtmp://', 'http://', 'https://'))
  56. # Directories
  57. save_dir = increment_path(Path(project) / name, exist_ok=exist_ok) # increment run
  58. (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
  59. # Initialize
  60. set_logging()
  61. device = select_device(device)
  62. half &= device.type != 'cpu' # half precision only supported on CUDA
  63. # Load model
  64. w = weights[0] if isinstance(weights, list) else weights
  65. classify, suffix, suffixes = False, Path(w).suffix.lower(), ['.pt', '.onnx', '.tflite', '.pb', '']
  66. check_suffix(w, suffixes) # check weights have acceptable suffix
  67. pt, onnx, tflite, pb, saved_model = (suffix == x for x in suffixes) # backend booleans
  68. stride, names = 64, [f'class{i}' for i in range(1000)] # assign defaults
  69. if pt:
  70. model = attempt_load(weights, map_location=device) # load FP32 model
  71. stride = int(model.stride.max()) # model stride
  72. names = model.module.names if hasattr(model, 'module') else model.names # get class names
  73. if half:
  74. model.half() # to FP16
  75. if classify: # second-stage classifier
  76. modelc = load_classifier(name='resnet50', n=2) # initialize
  77. modelc.load_state_dict(torch.load('resnet50.pt', map_location=device)['model']).to(device).eval()
  78. elif onnx:
  79. check_requirements(('onnx', 'onnxruntime'))
  80. import onnxruntime
  81. session = onnxruntime.InferenceSession(w, None)
  82. else: # TensorFlow models
  83. check_requirements(('tensorflow>=2.4.1',))
  84. import tensorflow as tf
  85. if pb: # https://www.tensorflow.org/guide/migrate#a_graphpb_or_graphpbtxt
  86. def wrap_frozen_graph(gd, inputs, outputs):
  87. x = tf.compat.v1.wrap_function(lambda: tf.compat.v1.import_graph_def(gd, name=""), []) # wrapped import
  88. return x.prune(tf.nest.map_structure(x.graph.as_graph_element, inputs),
  89. tf.nest.map_structure(x.graph.as_graph_element, outputs))
  90. graph_def = tf.Graph().as_graph_def()
  91. graph_def.ParseFromString(open(w, 'rb').read())
  92. frozen_func = wrap_frozen_graph(gd=graph_def, inputs="x:0", outputs="Identity:0")
  93. elif saved_model:
  94. model = tf.keras.models.load_model(w)
  95. elif tflite:
  96. interpreter = tf.lite.Interpreter(model_path=w) # load TFLite model
  97. interpreter.allocate_tensors() # allocate
  98. input_details = interpreter.get_input_details() # inputs
  99. output_details = interpreter.get_output_details() # outputs
  100. int8 = input_details[0]['dtype'] == np.uint8 # is TFLite quantized uint8 model
  101. imgsz = check_img_size(imgsz, s=stride) # check image size
  102. # Dataloader
  103. if webcam:
  104. view_img = check_imshow()
  105. cudnn.benchmark = True # set True to speed up constant image size inference
  106. dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt)
  107. bs = len(dataset) # batch_size
  108. else:
  109. dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt)
  110. bs = 1 # batch_size
  111. vid_path, vid_writer = [None] * bs, [None] * bs
  112. # Run inference
  113. if pt and device.type != 'cpu':
  114. model(torch.zeros(1, 3, *imgsz).to(device).type_as(next(model.parameters()))) # run once
  115. dt, seen = [0.0, 0.0, 0.0], 0
  116. for path, img, im0s, vid_cap in dataset:
  117. t1 = time_sync()
  118. if onnx:
  119. img = img.astype('float32')
  120. else:
  121. img = torch.from_numpy(img).to(device)
  122. img = img.half() if half else img.float() # uint8 to fp16/32
  123. img = img / 255.0 # 0 - 255 to 0.0 - 1.0
  124. if len(img.shape) == 3:
  125. img = img[None] # expand for batch dim
  126. t2 = time_sync()
  127. dt[0] += t2 - t1
  128. # Inference
  129. if pt:
  130. visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False
  131. pred = model(img, augment=augment, visualize=visualize)[0]
  132. elif onnx:
  133. pred = torch.tensor(session.run([session.get_outputs()[0].name], {session.get_inputs()[0].name: img}))
  134. else: # tensorflow model (tflite, pb, saved_model)
  135. imn = img.permute(0, 2, 3, 1).cpu().numpy() # image in numpy
  136. if pb:
  137. pred = frozen_func(x=tf.constant(imn)).numpy()
  138. elif saved_model:
  139. pred = model(imn, training=False).numpy()
  140. elif tflite:
  141. if int8:
  142. scale, zero_point = input_details[0]['quantization']
  143. imn = (imn / scale + zero_point).astype(np.uint8) # de-scale
  144. interpreter.set_tensor(input_details[0]['index'], imn)
  145. interpreter.invoke()
  146. pred = interpreter.get_tensor(output_details[0]['index'])
  147. if int8:
  148. scale, zero_point = output_details[0]['quantization']
  149. pred = (pred.astype(np.float32) - zero_point) * scale # re-scale
  150. pred[..., 0] *= imgsz[1] # x
  151. pred[..., 1] *= imgsz[0] # y
  152. pred[..., 2] *= imgsz[1] # w
  153. pred[..., 3] *= imgsz[0] # h
  154. pred = torch.tensor(pred)
  155. t3 = time_sync()
  156. dt[1] += t3 - t2
  157. # NMS
  158. pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
  159. dt[2] += time_sync() - t3
  160. # Second-stage classifier (optional)
  161. if classify:
  162. pred = apply_classifier(pred, modelc, img, im0s)
  163. # Process predictions
  164. for i, det in enumerate(pred): # per image
  165. seen += 1
  166. if webcam: # batch_size >= 1
  167. p, s, im0, frame = path[i], f'{i}: ', im0s[i].copy(), dataset.count
  168. else:
  169. p, s, im0, frame = path, '', im0s.copy(), getattr(dataset, 'frame', 0)
  170. p = Path(p) # to Path
  171. save_path = str(save_dir / p.name) # img.jpg
  172. txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}') # img.txt
  173. s += '%gx%g ' % img.shape[2:] # print string
  174. gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh
  175. imc = im0.copy() if save_crop else im0 # for save_crop
  176. annotator = Annotator(im0, line_width=line_thickness, example=str(names))
  177. if len(det):
  178. # Rescale boxes from img_size to im0 size
  179. det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
  180. # Print results
  181. for c in det[:, -1].unique():
  182. n = (det[:, -1] == c).sum() # detections per class
  183. s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string
  184. # Write results
  185. for *xyxy, conf, cls in reversed(det):
  186. if save_txt: # Write to file
  187. xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
  188. line = (cls, *xywh, conf) if save_conf else (cls, *xywh) # label format
  189. with open(txt_path + '.txt', 'a') as f:
  190. f.write(('%g ' * len(line)).rstrip() % line + '\n')
  191. if save_img or save_crop or view_img: # Add bbox to image
  192. c = int(cls) # integer class
  193. label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')
  194. annotator.box_label(xyxy, label, color=colors(c, True))
  195. if save_crop:
  196. save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)
  197. # Print time (inference-only)
  198. print(f'{s}Done. ({t3 - t2:.3f}s)')
  199. # Stream results
  200. im0 = annotator.result()
  201. if view_img:
  202. cv2.imshow(str(p), im0)
  203. cv2.waitKey(1) # 1 millisecond
  204. # Save results (image with detections)
  205. if save_img:
  206. if dataset.mode == 'image':
  207. cv2.imwrite(save_path, im0)
  208. else: # 'video' or 'stream'
  209. if vid_path[i] != save_path: # new video
  210. vid_path[i] = save_path
  211. if isinstance(vid_writer[i], cv2.VideoWriter):
  212. vid_writer[i].release() # release previous video writer
  213. if vid_cap: # video
  214. fps = vid_cap.get(cv2.CAP_PROP_FPS)
  215. w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
  216. h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
  217. else: # stream
  218. fps, w, h = 30, im0.shape[1], im0.shape[0]
  219. save_path += '.mp4'
  220. vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
  221. vid_writer[i].write(im0)
  222. # Print results
  223. t = tuple(x / seen * 1E3 for x in dt) # speeds per image
  224. print(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}' % t)
  225. if save_txt or save_img:
  226. s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
  227. print(f"Results saved to {colorstr('bold', save_dir)}{s}")
  228. if update:
  229. strip_optimizer(weights) # update model (to fix SourceChangeWarning)
  230. def parse_opt():
  231. parser = argparse.ArgumentParser()
  232. parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model path(s)')
  233. parser.add_argument('--source', type=str, default=ROOT / 'data/images', help='file/dir/URL/glob, 0 for webcam')
  234. parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='inference size h,w')
  235. parser.add_argument('--conf-thres', type=float, default=0.25, help='confidence threshold')
  236. parser.add_argument('--iou-thres', type=float, default=0.45, help='NMS IoU threshold')
  237. parser.add_argument('--max-det', type=int, default=1000, help='maximum detections per image')
  238. parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
  239. parser.add_argument('--view-img', action='store_true', help='show results')
  240. parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
  241. parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
  242. parser.add_argument('--save-crop', action='store_true', help='save cropped prediction boxes')
  243. parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
  244. parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --classes 0, or --classes 0 2 3')
  245. parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
  246. parser.add_argument('--augment', action='store_true', help='augmented inference')
  247. parser.add_argument('--visualize', action='store_true', help='visualize features')
  248. parser.add_argument('--update', action='store_true', help='update all models')
  249. parser.add_argument('--project', default=ROOT / 'runs/detect', help='save results to project/name')
  250. parser.add_argument('--name', default='exp', help='save results to project/name')
  251. parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
  252. parser.add_argument('--line-thickness', default=3, type=int, help='bounding box thickness (pixels)')
  253. parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels')
  254. parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences')
  255. parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
  256. opt = parser.parse_args()
  257. opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1 # expand
  258. print_args(FILE.stem, opt)
  259. return opt
  260. def main(opt):
  261. check_requirements(exclude=('tensorboard', 'thop'))
  262. run(**vars(opt))
  263. if __name__ == "__main__":
  264. opt = parse_opt()
  265. main(opt)