Roger's color ASCII-art thing

pyaa.py

#!/usr/bin/env python

import argparse
import string
from PIL import Image, ImageFont, ImageDraw
from operator import itemgetter

def image(string):
    '''catch the IOError if Image.open() goes screwy and raise a nicely formatted
    argparse error instead'''
    raise_error = False
    try:
        im = Image.open(string)
    except IOError:
        raise_error = True
    if raise_error:
        msg = "'%s' either does not exist or is not an image file" % string
        raise argparse.ArgumentTypeError(msg)
    return im

def build_letter_cache(font='FreeMonoBold.ttf', size=20):
    '''returns a list of tuples with letter images and data objects correlating to the images'''
    letter_cache = []
    letters = string.digits + string.ascii_letters + string.punctuation
    usefont = ImageFont.truetype(font, size)
    color = 0
    i = 0
    while i < len(letters):
        letter = letters[i]
        im = Image.new(size=(int(size*0.8), size), mode='1', color=1)
        draw = ImageDraw.Draw(im)
        draw.text((size*.1, 0), letter, color, font=usefont)
        letter_cache.append((im, im.getdata()))
        i += 1
    return letter_cache

def evaluate_image(im, size=20):
    '''returns a list of tuples containing 1-bit image data object, corresponding dominant color, 
    and box for re-assembly'''
    #Speedups needed - possible without concurrency?
    image_blocks = []
    width, height = im.size
    box_size = (int(size*0.8), size)
    h = 0
    while h < height+box_size[1]:
        w = 0
        while w < width+box_size[0]:
            maxcolors = 256l
            colors = None
            box = (w, h, w+box_size[0], h+box_size[1])
            lil_im = im.crop(box=box)
            while not colors:
                # Image.getcolors() returns None if there are more colors than maxcolors
                # but most blocks aren't going to have more than 256 colors
                # this loop compensates for the few that do
                colors = lil_im.getcolors(maxcolors=maxcolors)
                maxcolors += 256
            # get the color with the highest count in the colors list
            # itemgetter is faster than looping or a lambda
            color = max(colors, key=itemgetter(0))[1]
            lil_im = lil_im.convert(mode='1')
            image_blocks.append((lil_im.getdata(), color, box))
            w += box_size[0]
        h += box_size[1]
    return image_blocks

def build_aa_image(size, letters, blocks):
    '''build ASCII art image by matching letter blocks with image blocks'''
    new_image = Image.new(mode='RGB', size=size, color=(255,255,255))
    for lil_data, color, box in blocks:
        diffs = []
        for letter, letterdata in letters:
            pairs = zip(lil_data, letterdata)
            diff = sum(abs(p1-p2) for p1, p2 in pairs)
            diffs.append(diff)
        letter = letters[diffs.index(min(diffs))][0]
        letter = letter.convert(mode='RGB')
        # replace color in the letter image, quick method
        pixdata = letter.load()
        for y in xrange(letter.size[1]):
            for x in xrange(letter.size[0]):
                if pixdata[x, y] == (0, 0, 0):
                    pixdata[x, y] = color
        new_image.paste(letter, box)
    return new_image

parser = argparse.ArgumentParser()
parser.add_argument('filename', type=image, help="Filename of an image to process")
parser.add_argument('-o', '--output', help="Output image to a file")
# TODO resizing, font selection, ???
args = parser.parse_args()

im = args.filename
size = 10
if im.size[1] > 1440:
    ratio = int(im.size[1] / 1440)
    size = size * ratio
if im.mode is not 'RGB':
    im = im.convert(mode='RGB')
letters = build_letter_cache(size=size)
image_blocks = evaluate_image(im, size=size)
new_image = build_aa_image(im.size, letters, image_blocks)
if args.output:
    # TODO does not validate if already exists
    output = args.output
    new_image.save(output)
else:
    new_image.show()