device = torch.device("cuda")
else:
device = torch.device("cpu")
# input size
input_size = [args.input_size, args.input_size]
# load net
if args.version == 'centernet':
from models.centernet import CenterNet
net = CenterNet(device,
input_size=input_size,
num_classes=num_classes,
backbone=args.backbone,
use_nms=args.use_nms)
# load net
net.load_state_dict(torch.load(args.trained_model, map_location='cuda'))
net.eval()
print('Finished loading model!')
net = net.to(device)
# evaluation
with torch.no_grad():
if args.dataset == 'voc':
voc_test(net, device, input_size)
elif args.dataset == 'coco-val':
coco_test(net, device, input_size, test=False)
elif args.dataset == 'coco-test':
coco_test(net, device, input_size, test=True)
print('ap50_95 : ', ap50_95)
if __name__ == '__main__':
global cfg
cfg = coco_cfg
if args.cuda:
print('use cuda')
torch.backends.cudnn.benchmark = True
device = torch.device("cuda")
else:
device = torch.device("cpu")
if args.version == 'centernet':
from models.centernet import CenterNet
model = CenterNet(device,
input_size=cfg['min_dim'],
num_classes=args.num_classes)
else:
print('Unknown Version !!!')
exit()
# load model
model.load_state_dict(torch.load(args.trained_model, map_location=device))
model.eval().to(device)
print('Finished loading model!')
test(model, device)