visualization/logger.py

import rerun as rr
import numpy as np
from typing import Dict, Any, List
from utils.geometry import vector3_to_numpy, euler_to_quaternion


def create_subject_box(subject: Dict) -> Dict[str, np.ndarray]:
    position = vector3_to_numpy(subject['position'])
    size = vector3_to_numpy(subject['size'])

    return {
        'center': position,
        'half_size': size / 2  # Boxes3D uses half-sizes
    }


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:
        if not instructions:
            return

        rr.log("metadata/instructions", rr.TextDocument(
            "\n".join([
                f"Instruction {i+1}:\n" +
                f"  Movement: {inst.get('cameraMovement', 'N/A')}\n" +
                f"  Easing: {inst.get('movementEasing', 'N/A')}\n" +
                f"  Frames: {inst.get('frameCount', 'N/A')}\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]]) -> None:
        if not subjects:
            return

        # Prepare batch data for all subjects
        centers = []
        half_sizes = []
        colors = []
        labels = []

        for subject in subjects:
            try:
                box_params = create_subject_box(subject)
                centers.append(box_params['center'])
                half_sizes.append(box_params['half_size'])
                colors.append([0.8, 0.2, 0.2, 1.0])  # Red color with alpha
                labels.append(subject.get('objectClass', 'Unknown'))
            except Exception as e:
                print(f"Error creating box parameters: {str(e)}")
                continue

        if centers:  # Only log if we have valid subjects
            rr.log(
                "world/subjects",
                rr.Boxes3D(
                    centers=np.array(centers),
                    half_sizes=np.array(half_sizes),
                    colors=np.array(colors),
                    labels=labels,
                    fill_mode="solid"  # Use solid fill for better visualization
                ),
                timeless=True
            )

    def log_camera_trajectory(self, camera_frames: List[Dict[str, Any]]) -> None:
        if not camera_frames:
            return

        try:
            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
            )
        except Exception as e:
            print(f"Error logging camera trajectory: {str(e)}")

    def log_camera_frames(self, camera_frames: List[Dict[str, Any]]) -> None:
        if not camera_frames:
            return

        for frame_idx, camera_frame in enumerate(camera_frames):
            try:
                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.get('focalLength', 50),
                        width=1920,
                        height=1080
                    )
                )

                rr.log(
                    "metadata/current_frame",
                    rr.TextDocument(
                        f"Frame: {frame_idx + 1}/{len(camera_frames)}"),
                )
            except Exception as e:
                print(f"Error logging camera frame {frame_idx}: {str(e)}")