def get_train_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]
datasets: List[EhpiDataset] = [
# Set 1
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),
# Set 2
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_training_set_both(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_POSEALGO_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),
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_sim_pose_algo_only(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_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_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),
]
for dataset in datasets:
dataset.print_label_statistics()
return ConcatDataset(datasets)
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(os.path.join(dataset_path, "JHMDB_ITSC-1/"),
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_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)
