Add optional z-axis locking feature

flow_lab/flow_lab.py

@@ -132,8 +132,10 @@ def setup_visualizer(state_manager, annotation_saver, frames):
     # Use , and . keys for going forward and backward through the frames
     vis.register_key_callback(ord(","), state_manager.backward_frame_press)
     vis.register_key_callback(ord("."), state_manager.forward_frame_press)
-    # Use 'V' to toggle postion adjustment mode
+    # Use 'V' to toggle postion adjustment mode, "L" to toggle z axis lock
     vis.register_key_callback(ord("V"), state_manager.toggle_propagate_with_velocity)
+    vis.register_key_callback(ord("L"), state_manager.toggle_lock_z_axis)
+    vis.register_key_callback(ord("T"), state_manager.toggle_trajectory_visibility)
 
     # Use the 'Ctrl+S' to save, 'S' to translate
     vis.register_key_action_callback(ord("S"), state_manager.key_S_actions)

flow_lab/vis_classes.py

@@ -13,6 +13,10 @@
 
 
 def _update_o3d_mesh_pose(mesh: o3d.geometry.TriangleMesh, start_pose: SE3, target_pose: SE3):
+    # # Compute the relative transformation from start_pose to target_pose
+    # relative_transform = target_pose.compose(start_pose.inverse())
+    # # Apply the combined transformation to the mesh
+    # mesh.transform(relative_transform.transform_matrix)
     global_translation = target_pose.translation - start_pose.translation
     global_rotation = target_pose.rotation_matrix @ np.linalg.inv(start_pose.rotation_matrix)
 
@@ -61,10 +65,10 @@ def compute_new_global_pose_from_inputs(
         return self.pose.sensor_to_global.compose(local_frame_offset_se3)
 
     def update_from_global(self, global_target_se3: SE3):
-        _update_o3d_mesh_pose(self.o3d_wireframe, self.pose.sensor_to_global, global_target_se3)
-        _update_o3d_mesh_pose(self.o3d_triangle_mesh, self.pose.sensor_to_global, global_target_se3)
         # We only edit sensor to ego; ego to global stays fixed for the box
         # because the ego vehicle position never changes.
+        _update_o3d_mesh_pose(self.o3d_wireframe, self.pose.sensor_to_global, global_target_se3)
+        _update_o3d_mesh_pose(self.o3d_triangle_mesh, self.pose.sensor_to_global, global_target_se3)
         self.pose.sensor_to_ego = self.pose.ego_to_global.inverse().compose(global_target_se3)
 
     def triangle_mesh_o3d(self) -> o3d.geometry.TriangleMesh:
@@ -137,13 +141,15 @@ def __init__(
         self.rolling_window_size = rolling_window_size
         self.trajectory_geometries: list[RenderableBox] = []
         # the two blow is used for zoom
-        self.is_zoomed = False
+        self.is_zoomed = True
         self.original_view = None
         self.zoom_level = 0.06  # Default zoom level
         self.camera_relative_position = None
         # Used for velocity
         self.propagate_with_velocity = False
         self.velocities: dict[str, SE3] = {}
+        self.lock_z_axis = False
+        self.show_trajectory = True
         # Used for toggle box
         self.current_box_index = -1
 
@@ -177,8 +183,8 @@ def _update_selection(
         vis.update_geometry(selected_mesh.wireframe_o3d())
         vis.update_geometry(self.selection_axes)
 
-        if self.is_zoomed:
-            self.zoom_to_box(vis)
+        # if self.is_zoomed:
+        #     self.zoom_to_box(vis)
         if self.propagate_with_velocity:
             self.apply_velocity(vis)
 
@@ -194,6 +200,7 @@ def forward_press(self, vis, action, mods):
                 self._update_selection(vis, forward=self.default_tuning_scale)
             if mods == ["shift"]:
                 self._update_selection(vis, forward=self.fine_tuning_scale)
+            print("forward_press")
 
     def backward_press(self, vis, action, mods):
         if self.selected_mesh_id is None:
@@ -207,7 +214,7 @@ def backward_press(self, vis, action, mods):
                 self._update_selection(vis, forward=-self.default_tuning_scale)
             if mods == ["shift"]:
                 self._update_selection(vis, forward=-self.fine_tuning_scale)
-        print("backward_press")
+            print("backward_press")
 
     def left_press(self, vis, action, mods):
         if self.selected_mesh_id is None:
@@ -221,7 +228,7 @@ def left_press(self, vis, action, mods):
                 self._update_selection(vis, left=self.default_tuning_scale)
             if mods == ["shift"]:
                 self._update_selection(vis, left=self.fine_tuning_scale)
-        print("left_press")
+            print("left_press")
 
     def right_press(self, vis, action, mods):
         if self.selected_mesh_id is None:
@@ -236,7 +243,7 @@ def right_press(self, vis, action, mods):
             if mods == ["shift"]:
                 self._update_selection(vis, left=-self.fine_tuning_scale)
 
-        print("right_press")
+            print("right_press")
 
     def up_press(self, vis, action, mods):
         if self.selected_mesh_id is None:
@@ -250,7 +257,7 @@ def up_press(self, vis, action, mods):
                 self._update_selection(vis, up=self.default_tuning_scale)
             if mods == ["shift"]:
                 self._update_selection(vis, up=self.fine_tuning_scale)
-        print("up_press")
+            print("up_press")
 
     def down_press(self, vis, action, mods):
         if self.selected_mesh_id is None:
@@ -264,7 +271,7 @@ def down_press(self, vis, action, mods):
                 self._update_selection(vis, up=-self.default_tuning_scale)
             if mods == ["shift"]:
                 self._update_selection(vis, up=-self.fine_tuning_scale)
-        print("down_press")
+            print("down_press")
 
     def yaw_clockwise_press(self, vis, action, mods):
         if self.selected_mesh_id is None:
@@ -278,7 +285,7 @@ def yaw_clockwise_press(self, vis, action, mods):
                 self._update_selection(vis, yaw=self.default_tuning_scale)
             if mods == ["shift"]:
                 self._update_selection(vis, yaw=self.fine_tuning_scale)
-        print("yaw_clockwise_press")
+            print("yaw_clockwise_press")
 
     def yaw_counterclockwise_press(self, vis, action, mods):
         if self.selected_mesh_id is None:
@@ -293,7 +300,7 @@ def yaw_counterclockwise_press(self, vis, action, mods):
             if mods == ["shift"]:
                 self._update_selection(vis, yaw=-self.fine_tuning_scale)
 
-        print("yaw_counterclockwise_press")
+            print("yaw_counterclockwise_press")
 
     def pitch_up_press(self, vis, action, mods):
         if self.selected_mesh_id is None:
@@ -307,7 +314,7 @@ def pitch_up_press(self, vis, action, mods):
                 self._update_selection(vis, pitch=self.default_tuning_scale)
             if mods == ["shift"]:
                 self._update_selection(vis, pitch=self.fine_tuning_scale)
-        print("pitch_up_press")
+            print("pitch_up_press")
 
     def pitch_down_press(self, vis, action, mods):
         if self.selected_mesh_id is None:
@@ -322,7 +329,7 @@ def pitch_down_press(self, vis, action, mods):
             if mods == ["shift"]:
                 self._update_selection(vis, pitch=-self.fine_tuning_scale)
 
-        print("pitch_down_press")
+            print("pitch_down_press")
 
     def roll_clockwise_press(self, vis, action, mods):
         if self.selected_mesh_id is None:
@@ -337,7 +344,7 @@ def roll_clockwise_press(self, vis, action, mods):
             if mods == ["shift"]:
                 self._update_selection(vis, roll=self.fine_tuning_scale)
 
-        print("roll_clockwise_press")
+            print("roll_clockwise_press")
 
     def roll_counterclockwise_press(self, vis, action, mods):
         if self.selected_mesh_id is None:
@@ -352,7 +359,7 @@ def roll_counterclockwise_press(self, vis, action, mods):
             if mods == ["shift"]:
                 self._update_selection(vis, roll=-self.fine_tuning_scale)
 
-        print("roll_counterclockwise_press")
+            print("roll_counterclockwise_press")
 
     def forward_frame_press(self, vis):
         # self.annotation_saver.save(self.frames)
@@ -458,9 +465,11 @@ def select_mesh(self, vis, mesh_id: str, reset_bounding_box: bool = False):
                     )
         self.selected_mesh_id = mesh_id
         # Show trajectory of selected mesh
-        self.render_selected_mesh_trajectory(vis)
-        # if self.is_zoomed == True:
-        self.zoom_to_box(vis)
+        if self.show_trajectory:
+            self.clear_trajectory_geometries(vis)
+            self.render_selected_mesh_trajectory(vis)
+        if self.is_zoomed is True:
+            self.zoom_to_box(vis)
         base_bounding_box = self.clickable_geometries[self.selected_mesh_id].base_box
         print(
             f"Selected mesh: {base_bounding_box.track_uuid}, category: {base_bounding_box.category}"
@@ -475,7 +484,7 @@ def deselect_mesh(self, vis, reset_bounding_box: bool = False):
         if self.is_zoomed:
             view_control = vis.get_view_control()
             # view_control.convert_from_pinhole_camera_parameters(self.original_view)
-            self.is_zoomed = False
+            # self.is_zoomed = False
 
         # Re-add all meshes to the visualization
         if self.selected_mesh_id is not None:
@@ -689,6 +698,7 @@ def zoom_press(self, vis, action, mods):
                 view_control.convert_from_pinhole_camera_parameters(self.original_view)
                 self.is_zoomed = False
             elif mods == [] and self.selected_mesh_id:
+                self.is_zoomed = True
                 self.zoom_to_box(vis)
 
     def zoom_to_box(self, vis):
@@ -750,6 +760,25 @@ def toggle_propagate_with_velocity(self, vis):
             self.apply_velocity(vis)
         self.update_window_title(vis)
 
+    def toggle_lock_z_axis(self, vis):
+        """
+        Toggle the lock_z_axis feature on and off.
+        """
+        self.lock_z_axis = not self.lock_z_axis
+        print(f"Lock z axis: {'On' if self.lock_z_axis else 'Off'}")
+        if self.propagate_with_velocity:
+            self.apply_velocity(vis)
+
+    def toggle_trajectory_visibility(self, vis):
+        """
+        Toggle the visibility of the trajectory.
+        """
+        self.show_trajectory = not self.show_trajectory
+        if not self.show_trajectory:
+            self.clear_trajectory_geometries(vis)
+        else:
+            self.render_selected_mesh_trajectory(vis)
+
     def compute_velocities(self) -> bool:
         """
         Compute velocities for all the meshes in current frame.
@@ -787,14 +816,13 @@ def calculate_velocity(self, pose1: PoseInfo, pose2: PoseInfo) -> SE3:
         """
         Calculate displacement between two poses (suppose time interval=1 and use it as velocity).
         """
-
-        translation_velocity = (
-            pose2.sensor_to_global.translation - pose1.sensor_to_global.translation
-        )
         relative_transform = pose1.sensor_to_global.inverse().compose(pose2.sensor_to_global)
 
-        rotation_velocity = relative_transform.rotation_matrix
-        velocity = SE3(rotation_matrix=rotation_velocity, translation=translation_velocity)
+        # velocity = SE3(SE3.identity().rotation_matrix, relative_transform.translation)
+        velocity = SE3(
+            self.normalize_rotation_matrix(relative_transform.rotation_matrix),
+            relative_transform.translation,
+        )
 
         return velocity
 
@@ -804,8 +832,9 @@ def apply_velocity(self, vis):
         """
         get_velocity = self.compute_velocities()
         if get_velocity:
-
-            for offset in range(1, 3):  # Apply velocity to the next two frames
+            half_window_size = self.rolling_window_size // 2
+            # Apply velocity to the future frames that will be displayed in the scene
+            for offset in range(1, half_window_size + 1):
                 target_frame_index = self.current_frame_index + offset
                 if target_frame_index >= len(self.frames):
                     break
@@ -816,17 +845,35 @@ def apply_velocity(self, vis):
                 for uuid, velocity in self.velocities.items():
                     current_pose = self.find_pose_in_frame(current_frame, uuid)
                     if current_pose:
+                        # Project the velocity onto the XY plane
+                        xy_velocity_translation = velocity.translation.copy()
+                        xy_velocity_translation[2] = 0  # Set the Z component to 0
+                        xy_velocity = SE3(
+                            rotation_matrix=velocity.rotation_matrix,
+                            translation=xy_velocity_translation,
+                        )
                         # Calculate the new pose by composing the current pose with the velocity SE3
                         new_global_pose = current_pose.sensor_to_global
+                        # velocity =
+                        print(velocity)
+                        print(SE3.identity())
                         for i in range(offset):
-                            new_global_pose = new_global_pose.compose(velocity)
-
+                            if self.lock_z_axis:
+                                new_global_pose = new_global_pose.compose(xy_velocity)
+                            else:
+                                new_global_pose = new_global_pose.compose(velocity)
                         # Update the pose in the target frame
                         for box, pose in target_frame.boxes.valid_boxes():
                             if box.track_uuid == uuid:
                                 pose.sensor_to_ego = pose.ego_to_global.inverse().compose(
                                     new_global_pose
                                 )
                                 break
-            self.clear_trajectory_geometries(vis)
-            self.render_selected_mesh_trajectory(vis)
+            if self.show_trajectory:
+                self.clear_trajectory_geometries(vis)
+                self.render_selected_mesh_trajectory(vis)
+
+    def normalize_rotation_matrix(self, matrix):
+        # Use Singular Value Decomposition (SVD) to orthogonalize the rotation matrix
+        U, _, Vt = np.linalg.svd(matrix)
+        return np.dot(U, Vt)