visualization/visualizer.py

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