Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,259 Bytes
ffbf761 44a5fb8 ffbf761 |
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 77 78 79 80 81 82 83 84 85 86 87 88 89 |
import rerun as rr
import numpy as np
from typing import Dict, Any, List
from utils.geometry import vector3_to_numpy, euler_to_quaternion
from visualization.mesh import create_subject_mesh
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:
rr.log("metadata/instructions", rr.TextDocument(
"\n".join([
f"Instruction {i+1}:\n" +
f" Movement: {inst['cameraMovement']}\n" +
f" Easing: {inst['movementEasing']}\n" +
f" Frames: {inst['frameCount']}\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]], selected_subject: int = None) -> None:
for idx, subject in enumerate(subjects):
vertices, faces = create_subject_mesh(subject)
subject_color = [0.8, 0.2, 0.2, 1.0] if idx == selected_subject else [
0.8, 0.8, 0.8, 1.0]
rr.log(
f"world/subject_{idx}",
rr.Mesh3D(
vertex_positions=vertices,
indices=faces,
colors=np.tile(subject_color, (len(vertices), 1))
),
timeless=True
)
rr.log(f"world/subject_{idx}/class",
rr.TextDocument(subject['objectClass']),
timeless=True)
def log_camera_trajectory(self, camera_frames: List[Dict[str, Any]]) -> None:
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
)
def log_camera_frames(self, camera_frames: List[Dict[str, Any]]) -> None:
for frame_idx, camera_frame in enumerate(camera_frames):
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['focalLength'],
width=1920,
height=1080
)
)
rr.log(
"metadata/current_frame",
rr.TextDocument(
f"Frame: {frame_idx + 1}/{len(camera_frames)}"),
)
|