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

335 lines
16KB

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