No puede seleccionar más de 25 temas Los temas deben comenzar con una letra o número, pueden incluir guiones ('-') y pueden tener hasta 35 caracteres de largo.

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