Du kan inte välja fler än 25 ämnen Ämnen måste starta med en bokstav eller siffra, kan innehålla bindestreck ('-') och vara max 35 tecken långa.

367 lines
18KB

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