Merge pull request #128 from decargroup/add_camera_pose_viz

Add a method to visualize camera poses in 3D

navlie/__init__.py

@@ -58,6 +58,7 @@
     plot_meas,
     plot_meas_by_model,
     plot_poses,
+    plot_camera_poses,
     set_axes_equal
 )
 

navlie/utils/__init__.py

@@ -20,6 +20,7 @@
     plot_meas,
     plot_meas_by_model,
     plot_poses,
+    plot_camera_poses,
     set_axes_equal
 )
 

navlie/utils/plot.py

@@ -307,6 +307,161 @@ def plot_meas_by_model(
     return fig, axs
 
 
+def plot_camera_poses(
+    poses,
+    ax: plt.Axes = None,
+    color: str = "tab:blue",
+    line_thickness: float = 1,
+    step: int = 1,
+    scale: float = 0.25,
+):
+    """
+    Plots camera poses along a 3D plot.
+
+    The camera poses should be elements of SE(3), with the z-axis pointing
+    forward through the optical axis.
+
+    Parameters
+    ----------
+    poses : List[SE3State]
+        A list objects containing a ``position`` property and an attitude
+        property, representing the rotation matrix :math:``\mathbf{C}_{ab}``.
+    ax : plt.Axes, optional
+        Axes to plot on, if none, 3D axes are created.
+    color : str, optional
+        Color of the plotted camera, by default "tab:blue"
+    line_thickness : float, optional
+        Thickness of the camera line, by default 1
+    step : int, optional
+        Step size in number of poses to plot, by default 1 which plots all poses
+    scale : float, optional
+        Scale of the camera, by default 0.25
+    Returns
+    -------
+    plt.Figure
+        Handle to figure.
+    List[plt.Axes]
+        Handle to axes that were drawn on.
+    """
+    if isinstance(poses, GaussianResultList):
+        poses = poses.state
+
+    if isinstance(poses, StateWithCovariance):
+        poses = [poses.state]
+
+    if isinstance(poses, np.ndarray):
+        poses = poses.tolist()
+
+    if not isinstance(poses, list):
+        poses = [poses]
+
+    # Check if provided axes are in 3D
+    if ax is None:
+        fig = plt.figure()
+        ax = plt.axes(projection="3d")
+    else:
+        fig = ax.get_figure()
+
+    # Plot the individual camera poses
+    cam_pose_viz = CameraPoseVisualizer(
+        line_thickness=line_thickness,
+        scale=scale,
+        color=color,
+    )
+    for i in range(0, len(poses), step):
+        fig, ax = cam_pose_viz.plot_pose(
+            poses[i].attitude, poses[i].position, ax=ax
+        )
+
+    set_axes_equal(ax)
+    return fig, ax
+
+
+class CameraPoseVisualizer:
+    """A class to plot camera poses in 3D using matplotlib."""
+    def __init__(
+        self,
+        line_thickness: float = 1,
+        scale: float = 0.25,
+        color: str = "tab:blue",
+    ):
+        self.line_thickness = line_thickness
+        self.scale = scale
+        self.color = color
+
+        # Define points resolved in the camera frame
+        cam_points: List[np.ndarray] = []
+        # Image plane corners
+        cam_points.append(np.array([-1.0, -1.0, 1.0]))  # left top
+        cam_points.append(np.array([1.0, -1.0, 1.0]))  # right top
+        cam_points.append(np.array([-1.0, 1.0, 1.0]))  # left bottom
+        cam_points.append(np.array([1.0, 1.0, 1.0]))  # right bottom
+        # Optical center
+        cam_points.append(np.array([0.0, 0.0, 0.0]))
+        for point in cam_points:
+            point *= scale
+        self.cam_points = cam_points
+
+    def plot_pose(self, C: np.ndarray, r: np.ndarray, ax: plt.Axes = None):
+        """Plots a camera pose in 3D.
+
+        Plots lines representing connection between the optical center and the
+        camera corners.
+
+        Parameters
+        ----------
+        C : np.ndarray
+            Rotation matrix representing the camera attitude :math:``\mathbf{C}_{ab}``.
+        r : np.ndarray
+            Position of the camera in the inertial frame.
+
+        Returns
+        -------
+        plt.Figure
+            Handle to figure.
+        plt.Axes
+            Handle to axes that were drawn on.
+        """
+        if ax is None:
+            fig = plt.figure()
+            ax = fig.add_subplot(111, projection="3d")
+        else:
+            fig = ax.get_figure()
+
+        r = r.ravel()
+        # Resolve the points in the inertial frame
+        cam_points = self.cam_points
+        inertial_points: List[np.ndarray] = []
+        for point in cam_points:
+            inertial_points.append(C @ point + r)
+
+        # Define the connections between the points
+        connections = [
+            (0, 1),
+            (2, 3),
+            (0, 2),
+            (1, 3),
+            (0, 4),
+            (1, 4),
+            (2, 4),
+            (3, 4),
+        ]
+
+        # Plot lines between each point defined in the connections list
+        for connection in connections:
+            p1 = inertial_points[connection[0]]
+            p2 = inertial_points[connection[1]]
+
+            ax.plot(
+                [p1[0], p2[0]],
+                [p1[1], p2[1]],
+                [p1[2], p2[2]],
+                linewidth=self.line_thickness,
+                color=self.color,
+            )
+        return fig, ax
+
+
 def plot_poses(
     poses,
     ax: plt.Axes = None,

tests/unit/test_plot_poses.py

@@ -2,6 +2,8 @@
 import navlie as nav
 import matplotlib.pyplot as plt
 import seaborn as sns
+from navlie.lib.camera import PinholeCamera
+from pymlg import SO3, SE3
 
 sns.set_style("whitegrid")
 
@@ -45,7 +47,26 @@ def test_plot_poses_2d():
     # Test plotting SE(2) poses
     fig, ax = nav.plot_poses(x_traj)
 
+def test_plot_camera_poses():
+    """Tests the camera pose visualization function."""
+    C_bc = PinholeCamera.get_cam_to_enu()
+    poses = []
+
+    for i in range(5):
+        C_ab = SO3.Exp([0.0, 0.0, i / 2.0])
+        C_ac = C_ab @ C_bc
+        pose = nav.lib.SE3State(
+            value=SE3.from_components(C_ac, np.array([i / 2, i / 4, 0.0]))
+        )
+        poses.append(pose)
+
+    fig, ax = nav.utils.plot_camera_poses(poses, scale=0.15, line_thickness=1.0, color="tab:blue")
+    ax.set_xlabel("x (m)")
+    ax.set_ylabel("y (m)")
+    ax.set_zlabel("z (m)")
+
 if __name__ == "__main__":
     test_plot_poses_3d()
     test_plot_poses_2d()
+    test_plot_camera_poses()
     plt.show()