|
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 |
|
|
|
|
|
raw_rot = matrix_trajectory[:, :3, :3] |
|
rot6d = raw_rot[:, :, :2].permute(0, 2, 1).reshape(-1, 6) |
|
|
|
|
|
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 = 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) |
|
|
