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.

337 lines
16KB

  1. import argparse
  2. import json
  3. import os
  4. from pathlib import Path
  5. from threading import Thread
  6. import numpy as np
  7. import torch
  8. import yaml
  9. from tqdm import tqdm
  10. from models.experimental import attempt_load
  11. from utils.datasets import create_dataloader
  12. from utils.general import coco80_to_coco91_class, check_dataset, check_file, check_img_size, check_requirements, \
  13. box_iou, non_max_suppression, scale_coords, xyxy2xywh, xywh2xyxy, set_logging, increment_path, colorstr
  14. from utils.metrics import ap_per_class, ConfusionMatrix
  15. from utils.plots import plot_images, output_to_target, plot_study_txt
  16. from utils.torch_utils import select_device, time_synchronized
  17. def test(data,
  18. weights=None,
  19. batch_size=32,
  20. imgsz=640,
  21. conf_thres=0.001,
  22. iou_thres=0.6, # for NMS
  23. save_json=False,
  24. single_cls=False,
  25. augment=False,
  26. verbose=False,
  27. model=None,
  28. dataloader=None,
  29. save_dir=Path(''), # for saving images
  30. save_txt=False, # for auto-labelling
  31. save_hybrid=False, # for hybrid auto-labelling
  32. save_conf=False, # save auto-label confidences
  33. plots=True,
  34. log_imgs=0, # number of logged images
  35. compute_loss=None):
  36. # Initialize/load model and set device
  37. training = model is not None
  38. if training: # called by train.py
  39. device = next(model.parameters()).device # get model device
  40. else: # called directly
  41. set_logging()
  42. device = select_device(opt.device, batch_size=batch_size)
  43. # Directories
  44. save_dir = Path(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)) # increment run
  45. (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
  46. # Load model
  47. model = attempt_load(weights, map_location=device) # load FP32 model
  48. imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
  49. # Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99
  50. # if device.type != 'cpu' and torch.cuda.device_count() > 1:
  51. # model = nn.DataParallel(model)
  52. # Half
  53. half = device.type != 'cpu' # half precision only supported on CUDA
  54. if half:
  55. model.half()
  56. # Configure
  57. model.eval()
  58. is_coco = data.endswith('coco.yaml') # is COCO dataset
  59. with open(data) as f:
  60. data = yaml.load(f, Loader=yaml.FullLoader) # model dict
  61. check_dataset(data) # check
  62. nc = 1 if single_cls else int(data['nc']) # number of classes
  63. iouv = torch.linspace(0.5, 0.95, 10).to(device) # iou vector for mAP@0.5:0.95
  64. niou = iouv.numel()
  65. # Logging
  66. log_imgs, wandb = min(log_imgs, 100), None # ceil
  67. try:
  68. import wandb # Weights & Biases
  69. except ImportError:
  70. log_imgs = 0
  71. # Dataloader
  72. if not training:
  73. img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
  74. _ = model(img.half() if half else img) if device.type != 'cpu' else None # run once
  75. path = data['test'] if opt.task == 'test' else data['val'] # path to val/test images
  76. dataloader = create_dataloader(path, imgsz, batch_size, model.stride.max(), opt, pad=0.5, rect=True,
  77. prefix=colorstr('test: ' if opt.task == 'test' else 'val: '))[0]
  78. seen = 0
  79. confusion_matrix = ConfusionMatrix(nc=nc)
  80. names = {k: v for k, v in enumerate(model.names if hasattr(model, 'names') else model.module.names)}
  81. coco91class = coco80_to_coco91_class()
  82. s = ('%20s' + '%12s' * 6) % ('Class', 'Images', 'Targets', 'P', 'R', 'mAP@.5', 'mAP@.5:.95')
  83. p, r, f1, mp, mr, map50, map, t0, t1 = 0., 0., 0., 0., 0., 0., 0., 0., 0.
  84. loss = torch.zeros(3, device=device)
  85. jdict, stats, ap, ap_class, wandb_images = [], [], [], [], []
  86. for batch_i, (img, targets, paths, shapes) in enumerate(tqdm(dataloader, desc=s)):
  87. img = img.to(device, non_blocking=True)
  88. img = img.half() if half else img.float() # uint8 to fp16/32
  89. img /= 255.0 # 0 - 255 to 0.0 - 1.0
  90. targets = targets.to(device)
  91. nb, _, height, width = img.shape # batch size, channels, height, width
  92. with torch.no_grad():
  93. # Run model
  94. t = time_synchronized()
  95. inf_out, train_out = model(img, augment=augment) # inference and training outputs
  96. t0 += time_synchronized() - t
  97. # Compute loss
  98. if compute_loss:
  99. loss += compute_loss([x.float() for x in train_out], targets)[1][:3] # box, obj, cls
  100. # Run NMS
  101. targets[:, 2:] *= torch.Tensor([width, height, width, height]).to(device) # to pixels
  102. lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)] if save_hybrid else [] # for autolabelling
  103. t = time_synchronized()
  104. output = non_max_suppression(inf_out, conf_thres=conf_thres, iou_thres=iou_thres, labels=lb)
  105. t1 += time_synchronized() - t
  106. # Statistics per image
  107. for si, pred in enumerate(output):
  108. labels = targets[targets[:, 0] == si, 1:]
  109. nl = len(labels)
  110. tcls = labels[:, 0].tolist() if nl else [] # target class
  111. path = Path(paths[si])
  112. seen += 1
  113. if len(pred) == 0:
  114. if nl:
  115. stats.append((torch.zeros(0, niou, dtype=torch.bool), torch.Tensor(), torch.Tensor(), tcls))
  116. continue
  117. # Predictions
  118. predn = pred.clone()
  119. scale_coords(img[si].shape[1:], predn[:, :4], shapes[si][0], shapes[si][1]) # native-space pred
  120. # Append to text file
  121. if save_txt:
  122. gn = torch.tensor(shapes[si][0])[[1, 0, 1, 0]] # normalization gain whwh
  123. for *xyxy, conf, cls in predn.tolist():
  124. xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
  125. line = (cls, *xywh, conf) if save_conf else (cls, *xywh) # label format
  126. with open(save_dir / 'labels' / (path.stem + '.txt'), 'a') as f:
  127. f.write(('%g ' * len(line)).rstrip() % line + '\n')
  128. # W&B logging
  129. if plots and len(wandb_images) < log_imgs:
  130. box_data = [{"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
  131. "class_id": int(cls),
  132. "box_caption": "%s %.3f" % (names[cls], conf),
  133. "scores": {"class_score": conf},
  134. "domain": "pixel"} for *xyxy, conf, cls in pred.tolist()]
  135. boxes = {"predictions": {"box_data": box_data, "class_labels": names}} # inference-space
  136. wandb_images.append(wandb.Image(img[si], boxes=boxes, caption=path.name))
  137. # Append to pycocotools JSON dictionary
  138. if save_json:
  139. # [{"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}, ...
  140. image_id = int(path.stem) if path.stem.isnumeric() else path.stem
  141. box = xyxy2xywh(predn[:, :4]) # xywh
  142. box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner
  143. for p, b in zip(pred.tolist(), box.tolist()):
  144. jdict.append({'image_id': image_id,
  145. 'category_id': coco91class[int(p[5])] if is_coco else int(p[5]),
  146. 'bbox': [round(x, 3) for x in b],
  147. 'score': round(p[4], 5)})
  148. # Assign all predictions as incorrect
  149. correct = torch.zeros(pred.shape[0], niou, dtype=torch.bool, device=device)
  150. if nl:
  151. detected = [] # target indices
  152. tcls_tensor = labels[:, 0]
  153. # target boxes
  154. tbox = xywh2xyxy(labels[:, 1:5])
  155. scale_coords(img[si].shape[1:], tbox, shapes[si][0], shapes[si][1]) # native-space labels
  156. if plots:
  157. confusion_matrix.process_batch(pred, torch.cat((labels[:, 0:1], tbox), 1))
  158. # Per target class
  159. for cls in torch.unique(tcls_tensor):
  160. ti = (cls == tcls_tensor).nonzero(as_tuple=False).view(-1) # prediction indices
  161. pi = (cls == pred[:, 5]).nonzero(as_tuple=False).view(-1) # target indices
  162. # Search for detections
  163. if pi.shape[0]:
  164. # Prediction to target ious
  165. ious, i = box_iou(predn[pi, :4], tbox[ti]).max(1) # best ious, indices
  166. # Append detections
  167. detected_set = set()
  168. for j in (ious > iouv[0]).nonzero(as_tuple=False):
  169. d = ti[i[j]] # detected target
  170. if d.item() not in detected_set:
  171. detected_set.add(d.item())
  172. detected.append(d)
  173. correct[pi[j]] = ious[j] > iouv # iou_thres is 1xn
  174. if len(detected) == nl: # all targets already located in image
  175. break
  176. # Append statistics (correct, conf, pcls, tcls)
  177. stats.append((correct.cpu(), pred[:, 4].cpu(), pred[:, 5].cpu(), tcls))
  178. # Plot images
  179. if plots and batch_i < 3:
  180. f = save_dir / f'test_batch{batch_i}_labels.jpg' # labels
  181. Thread(target=plot_images, args=(img, targets, paths, f, names), daemon=True).start()
  182. f = save_dir / f'test_batch{batch_i}_pred.jpg' # predictions
  183. Thread(target=plot_images, args=(img, output_to_target(output), paths, f, names), daemon=True).start()
  184. # Compute statistics
  185. stats = [np.concatenate(x, 0) for x in zip(*stats)] # to numpy
  186. if len(stats) and stats[0].any():
  187. p, r, ap, f1, ap_class = ap_per_class(*stats, plot=plots, save_dir=save_dir, names=names)
  188. p, r, ap50, ap = p[:, 0], r[:, 0], ap[:, 0], ap.mean(1) # [P, R, AP@0.5, AP@0.5:0.95]
  189. mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean()
  190. nt = np.bincount(stats[3].astype(np.int64), minlength=nc) # number of targets per class
  191. else:
  192. nt = torch.zeros(1)
  193. # Print results
  194. pf = '%20s' + '%12.3g' * 6 # print format
  195. print(pf % ('all', seen, nt.sum(), mp, mr, map50, map))
  196. # Print results per class
  197. if (verbose or (nc <= 20 and not training)) and nc > 1 and len(stats):
  198. for i, c in enumerate(ap_class):
  199. print(pf % (names[c], seen, nt[c], p[i], r[i], ap50[i], ap[i]))
  200. # Print speeds
  201. t = tuple(x / seen * 1E3 for x in (t0, t1, t0 + t1)) + (imgsz, imgsz, batch_size) # tuple
  202. if not training:
  203. print('Speed: %.1f/%.1f/%.1f ms inference/NMS/total per %gx%g image at batch-size %g' % t)
  204. # Plots
  205. if plots:
  206. confusion_matrix.plot(save_dir=save_dir, names=list(names.values()))
  207. if wandb and wandb.run:
  208. wandb.log({"Images": wandb_images})
  209. wandb.log({"Validation": [wandb.Image(str(f), caption=f.name) for f in sorted(save_dir.glob('test*.jpg'))]})
  210. # Save JSON
  211. if save_json and len(jdict):
  212. w = Path(weights[0] if isinstance(weights, list) else weights).stem if weights is not None else '' # weights
  213. anno_json = '../coco/annotations/instances_val2017.json' # annotations json
  214. pred_json = str(save_dir / f"{w}_predictions.json") # predictions json
  215. print('\nEvaluating pycocotools mAP... saving %s...' % pred_json)
  216. with open(pred_json, 'w') as f:
  217. json.dump(jdict, f)
  218. try: # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb
  219. from pycocotools.coco import COCO
  220. from pycocotools.cocoeval import COCOeval
  221. anno = COCO(anno_json) # init annotations api
  222. pred = anno.loadRes(pred_json) # init predictions api
  223. eval = COCOeval(anno, pred, 'bbox')
  224. if is_coco:
  225. eval.params.imgIds = [int(Path(x).stem) for x in dataloader.dataset.img_files] # image IDs to evaluate
  226. eval.evaluate()
  227. eval.accumulate()
  228. eval.summarize()
  229. map, map50 = eval.stats[:2] # update results (mAP@0.5:0.95, mAP@0.5)
  230. except Exception as e:
  231. print(f'pycocotools unable to run: {e}')
  232. # Return results
  233. if not training:
  234. s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
  235. print(f"Results saved to {save_dir}{s}")
  236. model.float() # for training
  237. maps = np.zeros(nc) + map
  238. for i, c in enumerate(ap_class):
  239. maps[c] = ap[i]
  240. return (mp, mr, map50, map, *(loss.cpu() / len(dataloader)).tolist()), maps, t
  241. if __name__ == '__main__':
  242. parser = argparse.ArgumentParser(prog='test.py')
  243. parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)')
  244. parser.add_argument('--data', type=str, default='data/coco128.yaml', help='*.data path')
  245. parser.add_argument('--batch-size', type=int, default=32, help='size of each image batch')
  246. parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
  247. parser.add_argument('--conf-thres', type=float, default=0.001, help='object confidence threshold')
  248. parser.add_argument('--iou-thres', type=float, default=0.6, help='IOU threshold for NMS')
  249. parser.add_argument('--task', default='val', help="'val', 'test', 'study'")
  250. parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
  251. parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset')
  252. parser.add_argument('--augment', action='store_true', help='augmented inference')
  253. parser.add_argument('--verbose', action='store_true', help='report mAP by class')
  254. parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
  255. parser.add_argument('--save-hybrid', action='store_true', help='save label+prediction hybrid results to *.txt')
  256. parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
  257. parser.add_argument('--save-json', action='store_true', help='save a cocoapi-compatible JSON results file')
  258. parser.add_argument('--project', default='runs/test', help='save to project/name')
  259. parser.add_argument('--name', default='exp', help='save to project/name')
  260. parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
  261. opt = parser.parse_args()
  262. opt.save_json |= opt.data.endswith('coco.yaml')
  263. opt.data = check_file(opt.data) # check file
  264. print(opt)
  265. check_requirements()
  266. if opt.task in ['val', 'test']: # run normally
  267. test(opt.data,
  268. opt.weights,
  269. opt.batch_size,
  270. opt.img_size,
  271. opt.conf_thres,
  272. opt.iou_thres,
  273. opt.save_json,
  274. opt.single_cls,
  275. opt.augment,
  276. opt.verbose,
  277. save_txt=opt.save_txt | opt.save_hybrid,
  278. save_hybrid=opt.save_hybrid,
  279. save_conf=opt.save_conf,
  280. )
  281. elif opt.task == 'study': # run over a range of settings and save/plot
  282. for weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']:
  283. f = 'study_%s_%s.txt' % (Path(opt.data).stem, Path(weights).stem) # filename to save to
  284. x = list(range(320, 800, 64)) # x axis
  285. y = [] # y axis
  286. for i in x: # img-size
  287. print('\nRunning %s point %s...' % (f, i))
  288. r, _, t = test(opt.data, weights, opt.batch_size, i, opt.conf_thres, opt.iou_thres, opt.save_json,
  289. plots=False)
  290. y.append(r + t) # results and times
  291. np.savetxt(f, y, fmt='%10.4g') # save
  292. os.system('zip -r study.zip study_*.txt')
  293. plot_study_txt(f, x) # plot