Du kannst nicht mehr als 25 Themen auswählen Themen müssen entweder mit einem Buchstaben oder einer Ziffer beginnen. Sie können Bindestriche („-“) enthalten und bis zu 35 Zeichen lang sein.

360 Zeilen
17KB

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