low_poly_art.py

import collections
import os
import pathlib
import typing as t

import pygame
import pygame.gfxdraw

import numpy as np

from scipy.ndimage import gaussian_filter
from scipy.spatial import Delaunay


class LowPolyArt:
    '''
    - References:
        - https://github.com/Ovilia/Polyvia
        - https://cosmiccoding.com.au/tutorials/lowpoly
    '''
    AnyPath = t.Union[str, bytes, os.PathLike]

    def __init__(self, image: np.ndarray, scale: int = 2) -> None:
        perceptual_weight = np.array([0.2126, 0.7152, 0.0722])
        #
        self._image, self._scale = image, scale
        self._gray = (image*perceptual_weight).sum(axis=-1)
        #
        self._diff = gaussian_filter(self._gray, 2, mode='reflect') - \
            gaussian_filter(self._gray, 30, mode='reflect')
        self._diff[self._diff < 0] /= 10
        self._diff = np.sqrt(np.abs(self._diff) / self._diff.max())
        #
        self._screen = pygame.Surface((self.width*scale, self.height*scale))
        self._screen.fill(image.mean(axis=(0, 1)))
        #
        corners = np.array([
            (0, 0), (0, self.height-1), (self.width-1, 0), (self.width-1, self.height-1)
        ])
        self._points = np.concatenate((corners, self._samples()))

    @classmethod
    def from_file(cls, path: AnyPath, **kwargs) -> 'LowPolyArt':
        image = pygame.surfarray.pixels3d(pygame.image.load(path))
        return cls(image, **kwargs)

    @property
    def width(self) -> int:
        return self._image.shape[0]

    @property
    def height(self) -> int:
        return self._image.shape[1]

    def to_surface(self, n: int) -> pygame.Surface:
        triangles = Delaunay(self._points[:n, :])
        screen = self._draw(triangles, self._colors(triangles))
        return pygame.transform.smoothscale(screen, (self.width, self.height))

    def to_image(self, n: int, path: AnyPath) -> None:
        pygame.image.save(self.to_surface(n), path)

    def _samples(self, n: int = 1_000_000) -> np.ndarray:
        # np.random.seed(0)
        xs = np.random.randint(0, self.width, size=n)
        ys = np.random.randint(0, self.height, size=n)
        accept = np.random.random(size=n) < self._diff[xs, ys]
        return np.array([xs[accept], ys[accept]]).T

    def _colors(self, triangles: Delaunay) -> t.Dict[int, np.ndarray]:
        colors = collections.defaultdict(list)
        for ith in range(self.width):
            for jth in range(self.height):
                # Gets the index of the triangle the point is in
                index = triangles.find_simplex((ith, jth))
                colors[int(index)].append(self._image[ith, jth, :])
        # For each triangle, find the average colour
        return {
            index: np.mean(array, axis=0)
            for index, array in colors.items()
        }

    def _draw(
        self,
        triangles: Delaunay, colors: t.Dict[int, np.ndarray],
    ) -> pygame.Surface:
        screen = self._screen.copy()
        for key, color in colors.items():
            t = triangles.points[triangles.simplices[key]]
            pygame.gfxdraw.filled_polygon(screen, t*self._scale, color)
            pygame.gfxdraw.polygon(screen, t*self._scale, color)
        return screen


if __name__ == '__main__':
    input = 'ddlc.jpg'
    output = 'output'

    lpa = LowPolyArt.from_file(input, scale=2)
    directory = pathlib.Path(output)
    directory.mkdir(parents=True, exist_ok=True)
    for ith in range(100):
        n = 2*ith**2 + ith + 5
        lpa.to_image(n, directory/f'{ith}.png')