dots/dots.py

#! /usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Makes one of those meme pics with the dots.
"""
import random
import time

from PIL import Image, ImageDraw, ImageChops, ImageStat
import numpy as np

class Dots():
	def __init__(self):
		self.maxRadius = 11
		self.tolerance = 70

		self.original = None
		self.image = None
		self.colors = None


	def differenceImage(self, img1, img2):
		"""
		Returns the difference between two images as a numpy array.
		"""
		img1, img2 = map(np.uint8, [img1, img2])
		a = img1 - img2
		b = np.uint8(img1 < img2) * 254 + 1
		return a * b


	def compareDotsV1(self, back, dot, origDot):
		"""
		Finds the pseudo-distance between each of the images and determines
		if the placed dot is better or worse than before it was placed.
		"""
		diff1 = ImageChops.difference(origDot, dot)
		diff2 = ImageChops.difference(origDot, back)
		return np.mean(diff1) < np.mean(diff2)


	def compareDotsV12(self, back, dot, origDot):
		"""
		Finds the pseudo-distance between each of the images and determines
		if the placed dot is better or worse than before it was placed.
		"""
		diff1 = ImageChops.difference(origDot, dot)
		diff2 = ImageChops.difference(origDot, back)
		median1 = ImageStat.Stat(diff1).median
		median2 = ImageStat.Stat(diff2).median 
		return sum(median1)/3 < sum(median2)/3


	def compareDotsV2(self, back, dot, origDot):
		"""
		Finds the pseudo-distance between each of the images and determines
		if the placed dot is better or worse than before it was placed.
		"""
		diff1 = self.differenceImage(origDot, dot)
		diff2 = self.differenceImage(origDot, back)
		return np.mean(diff1) < np.mean(diff2)


	def placeDot(self, draw):
		"""
		Places random dot, then compares pixels.
		"""
		center = tuple(map(random.randrange, self.original.size))
		radius = random.randrange(1, self.maxRadius)
		boundBox = [center[0] - radius, center[1] - radius,
					center[0] + radius, center[1] + radius]
		boundBoxCir = [boundBox[0], boundBox[1], boundBox[2]-1, boundBox[3]-1]
		# color = tuple(map(random.randrange, (256, 256, 256)))
		color = random.choice(self.colors)

		# backup of the original
		back = self.image.crop(boundBox)
		draw.ellipse(boundBoxCir, fill=color)
		dot = self.image.crop(boundBox)
		origDot = self.original.crop(boundBox)

		if self.compareDots(back, dot, origDot):
			return True
		else:
			self.image.paste(back, boundBox)
			return False


	def compareImage(self):
		"""
		Compares the original image to the new image. If the mean is less
		than the tolerance, the new image is close enough to be done.
		"""
		# pass
		diff = self.differenceImage(self.original, self.image)
		mean = np.mean(diff)
		# print(mean)
		return mean < self.tolerance


	def buildDots(self, imagePath, output):
		"""
		Places randomly sized, randomly colored and radomly placed dots onto
		a blank canvas the same size as the original. The dot is then compared
		to the original: if the color of the effected pixels are closer to that
		of the original, the dot is kept. otherwise, it is removed.
		"""
		# TODO: handle RGBA
		self.original = Image.open(imagePath)
		self.colors = list(set(self.original.getdata()))
		self.image = Image.new("RGB", self.original.size, (255,255,255))
		draw = ImageDraw.Draw(self.image)

		# while True:
		for n in range(500000):
			ret = self.placeDot(draw)
			# if ret:
			# 	if self.compareImage():
			# 		break
		print(n)
		self.image.save(output, "PNG")


def testCompareDiffernece(size, func, nargs):
	"""
	Test for Dots.compareDots() or Dots.differenceImage(). func should be the
	function you're testing for, and nargs should be the number of args it
	takes.
	"""
	images = []
	for n in range(nargs):
		color = (n*100,) * 3
		image = Image.new("RGB", (size, size), color)
		images.append(image)

	then = time.time()
	func(*images)
	return time.time() - then


def printResults(times):
	"""
	Format and print the results of our tests.
	"""
	mean = format(np.mean(times), '.8f')
	minT = format(min(times), '.8f')
	maxT = format(max(times), '.8f')
	return f"Mean: {mean}  Min: {minT}  Max: {maxT}"


def test():
	"""
	Testing shit for speed.
	"""
	dots = Dots()
	sizes = [1, 5, 10, 20, 30, 40, 50, 100, 500, 1000]
	# print("------- Testing Dots.differenceImage() -------")
	# for size in sizes:
	# 	times = []
	# 	for n in range(1000):
	# 		t = testCompareDiffernece(size, dots.differenceImage, 2)
	# 		times.append(t)
	# 	print(f"Size: {size}x{size}\t{printResults(times)}")

	print("\n------- Testing Dots.compareDotsV1() (PIL-NumPy) -------")
	for size in sizes:
		times = []
		for n in range(1000):
			t = testCompareDiffernece(size, dots.compareDotsV1, 3)
			times.append(t)
		print(f"Size: {size}x{size}\t{printResults(times)}")

	print("\n------- Testing Dots.compareDotsV12() (Pure PIL) -------")
	for size in sizes:
		times = []
		for n in range(1000):
			t = testCompareDiffernece(size, dots.compareDotsV12, 3)
			times.append(t)
		print(f"Size: {size}x{size}\t{printResults(times)}")

	print("\n------- Testing Dots.compareDotsV2() (NumPy) -------")
	for size in sizes:
		times = []
		for n in range(1000):
			t = testCompareDiffernece(size, dots.compareDotsV2, 3)
			times.append(t)
		print(f"Size: {size}x{size}\t{printResults(times)}")


if __name__ == "__main__":
	import argparse

	parser = argparse.ArgumentParser(
		description="Makes one of those meme images with the dots.")
	parser.add_argument(
		"imagePath",
		nargs="?",
		help="Path to the original image.")
	parser.add_argument(
		"output",
		nargs="?",
		help="Path to output image.")
	parser.add_argument(
		"--test",
		action="store_true",
		help="Speed testing.")
	args = parser.parse_args()

	if args.test:
		test()
		exit()

	dots = Dots()
	then = time.time()
	dots.buildDots(args.imagePath, args.output)
	print(time.time() - then)