hacks/painter/painter.py

"""Converts a bitmap image into a minimal number of covering
rectangles.

Mainly plays with different methods.
"""

# python painter.py images/*.png

from __future__ import print_function

import sys
import collections
import time

import numpy
import scipy.misc

# A rectangle.
Rect = collections.namedtuple('Rect', 'x y w h')

# Chracters to use when printing the image to the screen.
EXT = range(ord('1'), ord('9') + 1) \
      + range(ord('a'), ord('z') + 1) \
      + range(ord('A'), ord('Z') + 1)
CHARS = [ord('.'), ord('#')] + EXT*20


def parse(filename):
    """Parse a image into a (0, 1) numpy array."""
    img = scipy.misc.imread(filename)
    img = numpy.clip(img, 0, 1)
    return img ^ 1


def render(shape, rects, with_ids=False):
    """Render rectangles into a new array of the given shape."""
    img = numpy.zeros(shape, numpy.uint16)

    for i, rect in enumerate(rects):
        for y in range(rect.y, rect.y + rect.h):
            for x in range(rect.x, rect.x + rect.w):
                if with_ids:
                    img[y][x] = i + 2
                else:
                    img[y][x] = 1
    return img

def show(img):
    """Print an image to the console."""
    for row in img:
        print(''.join(chr(CHARS[x]) for x in row))


def pointwise(img):
    """Method where each point is a rectangle."""
    for y, row in enumerate(img):
        for x, pix in enumerate(row):
            if pix:
                yield Rect(x, y, 1, 1)

def rowwise(img):
    """Method where rectangles are one row high."""
    for y, row in enumerate(img):
        start = None
        for x, pix in enumerate(row):
            if pix:
                if start is None:
                    start = x
            else:
                if start is not None:
                    yield Rect(start, y, x - start, 1)
                    start = None
        if start is not None:
            yield Rect(start, y, x - start + 1, 1)

def rotate(img, fun):
    """Rotate an image and rotate the final rectangles."""
    h, w = img.shape
    for rect in fun(numpy.rot90(img)):
        yield Rect(w - rect.y - 1, rect.x, rect.h, rect.w)

def colwise(img):
    """Method where rectangles are one column wide."""
    return list(rotate(img, rowwise))

def mergeup(rects):
    """Merge a rectangle into the identical width one below it."""
    # TODO(michaelh): ugly.
    while True:
        out = []
        merged = set()
        rects.sort(key=lambda x: x.y)

        for rect in rects:
            if rect in merged:
                continue

            top = rect.y + rect.h
            below = [x for x in rects if x.y == top]
            matches = [x for x in below if x.x == rect.x and x.w == rect.w]

            if matches:
                match = matches[0]
                merged.add(match)
                out.append(Rect(rect.x, rect.y, rect.w, rect.h + match.h))
            else:
                out.append(rect)
        if not merged:
            return out
        rects = out

def mergeleft(rects):
    """Merge a rectangle into the identical height one beside it."""
    # TODO(michaelh): merge with mergeup().

    while True:
        out = []
        merged = set()
        rects.sort(key=lambda x: x.x)

        for rect in rects:
            if rect in merged:
                continue

            edge = rect.x + rect.w
            adjacent = [x for x in rects if x.x == edge]
            matches = [x for x in adjacent if x.y == rect.y and x.h == rect.h]

            if matches:
                match = matches[0]
                merged.add(match)
                out.append(Rect(rect.x, rect.y, rect.w + match.w, rect.h))
            else:
                out.append(rect)
        if not merged:
            return out
        rects = out

            
def rowcolwise(img):
    """Method that picks the best of rowwise and colwise."""
    rows = list(rowwise(img))
    cols = list(colwise(img))
    return rows if len(rows) < len(cols) else cols

def mergerowcolwise(img):
    """Method that picks the best of the merged rowwise and colwise."""
    rows = mergeup(list(rowwise(img)))
    cols = mergeleft(list(colwise(img)))
    return rows if len(rows) < len(cols) else cols

def bothwise(img):
    """Method that combines row and col wise, then drops rectangles
    that are completly covered by larger rectangles.
    """
    rows = list(rowwise(img))
    cols = list(colwise(img))
    rects = rows + cols
    # Sort by area.  Smaller are less likely to contribute.
    rects.sort(key=lambda x: x.w*x.h)
    drop = set()

    for rect in rects:
        others = [x for x in rects if x != rect and x not in drop]
        got = render(img.shape, others)

        if not (got ^ img).any():
            # Dropping this rect didn't change the image.
            drop.add(rect)

    rects = [x for x in rects if x not in drop]
    return mergeleft(mergeup(rects))

def main():
    # All methods to try.
    methods = [
        pointwise,
        rowwise,
        colwise,
        rowcolwise,
        mergerowcolwise,
        bothwise,
    ]

    for filename in sys.argv[1:]:
        img = parse(filename)
        best = None

        print('{}:'.format(filename))
        for method in methods:
            start = time.clock()
            scan = list(method(img))
            elapsed = time.clock() - start

            name = method.__name__
            print('  {} -> {} rects in {:.3f} s'.format(name, len(scan), elapsed))

            rendered = render(img.shape, scan)
            diff = img ^ rendered

            if diff.any():
                print('Render error!')
            else:
                if best is None or len(scan) <= len(best):
                    best = scan
        # Show the best.
        show(render(img.shape, best, True))

if __name__ == '__main__':
    main()