def prepare_dataset(opt): """ Prepare PyTorch dataset objects used for training 2D-to-3D deep network Args opt: experiment options Returns train_dataset: training dataset as PyTorch dataset object eval_dataset: evaluation dataset as PyTorch dataset object data_stats: dataset statistics computed from the training dataset action_eval_list: a list of evaluation dataset objects where each corresponds to one action """ # get relevant paths data_dir = opt.data_dir cameras_path = os.path.join(data_dir, 'cameras.npy') # By default, all actions are used actions = define_actions(opt.actions) # load camera parameters to project 3D skeleton rcams = np.load(cameras_path).item() # produce more camera views by adding virtual cameras if needed if opt.virtual_cams: rcams = add_virtual_cams(rcams) # first prepare Python dictionary containing 2D and 3D data data_dic, data_stats = prepare_data_dict(rcams, opt, predict_14=False) input_size = len(data_stats['dim_use_2d']) output_size = len(data_stats['dim_use_3d']) # convert Python dictionary to numpy array train_input, train_output = get_all_data(data_dic['train_set_2d'], data_dic['train_set_3d'], camera_frame, norm_twoD=opt.norm_twoD, input_size=input_size, output_size=output_size) eval_input, eval_output = get_all_data(data_dic['test_set_2d'], data_dic['test_set_3d'], camera_frame, norm_twoD=opt.norm_twoD, input_size=input_size, output_size=output_size) # The Numpy arrays are finally used to initialize the dataset objects train_dataset = dataset.PoseDataset(train_input, train_output, 'train', refine_3d=opt.refine_3d) eval_dataset = dataset.PoseDataset(eval_input, eval_output, 'eval', refine_3d=opt.refine_3d) # Create a list of dataset objects for action-wise evaluation action_eval_list = split_action(data_dic['test_set_2d'], data_dic['test_set_3d'], actions, camera_frame, opt, input_size=input_size, output_size=output_size) return train_dataset, eval_dataset, data_stats, action_eval_list
def split_action(dic_2d, dic_3d, actions, camera_frame, opt, input_size, output_size): """ Generate a list of datasets for each action. Args dic_2d: dictionary containing 2d poses dic_3d: dictionary containing 3d poses actions: list of defined actions camera_frame: use camera coordinate system opt: experiment options input_size: input vector length output_size: output vector length Returns action_dataset_list: a list of datasets where each element correspond to one action """ action_dataset_list = [] for act_id in range(len(actions)): action = actions[act_id] dic_2d_action, dic_3d_action = select_action(dic_2d, dic_3d, action, opt.twoD_source) eval_input, eval_output = get_all_data(dic_2d_action, dic_3d_action, camera_frame, norm_twoD=opt.norm_twoD, input_size=input_size, output_size=output_size) action_dataset = dataset.PoseDataset(eval_input, eval_output, 'eval', action_name=action, refine_3d=opt.refine_3d) action_dataset_list.append(action_dataset) return action_dataset_list
def split_action(dic_2d, dic_3d, actions, camera_frame, opt, input_size, output_size): """ Generate a list of datasets for each action. """ action_dataset_list = [] for act_id in range(len(actions)): action = actions[act_id] dic_2d_action, dic_3d_action = select_action(dic_2d, dic_3d, action, opt.twoD_source) eval_input, eval_output = get_all_data(dic_2d_action, dic_3d_action, camera_frame, norm_twoD=opt.norm_twoD, input_size=input_size, output_size=output_size) action_dataset = dataset.PoseDataset(eval_input, eval_output, 'eval', action_name=action, refine_3d=opt.refine_3d) action_dataset_list.append(action_dataset) return action_dataset_list