diff --git a/mosiac.py b/mosiac.py index edbff24..10af902 100755 --- a/mosiac.py +++ b/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.bigImage.size[0] // mat[0], + self.tileSize = ( + 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 = [] - num = len(files) // self.numThreads - for n in range(self.numThreads): - t_files = [file for file in files[n*num:n*num+num]] - 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() + pool = ThreadPool() + pool.map(self.initImageBankWorker, files) + pool.close() + pool.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(root, "images", file)) + image = Image.open(os.path.join(self.tilesDir, file)) - 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 = image.convert("RGBA") + 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 = image.resize(self.tileSize, Image.ANTIALIAS) - 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()))) + image = image.resize(self.tileSize, Image.ANTIALIAS) + mean = tuple(np.mean(image, axis=(0,1))) - 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 j in range(0, 257, 16): - for k in range(0, 257, 16): + for i in range(0, 256, 15): + for j in range(0, 256, 15): + 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") - comp = Image.new("RGBA", image.size, pixel + (30,)) - image = Image.alpha_composite(image, comp).convert("RGB") + if tileAlpha: + image = image.convert("RGBA") + 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") - self.bigImage.putalpha(50) - image = Image.alpha_composite(image, self.bigImage).convert("RGB") - # self.bigImage.save(os.path.join(root, "mosiac.png"), "PNG") - image.save(os.path.join(root, "mosiac.jpg"), "JPEG", optimize=True, quality=90) + if bigAlpha: + image = image.convert("RGBA") + self.bigImage.putalpha(bigAlpha) + image = Image.alpha_composite(image, self.bigImage).convert("RGB") + # 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", - help="Root directory to work in.") - # parser.add_argument( - # "-s", - # "--size", - # default=(200,200), - # help="Size each image should be.") + "bigImagePath", + help="The big image that will be used as the 'guide'.") + parser.add_argument( + "tilesDir", + help="Directory full of images to be used as the tiles.") + parser.add_argument( + "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.initImageBank(args.root) + mosiac.openBigImage(args.bigImagePath, args.matrixSize) + mosiac.initImageBank(args.tilesDir) # mosiac.debugImageBank() mosiac.initClusters() - mosiac.buildMatrix() - mosiac.buildMosiac(args.root) + mosiac.buildMatrix(args.tileAlpha) + mosiac.buildMosiac(args.output, args.bigAlpha)