Console ANSI Art Generator

ansi_art_converter.py

#! /usr/bin/env python2
# -*- coding: utf-8 -*-
#
# This routine is adapted from: https://gist.github.com/jdiaz5513/9218791
#
# Things I changed:
# * Cache the results of color_distance() lookups, for a big speed-up.
# * Adjusted the RGB values for ANSI_COLORS to match original CGA values
# * Changed default fill character to a PC-ANSI shaded block character
# * Added some timer code to help with optimizing the conversion routine

from PIL import Image, ImageChops
from colormath.color_conversions import convert_color
from colormath.color_objects import LabColor
from colormath.color_objects import sRGBColor as RGBColor
from colormath.color_diff import delta_e_cmc as cmc
import argparse
import sys
import time
import functools


def timeit(func):
	@functools.wraps(func)
	def newfunc(*args, **kwargs):
		startTime = time.time()
		func(*args, **kwargs)
		elapsedTime = time.time() - startTime
		print('function [{}] finished in {} ms'.format(
			func.__name__, int(elapsedTime * 1000)))
	return newfunc

ANSI_SHADED_BLOCKS = (
	chr(176),
	chr(177),
	chr(178),
	chr(219),
)

ANSI_CODES = (
	'\033[00;30m',	 # black
	'\033[00;31m',	 # red
	'\033[00;32m',	 # green
	'\033[00;33m',	 # yellow
	'\033[00;34m',	 # blue
	'\033[00;35m',	 # magenta
	'\033[00;36m',	 # cyan
	'\033[00;37m',	 # gray
	'\033[01;30m',   # dark gray
	'\033[01;31m',   # bright red
	'\033[01;32m',   # bright green
	'\033[01;33m',   # bright yellow
	'\033[01;34m',   # bright blue
	'\033[01;35m',   # bright magenta
	'\033[01;36m',   # bright cyan
	'\033[01;37m',   # white
)
ANSI_COLORS = (
	RGBColor(0, 0, 0),		  # black
	RGBColor(170, 0, 0),		# red
	RGBColor(0, 170, 0),		# green
	RGBColor(170, 85, 0),	  # yellow
	RGBColor(0, 0, 170),		# blue
	RGBColor(170, 0, 170),	  # magenta
	RGBColor(0, 170, 170),	  # cyan
	RGBColor(170, 170, 170),	# gray
	RGBColor(85, 85, 85),	# dark gray
	RGBColor(255, 85, 85),		# bright red
	RGBColor(85, 255, 85),		# bright green
	RGBColor(255, 255, 85),	  # bright yellow
	RGBColor(85, 85, 255),	  # bright blue
	RGBColor(255, 85, 255),	  # bright magenta
	RGBColor(85, 255, 255),	  # bright cyan
	RGBColor(255, 255, 255),	# white
)
ANSI_RESET = '\033[0m'
INFINITY = float('inf')

COLOR_CACHE = {}

def closest_ansi_color(color):
	# Change RGB color value into a string we can use as a dict key
	color_id = ''.join( map(str, color) )
	# If we've calculated color_distance for this color before, it will be cached.
	# Use cached value instead of performing color_distance again.
	if color_id in COLOR_CACHE:
		return ANSI_CODES[ COLOR_CACHE[color_id] ]

	# Look up the closest ANSI color
	else:
		color = RGBColor(*color[:3])
		closest_dist = INFINITY
		closest_color_index = 0
		for i, c in enumerate(ANSI_COLORS):
			d = color_distance(c, color)
			if d < closest_dist:
				closest_dist = d
				closest_color_index = i
		# Add this index to our color cache so we don't have to look it up again
		COLOR_CACHE[color_id] = closest_color_index
		return ANSI_CODES[closest_color_index]


def color_distance(c1, c2):
	# return a value representing a relative distance between two RGB
	# color values, weighted for human eye sensitivity
	cl1 = convert_color(c1, LabColor)
	cl2 = convert_color(c2, LabColor)
	return cmc(cl1, cl2, pl=1, pc=1)

@timeit
def convert_image(filename, output_file, fill_char=ANSI_SHADED_BLOCKS[2]):
	# render an image as ASCII by converting it to RGBA then using the
	# color lookup table to find the closest colors, then filling with 
	# fill_char
	# TODO: use a set of fill characters and choose among them based on
	# color value
	im = Image.open(filename)
	if im.mode != 'RGBA':
		im = im.convert('RGBA')

	# Shrink the image to 79px wide, to match width of ANSI art
	basewidth = 79
	aspect_ratio = 0.5  #This value shrinks image vertically to accomodate tall characterset
	wpercent = (basewidth/float(im.size[0]))
	hsize = int( (float(im.size[1])*float(wpercent)) * float(aspect_ratio) )
	im = im.resize((basewidth,hsize))
	im.save('converted.png') # Save copy of shrunk image for debugging

	w = im.size[0]
	o = ''
	last_color = None
	for i, p in enumerate(im.getdata()):
		if i % w == 0:
			o += '\n'
		if im.mode == 'RGBA' and p[3] == 0:
			o += ' ' * len(fill_char)
		else:
			c = closest_ansi_color(p)
			if last_color != c:
				o += c
				last_color = c
			o += fill_char
	o += ANSI_RESET + '\n\n'
	if output_file is not sys.stdout:
		output_file = open(output_file, 'wb')
	output_file.write(o)
	output_file.close()


if __name__ == '__main__':
	parser = argparse.ArgumentParser()
	parser.add_argument('filename', help='File to convert to ASCII art')
	parser.add_argument('-o', '--output_file', nargs='?', default=sys.stdout,
						help='Path to the output file, defaults to stdout')
	parser.add_argument('-f', '--fill_char', nargs='?', default=ANSI_SHADED_BLOCKS[2],
						help='Character to use for solid pixels in the image')
	args = parser.parse_args()
	convert_image(args.filename, args.output_file, fill_char=args.fill_char)