You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

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