remove padding

visualization/et_visualizer.py

@@ -12,37 +12,18 @@ from visualization.logger import SimulationLogger
 from scipy.spatial.transform import Rotation
 
 
-def load_trajectory_data(traj_file: str, char_file: str, num_cams: int = 30) -> Dict:
+def load_trajectory_data(traj_file: str, char_file: str) -> Dict:
     trajectory = read_kitti_poses_file(traj_file)
     matrix_trajectory = torch.from_numpy(
         np.array(trajectory.poses_se3)).to(torch.float32)
 
-    raw_trans = torch.clone(matrix_trajectory[:, :3, 3])
-    raw_rot = matrix_trajectory[:, :3, :3]
-
-    rot6d = raw_rot[:, :, :2].permute(0, 2, 1).reshape(-1, 6)
-    trajectory_feature = torch.hstack([rot6d, raw_trans]).permute(1, 0)
-
-    padded_trajectory_feature = F.pad(
-        trajectory_feature,
-        (0, num_cams - trajectory_feature.shape[1])
-    )
-
-    padding_mask = torch.ones((num_cams))
-    padding_mask[trajectory_feature.shape[1]:] = 0
-
     char_feature = torch.from_numpy(np.load(char_file)).to(torch.float32)
-    padding_size = num_cams - char_feature.shape[0]
-    padded_char_feature = F.pad(
-        char_feature, (0, 0, 0, padding_size)).permute(1, 0)
 
     return {
         "traj_filename": Path(traj_file).name,
         "char_filename": Path(char_file).name,
-        "traj_feat": padded_trajectory_feature,
-        "char_feat": padded_char_feature,
-        "padding_mask": padding_mask,
-        "raw_matrix_trajectory": matrix_trajectory
+        "char_feat": char_feature,
+        "matrix_trajectory": matrix_trajectory
     }
 
 
@@ -58,11 +39,8 @@ class ETLogger(SimulationLogger):
             [0, 0, 1]
         ])
 
-    def log_trajectory(self, trajectory: np.ndarray, padding_mask: np.ndarray):
-        valid_frames = int(padding_mask.sum())
-        valid_trajectory = trajectory[:valid_frames]
-
-        positions = valid_trajectory[:, :3, 3]
+    def log_trajectory(self, trajectory: np.ndarray):
+        positions = trajectory[:, :3, 3]
         rr.log(
             "world/trajectory/points",
             rr.Points3D(
@@ -83,13 +61,12 @@ class ETLogger(SimulationLogger):
                 timeless=True
             )
 
-        for k in range(valid_frames):
-
+        for k in range(len(trajectory)):
             rr.set_time_sequence("frame_idx", k)
 
-            translation = valid_trajectory[k, :3, 3]
+            translation = trajectory[k, :3, 3]
             rotation_q = Rotation.from_matrix(
-                valid_trajectory[k, :3, :3]).as_quat()
+                trajectory[k, :3, :3]).as_quat()
 
             rr.log(
                 f"world/camera",
@@ -108,20 +85,16 @@ class ETLogger(SimulationLogger):
                 ),
             )
 
-    def log_character(self, char_feature: np.ndarray, padding_mask: np.ndarray):
-        valid_frames = int(padding_mask.sum())
-        valid_char = char_feature[:, :valid_frames]
-
-        if valid_char.shape[0] > 0:
-            rr.log(
-                "world/character",
-                rr.Points3D(
-                    valid_char.reshape(-1, 3),
-                    colors=np.full(
-                        (valid_char.reshape(-1, 3).shape[0], 4), [0.8, 0.2, 0.2, 1.0])
-                ),
-                timeless=True
-            )
+    def log_character(self, char_feature: np.ndarray):
+        rr.log(
+            "world/character",
+            rr.Points3D(
+                char_feature.reshape(-1, 3),
+                colors=np.full(
+                    (char_feature.reshape(-1, 3).shape[0], 4), [0.8, 0.2, 0.2, 1.0])
+            ),
+            timeless=True
+        )
 
 
 @spaces.GPU
@@ -134,14 +107,8 @@ def visualize_et_data(traj_file: str, char_file: str) -> Optional[str]:
         rrd_path = os.path.join(temp_dir, "et_visualization.rrd")
 
         logger = ETLogger()
-        logger.log_trajectory(
-            data["raw_matrix_trajectory"].numpy(),
-            data["padding_mask"].numpy()
-        )
-        logger.log_character(
-            data["char_feat"].numpy(),
-            data["padding_mask"].numpy()
-        )
+        logger.log_trajectory(data["matrix_trajectory"].numpy())
+        logger.log_character(data["char_feat"].numpy())
 
         rr.save(rrd_path)
         return rrd_path