| 72 | |
| 73 | |
| 74 | class PKLMotionLoader: |
| 75 | def __init__( |
| 76 | self, |
| 77 | motion_data: dict, |
| 78 | input_fps: int, |
| 79 | output_fps: int, |
| 80 | device: torch.device, |
| 81 | ): |
| 82 | self.motion_data = motion_data |
| 83 | self.input_fps = input_fps |
| 84 | self.output_fps = output_fps |
| 85 | self.input_dt = 1.0 / self.input_fps |
| 86 | self.output_dt = 1.0 / self.output_fps |
| 87 | self.current_idx = 0 |
| 88 | self.device = device |
| 89 | self._load_motion() |
| 90 | self._interpolate_motion() |
| 91 | self._compute_velocities() |
| 92 | |
| 93 | def _load_motion(self): |
| 94 | """Loads the motion from the motion data dict.""" |
| 95 | # assert self.motion_data["fps"] == self.input_fps, f"Input FPS Mismatch!, {self.motion_data['fps']} != {self.input_fps}" |
| 96 | |
| 97 | self.motion_base_poss_input = torch.from_numpy(self.motion_data["root_trans_offset"]).to(torch.float32).to(self.device) |
| 98 | |
| 99 | self.motion_base_rots_input = torch.from_numpy(self.motion_data["root_rot"]).to(torch.float32).to(self.device) |
| 100 | # motion_data['root_rot']: XYZW |
| 101 | self.motion_base_rots_input = self.motion_base_rots_input[:, [3, 0, 1, 2]] # convert xyzw to wxyz |
| 102 | |
| 103 | num_dof = self.motion_data["dof"].shape[1] |
| 104 | self.motion_dof_poss_input = torch.from_numpy(self.motion_data["dof"]).to(torch.float32).to(self.device) |
| 105 | if num_dof == 29: |
| 106 | # Everything good |
| 107 | pass |
| 108 | elif num_dof == 23: |
| 109 | # 29 = 15 + ( 4 + __3__)*2 |
| 110 | # = 19 + 3 + 4 + 3 |
| 111 | self.motion_dof_poss_input = torch.cat( |
| 112 | [ |
| 113 | self.motion_dof_poss_input[:, :19], |
| 114 | torch.zeros_like(self.motion_dof_poss_input[:, :3]), |
| 115 | self.motion_dof_poss_input[:, 19:23], |
| 116 | torch.zeros_like(self.motion_dof_poss_input[:, :3]), |
| 117 | ], |
| 118 | dim=1, |
| 119 | ) |
| 120 | assert self.motion_dof_poss_input.shape[1] == 29 |
| 121 | else: |
| 122 | raise NotImplementedError(f"Unsupported DOF count: {num_dof}") |
| 123 | |
| 124 | self.input_frames = self.motion_dof_poss_input.shape[0] |
| 125 | self.duration = (self.input_frames) * self.input_dt |
| 126 | print(f"Motion loaded, duration: {self.duration} sec, frames: {self.input_frames}") |
| 127 | |
| 128 | def _interpolate_motion(self): |
| 129 | """Interpolates the motion to the output fps.""" |
| 130 | times = torch.arange(0, self.duration, self.output_dt, device=self.device, dtype=torch.float32) |
| 131 | self.output_frames = times.shape[0] |