better multi-threading and ui

mosiac.py

@@ -6,7 +6,9 @@ Does mosiacs.
 import os
 import time
 import colorsys
-import threading
+# import threading
+from multiprocessing import Pool
+from multiprocessing.dummy import Pool as ThreadPool
 
 import numpy as np
 from PIL import Image
@@ -18,6 +20,7 @@ class Mosiac():
 		self.numThreads = 4
 
 		self.imageBank = {}
+		self.tilesDir = None
 		self.clusters = []
 		self.clusterMeans = []
 
@@ -29,66 +32,63 @@ class Mosiac():
 		self.matrixSize = None
 
 
-	def openBigImage(self, imagePath, mat=(40,20)):
+	def openBigImage(self, bigImagePath, mat=None):
 		"""
 		Opens and initializes the big image.
 		"""
+		self.bigImage = Image.open(bigImagePath).convert("RGB")
+		# TODO: find a better way to do this
+		if not mat:
+			if self.bigImage.size[0] > self.bigImage.size[1]:
+				mat = (40, 20)
+			elif self.bigImage.size[0] < self.bigImage.size[1]:
+				mat = (20, 40)
+			else:
+				mat = (20, 20)
+
 		self.matrixSize = mat
-		self.bigImage = Image.open(imagePath).convert("RGB")
+		self.tileSize = (
-		self.tileSize = (self.bigImage.size[0] // mat[0],
+			self.bigImage.size[0] // mat[0],
 			self.bigImage.size[1] // mat[1])
 		self.smallImage = self.bigImage.resize(mat)
 		self.bigImageSize = (self.tileSize[0]*mat[0], self.tileSize[1]*mat[1])
+		self.bigImage = self.bigImage.crop(
+			(0, 0, self.bigImageSize[0], self.bigImageSize[1]))
 
 
-	def initImageBank(self, root):
+	def initImageBank(self, tilesDir):
 		"""
 		Calculates the average pixel value of all the images in the directory.
 		This is the thread controller.
 		"""
 		then = time.time()
-		files = os.listdir(os.path.join(root, "images"))
+		self.tilesDir = tilesDir
+		files = os.listdir(self.tilesDir)
 
-		thread_files = []
+		pool = ThreadPool()
-		num = len(files) // self.numThreads
+		pool.map(self.initImageBankWorker, files)
-		for n in range(self.numThreads):
+		pool.close()
-			t_files = [file for file in files[n*num:n*num+num]]
+		pool.join()
-			thread_files.append(t_files)
-		thread_files[-1] += [file for file in files[-(len(files) % num):]]
-
-		threads = []
-		for t_files in thread_files:
-			t = threading.Thread(target=self.initImageBankWorker,
-				args=((root, t_files)))
-			threads.append(t)
-			t.start()
-		for t in threads:
-			t.join()
 		print(f"initImageBank took: {time.time()-then} seconds.")
 
 
-	def initImageBankWorker(self, root, files):
+	def initImageBankWorker(self, file):
 		"""
 		Thread worker.
 		"""
-		for file in files:
+		image = Image.open(os.path.join(self.tilesDir, file))
-			image = Image.open(os.path.join(root, "images", file))
 
-			if image.mode == "P":
+		if image.mode == "P":
-				image = image.convert(image.palette.mode)
+			image = image.convert(image.palette.mode)
-			if image.mode == "RGBA":
+		if image.mode == "RGBA":
-				image = alpha_composite(image).convert("RGB")
+			image = alpha_composite(image).convert("RGB")
-			if image.mode == "L":
+		if image.mode == "L":
-				image = image.convert("RGB")
+			image = image.convert("RGB")
-			# image = image.convert("RGBA")
 
-			image = image.resize(self.tileSize, Image.ANTIALIAS)
+		image = image.resize(self.tileSize, Image.ANTIALIAS)
-			mean = tuple(np.mean(image, axis=(0,1)))
+		mean = tuple(np.mean(image, axis=(0,1)))
-			# img = Image.new("RGBA", image.size)
-			# img.putdata(list(map(lambda pixel: (255,255,255,0) if pixel == \
-			# 	(255,255,255,255) else pixel, image.getdata())))
 
-			self.imageBank[mean] = image
+		self.imageBank[mean] = image
 
 
 	def debugImageBank(self):
@@ -96,9 +96,9 @@ class Mosiac():
 		Create a bank of test images.
 		"""
 		then = time.time()
-		for i in range(0, 257, 16):
+		for i in range(0, 256, 15):
-			for j in range(0, 257, 16):
+			for j in range(0, 256, 15):
-				for k in range(0, 257, 16):
+				for k in range(0, 256, 15):
 					image = Image.new("RGB", self.tileSize, (i,j,k, 255))
 					mean = tuple(np.mean(image, axis=(0,1)))
 					self.imageBank[mean] = image
@@ -150,7 +150,7 @@ class Mosiac():
 		return cluster[tuple(nodes[choice])]
 
 
-	def buildMatrix(self):
+	def buildMatrix(self, tileAlpha):
 		"""
 		Build the image matrix.
 		"""
@@ -158,31 +158,32 @@ class Mosiac():
 		pixels = list(self.smallImage.getdata())
 		for pixel in pixels:
 			image = self.nearestImage(pixel)
-			image = image.convert("RGBA")
+			if tileAlpha:
-			comp = Image.new("RGBA", image.size, pixel + (30,))
+				image = image.convert("RGBA")
-			image = Image.alpha_composite(image, comp).convert("RGB")
+				comp = Image.new("RGBA", image.size, pixel + (tileAlpha,))
+				image = Image.alpha_composite(image, comp).convert("RGB")
 			self.imageMatrix.append(image)
 		print(f"buildMatrix took: {time.time()-then} seconds.")
 
 
-	def buildMosiac(self, root):
+	def buildMosiac(self, output, bigAlpha):
 		"""
 		Builds the final mosiac image.
 		"""
 		then = time.time()
 		image = Image.new("RGB", self.bigImageSize, (255,255,255))
-		self.bigImage = self.bigImage.crop(
+
-			(0, 0, self.bigImageSize[0], self.bigImageSize[1]))
 		n = 0
 		for y in range(0, self.bigImageSize[1], self.tileSize[1]):
 			for x in range(0, self.bigImageSize[0], self.tileSize[0]):
 				image.paste(self.imageMatrix[n], box=(x,y))
 				n += 1
-		image = image.convert("RGBA")
+		if bigAlpha:
-		self.bigImage.putalpha(50)
+			image = image.convert("RGBA")
-		image = Image.alpha_composite(image, self.bigImage).convert("RGB")
+			self.bigImage.putalpha(bigAlpha)
-		# self.bigImage.save(os.path.join(root, "mosiac.png"), "PNG")
+			image = Image.alpha_composite(image, self.bigImage).convert("RGB")
-		image.save(os.path.join(root, "mosiac.jpg"), "JPEG", optimize=True, quality=90)
+		# self.bigImage.save(output, "PNG")
+		image.save(output, "JPEG", optimize=True, quality=90)
 		print(f"buildMosiac took: {time.time()-then} seconds.")
 
 
@@ -213,21 +214,44 @@ if __name__ == "__main__":
 	import argparse
 
 	parser = argparse.ArgumentParser(
-		description="Generates a mosiac of images.")
+		description="Stiches together a series of smaller images in the \
+		likeliness of a larger image. The 'big image' should be quite large.")
 	parser.add_argument(
-		"root",
+		"bigImagePath",
-		help="Root directory to work in.")
+		help="The big image that will be used as the 'guide'.")
-	# parser.add_argument(
+	parser.add_argument(
-	# 	"-s",
+		"tilesDir",
-	# 	"--size",
+		help="Directory full of images to be used as the tiles.")
-	# 	default=(200,200),
+	parser.add_argument(
-	# 	help="Size each image should be.")
+		"output",
+		help="File to be outputed.")
+	parser.add_argument(
+		"--matrix-size",
+		dest="matrixSize",
+		nargs=2,
+		type=int,
+		help="Size of the tile matrix. A 40x20 matrix would be '40 20'")
+	parser.add_argument(
+		"--tile-alpha",
+		dest="tileAlpha",
+		default=50,
+		help="Alpha channel value of the color filter that is applied to each \
+		tile. Range should be 0-255.")
+	parser.add_argument(
+		"--big-alpha",
+		dest="bigAlpha",
+		default=50,
+		help="Alpha channel value of big image when it is transposed onto the \
+		matrix. Range should be 0-255")
 	args = parser.parse_args()
 
+	if args.matrixSize:
+		args.matrixSize = tuple(args.matrixSize)
+
 	mosiac = Mosiac()
-	mosiac.openBigImage(os.path.join(args.root, "big.jpg"))
+	mosiac.openBigImage(args.bigImagePath, args.matrixSize)
-	mosiac.initImageBank(args.root)
+	mosiac.initImageBank(args.tilesDir)
 	# mosiac.debugImageBank()
 	mosiac.initClusters()
-	mosiac.buildMatrix()
+	mosiac.buildMatrix(args.tileAlpha)
-	mosiac.buildMosiac(args.root)
+	mosiac.buildMosiac(args.output, args.bigAlpha)