Merge pull request #255 from borglab/feature/python-plotting

Python plotting upgrades

cython/gtsam/examples/SFMExample.py

@@ -10,13 +10,17 @@ A structure-from-motion problem on a simulated dataset
 """
 from __future__ import print_function
 
-import gtsam
+import matplotlib.pyplot as plt
 import numpy as np
+
+import gtsam
 from gtsam.examples import SFMdata
 from gtsam.gtsam import (Cal3_S2, DoglegOptimizer,
-                         GenericProjectionFactorCal3_S2, NonlinearFactorGraph,
+                         GenericProjectionFactorCal3_S2, Marginals,
-                         Point3, Pose3, PriorFactorPoint3, PriorFactorPose3,
+                         NonlinearFactorGraph, Point3, Pose3,
-                         Rot3, PinholeCameraCal3_S2, Values)
+                         PriorFactorPoint3, PriorFactorPose3, Rot3,
+                         SimpleCamera, Values)
+from gtsam.utils import plot
 
 
 def symbol(name: str, index: int) -> int:
@@ -94,12 +98,10 @@ def main():
     # Intentionally initialize the variables off from the ground truth
     initial_estimate = Values()
     for i, pose in enumerate(poses):
-        r = Rot3.Rodrigues(-0.1, 0.2, 0.25)
+        transformed_pose = pose.retract(0.1*np.random.randn(6,1))
-        t = Point3(0.05, -0.10, 0.20)
-        transformed_pose = pose.compose(Pose3(r, t))
         initial_estimate.insert(symbol('x', i), transformed_pose)
     for j, point in enumerate(points):
-        transformed_point = Point3(point.vector() + np.array([-0.25, 0.20, 0.15]))
+        transformed_point = Point3(point.vector() + 0.1*np.random.randn(3))
         initial_estimate.insert(symbol('l', j), transformed_point)
     initial_estimate.print_('Initial Estimates:\n')
 
@@ -113,6 +115,11 @@ def main():
     print('initial error = {}'.format(graph.error(initial_estimate)))
     print('final error = {}'.format(graph.error(result)))
 
+    marginals = Marginals(graph, result)
+    plot.plot_3d_points(1, result, marginals=marginals)
+    plot.plot_trajectory(1, result, marginals=marginals, scale=8)
+    plot.set_axes_equal(1)
+    plt.show()
 
 if __name__ == '__main__':
     main()

cython/gtsam/examples/SFMdata.py

@@ -25,14 +25,15 @@ def createPoints():
 
 def createPoses(K):
     # Create the set of ground-truth poses
-    radius = 30.0
+    radius = 40.0
+    height = 10.0
     angles = np.linspace(0, 2*np.pi, 8, endpoint=False)
     up = gtsam.Point3(0, 0, 1)
     target = gtsam.Point3(0, 0, 0)
     poses = []
     for theta in angles:
         position = gtsam.Point3(radius*np.cos(theta),
-                                radius*np.sin(theta), 0.0)
+                                radius*np.sin(theta), height)
         camera = gtsam.PinholeCameraCal3_S2.Lookat(position, target, up, K)
         poses.append(camera.pose())
     return poses

cython/gtsam/utils/plot.py

@@ -3,6 +3,74 @@
 import numpy as np
 import matplotlib.pyplot as plt
 from matplotlib import patches
+from mpl_toolkits.mplot3d import Axes3D
+
+import gtsam
+
+
+def set_axes_equal(fignum):
+    """
+    Make axes of 3D plot have equal scale so that spheres appear as spheres,
+    cubes as cubes, etc..  This is one possible solution to Matplotlib's
+    ax.set_aspect('equal') and ax.axis('equal') not working for 3D.
+    Input
+      ax: a matplotlib axis, e.g., as output from plt.gca().
+    """
+    fig = plt.figure(fignum)
+    ax = fig.gca(projection='3d')
+
+    limits = np.array([
+        ax.get_xlim3d(),
+        ax.get_ylim3d(),
+        ax.get_zlim3d(),
+    ])
+
+    origin = np.mean(limits, axis=1)
+    radius = 0.5 * np.max(np.abs(limits[:, 1] - limits[:, 0]))
+
+    ax.set_xlim3d([origin[0] - radius, origin[0] + radius])
+    ax.set_ylim3d([origin[1] - radius, origin[1] + radius])
+    ax.set_zlim3d([origin[2] - radius, origin[2] + radius])
+
+
+def ellipsoid(xc, yc, zc, rx, ry, rz, n):
+    """Numpy equivalent of Matlab's ellipsoid function"""
+    u = np.linspace(0, 2*np.pi, n+1)
+    v = np.linspace(0, np.pi, n+1)
+    x = -rx * np.outer(np.cos(u), np.sin(v)).T
+    y = -ry * np.outer(np.sin(u), np.sin(v)).T
+    z = -rz * np.outer(np.ones_like(u), np.cos(v)).T
+
+    return x, y, z
+
+
+def plot_covariance_ellipse_3d(axes, origin, P, scale=1, n=8, alpha=0.5):
+    """
+    Plots a Gaussian as an uncertainty ellipse
+
+    Based on Maybeck Vol 1, page 366
+    k=2.296 corresponds to 1 std, 68.26% of all probability
+    k=11.82 corresponds to 3 std, 99.74% of all probability
+    """
+    k = 11.82
+    U, S, _ = np.linalg.svd(P)
+
+    radii = k * np.sqrt(S)
+    radii = radii * scale
+    rx, ry, rz = radii
+
+    # generate data for "unrotated" ellipsoid
+    xc, yc, zc = ellipsoid(0, 0, 0, rx, ry, rz, n)
+
+    # rotate data with orientation matrix U and center c
+    data = np.kron(U[:, 0:1], xc) + np.kron(U[:, 1:2], yc) + \
+        np.kron(U[:, 2:3], zc)
+    n = data.shape[1]
+    x = data[0:n, :] + origin[0]
+    y = data[n:2*n, :] + origin[1]
+    z = data[2*n:, :] + origin[2]
+
+    axes.plot_surface(x, y, z, alpha=alpha, cmap='hot')
 
 
 def plot_pose2_on_axes(axes, pose, axis_length=0.1, covariance=None):
@@ -35,6 +103,7 @@ def plot_pose2_on_axes(axes, pose, axis_length=0.1, covariance=None):
                              np.rad2deg(angle), fill=False)
         axes.add_patch(e1)
 
+
 def plot_pose2(fignum, pose, axis_length=0.1, covariance=None):
     """Plot a 2D pose on given figure with given 'axis_length'."""
     # get figure object
@@ -43,19 +112,21 @@ def plot_pose2(fignum, pose, axis_length=0.1, covariance=None):
     plot_pose2_on_axes(axes, pose, axis_length, covariance)
 
 
-def plot_point3_on_axes(axes, point, linespec):
+def plot_point3_on_axes(axes, point, linespec, P=None):
     """Plot a 3D point on given axis 'axes' with given 'linespec'."""
     axes.plot([point.x()], [point.y()], [point.z()], linespec)
+    if P is not None:
+        plot_covariance_ellipse_3d(axes, point.vector(), P)
 
 
-def plot_point3(fignum, point, linespec):
+def plot_point3(fignum, point, linespec, P=None):
     """Plot a 3D point on given figure with given 'linespec'."""
     fig = plt.figure(fignum)
     axes = fig.gca(projection='3d')
-    plot_point3_on_axes(axes, point, linespec)
+    plot_point3_on_axes(axes, point, linespec, P)
 
 
-def plot_3d_points(fignum, values, linespec, marginals=None):
+def plot_3d_points(fignum, values, linespec="g*", marginals=None):
     """
     Plots the Point3s in 'values', with optional covariances.
     Finds all the Point3 objects in the given Values object and plots them.
@@ -68,23 +139,25 @@ def plot_3d_points(fignum, values, linespec, marginals=None):
     # Plot points and covariance matrices
     for i in range(keys.size()):
         try:
-            p = values.atPoint3(keys.at(i))
+            key = keys.at(i)
-            # if haveMarginals
+            point = values.atPoint3(key)
-            #     P = marginals.marginalCovariance(key);
+            if marginals is not None:
-            #     gtsam.plot_point3(p, linespec, P);
+                P = marginals.marginalCovariance(key);
-            # else
+            else:
-            plot_point3(fignum, p, linespec)
+                P = None
+
+            plot_point3(fignum, point, linespec, P)
+
         except RuntimeError:
             continue
             # I guess it's not a Point3
 
 
-def plot_pose3_on_axes(axes, pose, axis_length=0.1):
+def plot_pose3_on_axes(axes, pose, P=None, scale=1, axis_length=0.1):
     """Plot a 3D pose on given axis 'axes' with given 'axis_length'."""
     # get rotation and translation (center)
     gRp = pose.rotation().matrix()  # rotation from pose to global
-    t = pose.translation()
+    origin = pose.translation().vector()
-    origin = np.array([t.x(), t.y(), t.z()])
 
     # draw the camera axes
     x_axis = origin + gRp[:, 0] * axis_length
@@ -100,17 +173,61 @@ def plot_pose3_on_axes(axes, pose, axis_length=0.1):
     axes.plot(line[:, 0], line[:, 1], line[:, 2], 'b-')
 
     # plot the covariance
-    # TODO (dellaert): make this work
+    if P is not None:
-    # if (nargin>2) && (~isempty(P))
+        # covariance matrix in pose coordinate frame
-    #     pPp = P(4:6,4:6); % covariance matrix in pose coordinate frame
+        pPp = P[3:6, 3:6]
-    #     gPp = gRp*pPp*gRp'; % convert the covariance matrix to global coordinate frame
+        # convert the covariance matrix to global coordinate frame
-    #     gtsam.covarianceEllipse3D(origin,gPp);
+        gPp = gRp @ pPp @ gRp.T
-    # end
+        plot_covariance_ellipse_3d(axes, origin, gPp)
 
 
-def plot_pose3(fignum, pose, axis_length=0.1):
+def plot_pose3(fignum, pose, P, axis_length=0.1):
     """Plot a 3D pose on given figure with given 'axis_length'."""
     # get figure object
     fig = plt.figure(fignum)
     axes = fig.gca(projection='3d')
-    plot_pose3_on_axes(axes, pose, axis_length)
+    plot_pose3_on_axes(axes, pose, P=P, axis_length=axis_length)
+
+def plot_trajectory(fignum, values, scale=1, marginals=None):
+    pose3Values = gtsam.allPose3s(values)
+    keys = gtsam.KeyVector(pose3Values.keys())
+    lastIndex = None
+
+    for i in range(keys.size()):
+        key = keys.at(i)
+        try:
+            pose = pose3Values.atPose3(key)
+        except:
+            print("Warning: no Pose3 at key: {0}".format(key))
+
+        if lastIndex is not None:
+            lastKey = keys.at(lastIndex)
+            try:
+                lastPose = pose3Values.atPose3(lastKey)
+            except:
+                print("Warning: no Pose3 at key: {0}".format(lastKey))
+                pass
+
+            if marginals:
+                P = marginals.marginalCovariance(lastKey)
+            else:
+                P = None
+
+            plot_pose3(fignum, lastPose, P, scale)
+
+        lastIndex = i
+
+    # Draw final pose
+    if lastIndex is not None:
+        lastKey = keys.at(lastIndex)
+        try:
+            lastPose = pose3Values.atPose3(lastKey)
+            if marginals:
+                P = marginals.marginalCovariance(lastKey)
+            else:
+                P = None
+
+            plot_pose3(fignum, lastPose, P, scale)
+
+        except:
+            pass