From 7591501e44a5e91d929b298254da6dc66f2c39e4 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Wed, 9 Jun 2021 21:36:10 +0200 Subject: [PATCH] Refactor test.py arguments (#3558) * remove opt from test() * pass kwargs * update comments * revert accidental default change * multiple --img options * add comments (cherry picked from commit a9553c04a7d32d5c8f29b0917fbeb6b1ef6cfe5f) --- detect.py | 2 +- test.py | 46 ++++++++++++++++++---------------------------- 2 files changed, 19 insertions(+), 29 deletions(-) diff --git a/detect.py b/detect.py index 8dbb656ed95f..5551824a4110 100644 --- a/detect.py +++ b/detect.py @@ -33,7 +33,7 @@ def detect(opt): # Load model model = attempt_load(weights, map_location=device) # load FP32 model stride = int(model.stride.max()) # model stride - imgsz = check_img_size(imgsz, s=stride) # check img_size + imgsz = check_img_size(imgsz, s=stride) # check image size names = model.module.names if hasattr(model, 'module') else model.names # get class names if half: model.half() # to FP16 diff --git a/test.py b/test.py index 971c4b005ca1..a10f0f88f8e6 100644 --- a/test.py +++ b/test.py @@ -22,9 +22,9 @@ def test(data, weights=None, batch_size=32, - imgsz=640, - conf_thres=0.001, - iou_thres=0.6, # for NMS + imgsz=640, # image size + conf_thres=0.001, # confidence threshold + iou_thres=0.6, # NMS IoU threshold save_json=False, single_cls=False, augment=False, @@ -38,8 +38,12 @@ def test(data, plots=True, wandb_logger=None, compute_loss=None, - half=True, - opt=None): + half=True, # FP16 half-precision inference + project='runs/test', + name='exp', + exist_ok=False, + task='val', + device=''): # Initialize/load model and set device training = model is not None if training: # called by train.py @@ -47,16 +51,16 @@ def test(data, else: # called directly set_logging() - device = select_device(opt.device, batch_size=batch_size) + device = select_device(device, batch_size=batch_size) # Directories - save_dir = increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok) # increment run + save_dir = increment_path(Path(project) / name, exist_ok=exist_ok) # increment run (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir # Load model model = attempt_load(weights, map_location=device) # load FP32 model gs = max(int(model.stride.max()), 32) # grid size (max stride) - imgsz = check_img_size(imgsz, s=gs) # check img_size + imgsz = check_img_size(imgsz, s=gs) # check image size # Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99 # if device.type != 'cpu' and torch.cuda.device_count() > 1: @@ -86,7 +90,7 @@ 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 - task = opt.task if opt.task in ('train', 'val', 'test') else 'val' # path to train/val/test images + task = task if task in ('train', 'val', 'test') else 'val' # path to train/val/test images dataloader = create_dataloader(data[task], imgsz, batch_size, gs, single_cls, pad=0.5, rect=True, prefix=colorstr(f'{task}: '))[0] @@ -294,7 +298,7 @@ def test(data, parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)') parser.add_argument('--data', type=str, default='data/coco128.yaml', help='*.data path') parser.add_argument('--batch-size', type=int, default=32, help='size of each image batch') - parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)') + parser.add_argument('--imgsz', '--img', '--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='train, val, test, speed or study') @@ -312,31 +316,17 @@ def test(data, parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference') opt = parser.parse_args() opt.save_json |= opt.data.endswith('coco.yaml') + opt.save_txt |= opt.save_hybrid opt.data = check_file(opt.data) # check file print(opt) check_requirements(exclude=('tensorboard', 'thop')) if opt.task in ('train', 'val', 'test'): # run normally - test(opt.data, - opt.weights, - opt.batch_size, - opt.img_size, - opt.conf_thres, - opt.iou_thres, - opt.save_json, - opt.single_cls, - opt.augment, - opt.verbose, - save_txt=opt.save_txt | opt.save_hybrid, - save_hybrid=opt.save_hybrid, - save_conf=opt.save_conf, - half=opt.half, - opt=opt - ) + test(**vars(opt)) elif opt.task == 'speed': # speed benchmarks for w in opt.weights if isinstance(opt.weights, list) else [opt.weights]: - test(opt.data, w, opt.batch_size, opt.img_size, 0.25, 0.45, save_json=False, plots=False, opt=opt) + test(opt.data, w, opt.batch_size, opt.imgsz, 0.25, 0.45, save_json=False, plots=False) elif opt.task == 'study': # run over a range of settings and save/plot # python test.py --task study --data coco.yaml --iou 0.7 --weights yolov5s.pt yolov5m.pt yolov5l.pt yolov5x.pt @@ -347,7 +337,7 @@ def test(data, for i in x: # img-size print(f'\nRunning {f} point {i}...') r, _, t = test(opt.data, w, opt.batch_size, i, opt.conf_thres, opt.iou_thres, opt.save_json, - plots=False, opt=opt) + plots=False) y.append(r + t) # results and times np.savetxt(f, y, fmt='%10.4g') # save os.system('zip -r study.zip study_*.txt')