sopel/modules/ascii.py

#! /usr/bin/env python3
# -*- coding: utf-8 -*-

from io import BytesIO
from PIL import Image
import requests
import module
from tools import web
from PIL import ImageFont
from PIL import ImageDraw

ASCII_CHARS = "$@%#*+=-:. "
headers = {'User-Agent': 'we wuz ascii and shiet'}


def scale_image(image, size=(100,100)):
	"""
	Resizes an image while preserving the aspect ratio. Chooses the
	dimension to scale by based on whichever is larger, ensuring that
	neither width or height is ever larger than the accepted size tuple.
	"""
	original_width, original_height = image.size
	original_width = original_width * 2 # because characters are generally
	if original_width > original_height:  # displayed as a 1:2 square
		if original_width > size[0]:
			new_width = 100
			aspect_ratio = original_height/float(original_width)
			new_height = int(aspect_ratio * new_width)
		else:
			new_width, new_height = image.size
	else:
		if original_height > size[1]:
			new_height = 100
			aspect_ratio = original_width/float(original_height)
			new_width = int(aspect_ratio * new_height)
		else:
			new_width, new_height = image.size
	image = image.resize((new_width, new_height))
	return image


def pixels_to_chars(image, reverse=False):
	"""
	Maps each pixel to an ascii char based on the range
	in which it lies.

	0-255 is divided into 11 ranges of 25 pixels each.
	"""
	range_width = int(255 / len(ASCII_CHARS)) + (255 % len(ASCII_CHARS) > 0)

	pixels_in_image = list(image.getdata())
	pixels_to_chars = []
	for pixel_value in pixels_in_image:
		if reverse:
			index = -int(pixel_value/range_width)-1
		else:
			index = int(pixel_value/range_width)
		pixels_to_chars.append(ASCII_CHARS[ index ])

	return "".join(pixels_to_chars)


def open_image(imagePath):
	"""
	Opens the image at the supplied file path in PIL. If an internet URL
	is supplied, it will download the image and then open it.
	"""
	try:
		if imagePath.startswith("http"):
			res = requests.get(imagePath, headers=headers, verify=True, timeout=20)
			res.raise_for_status()
			image = Image.open(BytesIO(res.content))
		else:
			image = Image.open(imagePath)
	except FileNotFoundError as e:
		return e
	except OSError:
		return e
	except Exception as e:
		return("Error opening image file: " + imagePath)

	return image


def image_to_ascii(image, reverse=False):
	"""
	Reads an image file and converts it to ascii art. Returns a
	newline-delineated string. If reverse is True, the ascii scale is
	reversed.
	"""
	image = scale_image(image)
	image = image.convert('L')

	chars = pixels_to_chars(image, reverse)

	image_ascii = []
	for index in range(0, len(chars), image.size[0]):
		image_ascii.append( chars[index: index + image.size[0]] )
	image.close()
	del image
	return "\n".join(image_ascii)


def ascii_to_image(image_ascii):
	"""
	Creates a plain image and draws text on it.
	"""
	width = len(image_ascii[:image_ascii.index("\n")]) * 8
	height = (image_ascii.count("\n")+1) * 12 + 4

	font = ImageFont.truetype("LiberationMono-Regular.ttf", 14)
	image = Image.new("RGB", (width, height), (255,255,255))
	draw = ImageDraw.Draw(image)
	draw.text((0,0), image_ascii, (0,0,0), font=font, spacing=0)
	return image


def handle_gif(output, reverse=False):
	image = open_image(args.imagePath)
	ascii_seq = []
	new_image = ascii_to_image(image_to_ascii(image, reverse))
	image.seek(1)
	while True:
		try:
			im = ascii_to_image(image_to_ascii(image, reverse))
			ascii_seq.append()
			image.seek(image.tell()+1)
		except EOFError:
			break # end of sequence

	new_image.save(args.output, save_all=True, append_images=ascii_seq, duration=60, loop=0, optimize=True)


@module.rate(user=60)
@module.require_chanmsg(message="It's impolite to whisper.")
@module.commands('ascii')
@module.example('.ascii [-r] https://www.freshports.org/images/freshports.jpg')
def ascii(bot, trigger):
	"""
	Downloads an image and converts it to ascii.
	"""
	reverse = False
	if trigger.group(3) == "-r":
		imagePath = trigger.group(4)
		reverse = True
	else:
		imagePath = trigger.group(2)

	if not web.secCheck(bot, imagePath):
		return bot.reply("Known malicious site. Ignoring.")

	if not imagePath.startswith("http"):
		bot.reply("Internet requests only.")
		return

	image = open_image(imagePath)
	image_ascii = image_to_ascii(image, reverse)
	bot.say(image_ascii)


if __name__=='__main__':
	import argparse

	# TODO: satisfy PEP8
	parser = argparse.ArgumentParser(description="Converts an image file to ascii art.")
	parser.add_argument("imagePath", help="The full path to the image file.")
	parser.add_argument("-r", "--reverse", action="store_true", help="Reverses the ascii scale.")
	parser.add_argument("-o", "--output", help="Outputs the ascii art into a file at the specified path.")
	parser.add_argument("-i", "--image", action="store_true", help="Outputs the ascii art as an image rather than plain text. Requires --output.")
	parser.add_argument("-a", "--animated", action="store_true", help="Handles animated GIFs. Includes --image.")
	parser.set_defaults(reverse=False, image=False, animated=False)
	args = parser.parse_args()

	if args.animated:	# --animated includes --image
		args.image = True
	if args.image:		# --image requires --output
		if not args.output:
			parser.error("--image requires --output")

	if args.animated:
		handle_gif(args.output, args.reverse)
	else:
		image = open_image(args.imagePath)
		image_ascii = image_to_ascii(image, args.reverse)
		if args.image:
			image = ascii_to_image(image_ascii)
			image.save(args.output, "PNG")
		elif args.output:
			with open(args.output, "w+") as file:
				file.write(image_ascii)
		else:
			print(image_ascii)