From d274dda340078ecaeeaf0b0c62caa8778d04f7cd Mon Sep 17 00:00:00 2001 From: Michail Kaseris Date: Mon, 15 Jan 2024 14:52:07 +0200 Subject: [PATCH 1/5] First take on Human3.6m. --- src/skelcast/data/dataset.py | 64 ++++++++++++++++++++++++++++++++++++ 1 file changed, 64 insertions(+) diff --git a/src/skelcast/data/dataset.py b/src/skelcast/data/dataset.py index 452cf25..38b0628 100644 --- a/src/skelcast/data/dataset.py +++ b/src/skelcast/data/dataset.py @@ -299,3 +299,67 @@ 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, actions, data_path, use_hourglass_detections=True, train=True) -> None: + self.actions = actions + self.data_path = data_path + self.use_hourglass_detections = use_hourglass_detections + self.train = train + + self.train_inputs, self.train_out, self.test_inputs, self.test_out = [], [], [], [] + self.train_meta, self.test_meta = [], [] + + 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 + num_f, _ = self.train_2d[k2d].shape + assert self.train_3d[k3d].shape[0] == self.train_2d[k2d].shape[0], f'(training) 3d and 2d shapes not matching' + for i in range(num_f): + self.train_inputs.append(self.train_2d[k2d][i]) + self.train_out.append(self.train_3d[k3d][i]) + 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 + # if act not in self.actions: + # continue + k3d = k2d + k3d = (sub, act, fname[:-3]) if fname.endswith('-sh') else k3d + num_f, _ = self.test_2d[k2d].shape + assert self.test_2d[k2d].shape[0] == self.test_3d[k3d].shape[0], '(test) 3d and 2d shapes not matching' + for i in range(num_f): + self.test_inputs.append(self.test_2d[k2d][i]) + self.test_out.append(self.test_3d[k3d][i]) + + def __getitem__(self, index) -> Any: + if self.train: + x = torch.from_numpy(self.train_inputs[index]).float() + y = torch.from_numpy(self.train_out[index]).float() + + else: + x = torch.from_numpy(self.test_inputs[index]).float() + y = torch.from_numpy(self.test_out[index]).float() + + return x, y + + def __len__(self): + if self.train: + return len(self.train_inputs) + else: + return len(self.test_inputs) \ No newline at end of file From 1fd09139a47348e167cfbc9bf1eb3777471f02bf Mon Sep 17 00:00:00 2001 From: Michail Kaseris Date: Mon, 15 Jan 2024 15:32:57 +0200 Subject: [PATCH 2/5] Return the samples as sequences --- src/skelcast/data/dataset.py | 30 +++++++++++++++++------------- 1 file changed, 17 insertions(+), 13 deletions(-) diff --git a/src/skelcast/data/dataset.py b/src/skelcast/data/dataset.py index 38b0628..be751b9 100644 --- a/src/skelcast/data/dataset.py +++ b/src/skelcast/data/dataset.py @@ -309,8 +309,8 @@ def __init__(self, actions, data_path, use_hourglass_detections=True, train=True self.use_hourglass_detections = use_hourglass_detections self.train = train - self.train_inputs, self.train_out, self.test_inputs, self.test_out = [], [], [], [] - self.train_meta, self.test_meta = [], [] + self.train_inputs, self.test_inputs = [], [] + # self.train_meta, self.test_meta = [], [] if self.use_hourglass_detections: train_2d_file = 'train_2d_ft.pth.tar' @@ -327,11 +327,13 @@ def __init__(self, actions, data_path, use_hourglass_detections=True, train=True (sub, act, fname) = k2d k3d = k2d k3d = (sub, act, fname[:-3]) if fname.endswith('-sh') else k3d - num_f, _ = self.train_2d[k2d].shape + # num_f, _ = self.train_2d[k2d].shape assert self.train_3d[k3d].shape[0] == self.train_2d[k2d].shape[0], f'(training) 3d and 2d shapes not matching' - for i in range(num_f): - self.train_inputs.append(self.train_2d[k2d][i]) - self.train_out.append(self.train_3d[k3d][i]) + # for i in range(num_f): + # self.train_inputs.append(self.train_2d[k2d][i]) + # self.train_out.append(self.train_3d[k3d][i]) + self.train_inputs.append(self.train_3d[k3d]) + 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)) @@ -341,22 +343,24 @@ def __init__(self, actions, data_path, use_hourglass_detections=True, train=True # continue k3d = k2d k3d = (sub, act, fname[:-3]) if fname.endswith('-sh') else k3d - num_f, _ = self.test_2d[k2d].shape + # num_f, _ = self.test_2d[k2d].shape assert self.test_2d[k2d].shape[0] == self.test_3d[k3d].shape[0], '(test) 3d and 2d shapes not matching' - for i in range(num_f): - self.test_inputs.append(self.test_2d[k2d][i]) - self.test_out.append(self.test_3d[k3d][i]) + self.test_inputs.append(self.test_3d[k3d]) + # for i in range(num_f): + # self.test_inputs.append(self.test_2d[k2d][i]) + # self.test_out.append(self.test_3d[k3d][i]) 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() - y = torch.from_numpy(self.train_out[index]).float() else: x = torch.from_numpy(self.test_inputs[index]).float() - y = torch.from_numpy(self.test_out[index]).float() + # y = torch.from_numpy(self.test_out[index]).float() - return x, y + return x.view(-1, 16, 3) def __len__(self): if self.train: From 061b30aeba8658e8faec9db6f021030d61e0730a Mon Sep 17 00:00:00 2001 From: Michail Kaseris Date: Mon, 15 Jan 2024 15:34:32 +0200 Subject: [PATCH 3/5] Cleanup. --- src/skelcast/data/dataset.py | 14 -------------- 1 file changed, 14 deletions(-) diff --git a/src/skelcast/data/dataset.py b/src/skelcast/data/dataset.py index be751b9..8e1b728 100644 --- a/src/skelcast/data/dataset.py +++ b/src/skelcast/data/dataset.py @@ -310,7 +310,6 @@ def __init__(self, actions, data_path, use_hourglass_detections=True, train=True self.train = train self.train_inputs, self.test_inputs = [], [] - # self.train_meta, self.test_meta = [], [] if self.use_hourglass_detections: train_2d_file = 'train_2d_ft.pth.tar' @@ -327,11 +326,7 @@ def __init__(self, actions, data_path, use_hourglass_detections=True, train=True (sub, act, fname) = k2d k3d = k2d k3d = (sub, act, fname[:-3]) if fname.endswith('-sh') else k3d - # num_f, _ = self.train_2d[k2d].shape assert self.train_3d[k3d].shape[0] == self.train_2d[k2d].shape[0], f'(training) 3d and 2d shapes not matching' - # for i in range(num_f): - # self.train_inputs.append(self.train_2d[k2d][i]) - # self.train_out.append(self.train_3d[k3d][i]) self.train_inputs.append(self.train_3d[k3d]) else: @@ -339,27 +334,18 @@ def __init__(self, actions, data_path, use_hourglass_detections=True, train=True self.test_2d = torch.load(os.path.join(data_path, test_2d_file)) for k2d in self.test_2d.keys(): (sub, act, fname) = k2d - # if act not in self.actions: - # continue k3d = k2d k3d = (sub, act, fname[:-3]) if fname.endswith('-sh') else k3d - # num_f, _ = self.test_2d[k2d].shape 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]) - # for i in range(num_f): - # self.test_inputs.append(self.test_2d[k2d][i]) - # self.test_out.append(self.test_3d[k3d][i]) 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() - # y = torch.from_numpy(self.test_out[index]).float() - return x.view(-1, 16, 3) def __len__(self): From c92b54614268a744e91ec1778669947b45ea0e32 Mon Sep 17 00:00:00 2001 From: Michail Kaseris Date: Mon, 15 Jan 2024 15:35:25 +0200 Subject: [PATCH 4/5] Do we actually care about the actions? We are doing forecasting --- src/skelcast/data/dataset.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/skelcast/data/dataset.py b/src/skelcast/data/dataset.py index 8e1b728..d26570c 100644 --- a/src/skelcast/data/dataset.py +++ b/src/skelcast/data/dataset.py @@ -303,8 +303,7 @@ def store_to_cache(self, cache_file: str) -> None: @DATASETS.register_module() class Human36mDataset(Dataset): - def __init__(self, actions, data_path, use_hourglass_detections=True, train=True) -> None: - self.actions = actions + 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 From cf2807f5f1736d99ebae55ec94f1116d12b7e69a Mon Sep 17 00:00:00 2001 From: Michail Kaseris Date: Mon, 15 Jan 2024 15:42:26 +0200 Subject: [PATCH 5/5] Return the sequence action --- src/skelcast/data/dataset.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/skelcast/data/dataset.py b/src/skelcast/data/dataset.py index d26570c..8a03201 100644 --- a/src/skelcast/data/dataset.py +++ b/src/skelcast/data/dataset.py @@ -309,6 +309,7 @@ def __init__(self, data_path, use_hourglass_detections=True, train=True) -> None self.train = train self.train_inputs, self.test_inputs = [], [] + self.act = [] if self.use_hourglass_detections: train_2d_file = 'train_2d_ft.pth.tar' @@ -327,6 +328,7 @@ def __init__(self, data_path, use_hourglass_detections=True, train=True) -> None 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')) @@ -337,6 +339,7 @@ def __init__(self, data_path, use_hourglass_detections=True, train=True) -> None 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: @@ -345,7 +348,7 @@ def __getitem__(self, index) -> Any: 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) + return x.view(-1, 16, 3), self.act[index] def __len__(self): if self.train: