Spaces:
Running
on
Zero
Running
on
Zero
| from pathlib import Path | |
| from evo.tools.file_interface import read_kitti_poses_file | |
| import numpy as np | |
| import torch | |
| from torch.utils.data import Dataset | |
| from torchtyping import TensorType | |
| import torch.nn.functional as F | |
| from typing import Tuple | |
| from utils.file_utils import load_txt | |
| from utils.rotation_utils import compute_rotation_matrix_from_ortho6d | |
| num_cams = None | |
| # ------------------------------------------------------------------------------------- # | |
| class TrajectoryDataset(Dataset): | |
| def __init__( | |
| self, | |
| name: str, | |
| set_name: str, | |
| dataset_dir: str, | |
| num_rawfeats: int, | |
| num_feats: int, | |
| num_cams: int, | |
| standardize: bool, | |
| **kwargs, | |
| ): | |
| super().__init__() | |
| self.name = name | |
| self.set_name = set_name | |
| self.dataset_dir = Path(dataset_dir) | |
| if name == "relative": | |
| self.data_dir = self.dataset_dir / "traj_raw" | |
| self.relative_dir = self.dataset_dir / "relative" | |
| else: | |
| self.data_dir = self.dataset_dir / "traj" | |
| self.intrinsics_dir = self.dataset_dir / "intrinsics" | |
| self.num_rawfeats = num_rawfeats | |
| self.num_feats = num_feats | |
| self.num_cams = num_cams | |
| self.augmentation = None | |
| self.standardize = standardize | |
| if self.standardize: | |
| mean_std = kwargs["standardization"] | |
| self.norm_mean = torch.Tensor(mean_std["norm_mean"]) | |
| self.norm_std = torch.Tensor(mean_std["norm_std"]) | |
| self.shift_mean = torch.Tensor(mean_std["shift_mean"]) | |
| self.shift_std = torch.Tensor(mean_std["shift_std"]) | |
| self.velocity = mean_std["velocity"] | |
| # --------------------------------------------------------------------------------- # | |
| def set_split(self, split: str, train_rate: float = 1.0): | |
| self.split = split | |
| split_path = Path(self.dataset_dir) / f"{split}_split.txt" | |
| split_traj = load_txt(split_path).split("\n") | |
| self.filenames = sorted(split_traj) | |
| return self | |
| # --------------------------------------------------------------------------------- # | |
| def get_feature( | |
| self, raw_matrix_trajectory: TensorType["num_cams", 4, 4] | |
| ) -> TensorType[9, "num_cams"]: | |
| matrix_trajectory = torch.clone(raw_matrix_trajectory) | |
| raw_trans = torch.clone(matrix_trajectory[:, :3, 3]) | |
| if self.velocity: | |
| velocity = raw_trans[1:] - raw_trans[:-1] | |
| raw_trans = torch.cat([raw_trans[0][None], velocity]) | |
| if self.standardize: | |
| raw_trans[0] -= self.shift_mean | |
| raw_trans[0] /= self.shift_std | |
| raw_trans[1:] -= self.norm_mean | |
| raw_trans[1:] /= self.norm_std | |
| # Compute the 6D continuous rotation | |
| raw_rot = matrix_trajectory[:, :3, :3] | |
| rot6d = raw_rot[:, :, :2].permute(0, 2, 1).reshape(-1, 6) | |
| # Stack rotation 6D and translation | |
| rot6d_trajectory = torch.hstack([rot6d, raw_trans]).permute(1, 0) | |
| return rot6d_trajectory | |
| def get_matrix( | |
| self, raw_rot6d_trajectory: TensorType[9, "num_cams"] | |
| ) -> TensorType["num_cams", 4, 4]: | |
| rot6d_trajectory = torch.clone(raw_rot6d_trajectory) | |
| device = rot6d_trajectory.device | |
| num_cams = rot6d_trajectory.shape[1] | |
| matrix_trajectory = torch.eye(4, device=device)[None].repeat(num_cams, 1, 1) | |
| raw_trans = rot6d_trajectory[6:].permute(1, 0) | |
| if self.standardize: | |
| raw_trans[0] *= self.shift_std.to(device) | |
| raw_trans[0] += self.shift_mean.to(device) | |
| raw_trans[1:] *= self.norm_std.to(device) | |
| raw_trans[1:] += self.norm_mean.to(device) | |
| if self.velocity: | |
| raw_trans = torch.cumsum(raw_trans, dim=0) | |
| matrix_trajectory[:, :3, 3] = raw_trans | |
| rot6d = rot6d_trajectory[:6].permute(1, 0) | |
| raw_rot = compute_rotation_matrix_from_ortho6d(rot6d) | |
| matrix_trajectory[:, :3, :3] = raw_rot | |
| return matrix_trajectory | |
| # --------------------------------------------------------------------------------- # | |
| def __getitem__(self, index: int) -> Tuple[str, TensorType["num_cams", 4, 4]]: | |
| filename = self.filenames[index] | |
| trajectory_filename = filename + ".txt" | |
| trajectory_path = self.data_dir / trajectory_filename | |
| trajectory = read_kitti_poses_file(trajectory_path) | |
| matrix_trajectory = torch.from_numpy(np.array(trajectory.poses_se3)).to( | |
| torch.float32 | |
| ) | |
| trajectory_feature = self.get_feature(matrix_trajectory) | |
| padded_trajectory_feature = F.pad( | |
| trajectory_feature, (0, self.num_cams - trajectory_feature.shape[1]) | |
| ) | |
| # Padding mask: 1 for valid cams, 0 for padded cams | |
| padding_mask = torch.ones((self.num_cams)) | |
| padding_mask[trajectory_feature.shape[1] :] = 0 | |
| intrinsics_filename = filename + ".npy" | |
| intrinsics_path = self.intrinsics_dir / intrinsics_filename | |
| intrinsics = np.load(intrinsics_path) | |
| return ( | |
| trajectory_filename, | |
| padded_trajectory_feature, | |
| padding_mask, | |
| intrinsics | |
| ) | |
| def __len__(self): | |
| return len(self.filenames) | |