From b2316888f61893aff1229c24002d903df79499cb Mon Sep 17 00:00:00 2001 From: natanielruiz <nataniel777@hotmail.com> Date: 星期三, 13 九月 2017 21:16:31 +0800 Subject: [PATCH] Removed unecessary testing scripts --- code/test_preangles.py | 15 +++++++++++++-- 1 files changed, 13 insertions(+), 2 deletions(-) diff --git a/code/test_preangles.py b/code/test_preangles.py index 4aedfd8..f2039f2 100644 --- a/code/test_preangles.py +++ b/code/test_preangles.py @@ -33,6 +33,7 @@ default=1, type=int) parser.add_argument('--save_viz', dest='save_viz', help='Save images with pose cube.', default=False, type=bool) + parser.add_argument('--dataset', dest='dataset', help='Dataset type.', default='AFLW2000', type=str) args = parser.parse_args() @@ -48,7 +49,7 @@ # ResNet101 with 3 outputs. # model = hopenet.Hopenet(torchvision.models.resnet.Bottleneck, [3, 4, 23, 3], 66) # ResNet50 - model = hopenet.Hopenet(torchvision.models.resnet.Bottleneck, [3, 4, 6, 3], 66) + model = hopenet.Hopenet(torchvision.models.resnet.Bottleneck, [3, 4, 6, 3], 66, 0) # ResNet18 # model = hopenet.Hopenet(torchvision.models.resnet.BasicBlock, [2, 2, 2, 2], 66) @@ -66,8 +67,18 @@ transforms.RandomCrop(224), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]) - pose_dataset = datasets.AFLW2000(args.data_dir, args.filename_list, + if args.dataset == 'AFLW2000': + pose_dataset = datasets.AFLW2000(args.data_dir, args.filename_list, transformations) + elif args.dataset == 'BIWI': + pose_dataset = datasets.BIWI(args.data_dir, args.filename_list, transformations) + elif args.dataset == 'AFLW': + pose_dataset = datasets.AFLW(args.data_dir, args.filename_list, transformations) + elif args.dataset == 'AFW': + pose_dataset = datasets.AFW(args.data_dir, args.filename_list, transformations) + else: + print 'Error: not a valid dataset name' + sys.exit() test_loader = torch.utils.data.DataLoader(dataset=pose_dataset, batch_size=args.batch_size, num_workers=2) -- Gitblit v1.8.0