|
|
@@ -22,9 +22,9 @@ from utils.torch_utils import select_device, time_synchronized |
|
|
|
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 @@ if __name__ == '__main__': |
|
|
|
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 @@ if __name__ == '__main__': |
|
|
|
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 @@ if __name__ == '__main__': |
|
|
|
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') |