| | |
|---|
| | | parser.add_argument('--save_viz', dest='save_viz', help='Save images with pose cube.', |
|---|
| | | default=False, type=bool) |
|---|
| | | parser.add_argument('--iter_ref', dest='iter_ref', default=1, type=int) |
|---|
| | | parser.add_argument('--dataset', dest='dataset', help='Dataset type.', default='AFLW2000', type=str) |
|---|
| | | |
|---|
| | | args = parser.parse_args() |
|---|
| | | |
|---|
| | |
|---|
| | | transforms.RandomCrop(224), transforms.ToTensor(), |
|---|
| | | transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]) |
|---|
| | | |
|---|
| | | 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) |
|---|
