File size: 2,456 Bytes
bdf752e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 |
import tempfile
import os
import spaces
from typing import Optional, List
from data.loader import load_simulation_data
from visualization.logger import SimulationLogger
import rerun as rr
def parse_frame_selection(selection_str: str, max_frame: int) -> List[int]:
if not selection_str:
return list(range(max_frame))
frames = set()
segments = selection_str.split(',')
for seg in segments:
seg = seg.strip()
if '-' in seg:
# range
start_str, end_str = seg.split('-', 1)
start, end = int(start_str), int(end_str)
# Convert to 0-based indexing
for f in range(start - 1, end):
if 0 <= f < max_frame:
frames.add(f)
else:
# single frame
f = int(seg) - 1
if 0 <= f < max_frame:
frames.add(f)
return sorted(frames)
@spaces.GPU
def visualize_simulation(file, simulation_index: Optional[int], frame_selection: str) -> 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
simulation = simulations[simulation_index]
camera_frames = simulation.get('cameraFrames', [])
max_frame = len(camera_frames)
selected_frames = parse_frame_selection(frame_selection, max_frame)
temp_dir = tempfile.mkdtemp()
rrd_path = os.path.join(temp_dir, "simulation.rrd")
logger = SimulationLogger()
logger.log_metadata(simulation['instructions'])
logger.log_subjects(simulation['subjects'])
if 'helper_keyframes' in simulation:
logger.log_helper_keyframes(simulation['helper_keyframes'])
# Filter camera trajectory to selected frames
selected_camera_frames = [camera_frames[i] for i in selected_frames]
# Log the trajectory and frames only for the selected frames
logger.log_camera_trajectory(selected_camera_frames)
logger.log_camera_frames(selected_camera_frames)
rr.save(rrd_path)
return rrd_path
except Exception as e:
print(f"Error processing simulation: {str(e)}")
return None
|