add support et visualization

app.py

@@ -1,77 +1,93 @@
 import gradio as gr
-import numpy as np
-import torch
-import torch.nn.functional as F
-from evo.tools.file_interface import read_kitti_poses_file
-from pathlib import Path
+from gradio_rerun import Rerun
+from data.loader import load_simulation_data
 from visualization.visualizer import visualize_simulation
+from visualization.et_visualizer import visualize_et_data
+from pathlib import Path
 
 
-def load_trajectory_data(traj_file, char_file, num_cams=30):
-    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])
+def update_simulation_dropdown(file):
+    simulations, descriptions = load_simulation_data(file)
+    return gr.Dropdown(
+        choices=descriptions if descriptions else [],
+        value=None,
+        allow_custom_value=False
     )
-    
-    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
-    }
-
-def create_trajectory_interface():
-    """Create Gradio interface for loading trajectory data"""
-    
-    def process_files(traj_file, char_file):
-        try:
-            result = load_trajectory_data(traj_file.name, char_file.name)
-            
-            # Convert tensors to numpy for display
-            info = {
-                "Trajectory filename": result["traj_filename"],
-                "Character filename": result["char_filename"],
-                "Trajectory shape": result["traj_feat"].shape,
-                "Character shape": result["char_feat"].shape,
-                "Valid frames": int(result["padding_mask"].sum().item())
-            }
-            
-            return str(info)
-            
-        except Exception as e:
-            return f"Error processing files: {str(e)}"
-
-    interface = gr.Interface(
-        fn=process_files,
-        inputs=[
-            gr.File(label="Trajectory File (.txt)"),
-            gr.File(label="Character File (.npy)")
-        ],
-        outputs=gr.Textbox(label="Results"),
-        title="Trajectory Data Loader",
-        description="Upload trajectory (.txt) and character (.npy) files to load and process them."
-    )
-    
-    return interface
+
+
+def create_app():
+    with gr.Blocks() as demo:
+        with gr.Tabs() as tabs:
+            # Camera Simulation Tab
+            with gr.Tab("Camera Simulation"):
+                gr.Markdown("""
+                # Camera Simulation Visualizer
+                Upload a JSON file containing camera simulation data and select a simulation to visualize.
+                """)
+
+                with gr.Row():
+                    file_input = gr.File(
+                        label="Upload Simulation JSON",
+                        file_types=[".json"]
+                    )
+                    simulation_dropdown = gr.Dropdown(
+                        label="Select Simulation",
+                        choices=[],
+                        type="index",
+                        scale=2
+                    )
+
+                with gr.Row():
+                    viewer = Rerun(streaming=False)
+
+                file_input.change(
+                    update_simulation_dropdown,
+                    inputs=[file_input],
+                    outputs=[simulation_dropdown]
+                )
+
+                simulation_dropdown.change(
+                    visualize_simulation,
+                    inputs=[file_input, simulation_dropdown],
+                    outputs=[viewer]
+                )
+
+            # E.T. Dataset Tab
+            with gr.Tab("E.T. Dataset"):
+                gr.Markdown("""
+                # E.T. Dataset Visualizer
+                Upload trajectory (.txt) and character (.npy) files to visualize the E.T. dataset.
+                """)
+
+                with gr.Row():
+                    traj_file = gr.File(
+                        label="Trajectory File (.txt)",
+                        file_types=[".txt"]
+                    )
+                    char_file = gr.File(
+                        label="Character File (.npy)",
+                        file_types=[".npy"]
+                    )
+
+                with gr.Row():
+                    et_viewer = Rerun(streaming=False)
+
+                def process_et_files(traj_file, char_file):
+                    if traj_file is None or char_file is None:
+                        return None
+                    return visualize_et_data(traj_file.name, char_file.name)
+
+                with gr.Row():
+                    visualize_btn = gr.Button("Visualize")
+                    visualize_btn.click(
+                        process_et_files,
+                        inputs=[traj_file, char_file],
+                        outputs=[et_viewer]
+                    )
+
+    return demo
+
 
 if __name__ == "__main__":
-    interface = create_trajectory_interface()
-    interface.launch()
+    demo = create_app()
+    demo.queue().launch(share=False)

visualization/et_visualizer.py

@@ -0,0 +1,131 @@
+import tempfile
+import os
+import spaces
+import numpy as np
+import torch
+import torch.nn.functional as F
+from evo.tools.file_interface import read_kitti_poses_file
+from pathlib import Path
+import rerun as rr
+from typing import Optional, Dict
+from visualization.logger import SimulationLogger
+
+
+def load_trajectory_data(traj_file: str, char_file: str, num_cams: int = 30) -> 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]
+
+    # Convert to 6D rotation representation
+    rot6d = raw_rot[:, :, :2].permute(0, 2, 1).reshape(-1, 6)
+    trajectory_feature = torch.hstack([rot6d, raw_trans]).permute(1, 0)
+
+    # Pad trajectory features
+    padded_trajectory_feature = F.pad(
+        trajectory_feature,
+        (0, num_cams - trajectory_feature.shape[1])
+    )
+
+    # Create padding mask
+    padding_mask = torch.ones((num_cams))
+    padding_mask[trajectory_feature.shape[1]:] = 0
+
+    # Load and pad character features
+    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
+    }
+
+
+class ETLogger(SimulationLogger):
+    def __init__(self):
+        super().__init__()
+        rr.init("et_visualization")
+        rr.log("world", rr.ViewCoordinates.RIGHT_HAND_Y_UP, timeless=True)
+
+    def log_trajectory(self, trajectory: np.ndarray, padding_mask: np.ndarray):
+        """Log camera trajectory."""
+        valid_frames = int(padding_mask.sum())
+        valid_trajectory = trajectory[:valid_frames]
+
+        # Log trajectory points
+        positions = valid_trajectory[:, :3, 3]
+        rr.log(
+            "world/trajectory/points",
+            rr.Points3D(
+                positions,
+                colors=np.full((len(positions), 4), [0.0, 0.8, 0.8, 1.0])
+            ),
+            timeless=True
+        )
+
+        # Log trajectory line
+        if len(positions) > 1:
+            lines = np.stack([positions[:-1], positions[1:]], axis=1)
+            rr.log(
+                "world/trajectory/line",
+                rr.LineStrips3D(
+                    lines,
+                    colors=[(0.0, 0.8, 0.8, 1.0)]
+                ),
+                timeless=True
+            )
+
+    def log_character(self, char_feature: np.ndarray, padding_mask: np.ndarray):
+        """Log character feature visualization."""
+        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
+            )
+
+
+@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