Puzzle Poster

puzzle-poster.py

from PIL import Image
import numpy as np
from pathlib import Path
from rich.progress import track
from scipy.spatial import KDTree

folder = Path('avatar')


def calc_all_colors():
    res = []
    for file in track(list(folder.glob('*.png'))):
        img = Image.open(file)
        mat = np.asarray(img)
        color = np.mean(mat, (0, 1))
        res.append((file.name, color))

    np.save('data.npy', np.array(res, dtype=object))


def main():
    colors = np.load('data.npy', allow_pickle=True)
    files = colors[:, 0]
    rgbs = np.array([c[1] for c in colors])
    tree = KDTree(rgbs)

    target = Image.open('avatar_poster.jpg')
    piece_width = 100
    resize_width = 100
    resize_height = int(target.height * resize_width / target.width)
    target = target.resize((resize_width, resize_height))
    mat = np.asarray(target)
    canvas = Image.new(
        'RGB', (resize_width * piece_width, resize_height * piece_width))

    for i in track(range(resize_height), total=resize_height):
        for j in range(resize_width):
            color = mat[i, j].astype(np.float32)
            dist, idx = tree.query(color, k=8)
            ind = np.random.choice(idx)
            dist = dist[idx.tolist().index(ind)]
            imgfile = folder / files[ind]
            img = Image.open(imgfile)
            img = img.resize((piece_width, piece_width))
            if dist > 30:
                offset = color - rgbs[ind]
                m = np.asarray(img, dtype=np.float32)
                m += offset
                m.clip(0, 255, out=m)
                img = Image.fromarray(m.astype(np.uint8))
            canvas.paste(img, (j * piece_width, i * piece_width))

    canvas.save('output.png')


if __name__ == '__main__':
    main()