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.

383 lines
18KB

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