Python script to convert OpenCV YAML to LaTeX bmatrix

yaml_to_latex.py

#!/usr/bin/python

"""
I'm expecting a YAML file from OpenCV that looks like:

    %YAML:1.0
    K: !!opencv-matrix
       rows: 3
       cols: 3
       dt: d
       data: [ 1.7331257894078371e+03, 0., 3.2615246280978710e+02, 0.,
           1.7348511763245624e+03, 2.4016376751013206e+02, 0., 0., 1. ]
    D: !!opencv-matrix
       rows: 1
       cols: 5
       dt: d
       data: [ -4.4987143758134934e-01, -4.0610929767103503e+00,
           1.3535676742197126e-03, -6.7574354060614061e-04,
           6.2389566493881304e+01 ]

Example:
$ ./yaml_to_latex.py ../build/calib_left.yaml K 1.7f
\begin{bmatrix}
1733.1257894 & 0.0000000 & 326.1524628 \\
0.0000000 & 1734.8511763 & 240.1637675 \\
0.0000000 & 0.0000000 & 1.0000000 \\
\end{bmatrix}

"""

import sys

import yaml
import numpy as np

"""
Since PyYAML doesn't understand OpenCV data types, you need to specifc a
constructor for each OpenCV dta type that you are trying to load.

See: http://stackoverflow.com/a/15942429
"""
def opencv_matrix(loader, node):
    mapping = loader.construct_mapping(node, deep=True)
    mat = np.array(mapping["data"])
    mat.resize(mapping["rows"], mapping["cols"])
    return mat
yaml.add_constructor(u"tag:yaml.org,2002:opencv-matrix", opencv_matrix)

"""
A yaml representer is for dumping structs into a yaml node.
So for an opencv_matrix type (to be compatible with c++'s FileStorage)
we save the rows, cols, type and flattened-data

See: http://stackoverflow.com/a/35758853
"""
def opencv_matrix_representer(dumper, mat):
    mapping = {'rows': mat.shape[0], 'cols': mat.shape[1], 'dt': 'd', 'data': mat.reshape(-1).tolist()}
    return dumper.represent_mapping(u"tag:yaml.org,2002:opencv-matrix", mapping)
yaml.add_representer(np.ndarray, opencv_matrix_representer)

def get_yaml(file):
    # Open the YAML file.
    f = open(file, 'r')

    # Ignore the first line: %YAML:1.0
    f.readline()

    # Read in the matrices
    return yaml.load(f)

# https://gist.github.com/malloc47/4665827
# http://stackoverflow.com/a/17131750
def to_matrix(a, frmt='{:1.2f}', arraytype='bmatrix'):
    """Returns a LaTeX array
    :a: numpy array
    :returns:
    :prints: LaTeX array

    Example:
    to_matrix(Krm, frmt = '{:1.2f}', arraytype = 'array')
    """
    if len(a.shape) > 2:
        raise ValueError('bmatrix can at most display two dimensions')
    lines = str(a).replace('[', '').replace(']', '').splitlines()

    # Create raw string literal vectors
    rv = [r'\begin{' + arraytype + '}']
    for l in lines:
        tmp = ""
        for (i, num) in enumerate(l.split()):
            tmp += frmt.format(float(num)) + r' '

            # Don't add a '&' if at the very end of a line
            if i<len(l.split())-1:
                tmp += r'& '

        # finish off the line
        rv += [tmp + r'\\']
    rv += [r'\end{' + arraytype + '}']
    return '\n'.join(rv)

def generate_latex(mat, frmt=None):

    frmt = '{:1.3f}' if frmt is None else '{:'+frmt+'}'

    # Since generate_latex depends on how numpy looks as a string,
    # give our self lots of precision to work with.
    np.set_printoptions(precision=10, linewidth=200)

    return to_matrix(mat, frmt=frmt, arraytype='bmatrix')


if __name__ == "__main__":

    if len(sys.argv) < 3:
        print("You don't know how to use this script... sorry :/")
        sys.exit(0)

    yaml_file = sys.argv[1] # location of the OpenCV YAML file
    mat_key = sys.argv[2]   # what is the name of the mat to LaTeX?
    # Format specifier (optional)
    frmt = sys.argv[3] if len(sys.argv) == 4 else None

    # Get a dict of the matrices in the YAML
    data = get_yaml(yaml_file)

    # Print out the generated LaTeX
    print(generate_latex(data[mat_key], frmt))