Update test.py --task train val study (#2453)
Browse files* Update test.py --task train val study
* update argparser --task
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 |
-
|
89 |
-
dataloader = create_dataloader(
|
90 |
-
prefix=colorstr('
|
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=
|
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
|
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,
|