sopel/modules/ascii.py

#!/usr/bin/env python3
"""
ASCII
"""
from io import BytesIO
import argparse

import requests
from PIL import Image, ImageFont, ImageDraw
#import imageio
import numpy as np
import numpngw

import module
from tools import web

ASCII_CHARS = "$@%#*+=-:. "
BRAIL_CHARS = "⠿⠾⠼⠸⠰⠠ "
HEADERS = {'User-Agent': 'Gimme ascii.'}


def scale_image(image, maxDim=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 maxDim argument.

	Because text characters are typically pretty close to a 1:2 rectangle,
	we weight the width twice as much.
	"""
	original_width, original_height = image.size
	original_width = original_width * 2
	if original_width <= maxDim and original_height <= maxDim:
		new_width, new_height = image.size
	elif original_width > original_height:
		new_width = maxDim
		aspect_ratio = original_height/float(original_width)
		new_height = int(aspect_ratio * new_width)
	else:
		new_height = maxDim
		aspect_ratio = original_width/float(original_height)
		new_width = int(aspect_ratio * new_height)
	image = image.resize((new_width, new_height), Image.ANTIALIAS)
	return image


def pixels_to_chars(image, scale="ascii", color_code=None):
	"""
	Maps each pixel to an ascii char based on where it falls in the range
	0-255 normalized to the length of the chosen scale.
	"""
	scales = {"ascii": ASCII_CHARS,
		"ascii_reverse": "".join(reversed(ASCII_CHARS)),
		"brail": BRAIL_CHARS,
		"brail_reverse": "".join(reversed(BRAIL_CHARS))}

	color_prefix = {"irc": "\03", "ansi":"\033"}

	range_width = int(255 / len(scales[scale])) + (255 % len(scales[scale]) > 0)

	pixels = list(image.getdata())
	pixels = [pixels[i:i + image.size[0]] for i in range(0, len(pixels),
		image.size[0])]

	chars = []
	for row in pixels:
		new_row = ""
		for pixel in row:
			R, G, B = pixel
			L = R * 299/1000 + G * 587/1000 + B * 114/1000
			index = int(L/range_width)
			char = scales[scale][index]
			if color_code and char is not " ":
				prefix = color_prefix[color_code] + char_color(pixel,color_code)
				char = prefix + char
			new_row += char
		chars.append(new_row)
	return "\n".join(chars)


def char_color(pixel, code="irc"):
	"""
	Maps a color to a character based on the original pixel color.
	Calculates the distance from the provided color to each of the 16
	colors in the IRC and ANSI color codes and selects the closest match.
	"""
	colors_index = ["white", "black", "blue", "green", "red", "brown", "purple",
		"orange", "yellow", "light green", "teal", "cyan", "light blue", "pink",
		"grey", "silver"]

	colors_rgb = {"white": (0,0,0), "black": (255,255,255), "blue": (0,0,255),
		"green": (0,255,0), "red": (255,0,0), "brown": (150,75,0),
		"purple": (128,0,128), "orange": (255,128,0), "yellow": (255,255,0),
		"light green": (191,255,0), "teal": (0,128,128), "cyan": (0,255,255),
		"light blue": (65,105,225), "pink":(255,192,203), "grey": (128,128,128),
		"silver": (192,192,192)}

	colors_irc = {"white": "0", "black": "1", "blue": "2", "green": "3",
		"red": "4", "brown": "5", "purple": "6", "orange": "7", "yellow": "8",
		"light green": "9", "teal": "10", "cyan": "11", "light blue": "12",
		"pink": "13", "grey": "14", "silver": "15"}

	colors_ansi = {"white": "[1;37m", "black": "[0;30m", "blue": "[0;34m",
		"green": "[0;32m", "red": "[0;31m", "brown": "[0;33m",
		"purple": "[0;35m", "orange": "[1;31m", "yellow": "[1;33m",
		"light green": "[1;32m", "teal": "[0;36m", "cyan": "[1;36m",
		"light blue": "[1;34m", "pink": "[1;35m", "grey": "[1;30m",
		"silver": "[0;37m"}

	dist = [(abs(pixel[0] - colors_rgb[color][0])**2
		+ abs(pixel[1] - colors_rgb[color][1])**2
		+ abs(pixel[2] - colors_rgb[color][2])**2)**0.5
		for color in colors_index]

	color = colors_index[dist.index(min(dist))]

	if code == "irc":
		return colors_irc[color]
	elif code == "ansi":
		return colors_ansi[color]


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. Returns a
	PIL image object.
	"""
	try:
		if imagePath.startswith("http"):
			res = requests.get(imagePath, headers=HEADERS, verify=True,
				timeout=20)
			if res.status_code == 404:
				return "404: file not found."
			res.raise_for_status()
			image = Image.open(BytesIO(res.content))
		else:
			image = Image.open(imagePath)
	except FileNotFoundError:
		return f"File not found: {imagePath}"
	except Exception as e:
		return(f"Error opening image: {imagePath}\n{e}")

	return image


def alpha_composite(image, color=(255, 255, 255)):
	"""
	Alpha composite an RGBA Image with a specified color.
	Source: http://stackoverflow.com/a/9166671/284318
	"""
	image.load()  # needed for split()
	background = Image.new('RGB', image.size, color)
	background.paste(image, mask=image.split()[3])  # 3 is the alpha channel
	return background


def image_to_ascii(image=None, reverse=False, brail=False, color=None,**kwargs):
	"""
	Reads an image file and converts it to ascii art. Returns a
	newline-delineated string. If reverse is True, the ascii scale is
	reversed.
	"""
	if not image:
		image = open_image(kwargs["imagePath"])

	if image.mode == "P":
		image = image.convert(image.palette.mode)
	if image.mode == "RGBA":
		image = alpha_composite(image).convert("RGB")
	if image.mode == "L":
		image = image.convert("RGB")

	image = scale_image(image)

	if brail:
		scale = "brail"
	else:
		scale = "ascii"
	if reverse:
		scale += "_reverse"

	chars = pixels_to_chars(image, scale, color)

	image.close()
	del(image)
	return chars


def ascii_to_image(image_ascii):
	"""
	Creates a plain image and draws text on it.
	"""
	# TODO: make font type, size and color non-fixed
	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(imagePath, **kwargs):
	"""
	Handle gifs seperately.
	"""
	image = open_image(imagePath)
	ascii_seq = []
	new_image = ascii_to_image(image_to_ascii(image, **kwargs))
	image.seek(1)
	while True:
		try:
			im = ascii_to_image(image_to_ascii(image, **kwargs))
			ascii_seq.append(im)
			image.seek(image.tell()+1)
		except EOFError:
			break # end of sequence

	# new_image.save(output, save_all=True, append_images=ascii_seq,
		# duration=60, loop=0, optimize=True)
	ascii_seq = [new_image] + ascii_seq
	np_ascii_seq = [np.array(im) for im in ascii_seq]
	with open(kwargs["output"], "wb") as file:
		numpngw.write_apng(file, np_ascii_seq)


@module.rate(user=60)
@module.require_chanmsg(message="It's impolite to whisper.")
@module.commands('ascii')
@module.example('.ascii [-rc] https://www.freshports.org/images/freshports.jpg')
def ascii(bot, trigger):
	"""
	Downloads an image and converts it to ascii.
	"""
	if not trigger.group(2):
		return bot.say()
	parser = argparse.ArgumentParser(add_help=False)
	parser.add_argument("imagePath")
	parser.add_argument("-r", "--reverse", action="store_true", help="Reverse.")
	parser.add_argument("-c", "--color", action="store_true")
	parser.add_argument("-b", "--brail", action="store_true")
	parser.add_argument("-a", "--animated", action="store_true")
	parser.add_argument("-h", "--help", action="store_true")
	args = parser.parse_args(trigger.group(2).split())

	if args.help:
		return bot.say(parser.print_help())	

	if args.color:
		args.color = "irc"

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

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

	if args.animated:
		args.output = "temp.png"
		handle_gif(**vars(args))
		file = {"file": open("temp.png", "rb")}
		res = requests.post("https://uguu.se/api.php?d=upload-tool", files=file)
		bot.say(res.text)
	else:
		image_ascii = image_to_ascii(None, **vars(args))
		bot.say(image_ascii)


if __name__=='__main__':
	import argparse

	parser = argparse.ArgumentParser(
		description="Converts an image file to ascii art.")
	parser.add_argument(
		"imagePath",
		help="The path to the image file. May be a local path or internet URL.")
	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",
		dest="drawImage",
		action="store_true",
		help="Outputs the ascii art as an image rather than text. Requires \
		--output.")
	parser.add_argument(
		"-a",
		"--animated",
		action="store_true",
		help="Handles animated GIFs. Includes --image.")
	parser.add_argument(
		"-c",
		"--color",
		type=str,
		help="Colorizes the ascii matrix. Currently supported modes are 'irc' \
		and 'ansi' for generating color codes compliant with those standards.")
	parser.add_argument(
		"--ansi",
		dest="color",
		action="store_const",
		const="ansi",
		help="Shortcut for '--color ansi'.")
	parser.add_argument(
		"--irc",
		dest="color",
		action="store_const",
		const="irc",
		help="Shortcut for '--color irc'.")
	parser.add_argument(
		"-b",
		"--brail",
		action="store_true",
		help="Uses brail unicode characters instead of ascii characters.")
	args = parser.parse_args()

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

	if args.animated:
		handle_gif(**vars(args))
		exit()

	image_ascii = image_to_ascii(None, **vars(args))
	if args.drawImage:
		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)