refactor and organize files

app.py

@@ -1,230 +1,50 @@
-import json
-import numpy as np
-import rerun as rr
-import spaces
 import gradio as gr
 from gradio_rerun import Rerun
-from scipy.spatial.transform import Rotation
-import tempfile
-import os
-from typing import Optional, Dict, Any, List, Tuple
+from data.loader import load_simulation_data
+from visualization.visualizer import visualize_simulation
 
 
-def vector3_to_numpy(vec):
-    """Convert Vector3 dictionary to numpy array"""
-    return np.array([vec['x'], vec['y'], vec['z']])
-
-
-def euler_to_quaternion(euler):
-    """Convert Euler angles dictionary to quaternion"""
-    return Rotation.from_euler('xyz', [euler['x'], euler['y'], euler['z']]).as_quat()
-
-
-def create_subject_mesh(subject):
-    """Create a simple cube mesh for a subject"""
-    position = vector3_to_numpy(subject['position'])
-    size = vector3_to_numpy(subject['size'])
-
-    # Create cube vertices
-    vertices = np.array([
-        [-0.5, -0.5, -0.5], [0.5, -0.5, -0.5], [0.5, 0.5, -0.5], [-0.5, 0.5, -0.5],
-        [-0.5, -0.5, 0.5], [0.5, -0.5, 0.5], [0.5, 0.5, 0.5], [-0.5, 0.5, 0.5]
-    ]) * size.reshape(1, 3) + position.reshape(1, 3)
-
-    # Create cube faces
-    faces = np.array([
-        [0, 1, 2], [0, 2, 3],  # front
-        [1, 5, 6], [1, 6, 2],  # right
-        [5, 4, 7], [5, 7, 6],  # back
-        [4, 0, 3], [4, 3, 7],  # left
-        [3, 2, 6], [3, 6, 7],  # top
-        [4, 5, 1], [4, 1, 0]   # bottom
-    ])
-
-    return vertices, faces
-
-
-def log_simulation(simulation_data: Dict[str, Any]) -> None:
-    """Log single simulation data to Rerun"""
-    rr.init("camera_simulation")
-
-    subjects = simulation_data['subjects']
-    camera_frames = simulation_data['cameraFrames']
-    instructions = simulation_data['instructions']
-
-    # Log simulation metadata
-    rr.log("metadata/instructions", rr.TextDocument(
-        "\n".join([
-            f"Instruction {i+1}:\n" +
-            f"  Movement: {inst['cameraMovement']}\n" +
-            f"  Easing: {inst['movementEasing']}\n" +
-            f"  Frames: {inst['frameCount']}\n" +
-            f"  Camera Angle: {inst.get('initialCameraAngle', 'N/A')}\n" +
-            f"  Shot Type: {inst.get('initialShotType', 'N/A')}\n" +
-            f"  Subject Index: {inst.get('subjectIndex', 'N/A')}"
-            for i, inst in enumerate(instructions)
-        ])
-    ), timeless=True)
-
-    # Set up world coordinate system
-    rr.log("world", rr.ViewCoordinates.RIGHT_HAND_Y_UP, timeless=True)
-
-    # Log subjects (as simple cubes)
-    for idx, subject in enumerate(subjects):
-        vertices, faces = create_subject_mesh(subject)
-        subject_color = [0.8, 0.2, 0.2, 1.0] if idx == simulation_data.get(
-            'selectedSubject') else [0.8, 0.8, 0.8, 1.0]
-
-        rr.log(
-            f"world/subject_{idx}",
-            rr.Mesh3D(
-                vertex_positions=vertices,
-                indices=faces,
-                # Apply color to all vertices
-                colors=np.tile(subject_color, (len(vertices), 1))
-            ),
-            timeless=True
-        )
-
-        # Log subject class
-        rr.log(f"world/subject_{idx}/class",
-               rr.TextDocument(subject['objectClass']),
-               timeless=True)
-
-    # Log camera trajectory
-    camera_positions = np.array(
-        [vector3_to_numpy(frame['position']) for frame in camera_frames])
-    rr.log(
-        "world/camera_trajectory",
-        rr.Points3D(
-            camera_positions,
-            # Cyan color for trajectory
-            colors=np.full((len(camera_positions), 4), [0.0, 0.8, 0.8, 1.0])
-        ),
-        timeless=True
-    )
+def update_simulation_dropdown(file):
+    _, descriptions = load_simulation_data(file)
+    return gr.Dropdown(choices=descriptions if descriptions else [], value=None)
 
-    # Log camera movement over time
-    for frame_idx, camera_frame in enumerate(camera_frames):
-        rr.set_time_sequence("frame", frame_idx)
 
-        position = vector3_to_numpy(camera_frame['position'])
-        rotation_q = euler_to_quaternion(camera_frame['angle'])
+def create_app():
+    with gr.Blocks() as demo:
+        gr.Markdown("""
+        # Camera Simulation Visualizer
+        Upload a JSON file containing camera simulation data and select a simulation to visualize.
+        """)
 
-        # Log camera transform
-        rr.log(
-            "world/camera",
-            rr.Transform3D(
-                translation=position,
-                rotation=rr.Quaternion(xyzw=rotation_q)
+        with gr.Row():
+            file_input = gr.File(
+                label="Upload Simulation JSON",
+                file_types=[".json"]
             )
-        )
-
-        # Log camera frustum
-        rr.log(
-            "world/camera/view",
-            rr.Pinhole(
-                focal_length=camera_frame['focalLength'],
-                width=1920,
-                height=1080
+            simulation_dropdown = gr.Dropdown(
+                label="Select Simulation",
+                choices=[],
+                type="index"
             )
-        )
 
-        # Log frame number
-        rr.log(
-            "metadata/current_frame",
-            rr.TextDocument(f"Frame: {frame_idx + 1}/{len(camera_frames)}"),
-        )
-
-
-def load_simulation_data(file) -> Tuple[Optional[List[Dict[str, Any]]], Optional[List[str]]]:
-    """Load simulation data from JSON file and return simulations with their descriptions"""
-    if file is None:
-        return None, None
-
-    try:
-        json_data = json.load(open(file.name))
-        simulations = json_data['simulations']
-
-        # Create descriptions for each simulation
-        descriptions = [
-            f"Simulation {i}: {len(sim['subjects'])} subjects, {len(sim['instructions'])} instructions"
-            for i, sim in enumerate(simulations)
-        ]
-
-        return simulations, descriptions
-    except Exception as e:
-        print(f"Error loading simulation data: {str(e)}")
-        return None, None
-
-
-@spaces.GPU
-def visualize_simulation(file, simulation_index: int) -> Optional[str]:
-    """Process selected simulation and create Rerun visualization"""
-    if file is None:
-        return None
-
-    try:
-        simulations, _ = load_simulation_data(file)
-        if simulations is None or simulation_index >= len(simulations):
-            return None
-
-        # Create temporary file for RRD
-        temp_dir = tempfile.mkdtemp()
-        rrd_path = os.path.join(temp_dir, "simulation.rrd")
-
-        # Log selected simulation
-        simulation = simulations[simulation_index]
-        log_simulation(simulation)
-        rr.save(rrd_path)
+        with gr.Row():
+            viewer = Rerun(streaming=False)
 
-        return rrd_path
-
-    except Exception as e:
-        print(f"Error processing simulation: {str(e)}")
-        return None
-
-
-def update_simulation_dropdown(file):
-    """Update simulation dropdown when file is uploaded"""
-    _, descriptions = load_simulation_data(file)
-    return gr.Dropdown(choices=descriptions if descriptions else [], value=None)
-
-
-# Create Gradio interface
-with gr.Blocks() as demo:
-    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"
+        file_input.change(
+            update_simulation_dropdown,
+            inputs=[file_input],
+            outputs=[simulation_dropdown]
         )
 
-    with gr.Row():
-        viewer = Rerun(streaming=False)
+        simulation_dropdown.change(
+            visualize_simulation,
+            inputs=[file_input, simulation_dropdown],
+            outputs=[viewer]
+        )
 
-    # Update dropdown when file is uploaded
-    file_input.change(
-        update_simulation_dropdown,
-        inputs=[file_input],
-        outputs=[simulation_dropdown]
-    )
+    return demo
 
-    # Visualize selected simulation
-    simulation_dropdown.change(
-        visualize_simulation,
-        inputs=[file_input, simulation_dropdown],
-        outputs=[viewer]
-    )
 
 if __name__ == "__main__":
+    demo = create_app()
     demo.queue().launch(share=False)

data/loader.py

@@ -0,0 +1,21 @@
+import json
+from typing import Optional, Dict, Any, List, Tuple
+
+
+def load_simulation_data(file) -> Tuple[Optional[List[Dict[str, Any]]], Optional[List[str]]]:
+    if file is None:
+        return None, None
+
+    try:
+        json_data = json.load(open(file.name))
+        simulations = json_data['simulations']
+
+        descriptions = [
+            f"Simulation {i}: {len(sim['subjects'])} subjects, {len(sim['instructions'])} instructions"
+            for i, sim in enumerate(simulations)
+        ]
+
+        return simulations, descriptions
+    except Exception as e:
+        print(f"Error loading simulation data: {str(e)}")
+        return None, None

utils/geometry.py

@@ -0,0 +1,11 @@
+import numpy as np
+from scipy.spatial.transform import Rotation
+from typing import Dict
+
+
+def vector3_to_numpy(vec: Dict[str, float]) -> np.ndarray:
+    return np.array([vec['x'], vec['y'], vec['z']])
+
+
+def euler_to_quaternion(euler: Dict[str, float]) -> np.ndarray:
+    return Rotation.from_euler('xyz', [euler['x'], euler['y'], euler['z']]).as_quat()

visualization/logger.py

@@ -0,0 +1,88 @@
+import rerun as rr
+import numpy as np
+from typing import Dict, Any, List
+from ..utils.geometry import vector3_to_numpy, euler_to_quaternion
+from .mesh import create_subject_mesh
+
+
+class SimulationLogger:
+    def __init__(self):
+        rr.init("camera_simulation")
+        rr.log("world", rr.ViewCoordinates.RIGHT_HAND_Y_UP, timeless=True)
+
+    def log_metadata(self, instructions: List[Dict[str, Any]]) -> None:
+        rr.log("metadata/instructions", rr.TextDocument(
+            "\n".join([
+                f"Instruction {i+1}:\n" +
+                f"  Movement: {inst['cameraMovement']}\n" +
+                f"  Easing: {inst['movementEasing']}\n" +
+                f"  Frames: {inst['frameCount']}\n" +
+                f"  Camera Angle: {inst.get('initialCameraAngle', 'N/A')}\n" +
+                f"  Shot Type: {inst.get('initialShotType', 'N/A')}\n" +
+                f"  Subject Index: {inst.get('subjectIndex', 'N/A')}"
+                for i, inst in enumerate(instructions)
+            ])
+        ), timeless=True)
+
+    def log_subjects(self, subjects: List[Dict[str, Any]], selected_subject: int = None) -> None:
+        for idx, subject in enumerate(subjects):
+            vertices, faces = create_subject_mesh(subject)
+            subject_color = [0.8, 0.2, 0.2, 1.0] if idx == selected_subject else [
+                0.8, 0.8, 0.8, 1.0]
+
+            rr.log(
+                f"world/subject_{idx}",
+                rr.Mesh3D(
+                    vertex_positions=vertices,
+                    indices=faces,
+                    colors=np.tile(subject_color, (len(vertices), 1))
+                ),
+                timeless=True
+            )
+
+            rr.log(f"world/subject_{idx}/class",
+                   rr.TextDocument(subject['objectClass']),
+                   timeless=True)
+
+    def log_camera_trajectory(self, camera_frames: List[Dict[str, Any]]) -> None:
+        camera_positions = np.array(
+            [vector3_to_numpy(frame['position']) for frame in camera_frames])
+        rr.log(
+            "world/camera_trajectory",
+            rr.Points3D(
+                camera_positions,
+                colors=np.full((len(camera_positions), 4),
+                               [0.0, 0.8, 0.8, 1.0])
+            ),
+            timeless=True
+        )
+
+    def log_camera_frames(self, camera_frames: List[Dict[str, Any]]) -> None:
+        for frame_idx, camera_frame in enumerate(camera_frames):
+            rr.set_time_sequence("frame", frame_idx)
+
+            position = vector3_to_numpy(camera_frame['position'])
+            rotation_q = euler_to_quaternion(camera_frame['angle'])
+
+            rr.log(
+                "world/camera",
+                rr.Transform3D(
+                    translation=position,
+                    rotation=rr.Quaternion(xyzw=rotation_q)
+                )
+            )
+
+            rr.log(
+                "world/camera/view",
+                rr.Pinhole(
+                    focal_length=camera_frame['focalLength'],
+                    width=1920,
+                    height=1080
+                )
+            )
+
+            rr.log(
+                "metadata/current_frame",
+                rr.TextDocument(
+                    f"Frame: {frame_idx + 1}/{len(camera_frames)}"),
+            )

visualization/mesh.py

@@ -0,0 +1,26 @@
+import numpy as np
+from typing import Dict, Tuple
+from ..utils.geometry import vector3_to_numpy
+
+
+def create_subject_mesh(subject: Dict) -> Tuple[np.ndarray, np.ndarray]:
+    position = vector3_to_numpy(subject['position'])
+    size = vector3_to_numpy(subject['size'])
+
+    # Create cube vertices
+    vertices = np.array([
+        [-0.5, -0.5, -0.5], [0.5, -0.5, -0.5], [0.5, 0.5, -0.5], [-0.5, 0.5, -0.5],
+        [-0.5, -0.5, 0.5], [0.5, -0.5, 0.5], [0.5, 0.5, 0.5], [-0.5, 0.5, 0.5]
+    ]) * size.reshape(1, 3) + position.reshape(1, 3)
+
+    # Create cube faces
+    faces = np.array([
+        [0, 1, 2], [0, 2, 3],  # front
+        [1, 5, 6], [1, 6, 2],  # right
+        [5, 4, 7], [5, 7, 6],  # back
+        [4, 0, 3], [4, 3, 7],  # left
+        [3, 2, 6], [3, 6, 7],  # top
+        [4, 5, 1], [4, 1, 0]   # bottom
+    ])
+
+    return vertices, faces

visualization/visualizer.py

@@ -0,0 +1,40 @@
+import tempfile
+import os
+import spaces
+from typing import Optional
+from ..data.loader import load_simulation_data
+from .logger import SimulationLogger
+import rerun as rr
+
+
+@spaces.GPU
+def visualize_simulation(file, simulation_index: int) -> Optional[str]:
+    if file is None:
+        return None
+
+    try:
+        simulations, _ = load_simulation_data(file)
+        if simulations is None or simulation_index >= len(simulations):
+            return None
+
+        # Create temporary file for RRD
+        temp_dir = tempfile.mkdtemp()
+        rrd_path = os.path.join(temp_dir, "simulation.rrd")
+
+        # Log selected simulation
+        simulation = simulations[simulation_index]
+
+        logger = SimulationLogger()
+        logger.log_metadata(simulation['instructions'])
+        logger.log_subjects(
+            simulation['subjects'], simulation.get('selectedSubject'))
+        logger.log_camera_trajectory(simulation['cameraFrames'])
+        logger.log_camera_frames(simulation['cameraFrames'])
+
+        rr.save(rrd_path)
+
+        return rrd_path
+
+    except Exception as e:
+        print(f"Error processing simulation: {str(e)}")
+        return None