Update visualization/logger.py

visualization/logger.py

@@ -4,20 +4,28 @@ from typing import Dict, Any, List
 from utils.geometry import vector3_to_numpy, euler_to_quaternion
 
 
-def create_subject_box(frame: Dict, dimensions: Dict) -> Dict[str, np.ndarray]:
+def create_subject_box(frame: Dict[str, Any], dimensions: Dict[str, float]) -> Dict[str, np.ndarray]:
+    """
+    Given a single frame (position, rotation) and subject's dimensions,
+    return the center and half_size for a box representation.
+    """
     position = vector3_to_numpy(frame['position'])
-    size = np.array([dimensions['width'], dimensions['height'], dimensions['depth']])
+    half_size = np.array([
+        dimensions['width'] / 2.0,
+        dimensions['height'] / 2.0,
+        dimensions['depth'] / 2.0
+    ], dtype=np.float32)
 
     return {
         'center': position,
-        'half_size': size / 2
+        'half_size': half_size
     }
 
 
 class SimulationLogger:
     def __init__(self):
         rr.init("camera_simulation")
-        rr.log("world", rr.ViewCoordinates.RIGHT_HAND_Y_UP)
+        rr.log("world", rr.ViewCoordinates.RIGHT_HAND_Y_UP, timeless=True)
 
         self.K = np.array([
             [500, 0, 960],
@@ -29,52 +37,72 @@ class SimulationLogger:
         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)
-            ])
-        ))
+        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]], frame_idx: int = 0) -> None:
+    def log_subjects(self, subjects: List[Dict[str, Any]]) -> None:
+        """
+        Updated to handle dynamic subjects:
+        - Each subject has multiple frames (position, rotation).
+        - Subject dimensions (width, height, depth) remain constant across frames.
+        """
         if not subjects:
             return
 
-        centers = []
-        half_sizes = []
-        colors = []
-        labels = []
+        # For each subject, log each frame at a different time index.
+        for subject_idx, subject in enumerate(subjects):
+            dimensions = subject.get("dimensions", {})
+            frames = subject.get("frames", [])
+            label = subject.get('objectClass', f"Subject_{subject_idx}")
 
-        for subject in subjects:
-            try:
-                # Use the frame at frame_idx, or the first frame if frame_idx is out of range
-                frame = subject['frames'][min(frame_idx, len(subject['frames']) - 1)]
-                box_params = create_subject_box(frame, subject['dimensions'])
-                centers.append(box_params['center'])
-                half_sizes.append(box_params['half_size'])
-                colors.append([0.8, 0.2, 0.2, 1.0])
-                labels.append(subject.get('objectClass', 'Unknown'))
-            except Exception as e:
-                print(f"Error creating box parameters: {str(e)}")
-                continue
+            for frame_idx, frame in enumerate(frames):
+                rr.set_time_sequence("frame_idx", frame_idx)
 
-        if centers:
-            rr.log(
-                "world/subjects",
-                rr.Boxes3D(
-                    centers=np.array(centers),
-                    half_sizes=np.array(half_sizes),
-                    colors=np.array(colors),
-                    show_labels=False,
-                    fill_mode="solid"
-                )
-            )
+                try:
+                    box_params = create_subject_box(frame, dimensions)
+                    center = box_params['center']
+                    half_size = box_params['half_size']
+
+                    # Convert Euler angles to quaternion for rotation
+                    rotation_q = euler_to_quaternion(frame['rotation'])
+
+                    # 1) Log a transform for the subject so the box is placed at 'center' with correct rotation
+                    rr.log(
+                        f"world/subjects/subject_{subject_idx}",
+                        rr.Transform3D(
+                            translation=center,
+                            rotation=rr.Quaternion(xyzw=rotation_q)
+                        )
+                    )
+
+                    # 2) Log a box at the origin of the subject's local transform
+                    rr.log(
+                        f"world/subjects/subject_{subject_idx}/box",
+                        rr.Boxes3D(
+                            centers=np.array([[0.0, 0.0, 0.0]], dtype=np.float32),
+                            half_sizes=np.array([half_size], dtype=np.float32),
+                            colors=np.array([[0.8, 0.2, 0.2, 1.0]], dtype=np.float32),
+                            labels=[label],
+                            show_labels=False,
+                            fill_mode="solid"
+                        )
+                    )
+                except Exception as e:
+                    print(f"Error creating box parameters for subject {subject_idx}, frame {frame_idx}: {str(e)}")
 
     def log_camera_trajectory(self, camera_frames: List[Dict[str, Any]]) -> None:
         if not camera_frames:
@@ -88,36 +116,29 @@ class SimulationLogger:
                 "world/camera_trajectory",
                 rr.Points3D(
                     camera_positions,
-                    colors=np.full((len(camera_positions), 4),
-                                   [0.0, 0.8, 0.8, 1.0])
-                )
+                    colors=np.full((len(camera_positions), 4), [0.0, 0.8, 0.8, 1.0])
+                ),
+                timeless=True
             )
 
             if len(camera_positions) > 1:
-                lines = np.stack(
-                    [camera_positions[:-1], camera_positions[1:]], axis=1)
+                lines = np.stack([camera_positions[:-1], camera_positions[1:]], axis=1)
                 rr.log(
                     "world/camera_trajectory/line",
                     rr.LineStrips3D(
                         lines,
                         colors=[(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]], subjects: List[Dict[str, Any]] = None) -> None:
-        """Log camera frames and optionally their corresponding subjects"""
-        if subjects is None:
-            subjects = []
+    def log_camera_frames(self, camera_frames: List[Dict[str, Any]]) -> None:
         if not camera_frames:
             return
 
-        # First log the camera trajectory
-        self.log_camera_trajectory(camera_frames)
-
-        # Then log each frame with its corresponding subject positions
         for frame_idx, camera_frame in enumerate(camera_frames):
             try:
                 rr.set_time_sequence("frame_idx", frame_idx)
@@ -142,9 +163,6 @@ class SimulationLogger:
                     )
                 )
 
-                # Log subjects for this frame
-                self.log_subjects(subjects, frame_idx)
-
             except Exception as e:
                 print(f"Error logging camera frame {frame_idx}: {str(e)}")
 
@@ -161,32 +179,10 @@ class SimulationLogger:
                 helper_positions,
                 radii=np.full(len(helper_positions), 0.03),
                 colors=np.full((len(helper_positions), 4), [1.0, 1.0, 0.0, 1.0]),
-            )
+            ),
+            timeless=True
         )
 
         for keyframe_idx, helper_keyframe in enumerate(helper_keyframes):
             try:
-                position = vector3_to_numpy(helper_keyframe['position'])
-                rotation_q = euler_to_quaternion(helper_keyframe['angle'])
-
-                rr.log(
-                    f"world/helper_camera_{keyframe_idx}",
-                    rr.Transform3D(
-                        translation=position,
-                        rotation=rr.Quaternion(xyzw=rotation_q),
-                        scale=(.5, .5, .5)
-                    )
-                )
-
-                rr.log(
-                    f"world/helper_camera_{keyframe_idx}/image",
-                    rr.Pinhole(
-                        image_from_camera=self.K,
-                        width=1920,
-                        height=1080,
-                    )
-                )
-
-            except Exception as e:
-                print(
-                    f"Error logging helper keyframe {keyframe_idx}: {str(e)}")
+