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_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_lab(dataset_path: str, image_size: ImageSize):
num_joints = 15
datasets = [
EhpiLSTMDataset(os.path.join(dataset_path, "JOURNAL_2019_03_TEST_VUE01_30FPS"),
transform=transforms.Compose([
RemoveJointsOutsideImgEhpi(image_size),
NormalizeEhpi(image_size)
]), num_joints=num_joints, dataset_part=DatasetPart.TEST),
EhpiLSTMDataset(os.path.join(dataset_path, "JOURNAL_2019_03_TEST_VUE02_30FPS"),
transform=transforms.Compose([
RemoveJointsOutsideImgEhpi(image_size),
NormalizeEhpi(image_size)
]), num_joints=num_joints, dataset_part=DatasetPart.TEST),
]
for dataset in datasets:
dataset.print_label_statistics()
return ConcatDataset(datasets)
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 __init__(self, model, feature_vec_producer: FeatureVecProducerEhpi,
image_size: ImageSize):
self.model = model
self.feature_vec_producer = feature_vec_producer
self.action_buffer: AlgorithmOutputBuffer = AlgorithmOutputBuffer(
buffer_size=32)
self.remove = RemoveJointsOutsideImgEhpi(image_size)
self.normalize = NormalizeEhpi(image_size)
model.cuda()
model.eval()
def get_test_set_office(dataset_path: str, image_size: ImageSize):
num_joints = 15
dataset = EhpiLSTMDataset(os.path.join(dataset_path, "JOURNAL_2019_04_TEST_EVAL2_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)
dataset.print_label_statistics()
return dataset
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)
# balances = [True, False]
for seed in seeds:
for batch_size in batch_sizes:
for lr in lrs:
for weight_decay in weight_decays:
print("Seed: {}, Batchsize: {}, LR: {}, Weight Decay: {}".
format(seed, batch_size, lr, weight_decay))
set_seed(0) # FIXED SEED FOR DATASET SPLIT!
# Load full dataset
train_full_set = get_training_set(
os.path.join(ehpi_dataset_path, "jhmdb"), image_size)
test_full_set = copy(train_full_set)
test_full_set.transform = transforms.Compose([
RemoveJointsOutsideImgEhpi(image_size),
NormalizeEhpi(image_size)
])
# Create train and validation splits
train_indices, val_indices = train_full_set.get_subsplit_indices(
validation_percentage=0.3)
# Train set
train_set = Subset(train_full_set, train_indices)
# Validation set
val_set = Subset(test_full_set, val_indices)
val_loader = DataLoader(val_set,
batch_size=1,
shuffle=False)
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)