Update visualization/logger.py

visualization/logger.py

@@ -4,9 +4,9 @@ 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'])
+def create_subject_box(frame: Dict, dimensions: Dict) -> Dict[str, np.ndarray]:
+    position = vector3_to_numpy(frame['position'])
+    size = np.array([dimensions['width'], dimensions['height'], dimensions['depth']])
 
     return {
         'center': position,
@@ -42,7 +42,7 @@ class SimulationLogger:
             ])
         ), timeless=True)
 
-    def log_subjects(self, subjects: List[Dict[str, Any]]) -> None:
+    def log_subjects(self, subjects: List[Dict[str, Any]], frame_idx: int = 0) -> None:
         if not subjects:
             return
 
@@ -53,7 +53,9 @@ class SimulationLogger:
 
         for subject in subjects:
             try:
-                box_params = create_subject_box(subject)
+                # 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])
@@ -71,8 +73,7 @@ class SimulationLogger:
                     colors=np.array(colors),
                     show_labels=False,
                     fill_mode="solid"
-                ),
-                timeless=True
+                )
             )
 
     def log_camera_trajectory(self, camera_frames: List[Dict[str, Any]]) -> None:
@@ -108,7 +109,7 @@ class SimulationLogger:
         except Exception as e:
             print(f"Error logging camera trajectory: {str(e)}")
 
-    def log_camera_frames(self, camera_frames: List[Dict[str, Any]]) -> None:
+    def log_camera_frames(self, camera_frames: List[Dict[str, Any]], subjects: List[Dict[str, Any]]) -> None:
         if not camera_frames:
             return
 
@@ -136,6 +137,9 @@ 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)}")
 
@@ -143,7 +147,6 @@ class SimulationLogger:
         if not helper_keyframes:
             return
 
-
         helper_positions = np.array([
             vector3_to_numpy(frame['position']) for frame in helper_keyframes
         ])
@@ -183,4 +186,4 @@ class SimulationLogger:
 
             except Exception as e:
                 print(
-                    f"Error logging helper keyframe {keyframe_idx}: {str(e)}")
+                    f"Error logging helper keyframe {keyframe_idx}: {str(e)}")