Ви не можете вибрати більше 25 тем Теми мають розпочинатися з літери або цифри, можуть містити дефіси (-) і не повинні перевищувати 35 символів.

334 lines
16KB

  1. import argparse
  2. import glob
  3. import json
  4. import os
  5. from pathlib import Path
  6. from threading import Thread
  7. import numpy as np
  8. import torch
  9. import yaml
  10. from tqdm import tqdm
  11. from models.experimental import attempt_load
  12. from utils.datasets import create_dataloader
  13. from utils.general import coco80_to_coco91_class, check_dataset, check_file, check_img_size, box_iou, \
  14. non_max_suppression, scale_coords, xyxy2xywh, xywh2xyxy, set_logging, increment_path
  15. from utils.loss import compute_loss
  16. from utils.metrics import ap_per_class, ConfusionMatrix
  17. from utils.plots import plot_images, output_to_target, plot_study_txt
  18. from utils.torch_utils import select_device, time_synchronized
  19. def test(data,
  20. weights=None,
  21. batch_size=32,
  22. imgsz=640,
  23. conf_thres=0.001,
  24. iou_thres=0.6, # for NMS
  25. save_json=False,
  26. single_cls=False,
  27. augment=False,
  28. verbose=False,
  29. model=None,
  30. dataloader=None,
  31. save_dir=Path(''), # for saving images
  32. save_txt=False, # for auto-labelling
  33. save_conf=False,
  34. plots=True,
  35. log_imgs=0): # number of logged images
  36. # Initialize/load model and set device
  37. training = model is not None
  38. if training: # called by train.py
  39. device = next(model.parameters()).device # get model device
  40. else: # called directly
  41. set_logging()
  42. device = select_device(opt.device, batch_size=batch_size)
  43. save_txt = opt.save_txt # save *.txt labels
  44. # Directories
  45. save_dir = Path(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)) # increment run
  46. (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
  47. # Load model
  48. model = attempt_load(weights, map_location=device) # load FP32 model
  49. imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
  50. # Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99
  51. # if device.type != 'cpu' and torch.cuda.device_count() > 1:
  52. # model = nn.DataParallel(model)
  53. # Half
  54. half = device.type != 'cpu' # half precision only supported on CUDA
  55. if half:
  56. model.half()
  57. # Configure
  58. model.eval()
  59. is_coco = data.endswith('coco.yaml') # is COCO dataset
  60. with open(data) as f:
  61. data = yaml.load(f, Loader=yaml.FullLoader) # model dict
  62. check_dataset(data) # check
  63. nc = 1 if single_cls else int(data['nc']) # number of classes
  64. iouv = torch.linspace(0.5, 0.95, 10).to(device) # iou vector for mAP@0.5:0.95
  65. niou = iouv.numel()
  66. # Logging
  67. log_imgs, wandb = min(log_imgs, 100), None # ceil
  68. try:
  69. import wandb # Weights & Biases
  70. except ImportError:
  71. log_imgs = 0
  72. # Dataloader
  73. if not training:
  74. img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
  75. _ = model(img.half() if half else img) if device.type != 'cpu' else None # run once
  76. path = data['test'] if opt.task == 'test' else data['val'] # path to val/test images
  77. dataloader = create_dataloader(path, imgsz, batch_size, model.stride.max(), opt, pad=0.5, rect=True)[0]
  78. seen = 0
  79. confusion_matrix = ConfusionMatrix(nc=nc)
  80. names = {k: v for k, v in enumerate(model.names if hasattr(model, 'names') else model.module.names)}
  81. coco91class = coco80_to_coco91_class()
  82. s = ('%20s' + '%12s' * 6) % ('Class', 'Images', 'Targets', 'P', 'R', 'mAP@.5', 'mAP@.5:.95')
  83. p, r, f1, mp, mr, map50, map, t0, t1 = 0., 0., 0., 0., 0., 0., 0., 0., 0.
  84. loss = torch.zeros(3, device=device)
  85. jdict, stats, ap, ap_class, wandb_images = [], [], [], [], []
  86. for batch_i, (img, targets, paths, shapes) in enumerate(tqdm(dataloader, desc=s)):
  87. img = img.to(device, non_blocking=True)
  88. img = img.half() if half else img.float() # uint8 to fp16/32
  89. img /= 255.0 # 0 - 255 to 0.0 - 1.0
  90. targets = targets.to(device)
  91. nb, _, height, width = img.shape # batch size, channels, height, width
  92. with torch.no_grad():
  93. # Run model
  94. t = time_synchronized()
  95. inf_out, train_out = model(img, augment=augment) # inference and training outputs
  96. t0 += time_synchronized() - t
  97. # Compute loss
  98. if training:
  99. loss += compute_loss([x.float() for x in train_out], targets, model)[1][:3] # box, obj, cls
  100. # Run NMS
  101. targets[:, 2:] *= torch.Tensor([width, height, width, height]).to(device) # to pixels
  102. lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)] if save_txt else [] # for autolabelling
  103. t = time_synchronized()
  104. output = non_max_suppression(inf_out, conf_thres=conf_thres, iou_thres=iou_thres, labels=lb)
  105. t1 += time_synchronized() - t
  106. # Statistics per image
  107. for si, pred in enumerate(output):
  108. labels = targets[targets[:, 0] == si, 1:]
  109. nl = len(labels)
  110. tcls = labels[:, 0].tolist() if nl else [] # target class
  111. path = Path(paths[si])
  112. seen += 1
  113. if len(pred) == 0:
  114. if nl:
  115. stats.append((torch.zeros(0, niou, dtype=torch.bool), torch.Tensor(), torch.Tensor(), tcls))
  116. continue
  117. # Predictions
  118. predn = pred.clone()
  119. scale_coords(img[si].shape[1:], predn[:, :4], shapes[si][0], shapes[si][1]) # native-space pred
  120. # Append to text file
  121. if save_txt:
  122. gn = torch.tensor(shapes[si][0])[[1, 0, 1, 0]] # normalization gain whwh
  123. for *xyxy, conf, cls in predn.tolist():
  124. xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
  125. line = (cls, *xywh, conf) if save_conf else (cls, *xywh) # label format
  126. with open(save_dir / 'labels' / (path.stem + '.txt'), 'a') as f:
  127. f.write(('%g ' * len(line)).rstrip() % line + '\n')
  128. # W&B logging
  129. if plots and len(wandb_images) < log_imgs:
  130. box_data = [{"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
  131. "class_id": int(cls),
  132. "box_caption": "%s %.3f" % (names[cls], conf),
  133. "scores": {"class_score": conf},
  134. "domain": "pixel"} for *xyxy, conf, cls in pred.tolist()]
  135. boxes = {"predictions": {"box_data": box_data, "class_labels": names}} # inference-space
  136. wandb_images.append(wandb.Image(img[si], boxes=boxes, caption=path.name))
  137. # Append to pycocotools JSON dictionary
  138. if save_json:
  139. # [{"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}, ...
  140. image_id = int(path.stem) if path.stem.isnumeric() else path.stem
  141. box = xyxy2xywh(predn[:, :4]) # xywh
  142. box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner
  143. for p, b in zip(pred.tolist(), box.tolist()):
  144. jdict.append({'image_id': image_id,
  145. 'category_id': coco91class[int(p[5])] if is_coco else int(p[5]),
  146. 'bbox': [round(x, 3) for x in b],
  147. 'score': round(p[4], 5)})
  148. # Assign all predictions as incorrect
  149. correct = torch.zeros(pred.shape[0], niou, dtype=torch.bool, device=device)
  150. if nl:
  151. detected = [] # target indices
  152. tcls_tensor = labels[:, 0]
  153. # target boxes
  154. tbox = xywh2xyxy(labels[:, 1:5])
  155. scale_coords(img[si].shape[1:], tbox, shapes[si][0], shapes[si][1]) # native-space labels
  156. if plots:
  157. confusion_matrix.process_batch(pred, torch.cat((labels[:, 0:1], tbox), 1))
  158. # Per target class
  159. for cls in torch.unique(tcls_tensor):
  160. ti = (cls == tcls_tensor).nonzero(as_tuple=False).view(-1) # prediction indices
  161. pi = (cls == pred[:, 5]).nonzero(as_tuple=False).view(-1) # target indices
  162. # Search for detections
  163. if pi.shape[0]:
  164. # Prediction to target ious
  165. ious, i = box_iou(predn[pi, :4], tbox[ti]).max(1) # best ious, indices
  166. # Append detections
  167. detected_set = set()
  168. for j in (ious > iouv[0]).nonzero(as_tuple=False):
  169. d = ti[i[j]] # detected target
  170. if d.item() not in detected_set:
  171. detected_set.add(d.item())
  172. detected.append(d)
  173. correct[pi[j]] = ious[j] > iouv # iou_thres is 1xn
  174. if len(detected) == nl: # all targets already located in image
  175. break
  176. # Append statistics (correct, conf, pcls, tcls)
  177. stats.append((correct.cpu(), pred[:, 4].cpu(), pred[:, 5].cpu(), tcls))
  178. # Plot images
  179. if plots and batch_i < 3:
  180. f = save_dir / f'test_batch{batch_i}_labels.jpg' # labels
  181. Thread(target=plot_images, args=(img, targets, paths, f, names), daemon=True).start()
  182. f = save_dir / f'test_batch{batch_i}_pred.jpg' # predictions
  183. Thread(target=plot_images, args=(img, output_to_target(output), paths, f, names), daemon=True).start()
  184. # Compute statistics
  185. stats = [np.concatenate(x, 0) for x in zip(*stats)] # to numpy
  186. if len(stats) and stats[0].any():
  187. p, r, ap, f1, ap_class = ap_per_class(*stats, plot=plots, save_dir=save_dir, names=names)
  188. p, r, ap50, ap = p[:, 0], r[:, 0], ap[:, 0], ap.mean(1) # [P, R, AP@0.5, AP@0.5:0.95]
  189. mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean()
  190. nt = np.bincount(stats[3].astype(np.int64), minlength=nc) # number of targets per class
  191. else:
  192. nt = torch.zeros(1)
  193. # Print results
  194. pf = '%20s' + '%12.3g' * 6 # print format
  195. print(pf % ('all', seen, nt.sum(), mp, mr, map50, map))
  196. # Print results per class
  197. if verbose and nc > 1 and len(stats):
  198. for i, c in enumerate(ap_class):
  199. print(pf % (names[c], seen, nt[c], p[i], r[i], ap50[i], ap[i]))
  200. # Print speeds
  201. t = tuple(x / seen * 1E3 for x in (t0, t1, t0 + t1)) + (imgsz, imgsz, batch_size) # tuple
  202. if not training:
  203. print('Speed: %.1f/%.1f/%.1f ms inference/NMS/total per %gx%g image at batch-size %g' % t)
  204. # Plots
  205. if plots:
  206. confusion_matrix.plot(save_dir=save_dir, names=list(names.values()))
  207. if wandb and wandb.run:
  208. wandb.log({"Images": wandb_images})
  209. wandb.log({"Validation": [wandb.Image(str(f), caption=f.name) for f in sorted(save_dir.glob('test*.jpg'))]})
  210. # Save JSON
  211. if save_json and len(jdict):
  212. w = Path(weights[0] if isinstance(weights, list) else weights).stem if weights is not None else '' # weights
  213. anno_json = glob.glob('../coco/annotations/instances_val*.json')[0] # annotations json
  214. pred_json = str(save_dir / f"{w}_predictions.json") # predictions json
  215. print('\nEvaluating pycocotools mAP... saving %s...' % pred_json)
  216. with open(pred_json, 'w') as f:
  217. json.dump(jdict, f)
  218. try: # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb
  219. from pycocotools.coco import COCO
  220. from pycocotools.cocoeval import COCOeval
  221. anno = COCO(anno_json) # init annotations api
  222. pred = anno.loadRes(pred_json) # init predictions api
  223. eval = COCOeval(anno, pred, 'bbox')
  224. if is_coco:
  225. eval.params.imgIds = [int(Path(x).stem) for x in dataloader.dataset.img_files] # image IDs to evaluate
  226. eval.evaluate()
  227. eval.accumulate()
  228. eval.summarize()
  229. map, map50 = eval.stats[:2] # update results (mAP@0.5:0.95, mAP@0.5)
  230. except Exception as e:
  231. print('ERROR: pycocotools unable to run: %s' % e)
  232. # Return results
  233. if not training:
  234. s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
  235. print(f"Results saved to {save_dir}{s}")
  236. model.float() # for training
  237. maps = np.zeros(nc) + map
  238. for i, c in enumerate(ap_class):
  239. maps[c] = ap[i]
  240. return (mp, mr, map50, map, *(loss.cpu() / len(dataloader)).tolist()), maps, t
  241. if __name__ == '__main__':
  242. parser = argparse.ArgumentParser(prog='test.py')
  243. parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)')
  244. parser.add_argument('--data', type=str, default='data/coco128.yaml', help='*.data path')
  245. parser.add_argument('--batch-size', type=int, default=32, help='size of each image batch')
  246. parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
  247. parser.add_argument('--conf-thres', type=float, default=0.001, help='object confidence threshold')
  248. parser.add_argument('--iou-thres', type=float, default=0.6, help='IOU threshold for NMS')
  249. parser.add_argument('--task', default='val', help="'val', 'test', 'study'")
  250. parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
  251. parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset')
  252. parser.add_argument('--augment', action='store_true', help='augmented inference')
  253. parser.add_argument('--verbose', action='store_true', help='report mAP by class')
  254. parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
  255. parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
  256. parser.add_argument('--save-json', action='store_true', help='save a cocoapi-compatible JSON results file')
  257. parser.add_argument('--project', default='runs/test', help='save to project/name')
  258. parser.add_argument('--name', default='exp', help='save to project/name')
  259. parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
  260. opt = parser.parse_args()
  261. opt.save_json |= opt.data.endswith('coco.yaml')
  262. opt.data = check_file(opt.data) # check file
  263. print(opt)
  264. if opt.task in ['val', 'test']: # run normally
  265. test(opt.data,
  266. opt.weights,
  267. opt.batch_size,
  268. opt.img_size,
  269. opt.conf_thres,
  270. opt.iou_thres,
  271. opt.save_json,
  272. opt.single_cls,
  273. opt.augment,
  274. opt.verbose,
  275. save_txt=opt.save_txt,
  276. save_conf=opt.save_conf,
  277. )
  278. elif opt.task == 'study': # run over a range of settings and save/plot
  279. for weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']:
  280. f = 'study_%s_%s.txt' % (Path(opt.data).stem, Path(weights).stem) # filename to save to
  281. x = list(range(320, 800, 64)) # x axis
  282. y = [] # y axis
  283. for i in x: # img-size
  284. print('\nRunning %s point %s...' % (f, i))
  285. r, _, t = test(opt.data, weights, opt.batch_size, i, opt.conf_thres, opt.iou_thres, opt.save_json,
  286. plots=False)
  287. y.append(r + t) # results and times
  288. np.savetxt(f, y, fmt='%10.4g') # save
  289. os.system('zip -r study.zip study_*.txt')
  290. plot_study_txt(f, x) # plot