diff --git a/navlie/__init__.py b/navlie/__init__.py index 2d805305..8050a63a 100644 --- a/navlie/__init__.py +++ b/navlie/__init__.py @@ -58,6 +58,7 @@ plot_meas, plot_meas_by_model, plot_poses, + plot_camera_poses, set_axes_equal ) diff --git a/navlie/utils/__init__.py b/navlie/utils/__init__.py index 91ac3632..eae2a0ec 100644 --- a/navlie/utils/__init__.py +++ b/navlie/utils/__init__.py @@ -20,6 +20,7 @@ plot_meas, plot_meas_by_model, plot_poses, + plot_camera_poses, set_axes_equal ) diff --git a/navlie/utils/plot.py b/navlie/utils/plot.py index 005997b7..61aa7672 100644 --- a/navlie/utils/plot.py +++ b/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, diff --git a/tests/unit/test_plot_poses.py b/tests/unit/test_plot_poses.py index f12d5204..9b12d127 100644 --- a/tests/unit/test_plot_poses.py +++ b/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()