Nelze vybrat více než 25 témat Téma musí začínat písmenem nebo číslem, může obsahovat pomlčky („-“) a může být dlouhé až 35 znaků.

384 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({'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 if pt else batch_size, 3, imgsz, imgsz), half=half) # warmup
  146. pad = 0.0 if task in ('speed', 'benchmark') else 0.5
  147. rect = False if task == 'benchmark' else pt # square inference for benchmarks
  148. task = task if task in ('train', 'val', 'test') else 'val' # path to train/val/test images
  149. dataloader = create_dataloader(data[task], imgsz, batch_size, stride, single_cls, pad=pad, rect=rect,
  150. workers=workers, prefix=colorstr(f'{task}: '))[0]
  151. seen = 0
  152. confusion_matrix = ConfusionMatrix(nc=nc)
  153. names = {k: v for k, v in enumerate(model.names if hasattr(model, 'names') else model.module.names)}
  154. class_map = coco80_to_coco91_class() if is_coco else list(range(1000))
  155. s = ('%20s' + '%11s' * 6) % ('Class', 'Images', 'Labels', 'P', 'R', 'mAP@.5', 'mAP@.5:.95')
  156. 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
  157. loss = torch.zeros(3, device=device)
  158. jdict, stats, ap, ap_class = [], [], [], []
  159. pbar = tqdm(dataloader, desc=s, bar_format='{l_bar}{bar:10}{r_bar}{bar:-10b}') # progress bar
  160. for batch_i, (im, targets, paths, shapes) in enumerate(pbar):
  161. t1 = time_sync()
  162. if pt or jit or engine:
  163. im = im.to(device, non_blocking=True)
  164. targets = targets.to(device)
  165. im = im.half() if half else im.float() # uint8 to fp16/32
  166. im /= 255 # 0 - 255 to 0.0 - 1.0
  167. nb, _, height, width = im.shape # batch size, channels, height, width
  168. t2 = time_sync()
  169. dt[0] += t2 - t1
  170. # Inference
  171. out, train_out = model(im) if training else model(im, augment=augment, val=True) # inference, loss outputs
  172. dt[1] += time_sync() - t2
  173. # Loss
  174. if compute_loss:
  175. loss += compute_loss([x.float() for x in train_out], targets)[1] # box, obj, cls
  176. # NMS
  177. targets[:, 2:] *= torch.Tensor([width, height, width, height]).to(device) # to pixels
  178. lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)] if save_hybrid else [] # for autolabelling
  179. t3 = time_sync()
  180. out = non_max_suppression(out, conf_thres, iou_thres, labels=lb, multi_label=True, agnostic=single_cls)
  181. dt[2] += time_sync() - t3
  182. # Metrics
  183. for si, pred in enumerate(out):
  184. labels = targets[targets[:, 0] == si, 1:]
  185. nl = len(labels)
  186. tcls = labels[:, 0].tolist() if nl else [] # target class
  187. path, shape = Path(paths[si]), shapes[si][0]
  188. seen += 1
  189. if len(pred) == 0:
  190. if nl:
  191. stats.append((torch.zeros(0, niou, dtype=torch.bool), torch.Tensor(), torch.Tensor(), tcls))
  192. continue
  193. # Predictions
  194. if single_cls:
  195. pred[:, 5] = 0
  196. predn = pred.clone()
  197. scale_coords(im[si].shape[1:], predn[:, :4], shape, shapes[si][1]) # native-space pred
  198. # Evaluate
  199. if nl:
  200. tbox = xywh2xyxy(labels[:, 1:5]) # target boxes
  201. scale_coords(im[si].shape[1:], tbox, shape, shapes[si][1]) # native-space labels
  202. labelsn = torch.cat((labels[:, 0:1], tbox), 1) # native-space labels
  203. correct = process_batch(predn, labelsn, iouv)
  204. if plots:
  205. confusion_matrix.process_batch(predn, labelsn)
  206. else:
  207. correct = torch.zeros(pred.shape[0], niou, dtype=torch.bool)
  208. stats.append((correct.cpu(), pred[:, 4].cpu(), pred[:, 5].cpu(), tcls)) # (correct, conf, pcls, tcls)
  209. # Save/log
  210. if save_txt:
  211. save_one_txt(predn, save_conf, shape, file=save_dir / 'labels' / (path.stem + '.txt'))
  212. if save_json:
  213. save_one_json(predn, jdict, path, class_map) # append to COCO-JSON dictionary
  214. callbacks.run('on_val_image_end', pred, predn, path, names, im[si])
  215. # Plot images
  216. if plots and batch_i < 3:
  217. f = save_dir / f'val_batch{batch_i}_labels.jpg' # labels
  218. Thread(target=plot_images, args=(im, targets, paths, f, names), daemon=True).start()
  219. f = save_dir / f'val_batch{batch_i}_pred.jpg' # predictions
  220. Thread(target=plot_images, args=(im, output_to_target(out), paths, f, names), daemon=True).start()
  221. # Compute metrics
  222. stats = [np.concatenate(x, 0) for x in zip(*stats)] # to numpy
  223. if len(stats) and stats[0].any():
  224. tp, fp, p, r, f1, ap, ap_class = ap_per_class(*stats, plot=plots, save_dir=save_dir, names=names)
  225. ap50, ap = ap[:, 0], ap.mean(1) # AP@0.5, AP@0.5:0.95
  226. mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean()
  227. nt = np.bincount(stats[3].astype(np.int64), minlength=nc) # number of targets per class
  228. else:
  229. nt = torch.zeros(1)
  230. # Print results
  231. pf = '%20s' + '%11i' * 2 + '%11.3g' * 4 # print format
  232. LOGGER.info(pf % ('all', seen, nt.sum(), mp, mr, map50, map))
  233. # Print results per class
  234. if (verbose or (nc < 50 and not training)) and nc > 1 and len(stats):
  235. for i, c in enumerate(ap_class):
  236. LOGGER.info(pf % (names[c], seen, nt[c], p[i], r[i], ap50[i], ap[i]))
  237. # Print speeds
  238. t = tuple(x / seen * 1E3 for x in dt) # speeds per image
  239. if not training:
  240. shape = (batch_size, 3, imgsz, imgsz)
  241. LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {shape}' % t)
  242. # Plots
  243. if plots:
  244. confusion_matrix.plot(save_dir=save_dir, names=list(names.values()))
  245. callbacks.run('on_val_end')
  246. # Save JSON
  247. if save_json and len(jdict):
  248. w = Path(weights[0] if isinstance(weights, list) else weights).stem if weights is not None else '' # weights
  249. anno_json = str(Path(data.get('path', '../coco')) / 'annotations/instances_val2017.json') # annotations json
  250. pred_json = str(save_dir / f"{w}_predictions.json") # predictions json
  251. LOGGER.info(f'\nEvaluating pycocotools mAP... saving {pred_json}...')
  252. with open(pred_json, 'w') as f:
  253. json.dump(jdict, f)
  254. try: # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb
  255. check_requirements(['pycocotools'])
  256. from pycocotools.coco import COCO
  257. from pycocotools.cocoeval import COCOeval
  258. anno = COCO(anno_json) # init annotations api
  259. pred = anno.loadRes(pred_json) # init predictions api
  260. eval = COCOeval(anno, pred, 'bbox')
  261. if is_coco:
  262. eval.params.imgIds = [int(Path(x).stem) for x in dataloader.dataset.img_files] # image IDs to evaluate
  263. eval.evaluate()
  264. eval.accumulate()
  265. eval.summarize()
  266. map, map50 = eval.stats[:2] # update results (mAP@0.5:0.95, mAP@0.5)
  267. except Exception as e:
  268. LOGGER.info(f'pycocotools unable to run: {e}')
  269. # Return results
  270. model.float() # for training
  271. if not training:
  272. s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
  273. LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
  274. maps = np.zeros(nc) + map
  275. for i, c in enumerate(ap_class):
  276. maps[c] = ap[i]
  277. return (mp, mr, map50, map, *(loss.cpu() / len(dataloader)).tolist()), maps, t
  278. def parse_opt():
  279. parser = argparse.ArgumentParser()
  280. parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
  281. parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model.pt path(s)')
  282. parser.add_argument('--batch-size', type=int, default=32, help='batch size')
  283. parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='inference size (pixels)')
  284. parser.add_argument('--conf-thres', type=float, default=0.001, help='confidence threshold')
  285. parser.add_argument('--iou-thres', type=float, default=0.6, help='NMS IoU threshold')
  286. parser.add_argument('--task', default='val', help='train, val, test, speed or study')
  287. parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
  288. parser.add_argument('--workers', type=int, default=8, help='max dataloader workers (per RANK in DDP mode)')
  289. parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset')
  290. parser.add_argument('--augment', action='store_true', help='augmented inference')
  291. parser.add_argument('--verbose', action='store_true', help='report mAP by class')
  292. parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
  293. parser.add_argument('--save-hybrid', action='store_true', help='save label+prediction hybrid results to *.txt')
  294. parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
  295. parser.add_argument('--save-json', action='store_true', help='save a COCO-JSON results file')
  296. parser.add_argument('--project', default=ROOT / 'runs/val', help='save to project/name')
  297. parser.add_argument('--name', default='exp', help='save to project/name')
  298. parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
  299. parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
  300. parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')
  301. opt = parser.parse_args()
  302. opt.data = check_yaml(opt.data) # check YAML
  303. opt.save_json |= opt.data.endswith('coco.yaml')
  304. opt.save_txt |= opt.save_hybrid
  305. print_args(FILE.stem, opt)
  306. return opt
  307. def main(opt):
  308. check_requirements(requirements=ROOT / 'requirements.txt', exclude=('tensorboard', 'thop'))
  309. if opt.task in ('train', 'val', 'test'): # run normally
  310. if opt.conf_thres > 0.001: # https://github.com/ultralytics/yolov5/issues/1466
  311. LOGGER.info(f'WARNING: confidence threshold {opt.conf_thres} >> 0.001 will produce invalid mAP values.')
  312. run(**vars(opt))
  313. else:
  314. weights = opt.weights if isinstance(opt.weights, list) else [opt.weights]
  315. opt.half = True # FP16 for fastest results
  316. if opt.task == 'speed': # speed benchmarks
  317. # python val.py --task speed --data coco.yaml --batch 1 --weights yolov5n.pt yolov5s.pt...
  318. opt.conf_thres, opt.iou_thres, opt.save_json = 0.25, 0.45, False
  319. for opt.weights in weights:
  320. run(**vars(opt), plots=False)
  321. elif opt.task == 'study': # speed vs mAP benchmarks
  322. # python val.py --task study --data coco.yaml --iou 0.7 --weights yolov5n.pt yolov5s.pt...
  323. for opt.weights in weights:
  324. f = f'study_{Path(opt.data).stem}_{Path(opt.weights).stem}.txt' # filename to save to
  325. x, y = list(range(256, 1536 + 128, 128)), [] # x axis (image sizes), y axis
  326. for opt.imgsz in x: # img-size
  327. LOGGER.info(f'\nRunning {f} --imgsz {opt.imgsz}...')
  328. r, _, t = run(**vars(opt), plots=False)
  329. y.append(r + t) # results and times
  330. np.savetxt(f, y, fmt='%10.4g') # save
  331. os.system('zip -r study.zip study_*.txt')
  332. plot_val_study(x=x) # plot
  333. if __name__ == "__main__":
  334. opt = parse_opt()
  335. main(opt)