rerun-viewer / visualization /visualizer.py
abreza's picture
enhance logging for helper keyframes by adding timeless parameter and reorganizing camera trajectory logging
15db7ce
import tempfile
import os
import spaces
from typing import Optional
from data.loader import load_simulation_data
from visualization.logger import SimulationLogger
import rerun as rr
@spaces.GPU
def visualize_simulation(file, simulation_index: Optional[int]) -> Optional[str]:
if file is None or simulation_index is None:
return None
try:
simulations, _ = load_simulation_data(file)
if not simulations or not isinstance(simulation_index, int) or simulation_index < 0 or simulation_index >= len(simulations):
print(f"Invalid simulation data or index: {simulation_index}")
return None
temp_dir = tempfile.mkdtemp()
rrd_path = os.path.join(temp_dir, "simulation.rrd")
simulation = simulations[simulation_index]
logger = SimulationLogger()
logger.log_metadata(simulation['instructions'])
logger.log_subjects(simulation['subjects'])
if 'helper_keyframes' in simulation:
logger.log_helper_keyframes(simulation['helper_keyframes'])
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