abreza/rerun-viewer

1import tempfile2import os3import spaces4from typing import Optional5from data.loader import load_simulation_data6from visualization.logger import SimulationLogger7import rerun as rr8 9 10@spaces.GPU11def visualize_simulation(file, simulation_index: Optional[int]) -> Optional[str]:12    if file is None or simulation_index is None:13        return None14 15    try:16        simulations, _ = load_simulation_data(file)17        if not simulations or not isinstance(simulation_index, int) or simulation_index < 0 or simulation_index >= len(simulations):18            print(f"Invalid simulation data or index: {simulation_index}")19            return None20 21        temp_dir = tempfile.mkdtemp()22        rrd_path = os.path.join(temp_dir, "simulation.rrd")23 24        simulation = simulations[simulation_index]25 26        logger = SimulationLogger()27 28        # Log metadata29        logger.log_metadata(simulation['instructions'])30 31        # Log subjects (individual frames)32        logger.log_subjects(simulation['subjects'])33 34        # Log subject trajectories (full center path)35        logger.log_subject_trajectories(simulation['subjects'])36 37        # Log camera trajectory (full path)38        logger.log_camera_trajectory(simulation['cameraFrames'])39 40        # Log individual camera frames41        logger.log_camera_frames(simulation['cameraFrames'])42 43        # Log helper keyframes if available44        if 'helper_keyframes' in simulation:45            logger.log_helper_keyframes(simulation['helper_keyframes'])46 47        # Save the visualization data48        rr.save(rrd_path)49 50        return rrd_path51 52    except Exception as e:53        print(f"Error processing simulation: {str(e)}")54        return None55