mosiac/mosiac.py

250 lines
6.9 KiB
Python
Executable File

#! /usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Does mosiacs.
"""
import os
import time
import random
import colorsys
import numpy as np
from PIL import Image
class Mosiac():
def __init__(self):
self.imageBank = {}
self.numClusters = 8
self.clusters = []
self.clusterMeans = []
self.bigImageSize = None
self.bigImage = None
self.smallImage = None
self.imageMatrix = []
self.tileSize = None
self.matrixSize = None
def openBigImage(self, imagePath, mat=(40,20)):
"""
Opens and initializes the big image.
"""
self.matrixSize = mat
self.bigImage = Image.open(imagePath).convert("RGB")
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])
def initImageBank(self, root, size=(200,200)):
"""
Calculates the average pixel value of all the images in the directory.
"""
then = time.time()
for file in os.listdir(os.path.join(root, "images")):
image = Image.open(os.path.join(root, "images", file))
if image.mode == "P":
image = image.convert(image.palette.mode)
if image.mode == "RGBA":
image = alpha_composite_with_color(image).convert("RGB")
if image.mode == "L":
image = image.convert("RGB")
# image = image.convert("RGBA")
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())))
self.imageBank[mean] = image
print(f"initImageBank took: {time.time()-then} seconds.")
def debugImageBank(self):
"""
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):
image = Image.new("RGB", self.tileSize, (i,j,k, 255))
mean = tuple(np.mean(image, axis=(0,1)))
self.imageBank[mean] = image
print(f"debugImageBank took: {time.time()-then} seconds.")
def initClusters(self):
"""
Divides the images into clusters.
"""
then = time.time()
sort = sorted(self.imageBank.keys(), key=lambda pixel: step(*pixel,8))
num = len(sort) // self.numClusters
for n in range(self.numClusters):
cluster = {pixel: self.imageBank[pixel] for pixel in \
sort[n*num:n*num+num]}
self.clusters.append(cluster)
self.clusters[-1].update({pixel: self.imageBank[pixel] for pixel in \
sort[-(len(sort) % num):]})
for cluster in self.clusters:
mean = tuple(np.mean(list(cluster.keys()), axis=(0)))
self.clusterMeans.append(mean)
"""
indexs = list(range(len(self.images)))
step = len(self.images) // self.numClusters
self.clusters = [indexs[i*step : (i+1)*step] for i in
range(self.numClusters)]
for cluster in self.clusters:
means = [self.means[i] for i in cluster]
mean = tuple(np.mean(means, axis=(0)))
self.clusterMeans.append(mean)
"""
print(f"initClusters took: {time.time()-then} seconds.")
def nearestImage(self, pixel):
"""
Finds the nearest image within the mosiac image bank to the provided
pixel.
"""
dists = []
for mean in self.clusterMeans:
dist = np.linalg.norm(np.array(pixel)-np.array(mean))
dists.append(dist)
clusterMean = self.clusterMeans[dists.index(min(dists))]
cluster = self.clusters[self.clusterMeans.index(clusterMean)]
nodes = np.array(list(cluster.keys()))
dist = np.sum((nodes - np.array(pixel))**2, axis=1)
# return np.argmin(dist) # closet value
dist.argsort()[:10]
return random.choice(list(dist))
"""
dists = []
for mean in cluster.keys():
dist = np.linalg.norm(np.array(pixel)-np.array(mean))
dists.append(dist)
"""
# choice = random.choice(sorted(dists)[:10])
# return cluster[dists.index(choice)]
def buildMatrix(self):
"""
Build the image matrix.
"""
then = time.time()
# for row in range(self.smallImage.size[1]):
# new_row = []
# for col in range(self.smallImage.size[0]):
# image = self.nearestImage(self.smallImage.getpixel((col, row)))
# new_row.append(image)
# self.imageMatrix.append(new_row)
pixels = list(self.smallImage.getdata())
for pixel in pixels:
image = self.nearestImage(pixel)
self.imageMatrix.append(image)
print(f"buildMatrix took: {time.time()-then} seconds.")
def buildMosiac(self, root):
"""
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
# self.bigImage.save(os.path.join(root, "mosiac.png"), "PNG")
image.save(os.path.join(root, "mosiac.jpg"), "JPEG", optimize=True, quality=90)
print(f"buildMosiac took: {time.time()-then} seconds.")
def step (r,g,b, repetitions=1):
lum = ( .241 * r + .691 * g + .068 * b )**0.5
h, s, v = colorsys.rgb_to_hsv(r,g,b)
h2 = int(h * repetitions)
lum2 = int(lum * repetitions)
v2 = int(v * repetitions)
return (h2, lum, v2)
def alpha_composite(front, back):
"""Alpha composite two RGBA images.
Source: http://stackoverflow.com/a/9166671/284318
Keyword Arguments:
front -- PIL RGBA Image object
back -- PIL RGBA Image object
"""
front = np.asarray(front)
back = np.asarray(back)
result = np.empty(front.shape, dtype='float')
alpha = np.index_exp[:, :, 3:]
rgb = np.index_exp[:, :, :3]
falpha = front[alpha] / 255.0
balpha = back[alpha] / 255.0
result[alpha] = falpha + balpha * (1 - falpha)
old_setting = np.seterr(invalid='ignore')
result[rgb] = (front[rgb] * falpha + back[rgb] * balpha * (1 - falpha)) / result[alpha]
np.seterr(**old_setting)
result[alpha] *= 255
np.clip(result, 0, 255)
# astype('uint8') maps np.nan and np.inf to 0
result = result.astype('uint8')
result = Image.fromarray(result, 'RGBA')
return result
def alpha_composite_with_color(image, color=(255, 255, 255)):
"""Alpha composite an RGBA image with a single color image of the
specified color and the same size as the original image.
Keyword Arguments:
image -- PIL RGBA Image object
color -- Tuple r, g, b (default 255, 255, 255)
"""
back = Image.new('RGBA', size=image.size, color=color + (255,))
return alpha_composite(image, back)
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(
description="Generates a mosiac of images.")
parser.add_argument(
"root",
help="Root directory to work in.")
# parser.add_argument(
# "-s",
# "--size",
# default=(200,200),
# help="Size each image should be.")
args = parser.parse_args()
mosiac = Mosiac()
mosiac.openBigImage(os.path.join(args.root, "big.jpg"))
mosiac.initImageBank(**vars(args))
# mosiac.debugImageBank()
mosiac.initClusters()
mosiac.buildMatrix()
mosiac.buildMosiac(args.root)