From 545033df0ebaa1ee073aa3506106404605a1fb81 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 12 Mar 2021 21:28:11 -0800 Subject: [PATCH 1/2] Update test.py --task train val study --- test.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/test.py b/test.py index 46288019a8bd..92f18fff7c24 100644 --- a/test.py +++ b/test.py @@ -85,9 +85,9 @@ def test(data, if not training: if device.type != 'cpu': model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters()))) # run once - path = data['test'] if opt.task == 'test' else data['val'] # path to val/test images - dataloader = create_dataloader(path, imgsz, batch_size, gs, opt, pad=0.5, rect=True, - prefix=colorstr('test: ' if opt.task == 'test' else 'val: '))[0] + task = opt.task if opt.task in ('train', 'val', 'test') else 'val' # path to train/val/test images + dataloader = create_dataloader(data[task], imgsz, batch_size, gs, opt, pad=0.5, rect=True, + prefix=colorstr(f'{task}: '))[0] seen = 0 confusion_matrix = ConfusionMatrix(nc=nc) @@ -305,7 +305,7 @@ def test(data, print(opt) check_requirements() - if opt.task in ['val', 'test']: # run normally + if opt.task in ('train', 'val', 'test'): # run normally test(opt.data, opt.weights, opt.batch_size, From 209514e8b726cdac0ddef3b8edbcbac0412b838e Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 12 Mar 2021 21:34:03 -0800 Subject: [PATCH 2/2] update argparser --task --- test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test.py b/test.py index 92f18fff7c24..39e0992264ec 100644 --- a/test.py +++ b/test.py @@ -287,7 +287,7 @@ def test(data, parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)') parser.add_argument('--conf-thres', type=float, default=0.001, help='object confidence threshold') parser.add_argument('--iou-thres', type=float, default=0.6, help='IOU threshold for NMS') - parser.add_argument('--task', default='val', help="'val', 'test', 'study'") + parser.add_argument('--task', default='val', help='train, val, test, speed or study') parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset') parser.add_argument('--augment', action='store_true', help='augmented inference')