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.

395 lines
19KB

  1. # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
  2. """
  3. Validate a trained YOLOv5 model accuracy on a custom dataset
  4. Usage:
  5. $ python path/to/val.py --weights yolov5s.pt --data coco128.yaml --img 640
  6. Usage - formats:
  7. $ python path/to/val.py --weights yolov5s.pt # PyTorch
  8. yolov5s.torchscript # TorchScript
  9. yolov5s.onnx # ONNX Runtime or OpenCV DNN with --dnn
  10. yolov5s.xml # OpenVINO
  11. yolov5s.engine # TensorRT
  12. yolov5s.mlmodel # CoreML (MacOS-only)
  13. yolov5s_saved_model # TensorFlow SavedModel
  14. yolov5s.pb # TensorFlow GraphDef
  15. yolov5s.tflite # TensorFlow Lite
  16. yolov5s_edgetpu.tflite # TensorFlow Edge TPU
  17. """
  18. import argparse
  19. import json
  20. import os
  21. import sys
  22. from pathlib import Path
  23. from threading import Thread
  24. import numpy as np
  25. import torch
  26. from tqdm import tqdm
  27. FILE = Path(__file__).resolve()
  28. ROOT = FILE.parents[0] # YOLOv5 root directory
  29. if str(ROOT) not in sys.path:
  30. sys.path.append(str(ROOT)) # add ROOT to PATH
  31. ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
  32. from models.common import DetectMultiBackend
  33. from utils.callbacks import Callbacks
  34. from utils.datasets import create_dataloader
  35. from utils.general import (LOGGER, box_iou, check_dataset, check_img_size, check_requirements, check_yaml,
  36. coco80_to_coco91_class, colorstr, increment_path, non_max_suppression, print_args,
  37. scale_coords, xywh2xyxy, xyxy2xywh)
  38. from utils.metrics import ConfusionMatrix, ap_per_class
  39. from utils.plots import output_to_target, plot_images, plot_val_study
  40. from utils.torch_utils import select_device, time_sync
  41. def save_one_txt(predn, save_conf, shape, file):
  42. # Save one txt result
  43. gn = torch.tensor(shape)[[1, 0, 1, 0]] # normalization gain whwh
  44. for *xyxy, conf, cls in predn.tolist():
  45. xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
  46. line = (cls, *xywh, conf) if save_conf else (cls, *xywh) # label format
  47. with open(file, 'a') as f:
  48. f.write(('%g ' * len(line)).rstrip() % line + '\n')
  49. def save_one_json(predn, jdict, path, class_map):
  50. # Save one JSON result {"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}
  51. image_id = int(path.stem) if path.stem.isnumeric() else path.stem
  52. box = xyxy2xywh(predn[:, :4]) # xywh
  53. box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner
  54. for p, b in zip(predn.tolist(), box.tolist()):
  55. jdict.append({
  56. 'image_id': image_id,
  57. 'category_id': class_map[int(p[5])],
  58. 'bbox': [round(x, 3) for x in b],
  59. 'score': round(p[4], 5)})
  60. def process_batch(detections, labels, iouv):
  61. """
  62. Return correct predictions matrix. Both sets of boxes are in (x1, y1, x2, y2) format.
  63. Arguments:
  64. detections (Array[N, 6]), x1, y1, x2, y2, conf, class
  65. labels (Array[M, 5]), class, x1, y1, x2, y2
  66. Returns:
  67. correct (Array[N, 10]), for 10 IoU levels
  68. """
  69. correct = torch.zeros(detections.shape[0], iouv.shape[0], dtype=torch.bool, device=iouv.device)
  70. iou = box_iou(labels[:, 1:], detections[:, :4])
  71. x = torch.where((iou >= iouv[0]) & (labels[:, 0:1] == detections[:, 5])) # IoU above threshold and classes match
  72. if x[0].shape[0]:
  73. matches = torch.cat((torch.stack(x, 1), iou[x[0], x[1]][:, None]), 1).cpu().numpy() # [label, detection, iou]
  74. if x[0].shape[0] > 1:
  75. matches = matches[matches[:, 2].argsort()[::-1]]
  76. matches = matches[np.unique(matches[:, 1], return_index=True)[1]]
  77. # matches = matches[matches[:, 2].argsort()[::-1]]
  78. matches = matches[np.unique(matches[:, 0], return_index=True)[1]]
  79. matches = torch.from_numpy(matches).to(iouv.device)
  80. correct[matches[:, 1].long()] = matches[:, 2:3] >= iouv
  81. return correct
  82. @torch.no_grad()
  83. def run(
  84. data,
  85. weights=None, # model.pt path(s)
  86. batch_size=32, # batch size
  87. imgsz=640, # inference size (pixels)
  88. conf_thres=0.001, # confidence threshold
  89. iou_thres=0.6, # NMS IoU threshold
  90. task='val', # train, val, test, speed or study
  91. device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
  92. workers=8, # max dataloader workers (per RANK in DDP mode)
  93. single_cls=False, # treat as single-class dataset
  94. augment=False, # augmented inference
  95. verbose=False, # verbose output
  96. save_txt=False, # save results to *.txt
  97. save_hybrid=False, # save label+prediction hybrid results to *.txt
  98. save_conf=False, # save confidences in --save-txt labels
  99. save_json=False, # save a COCO-JSON results file
  100. project=ROOT / 'runs/val', # save to project/name
  101. name='exp', # save to project/name
  102. exist_ok=False, # existing project/name ok, do not increment
  103. half=True, # use FP16 half-precision inference
  104. dnn=False, # use OpenCV DNN for ONNX inference
  105. model=None,
  106. dataloader=None,
  107. save_dir=Path(''),
  108. plots=True,
  109. callbacks=Callbacks(),
  110. compute_loss=None,
  111. ):
  112. # Initialize/load model and set device
  113. training = model is not None
  114. if training: # called by train.py
  115. device, pt, jit, engine = next(model.parameters()).device, True, False, False # get model device, PyTorch model
  116. half &= device.type != 'cpu' # half precision only supported on CUDA
  117. model.half() if half else model.float()
  118. else: # called directly
  119. device = select_device(device, batch_size=batch_size)
  120. # Directories
  121. save_dir = increment_path(Path(project) / name, exist_ok=exist_ok) # increment run
  122. (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
  123. # Load model
  124. model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half)
  125. stride, pt, jit, engine = model.stride, model.pt, model.jit, model.engine
  126. imgsz = check_img_size(imgsz, s=stride) # check image size
  127. half = model.fp16 # FP16 supported on limited backends with CUDA
  128. if engine:
  129. batch_size = model.batch_size
  130. else:
  131. device = model.device
  132. if not (pt or jit):
  133. batch_size = 1 # export.py models default to batch-size 1
  134. LOGGER.info(f'Forcing --batch-size 1 square inference (1,3,{imgsz},{imgsz}) for non-PyTorch models')
  135. # Data
  136. data = check_dataset(data) # check
  137. # Configure
  138. model.eval()
  139. cuda = device.type != 'cpu'
  140. is_coco = isinstance(data.get('val'), str) and data['val'].endswith('coco/val2017.txt') # COCO dataset
  141. nc = 1 if single_cls else int(data['nc']) # number of classes
  142. iouv = torch.linspace(0.5, 0.95, 10, device=device) # iou vector for mAP@0.5:0.95
  143. niou = iouv.numel()
  144. # Dataloader
  145. if not training:
  146. if pt and not single_cls: # check --weights are trained on --data
  147. ncm = model.model.yaml['nc']
  148. assert ncm == nc, f'{weights[0]} ({ncm} classes) trained on different --data than what you passed ({nc} ' \
  149. f'classes). Pass correct combination of --weights and --data that are trained together.'
  150. model.warmup(imgsz=(1 if pt else batch_size, 3, imgsz, imgsz)) # warmup
  151. pad = 0.0 if task in ('speed', 'benchmark') else 0.5
  152. rect = False if task == 'benchmark' else pt # square inference for benchmarks
  153. task = task if task in ('train', 'val', 'test') else 'val' # path to train/val/test images
  154. dataloader = create_dataloader(data[task],
  155. imgsz,
  156. batch_size,
  157. stride,
  158. single_cls,
  159. pad=pad,
  160. rect=rect,
  161. workers=workers,
  162. prefix=colorstr(f'{task}: '))[0]
  163. seen = 0
  164. confusion_matrix = ConfusionMatrix(nc=nc)
  165. names = {k: v for k, v in enumerate(model.names if hasattr(model, 'names') else model.module.names)}
  166. class_map = coco80_to_coco91_class() if is_coco else list(range(1000))
  167. s = ('%20s' + '%11s' * 6) % ('Class', 'Images', 'Labels', 'P', 'R', 'mAP@.5', 'mAP@.5:.95')
  168. dt, p, r, f1, mp, mr, map50, map = [0.0, 0.0, 0.0], 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
  169. loss = torch.zeros(3, device=device)
  170. jdict, stats, ap, ap_class = [], [], [], []
  171. pbar = tqdm(dataloader, desc=s, bar_format='{l_bar}{bar:10}{r_bar}{bar:-10b}') # progress bar
  172. for batch_i, (im, targets, paths, shapes) in enumerate(pbar):
  173. t1 = time_sync()
  174. if cuda:
  175. im = im.to(device, non_blocking=True)
  176. targets = targets.to(device)
  177. im = im.half() if half else im.float() # uint8 to fp16/32
  178. im /= 255 # 0 - 255 to 0.0 - 1.0
  179. nb, _, height, width = im.shape # batch size, channels, height, width
  180. t2 = time_sync()
  181. dt[0] += t2 - t1
  182. # Inference
  183. out, train_out = model(im) if training else model(im, augment=augment, val=True) # inference, loss outputs
  184. dt[1] += time_sync() - t2
  185. # Loss
  186. if compute_loss:
  187. loss += compute_loss([x.float() for x in train_out], targets)[1] # box, obj, cls
  188. # NMS
  189. targets[:, 2:] *= torch.tensor((width, height, width, height), device=device) # to pixels
  190. lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)] if save_hybrid else [] # for autolabelling
  191. t3 = time_sync()
  192. out = non_max_suppression(out, conf_thres, iou_thres, labels=lb, multi_label=True, agnostic=single_cls)
  193. dt[2] += time_sync() - t3
  194. # Metrics
  195. for si, pred in enumerate(out):
  196. labels = targets[targets[:, 0] == si, 1:]
  197. nl = len(labels)
  198. tcls = labels[:, 0].tolist() if nl else [] # target class
  199. path, shape = Path(paths[si]), shapes[si][0]
  200. seen += 1
  201. if len(pred) == 0:
  202. if nl:
  203. stats.append((torch.zeros(0, niou, dtype=torch.bool), torch.Tensor(), torch.Tensor(), tcls))
  204. continue
  205. # Predictions
  206. if single_cls:
  207. pred[:, 5] = 0
  208. predn = pred.clone()
  209. scale_coords(im[si].shape[1:], predn[:, :4], shape, shapes[si][1]) # native-space pred
  210. # Evaluate
  211. if nl:
  212. tbox = xywh2xyxy(labels[:, 1:5]) # target boxes
  213. scale_coords(im[si].shape[1:], tbox, shape, shapes[si][1]) # native-space labels
  214. labelsn = torch.cat((labels[:, 0:1], tbox), 1) # native-space labels
  215. correct = process_batch(predn, labelsn, iouv)
  216. if plots:
  217. confusion_matrix.process_batch(predn, labelsn)
  218. else:
  219. correct = torch.zeros(pred.shape[0], niou, dtype=torch.bool)
  220. stats.append((correct.cpu(), pred[:, 4].cpu(), pred[:, 5].cpu(), tcls)) # (correct, conf, pcls, tcls)
  221. # Save/log
  222. if save_txt:
  223. save_one_txt(predn, save_conf, shape, file=save_dir / 'labels' / (path.stem + '.txt'))
  224. if save_json:
  225. save_one_json(predn, jdict, path, class_map) # append to COCO-JSON dictionary
  226. callbacks.run('on_val_image_end', pred, predn, path, names, im[si])
  227. # Plot images
  228. if plots and batch_i < 3:
  229. f = save_dir / f'val_batch{batch_i}_labels.jpg' # labels
  230. Thread(target=plot_images, args=(im, targets, paths, f, names), daemon=True).start()
  231. f = save_dir / f'val_batch{batch_i}_pred.jpg' # predictions
  232. Thread(target=plot_images, args=(im, output_to_target(out), paths, f, names), daemon=True).start()
  233. # Compute metrics
  234. stats = [np.concatenate(x, 0) for x in zip(*stats)] # to numpy
  235. if len(stats) and stats[0].any():
  236. tp, fp, p, r, f1, ap, ap_class = ap_per_class(*stats, plot=plots, save_dir=save_dir, names=names)
  237. ap50, ap = ap[:, 0], ap.mean(1) # AP@0.5, AP@0.5:0.95
  238. mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean()
  239. nt = np.bincount(stats[3].astype(np.int64), minlength=nc) # number of targets per class
  240. else:
  241. nt = torch.zeros(1)
  242. # Print results
  243. pf = '%20s' + '%11i' * 2 + '%11.3g' * 4 # print format
  244. LOGGER.info(pf % ('all', seen, nt.sum(), mp, mr, map50, map))
  245. # Print results per class
  246. if (verbose or (nc < 50 and not training)) and nc > 1 and len(stats):
  247. for i, c in enumerate(ap_class):
  248. LOGGER.info(pf % (names[c], seen, nt[c], p[i], r[i], ap50[i], ap[i]))
  249. # Print speeds
  250. t = tuple(x / seen * 1E3 for x in dt) # speeds per image
  251. if not training:
  252. shape = (batch_size, 3, imgsz, imgsz)
  253. LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {shape}' % t)
  254. # Plots
  255. if plots:
  256. confusion_matrix.plot(save_dir=save_dir, names=list(names.values()))
  257. callbacks.run('on_val_end')
  258. # Save JSON
  259. if save_json and len(jdict):
  260. w = Path(weights[0] if isinstance(weights, list) else weights).stem if weights is not None else '' # weights
  261. anno_json = str(Path(data.get('path', '../coco')) / 'annotations/instances_val2017.json') # annotations json
  262. pred_json = str(save_dir / f"{w}_predictions.json") # predictions json
  263. LOGGER.info(f'\nEvaluating pycocotools mAP... saving {pred_json}...')
  264. with open(pred_json, 'w') as f:
  265. json.dump(jdict, f)
  266. try: # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb
  267. check_requirements(['pycocotools'])
  268. from pycocotools.coco import COCO
  269. from pycocotools.cocoeval import COCOeval
  270. anno = COCO(anno_json) # init annotations api
  271. pred = anno.loadRes(pred_json) # init predictions api
  272. eval = COCOeval(anno, pred, 'bbox')
  273. if is_coco:
  274. eval.params.imgIds = [int(Path(x).stem) for x in dataloader.dataset.im_files] # image IDs to evaluate
  275. eval.evaluate()
  276. eval.accumulate()
  277. eval.summarize()
  278. map, map50 = eval.stats[:2] # update results (mAP@0.5:0.95, mAP@0.5)
  279. except Exception as e:
  280. LOGGER.info(f'pycocotools unable to run: {e}')
  281. # Return results
  282. model.float() # for training
  283. if not training:
  284. s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
  285. LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
  286. maps = np.zeros(nc) + map
  287. for i, c in enumerate(ap_class):
  288. maps[c] = ap[i]
  289. return (mp, mr, map50, map, *(loss.cpu() / len(dataloader)).tolist()), maps, t
  290. def parse_opt():
  291. parser = argparse.ArgumentParser()
  292. parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
  293. parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model.pt path(s)')
  294. parser.add_argument('--batch-size', type=int, default=32, help='batch size')
  295. parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='inference size (pixels)')
  296. parser.add_argument('--conf-thres', type=float, default=0.001, help='confidence threshold')
  297. parser.add_argument('--iou-thres', type=float, default=0.6, help='NMS IoU threshold')
  298. parser.add_argument('--task', default='val', help='train, val, test, speed or study')
  299. parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
  300. parser.add_argument('--workers', type=int, default=8, help='max dataloader workers (per RANK in DDP mode)')
  301. parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset')
  302. parser.add_argument('--augment', action='store_true', help='augmented inference')
  303. parser.add_argument('--verbose', action='store_true', help='report mAP by class')
  304. parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
  305. parser.add_argument('--save-hybrid', action='store_true', help='save label+prediction hybrid results to *.txt')
  306. parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
  307. parser.add_argument('--save-json', action='store_true', help='save a COCO-JSON results file')
  308. parser.add_argument('--project', default=ROOT / 'runs/val', help='save to project/name')
  309. parser.add_argument('--name', default='exp', help='save to project/name')
  310. parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
  311. parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
  312. parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')
  313. opt = parser.parse_args()
  314. opt.data = check_yaml(opt.data) # check YAML
  315. opt.save_json |= opt.data.endswith('coco.yaml')
  316. opt.save_txt |= opt.save_hybrid
  317. print_args(vars(opt))
  318. return opt
  319. def main(opt):
  320. check_requirements(requirements=ROOT / 'requirements.txt', exclude=('tensorboard', 'thop'))
  321. if opt.task in ('train', 'val', 'test'): # run normally
  322. if opt.conf_thres > 0.001: # https://github.com/ultralytics/yolov5/issues/1466
  323. LOGGER.info(f'WARNING: confidence threshold {opt.conf_thres} >> 0.001 will produce invalid mAP values.')
  324. run(**vars(opt))
  325. else:
  326. weights = opt.weights if isinstance(opt.weights, list) else [opt.weights]
  327. opt.half = True # FP16 for fastest results
  328. if opt.task == 'speed': # speed benchmarks
  329. # python val.py --task speed --data coco.yaml --batch 1 --weights yolov5n.pt yolov5s.pt...
  330. opt.conf_thres, opt.iou_thres, opt.save_json = 0.25, 0.45, False
  331. for opt.weights in weights:
  332. run(**vars(opt), plots=False)
  333. elif opt.task == 'study': # speed vs mAP benchmarks
  334. # python val.py --task study --data coco.yaml --iou 0.7 --weights yolov5n.pt yolov5s.pt...
  335. for opt.weights in weights:
  336. f = f'study_{Path(opt.data).stem}_{Path(opt.weights).stem}.txt' # filename to save to
  337. x, y = list(range(256, 1536 + 128, 128)), [] # x axis (image sizes), y axis
  338. for opt.imgsz in x: # img-size
  339. LOGGER.info(f'\nRunning {f} --imgsz {opt.imgsz}...')
  340. r, _, t = run(**vars(opt), plots=False)
  341. y.append(r + t) # results and times
  342. np.savetxt(f, y, fmt='%10.4g') # save
  343. os.system('zip -r study.zip study_*.txt')
  344. plot_val_study(x=x) # plot
  345. if __name__ == "__main__":
  346. opt = parse_opt()
  347. main(opt)