Merge pull request #67 from kaseris/datasets/human36m

Human 3.6m Dataset

src/skelcast/data/dataset.py

@@ -299,3 +299,59 @@ def store_to_cache(self, cache_file: str) -> None:
  with open(cache_file, 'wb') as f:
  pickle.dump(self.skeleton_files_clean, f)
  logging.info(f"Stored {len(self.skeleton_files_clean)} files to cache file {cache_file}.")
+
+
+@DATASETS.register_module()
+class Human36mDataset(Dataset):
+ def __init__(self, data_path, use_hourglass_detections=True, train=True) -> None:
+ self.data_path = data_path
+ self.use_hourglass_detections = use_hourglass_detections
+ self.train = train
+
+ self.train_inputs, self.test_inputs = [], []
+ self.act = []
+
+ if self.use_hourglass_detections:
+ train_2d_file = 'train_2d_ft.pth.tar'
+ test_2d_file = 'test_2d_ft.pth.tar'
+ else:
+ train_2d_file = 'train_2d.pth.tar'
+ train_2d_file = 'test_2d.pth.tar'
+
+ if self.train:
+ self.train_3d = torch.load(os.path.join(data_path, 'train_3d.pth.tar'))
+ self.train_2d = torch.load(os.path.join(data_path, train_2d_file))
+
+ for k2d in self.train_2d.keys():
+ (sub, act, fname) = k2d
+ k3d = k2d
+ k3d = (sub, act, fname[:-3]) if fname.endswith('-sh') else k3d
+ assert self.train_3d[k3d].shape[0] == self.train_2d[k2d].shape[0], f'(training) 3d and 2d shapes not matching'
+ self.train_inputs.append(self.train_3d[k3d])
+ self.act.append(act)
+
+ else:
+ self.test_3d = torch.load(os.path.join(data_path, 'test_3d.pth.tar'))
+ self.test_2d = torch.load(os.path.join(data_path, test_2d_file))
+ for k2d in self.test_2d.keys():
+ (sub, act, fname) = k2d
+ k3d = k2d
+ k3d = (sub, act, fname[:-3]) if fname.endswith('-sh') else k3d
+ assert self.test_2d[k2d].shape[0] == self.test_3d[k3d].shape[0], '(test) 3d and 2d shapes not matching'
+ self.test_inputs.append(self.test_3d[k3d])
+ self.act.append(act)
+
+ def __getitem__(self, index) -> Any:
+ if self.train:
+ # We want the sampeles to be returned as sequences
+ # i.e.: [seq_len, n_joints, 3]
+ x = torch.from_numpy(self.train_inputs[index]).float()
+ else:
+ x = torch.from_numpy(self.test_inputs[index]).float()
+ return x.view(-1, 16, 3), self.act[index]
+
+ def __len__(self):
+ if self.train:
+ return len(self.train_inputs)
+ else:
+ return len(self.test_inputs)