Bayer Matrix Generator

bayer_matrix.py

"""
Bayer Matrix

Generate a PNG Bayer Matrix

Usage:
  bayer_matrix <order> [--output=<file>]

Options:
  -h --help                     Show this screen
  -o <file>, --output=<file>    Output to the named file
"""
import sys

import docopt
import png
import numpy


BITDEPTH = 16


def bayer_matrix_unnormalized(n):
    if n == 1:
        return numpy.array([[0, 2], [3, 1]], dtype=numpy.int)

    mat = bayer_matrix_unnormalized(n - 1)
    return numpy.block([[4 * mat, 4 * mat + 2], [4 * mat + 3, 4 * mat + 1]])


def bayer_matrix_image(n):
    assert 1 <= n <= BITDEPTH

    matrix = (2**BITDEPTH) * bayer_matrix_unnormalized(n) >> 2*n
    matrix = matrix.astype(f'u{BITDEPTH // 8}')

    return png.from_array(
        matrix,
        mode='L',
        info={'bitdepth': BITDEPTH},
    )

def main(order, filename):
    image = bayer_matrix_image(order)
    image.save(filename)


if __name__ == '__main__':
    args = docopt.docopt(__doc__)
    try:
        order = int(args['<order>'])
    except ValueError:
        print('order must be an integer', file=sys.stderr)
        sys.exit(1)
    if not 1 <= order <= 16:
        print('order must be between 1 and 16', file=sys.stderr)
        sys.exit(1)
    filename = args['--output'] or f'bayer_{order}.png'

    main(order, filename)