fix et issue

visualization/et_visualizer.py

@@ -1,3 +1,6 @@
+import tempfile
+import os
+import spaces
 import numpy as np
 import torch
 import torch.nn.functional as F
@@ -7,7 +10,6 @@ import rerun as rr
 from typing import Optional, Dict
 from visualization.logger import SimulationLogger
 from scipy.spatial.transform import Rotation
-from rerun.components import Material
 
 
 def load_trajectory_data(traj_file: str, char_file: str, num_cams: int = 30) -> Dict:
@@ -53,137 +55,113 @@ class ETLogger(SimulationLogger):
         super().__init__()
         rr.init("et_visualization")
         rr.log("world", rr.ViewCoordinates.RIGHT_HAND_Y_UP, timeless=True)
-
-        # Define default camera parameters
-        self.camera_width = 640  # default width
-        self.camera_height = 480  # default height
-        self.focal_length = 500  # default focal length
+        # Default camera intrinsics
         self.K = np.array([
-            [self.focal_length, 0, self.camera_width/2],
-            [0, self.focal_length, self.camera_height/2],
+            [500, 0, 320],
+            [0, 500, 240],
             [0, 0, 1]
         ])
 
     def log_trajectory(self, trajectory: np.ndarray, padding_mask: np.ndarray):
-        """Log camera trajectory with enhanced visualization."""
+        """Log camera trajectory."""
         valid_frames = int(padding_mask.sum())
         valid_trajectory = trajectory[:valid_frames]
 
-        # Log trajectory points with rainbow coloring
+        # Log trajectory points
         positions = valid_trajectory[:, :3, 3]
-        colors = np.zeros((len(positions), 4))
-        colors[:, :3] = plt.cm.rainbow(
-            np.linspace(0, 1, len(positions)))[:, :3]
-        colors[:, 3] = 1.0  # Set alpha to 1
-
         rr.log(
             "world/trajectory/points",
             rr.Points3D(
                 positions,
-                colors=colors
+                colors=np.full((len(positions), 4), [0.0, 0.8, 0.8, 1.0])
             ),
             timeless=True
         )
 
-        # Log trajectory line with gradient color
+        # Log trajectory line
         if len(positions) > 1:
             lines = np.stack([positions[:-1], positions[1:]], axis=1)
-            line_colors = np.zeros((len(lines), 4))
-            line_colors[:, :3] = plt.cm.rainbow(
-                np.linspace(0, 1, len(lines)))[:, :3]
-            line_colors[:, 3] = 1.0
-
             rr.log(
                 "world/trajectory/line",
                 rr.LineStrips3D(
                     lines,
-                    colors=line_colors
+                    colors=[(0.0, 0.8, 0.8, 1.0)]
                 ),
                 timeless=True
             )
 
-        # Log camera frustums
-        for i in range(valid_frames):
-            # Get camera position and rotation
-            translation = valid_trajectory[i, :3, 3]
-            rotation_matrix = valid_trajectory[i, :3, :3]
-            rotation_quat = Rotation.from_matrix(rotation_matrix).as_quat()
+        # Log cameras
+        for k in range(valid_frames):
+            # Set time sequence
+            rr.set_time_sequence("frame_idx", k)
 
-            # Set time sequence for animation
-            rr.set_time_sequence("frame_idx", i)
+            # Get camera pose
+            translation = valid_trajectory[k, :3, 3]
+            rotation_q = Rotation.from_matrix(
+                valid_trajectory[k, :3, :3]).as_quat()
 
-            # Log camera frustum
+            # Log camera transform
             rr.log(
-                f"world/cameras/camera_{i}",
+                f"world/camera",
                 rr.Transform3D(
                     translation=translation,
-                    rotation=rr.Quaternion(xyzw=rotation_quat),
-                )
+                    rotation=rr.Quaternion(xyzw=rotation_q),
+                ),
             )
 
-            # Add camera visualization
+            # Log camera frustum
             rr.log(
-                f"world/cameras/camera_{i}/frustum",
+                f"world/camera/image",
                 rr.Pinhole(
                     image_from_camera=self.K,
-                    width=self.camera_width,
-                    height=self.camera_height,
-                    focal_length=self.focal_length,
+                    width=640,
+                    height=480,
                 ),
             )
 
-            # Add coordinate axes for each camera
-            rr.log(
-                f"world/cameras/camera_{i}/axes",
-                rr.Arrows3D(
-                    origins=np.zeros((3, 3)),
-                    vectors=np.eye(3) * 0.5,  # 0.5 meter long axes
-                    colors=[[1, 0, 0, 1], [0, 1, 0, 1], [
-                        0, 0, 1, 1]]  # RGB colors for XYZ
-                )
-            )
-
     def log_character(self, char_feature: np.ndarray, padding_mask: np.ndarray):
-        """Log character feature visualization with enhanced appearance."""
+        """Log character feature visualization."""
         valid_frames = int(padding_mask.sum())
         valid_char = char_feature[:, :valid_frames]
 
         if valid_char.shape[0] > 0:
-            # Create gradient colors for character points
-            num_points = valid_char.reshape(-1, 3).shape[0]
-            colors = np.zeros((num_points, 4))
-            colors[:, 0] = 0.8  # Red component
-            colors[:, 1] = 0.2  # Green component
-            colors[:, 2] = np.linspace(0.2, 0.8, num_points)  # Blue gradient
-            colors[:, 3] = 1.0  # Alpha
-
             rr.log(
                 "world/character",
                 rr.Points3D(
                     valid_char.reshape(-1, 3),
-                    colors=colors,
-                    radii=0.05  # Add point size for better visibility
+                    colors=np.full(
+                        (valid_char.reshape(-1, 3).shape[0], 4), [0.8, 0.2, 0.2, 1.0])
                 ),
                 timeless=True
             )
 
-            # Add a semi-transparent hull around character points
-            try:
-                from scipy.spatial import ConvexHull
-                points = valid_char.reshape(-1, 3)
-                hull = ConvexHull(points)
-
-                rr.log(
-                    "world/character/hull",
-                    rr.Mesh3D(
-                        vertex_positions=points[hull.vertices],
-                        indices=hull.simplices,
-                        mesh_material=Material(
-                            # Semi-transparent red
-                            albedo_factor=[0.8, 0.2, 0.2, 0.3]
-                        )
-                    ),
-                    timeless=True
-                )
-            except Exception:
-                pass  # Skip hull visualization if it fails
+
+@spaces.GPU
+def visualize_et_data(traj_file: str, char_file: str) -> Optional[str]:
+    """Visualize E.T. dataset using Rerun."""
+    try:
+        # Load data
+        data = load_trajectory_data(traj_file, char_file)
+
+        # Create temporary file for RRD
+        temp_dir = tempfile.mkdtemp()
+        rrd_path = os.path.join(temp_dir, "et_visualization.rrd")
+
+        # Initialize logger and log data
+        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()
+        )
+
+        # Save visualization
+        rr.save(rrd_path)
+        return rrd_path
+
+    except Exception as e:
+        print(f"Error visualizing E.T. data: {str(e)}")
+        return None