merged ascii repo and sopel/ascii

modules/ascii.py

@@ -4,7 +4,6 @@
 ASCII
 """
 from io import BytesIO
-import argparse
 
 import requests
 from PIL import Image, ImageFont, ImageDraw
@@ -17,58 +16,65 @@ from tools import web
 
 ASCII_CHARS = "$@%#*+=-:. "
 BRAIL_CHARS = "⠿⠾⠼⠸⠰⠠ "
-HEADERS = {'User-Agent': 'Gimme the ascii.'}
+HEADERS = {'User-Agent': 'Gimme ascii.'}
 
 
-def scale_image(image, size=(100,100)):
+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 accepted size tuple.
+	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 > original_height:
+	if original_width <= maxDim and original_height <= maxDim:
-		if original_width > size[0]:
+		new_width, new_height = image.size
-			new_width = size[0]
+	elif original_width > original_height:
+		new_width = maxDim
 		aspect_ratio = original_height/float(original_width)
 		new_height = int(aspect_ratio * new_width)
 	else:
-			new_width, new_height = image.size
+		new_height = maxDim
-	else:
-		if original_height > size[1]:
-			new_height = size[1]
 		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, scale="ascii"):
+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 ASCII_CHARS.
+	0-255 normalized to the length of the chosen scale.
 	"""
 	scales = {"ascii": ASCII_CHARS,
-		"ascii_reverse": "".join(list(reversed(ASCII_CHARS))),
+		"ascii_reverse": "".join(reversed(ASCII_CHARS)),
 		"brail": BRAIL_CHARS,
-		"brail_reverse": "".join(list(reversed(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())
-	chars = []
+	pixels = [pixels[i:i + image.size[0]] for i in range(0, len(pixels),
-	for pixel in pixels:
-		index = int(pixel/range_width)
-		chars.append(scales[scale][index])
-
-	chars = "".join(chars)
-	chars = [chars[i:i + image.size[0]] for i in range(0, len(chars),
 		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)
 
 
@@ -124,11 +130,13 @@ def open_image(imagePath):
 		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 as e:
+	except FileNotFoundError:
 		return f"File not found: {imagePath}"
 	except Exception as e:
 		return(f"Error opening image: {imagePath}\n{e}")
@@ -136,21 +144,6 @@ def open_image(imagePath):
 	return image
 
 
-def colorize(chars, image, code):
-	"""
-	Colorizes the ascii matrix.
-	"""
-	prefix = {"irc": "\03", "ansi":"\033"}
-	chars = chars.split("\n")
-	for j in range(0, image.size[1]):
-		new_row = ""
-		for k in range(0, image.size[0]):
-			new_row += prefix[code] + char_color(image.getpixel((k,j)), code)
-			new_row += chars[j][k]
-		chars[j] = new_row
-	return "\n".join(chars)
-
-
 def alpha_composite(front, back):
 	"""Alpha composite two RGBA images.
 
@@ -193,36 +186,33 @@ def alpha_composite_with_color(image, color=(255, 255, 255)):
 	return alpha_composite(image, back)
 
 
-def image_to_ascii(image, reverse=False, colors=None, brail=False):
+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_with_color(image).convert("RGB")
 
 	image = scale_image(image)
-	image_grey = image.convert('L')	# convert to grayscale
 
-	if reverse:
-		if brail:
-			scale = "brail_reverse"
-		else:
-			scale = "ascii_reverse"
-	else:
 	if brail:
 		scale = "brail"
 	else:
 		scale = "ascii"
-	chars = pixels_to_chars(image_grey, scale)
+	if reverse:
+		scale += "_reverse"
+
+	chars = pixels_to_chars(image, scale, color)
 
-	if colors:
-		chars = colorize(chars, image, colors)
 	image.close()
-	image_grey.close()
 	del(image)
-	del(image_grey)
 	return chars
 
 
@@ -230,7 +220,7 @@ def ascii_to_image(image_ascii):
 	"""
 	Creates a plain image and draws text on it.
 	"""
-	# TODO: make font type and size non-fixed
+	# 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
 
@@ -241,27 +231,27 @@ def ascii_to_image(image_ascii):
 	return image
 
 
-def handle_gif(image, output, reverse=False):
+def handle_gif(imagePath, **kwargs):
 	"""
 	Handle gifs seperately.
 	"""
-	# image = open_image(args.imagePath)
+	image = open_image(imagePath)
 	ascii_seq = []
-	new_image = ascii_to_image(image_to_ascii(image, reverse))
+	new_image = ascii_to_image(image_to_ascii(image, **kwargs))
 	image.seek(1)
 	while True:
 		try:
-			im = ascii_to_image(image_to_ascii(image, reverse))
+			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,
+	# 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(output, "wb") as file:
+	with open(kwargs["output"], "wb") as file:
 		numpngw.write_apng(file, np_ascii_seq)
 
 
@@ -275,14 +265,18 @@ def ascii(bot, trigger):
 	"""
 	if not trigger.group(2):
 		return bot.say()
-	parser = argparse.ArgumentParser()
+	parser = argparse.ArgumentParser(add_help=False)
 	parser.add_argument("imagePath")
-	parser.add_argument("-r", "--reverse", action="store_true")
+	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"
 
@@ -293,24 +287,25 @@ def ascii(bot, trigger):
 		bot.reply("Internet requests only.")
 		return
 
-	image = open_image(args.imagePath)
 	if args.animated:
-		handle_gif(image, "temp.png", args.reverse)
+		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)
-		# print(res.text)
 		bot.say(res.text)
 	else:
-		image_ascii = image_to_ascii(image, args.reverse, args.color, args.brail)
+		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 full path to the image file.")
+		help="The path to the image file. May be a local path or internet URL.")
 	parser.add_argument(
 		"-r",
 		"--reverse",
@@ -323,9 +318,10 @@ if __name__=='__main__':
 	parser.add_argument(
 		"-i",
 		"--image",
+		dest="drawImage",
 		action="store_true",
-		help="Outputs the ascii art as an image rather than plain text. \
+		help="Outputs the ascii art as an image rather than text. Requires \
-		Requires --output.")
+		--output.")
 	parser.add_argument(
 		"-a",
 		"--animated",
@@ -335,23 +331,39 @@ if __name__=='__main__':
 		"-c",
 		"--color",
 		type=str,
-		help="Colorizes the ascii matrix.")
+		help="Colorizes the ascii matrix. Currently supported modes are 'irc' \
-	parser.set_defaults(reverse=False, image=False, animated=False)
+		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.image = True
+		args.drawImage = True
-	if args.image:		# --image requires --output
+	if args.drawImage:		# --image requires --output
 		if not args.output:
 			parser.error("--image requires --output")
 
-	image = open_image(args.imagePath)
 	if args.animated:
-		handle_gif(image, args.output, args.reverse)
+		handle_gif(**vars(args))
 		exit()
 
-	image_ascii = image_to_ascii(image, args.reverse, args.color)
+	image_ascii = image_to_ascii(None, **vars(args))
-	if args.image:
+	if args.drawImage:
 		image = ascii_to_image(image_ascii)
 		image.save(args.output, "PNG")
 	elif args.output: