glenn-jocher commited on
Commit
f419721
·
unverified ·
1 Parent(s): f01f322

Update test.py --task train val study (#2453)

Browse files

* Update test.py --task train val study

* update argparser --task

Files changed (1) hide show
  1. test.py +5 -5
test.py CHANGED
@@ -85,9 +85,9 @@ def test(data,
85
  if not training:
86
  if device.type != 'cpu':
87
  model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters()))) # run once
88
- path = data['test'] if opt.task == 'test' else data['val'] # path to val/test images
89
- dataloader = create_dataloader(path, imgsz, batch_size, gs, opt, pad=0.5, rect=True,
90
- prefix=colorstr('test: ' if opt.task == 'test' else 'val: '))[0]
91
 
92
  seen = 0
93
  confusion_matrix = ConfusionMatrix(nc=nc)
@@ -287,7 +287,7 @@ if __name__ == '__main__':
287
  parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
288
  parser.add_argument('--conf-thres', type=float, default=0.001, help='object confidence threshold')
289
  parser.add_argument('--iou-thres', type=float, default=0.6, help='IOU threshold for NMS')
290
- parser.add_argument('--task', default='val', help="'val', 'test', 'study'")
291
  parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
292
  parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset')
293
  parser.add_argument('--augment', action='store_true', help='augmented inference')
@@ -305,7 +305,7 @@ if __name__ == '__main__':
305
  print(opt)
306
  check_requirements()
307
 
308
- if opt.task in ['val', 'test']: # run normally
309
  test(opt.data,
310
  opt.weights,
311
  opt.batch_size,
 
85
  if not training:
86
  if device.type != 'cpu':
87
  model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters()))) # run once
88
+ task = opt.task if opt.task in ('train', 'val', 'test') else 'val' # path to train/val/test images
89
+ dataloader = create_dataloader(data[task], imgsz, batch_size, gs, opt, pad=0.5, rect=True,
90
+ prefix=colorstr(f'{task}: '))[0]
91
 
92
  seen = 0
93
  confusion_matrix = ConfusionMatrix(nc=nc)
 
287
  parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
288
  parser.add_argument('--conf-thres', type=float, default=0.001, help='object confidence threshold')
289
  parser.add_argument('--iou-thres', type=float, default=0.6, help='IOU threshold for NMS')
290
+ parser.add_argument('--task', default='val', help='train, val, test, speed or study')
291
  parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
292
  parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset')
293
  parser.add_argument('--augment', action='store_true', help='augmented inference')
 
305
  print(opt)
306
  check_requirements()
307
 
308
+ if opt.task in ('train', 'val', 'test'): # run normally
309
  test(opt.data,
310
  opt.weights,
311
  opt.batch_size,