update et visualization

visualization/et_visualizer.py

@@ -1,6 +1,3 @@
-import tempfile
-import os
-import spaces
 import numpy as np
 import torch
 import torch.nn.functional as F
@@ -9,6 +6,8 @@ from pathlib import Path
 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:
@@ -55,77 +54,136 @@ class ETLogger(SimulationLogger):
         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
+        self.K = np.array([
+            [self.focal_length, 0, self.camera_width/2],
+            [0, self.focal_length, self.camera_height/2],
+            [0, 0, 1]
+        ])
+
     def log_trajectory(self, trajectory: np.ndarray, padding_mask: np.ndarray):
-        """Log camera trajectory."""
+        """Log camera trajectory with enhanced visualization."""
         valid_frames = int(padding_mask.sum())
         valid_trajectory = trajectory[:valid_frames]
 
-        # Log trajectory points
+        # Log trajectory points with rainbow coloring
         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=np.full((len(positions), 4), [0.0, 0.8, 0.8, 1.0])
+                colors=colors
             ),
             timeless=True
         )
 
-        # Log trajectory line
+        # Log trajectory line with gradient color
         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=[(0.0, 0.8, 0.8, 1.0)]
+                    colors=line_colors
                 ),
                 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()
+
+            # Set time sequence for animation
+            rr.set_time_sequence("frame_idx", i)
+
+            # Log camera frustum
+            rr.log(
+                f"world/cameras/camera_{i}",
+                rr.Transform3D(
+                    translation=translation,
+                    rotation=rr.Quaternion(xyzw=rotation_quat),
+                )
+            )
+
+            # Add camera visualization
+            rr.log(
+                f"world/cameras/camera_{i}/frustum",
+                rr.Pinhole(
+                    image_from_camera=self.K,
+                    width=self.camera_width,
+                    height=self.camera_height,
+                    focal_length=self.focal_length,
+                ),
+            )
+
+            # 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."""
+        """Log character feature visualization with enhanced appearance."""
         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=np.full(
-                        (valid_char.reshape(-1, 3).shape[0], 4), [0.8, 0.2, 0.2, 1.0])
+                    colors=colors,
+                    radii=0.05  # Add point size for better visibility
                 ),
                 timeless=True
             )
 
-
-@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
+            # 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