some plotting utils for displaying camera calibration results

plot.py

import matplotlib.pyplot as plt
import numpy as np
import cv2

"""
Some utilities fr plotting camera calibrations:

Example:
fig = plt.figure(figsize=[10,10])
ax = fig.add_subplot(111, projection='3d')
draw_target(ax, pts, shape=[9, 5])

"""

def axes_equal(ax):
    """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.
    
    From https://stackoverflow.com/a/31364297
    
    Input
      ax: a matplotlib axis, e.g., as output from plt.gca().
    """
    x_limits = ax.get_xlim3d()
    y_limits = ax.get_ylim3d()
    z_limits = ax.get_zlim3d()

    x_range = abs(x_limits[1] - x_limits[0])
    x_middle = np.mean(x_limits)
    y_range = abs(y_limits[1] - y_limits[0])
    y_middle = np.mean(y_limits)
    z_range = abs(z_limits[1] - z_limits[0])
    z_middle = np.mean(z_limits)

    # The plot bounding box is a sphere in the sense of the infinity
    # norm, hence I call half the max range the plot radius.
    plot_radius = 0.5*max([x_range, y_range, z_range])

    ax.set_xlim3d([x_middle - plot_radius, x_middle + plot_radius])
    ax.set_ylim3d([y_middle - plot_radius, y_middle + plot_radius])
    ax.set_zlim3d([z_middle - plot_radius, z_middle + plot_radius])


def repro_error(ax, objpoints, imgpoints, rvecs, tvecs, mtx, dist):
    """ plot the reprojection error of a camera calibratioin"""
    lim = 0
    for i, op in enumerate(objpoints):
        ip2 = cv2.projectPoints(op, rvecs[i], tvecs[i], mtx, dist)[0]
        p = (ip2 - imgpoints[i]).reshape(2, -1)
        lim = lim if lim > np.abs(p).max() else np.abs(p).max()
        ax.plot(p[0], p[1], 'x', lw=0, ms=8)
    lim += (lim * 0.1)
    ax.set_xlim([-lim, lim])
    ax.set_ylim([-lim, lim])
    ax.set_xlabel('X Error (px)')
    ax.set_ylabel('Y Error (px)')
    ax.set_title('Reprojection Error')


def two(ax, left, right, lines=False):
    """
    plot two images side by side.

    if lines == True, plot horizontal  lines over each image.
    this is to check correspondence for stereo rectified images.
    """
    ax[0].imshow(left, cmap='gray')
    ax[1].imshow(right, cmap='gray')
    if lines:
        [ylim, xlim], dist = left.shape, 100
        ax[0].autoscale(False)
        ax[1].autoscale(False)
        for y in range(dist, ylim, dist):
            ax[0].plot([0, xlim], [y, y], 'r', lw=2.0)
            ax[1].plot([0, xlim], [y, y], 'r', lw=2.0)

            
def target(ax, pts, shape=[9, 5]):
    """ draw an opencv chessboard target"""
    [m, n], x, y, z = shape, [], [], []
    for i in range(n):
        j = i * m
        x += [pts[j][0], pts[j + m - 1][0], np.nan]
        y += [pts[j][1], pts[j + m - 1][1], np.nan]
        z += [pts[j][2], pts[j + m - 1][2], np.nan]
    for i in range(m):
        k = (n-1)*m
        x += [pts[i][0], pts[i + k][0], np.nan]
        y += [pts[i][1], pts[i + k][1], np.nan]
        z += [pts[i][2], pts[i + k][2], np.nan]
    ax.plot(x, y, z, lw=0.7)


def camera(ax, size=1.0, R=None, T=None):
    """draw a camera at the origin."""
    cam = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0],  [0, 0, 1]]) * size
    fst = np.array([[0, 0, 0], [-1, -1, 4], [-1, 1, 4],
                    [1, 1, 4], [1, -1, 4]]) * 0.5 * size

    if R is not None and T is not None:
        cam = (R.dot(cam.T) + T).T
        fst = (R.dot(fst.T) + T).T

    # cam
    ax.plot(*zip(cam[0], cam[1]), color='r')
    ax.plot(*zip(cam[0], cam[2]), color='g')
    ax.plot(*zip(cam[0], cam[3]), color='b')
    # frustum
    ax.plot(*zip(fst[0], fst[1]), color='gray')
    ax.plot(*zip(fst[0], fst[2]), color='gray')
    ax.plot(*zip(fst[0], fst[3]), color='gray')
    ax.plot(*zip(fst[0], fst[4]), color='gray')
    ax.plot(*zip(fst[1], fst[2], fst[3], fst[4], fst[1]), color='gray')

set_axes_equal is from the SO answer:
https://stackoverflow.com/a/31364297