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_test_set(image_size: ImageSize):
num_joints = 15
return EhpiDataset("/media/disks/beta/datasets/ehpi/2019_03_13_Freilichtmuseum_30FPS",
transform=transforms.Compose([
RemoveJointsOutsideImgEhpi(image_size),
NormalizeEhpi(image_size)
]), dataset_part=DatasetPart.TEST, num_joints=num_joints)
def get_test_set(dataset_path: str, image_size: ImageSize):
num_joints = 15
return EhpiDataset(os.path.join(dataset_path, "2019_03_13_Freilichtmuseum_30FPS"),
transform=transforms.Compose([
RemoveJointsOutsideImgEhpi(image_size),
NormalizeEhpi(image_size)
]), dataset_part=DatasetPart.TEST, num_joints=num_joints)
def get_test_set(dataset_path: str, image_size: ImageSize):
num_joints = 15
return EhpiDataset(dataset_path,
transform=transforms.Compose([
RemoveJointsOutsideImgEhpi(image_size),
NormalizeEhpi(image_size)
]),
dataset_part=DatasetPart.TEST,
num_joints=num_joints)
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)
import cv2
import numpy as np
from nobos_commons.data_structures.constants.dataset_part import DatasetPart
from nobos_torch_lib.datasets.action_recognition_datasets.ehpi_dataset import EhpiDataset
set_to_test = EhpiDataset("/media/disks/beta/datasets/ehpi/JHMDB_ITSC-1-GT/",
dataset_part=DatasetPart.TRAIN, num_joints=15)
for ehpi in set_to_test:
x = ehpi["x"]
y = ehpi["y"]
x = np.transpose(x, (1, 2, 0))
x = cv2.resize(x, (x.shape[0] * 10, x.shape[1] * 10))
cv2.imshow("preview", x)
cv2.waitKey(0)