def get_training_set(dataset_path: str, image_size: ImageSize): num_joints = 15 left_indexes: List[int] = [3, 4, 5, 9, 10, 11] right_indexes: List[int] = [6, 7, 8, 12, 13, 14] return EhpiDataset(dataset_path, transform=transforms.Compose([ RemoveJointsOutsideImgEhpi(image_size), RemoveJointsEhpi(indexes_to_remove=foot_indexes, indexes_to_remove_2=knee_indexes, probability=0.25), ScaleEhpi(image_size), TranslateEhpi(image_size), FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes), NormalizeEhpi(image_size) ]), num_joints=num_joints)
def get_training_set_gt(dataset_path: str, image_size: ImageSize): num_joints = 15 left_indexes: List[int] = [3, 4, 5, 9, 10, 11] right_indexes: List[int] = [6, 7, 8, 12, 13, 14] datasets: List[EhpiLSTMDataset] = [ EhpiLSTMDataset(os.path.join(dataset_path, "JOURNAL_2019_03_GT_30fps"), transform=transforms.Compose([ RemoveJointsOutsideImgEhpi(image_size), ScaleEhpi(image_size), TranslateEhpi(image_size), FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes), NormalizeEhpi(image_size) ]), num_joints=num_joints), ] for dataset in datasets: dataset.print_label_statistics() return ConcatDataset(datasets)
def get_set_wo_sim(dataset_path: str, image_size: ImageSize): num_joints = 15 left_indexes: List[int] = [3, 4, 5, 9, 10, 11] right_indexes: List[int] = [6, 7, 8, 12, 13, 14] datasets: List[EhpiDataset] = [ EhpiDataset(os.path.join(dataset_path, "ofp_webcam"), transform=transforms.Compose([ RemoveJointsOutsideImgEhpi(image_size), ScaleEhpi(image_size), TranslateEhpi(image_size), FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes), NormalizeEhpi(image_size) ]), num_joints=num_joints), EhpiDataset(os.path.join(dataset_path, "ofp_record_2019_03_11_30FPS"), transform=transforms.Compose([ RemoveJointsOutsideImgEhpi(image_size), ScaleEhpi(image_size), TranslateEhpi(image_size), FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes), NormalizeEhpi(image_size) ]), num_joints=num_joints), EhpiDataset(os.path.join(dataset_path, "ofp_record_2019_03_11_HSRT_30FPS"), transform=transforms.Compose([ RemoveJointsOutsideImgEhpi(image_size), ScaleEhpi(image_size), TranslateEhpi(image_size), FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes), NormalizeEhpi(image_size) ]), num_joints=num_joints, dataset_part=DatasetPart.TEST), EhpiDataset(os.path.join(dataset_path, "ofp_record_2019_03_11_HELLA_30FPS"), transform=transforms.Compose([ RemoveJointsOutsideImgEhpi(image_size), ScaleEhpi(image_size), TranslateEhpi(image_size), FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes), NormalizeEhpi(image_size) ]), num_joints=num_joints, dataset_part=DatasetPart.TRAIN), # Freilichtmuseum EhpiDataset(os.path.join(dataset_path, "2019_03_13_Freilichtmuseum_30FPS"), transform=transforms.Compose([ RemoveJointsOutsideImgEhpi(image_size), ScaleEhpi(image_size), TranslateEhpi(image_size), FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes), NormalizeEhpi(image_size) ]), num_joints=num_joints, dataset_part=DatasetPart.TRAIN), ] for dataset in datasets: dataset.print_label_statistics() return ConcatDataset(datasets)
def get_full(dataset_path: str, image_size: ImageSize): num_joints = 15 left_indexes: List[int] = [3, 4, 5, 9, 10, 11] right_indexes: List[int] = [6, 7, 8, 12, 13, 14] datasets: List[EhpiDataset] = [ # Real EhpiDataset(os.path.join(dataset_path, "ofp_webcam"), transform=transforms.Compose([ RemoveJointsOutsideImgEhpi(image_size), ScaleEhpi(image_size), TranslateEhpi(image_size), FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes), NormalizeEhpi(image_size) ]), num_joints=num_joints), EhpiDataset(os.path.join(dataset_path, "ofp_record_2019_03_11_30FPS"), transform=transforms.Compose([ RemoveJointsOutsideImgEhpi(image_size), ScaleEhpi(image_size), TranslateEhpi(image_size), FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes), NormalizeEhpi(image_size) ]), num_joints=num_joints), EhpiDataset(os.path.join(dataset_path, "ofp_record_2019_03_11_HSRT_30FPS"), transform=transforms.Compose([ RemoveJointsOutsideImgEhpi(image_size), ScaleEhpi(image_size), TranslateEhpi(image_size), FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes), NormalizeEhpi(image_size) ]), num_joints=num_joints, dataset_part=DatasetPart.TEST), EhpiDataset(os.path.join(dataset_path, "ofp_record_2019_03_11_HELLA_30FPS"), transform=transforms.Compose([ RemoveJointsOutsideImgEhpi(image_size), ScaleEhpi(image_size), TranslateEhpi(image_size), FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes), NormalizeEhpi(image_size) ]), num_joints=num_joints, dataset_part=DatasetPart.TRAIN), # Freilichtmuseum EhpiDataset(os.path.join(dataset_path, "2019_03_13_Freilichtmuseum_30FPS"), transform=transforms.Compose([ RemoveJointsOutsideImgEhpi(image_size), ScaleEhpi(image_size), TranslateEhpi(image_size), FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes), NormalizeEhpi(image_size) ]), num_joints=num_joints, dataset_part=DatasetPart.TRAIN), # Simulated EhpiDataset(os.path.join(dataset_path, "ofp_from_mocap_30fps/"), transform=transforms.Compose([ RemoveJointsOutsideImgEhpi(image_size), RemoveJointsEhpi(indexes_to_remove=foot_indexes, indexes_to_remove_2=knee_indexes, probability=0.25), ScaleEhpi(image_size), TranslateEhpi(image_size), FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes), NormalizeEhpi(image_size) ]), num_joints=num_joints), EhpiDataset(os.path.join(dataset_path, "ofp_sim_pose_algo_equal_30fps"), transform=transforms.Compose([ RemoveJointsOutsideImgEhpi(image_size), RemoveJointsEhpi(indexes_to_remove=foot_indexes, indexes_to_remove_2=knee_indexes, probability=0.25), ScaleEhpi(image_size), TranslateEhpi(image_size), FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes), NormalizeEhpi(image_size) ]), num_joints=num_joints), EhpiDataset(os.path.join(dataset_path, "ofp_sim_gt_equal_30fps"), transform=transforms.Compose([ RemoveJointsOutsideImgEhpi(image_size), RemoveJointsEhpi(indexes_to_remove=foot_indexes, indexes_to_remove_2=knee_indexes, probability=0.25), ScaleEhpi(image_size), TranslateEhpi(image_size), FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes), NormalizeEhpi(image_size) ]), num_joints=num_joints), EhpiDataset(os.path.join(dataset_path, "ofp_from_mocap_gt_30fps"), transform=transforms.Compose([ RemoveJointsOutsideImgEhpi(image_size), RemoveJointsEhpi(indexes_to_remove=foot_indexes, indexes_to_remove_2=knee_indexes, probability=0.25), ScaleEhpi(image_size), TranslateEhpi(image_size), FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes), NormalizeEhpi(image_size) ]), num_joints=num_joints), ] for dataset in datasets: dataset.print_label_statistics() return ConcatDataset(datasets)
def get_sim(image_size: ImageSize): num_joints = 15 left_indexes: List[int] = [3, 4, 5, 9, 10, 11] right_indexes: List[int] = [6, 7, 8, 12, 13, 14] datasets: List[EhpiDataset] = [ EhpiDataset( "/media/disks/beta/datasets/ehpi/ofp_sim_pose_algo_equal_30fps", transform=transforms.Compose([ RemoveJointsOutsideImgEhpi(image_size), RemoveJointsEhpi(indexes_to_remove=foot_indexes, indexes_to_remove_2=knee_indexes, probability=0.25), ScaleEhpi(image_size), TranslateEhpi(image_size), FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes), NormalizeEhpi(image_size) ]), num_joints=num_joints), EhpiDataset( "/media/disks/beta/datasets/ehpi/ofp_from_mocap_pose_algo_30fps", transform=transforms.Compose([ RemoveJointsOutsideImgEhpi(image_size), RemoveJointsEhpi(indexes_to_remove=foot_indexes, indexes_to_remove_2=knee_indexes, probability=0.25), ScaleEhpi(image_size), TranslateEhpi(image_size), FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes), NormalizeEhpi(image_size) ]), num_joints=num_joints), EhpiDataset("/media/disks/beta/datasets/ehpi/ofp_sim_gt_equal_30fps", transform=transforms.Compose([ RemoveJointsOutsideImgEhpi(image_size), RemoveJointsEhpi(indexes_to_remove=foot_indexes, indexes_to_remove_2=knee_indexes, probability=0.25), ScaleEhpi(image_size), TranslateEhpi(image_size), FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes), NormalizeEhpi(image_size) ]), num_joints=num_joints), EhpiDataset("/media/disks/beta/datasets/ehpi/ofp_from_mocap_gt_30fps", transform=transforms.Compose([ RemoveJointsOutsideImgEhpi(image_size), RemoveJointsEhpi(indexes_to_remove=foot_indexes, indexes_to_remove_2=knee_indexes, probability=0.25), ScaleEhpi(image_size), TranslateEhpi(image_size), FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes), NormalizeEhpi(image_size) ]), num_joints=num_joints) ] for dataset in datasets: dataset.print_label_statistics() return ConcatDataset(datasets)