A prototype color by number template generator

generate.py

import cv2
import numpy as np
import random

class SVG:
    def __init__(self, width, height):
        self.width = width
        self.height = height
        self.elements = []

    def add_text(self, text, x, y):
        self.elements.append(f'<text x="{x}" y="{y}" font-size="10px">{text}</text>')

    def add_path(self, points):
        path = "M" # start path
        for point in points:
            path += f" {point[0]},{point[1]} "
        path += "Z" # end path

        fmt_string = f'<path d="{path}" fill="white" stroke="black" stroke-width="1"/>'
        self.elements.append(fmt_string)

    def write(self, output_path):
        content = '<?xml version="1.0" encoding="UTF-8"?>'
        content += f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {self.width} {self.height}">'
        content += f'<rect width="{self.width}" height="{self.height}" fill="white"/>' # Background
        for element in self.elements:
            content += element
        content += "</svg>"

        with open(output_path, "w") as file:
            file.write(content)

# Resize an image while preserving its aspect ratio
def resize(image, new_width):
    height, width, _ = image.shape
    new_height = round(height / width * new_width)
    return cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_LINEAR)

# Resize, denoise and flatten the input image
def preprocess_image(image_path):
    original = cv2.imread(image_path)
    resized = resize(original, 600)
    blurred = cv2.bilateralFilter(resized, 15, 75, 75)

    height, width, channels = blurred.shape
    flattened = blurred.reshape(width * height, channels)
    return flattened.astype(np.float32), width, height, channels

# Perform K-means clustering on a set of points
def kmeans(k, points, num_steps):
    # Randomly sample k centroids
    centroids = [random.choice(points) for _ in range(k)]

    # List of point indexes correspnding to each cluster centroid
    clusters = [[] for _ in range(k)]

    inertia = 0 # Measure of how dense the clusters are

    for _ in range(num_steps):
        clusters = [[] for _ in range(k)]

        # Add each point to the cluster it's closest to
        for i, point in enumerate(points):
            distances = []
            for centroid in centroids:
                distance = np.sum((point - centroid) ** 2)
                distances.append(distance)
            # Assign to centoid with the smallest euclidean distance
            closest = np.argmin(distances)
            clusters[closest].append(i)

        # Recompute the centroid based on the average of all the points in the cluster
        # Also compute the inertia of the clustering. Inertia is the sum of all the
        # distances of all the points in the cluster to the centroid
        inertia = 0
        for i, point_indexes in enumerate(clusters):
            if len(point_indexes) == 0:
                continue # Leave the centroid as is

            cluster = points[point_indexes]
            centroids[i] = np.mean(cluster, axis=0)
            for point in cluster:
                inertia += np.sum((point - centroids[i]) ** 2)

    return centroids, clusters, inertia

# Get the image that is best reduced to `num_colors` colors
# `iterations` specifies how many times we'll run our clustering algorithm
# and `attempts` specifies how many different quantized images we'll consider
def reduce_colors(image, num_colors, iterations, attempts):
    min_inertia = float("inf")
    new_image = image.copy()
    colors_used = []

    for _ in range(attempts):
        centroids, clusters, inertia = kmeans(num_colors, image, iterations)
        if inertia < min_inertia: # TODO: is there a better way to judge quality??
            colors_used = centroids
            # Update the pixels in the image
            for i, color in enumerate(centroids):
                pixel_indexes = clusters[i]
                new_image[pixel_indexes] = color

    return new_image, colors_used

def convert_to_template(image, colors):
    height, width, _ = image.shape
    min_area = round(width * height * 0.001)
    morph_kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
    svg = SVG(width, height)

    for color in colors:
        # Apply binary thresholding.
        # Each pixel that is equal to our color is [255, 255, 255], else it's [0, 0, 0]
        mask = (image == color).astype(np.uint8) * 255
        grayscale = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)

        # Find groups of same colored pixels
        result = cv2.connectedComponentsWithStats(image=grayscale, connectivity=8, ltype=cv2.CV_32S)
        (num_groups, matrix, stats, _) = result

        # Draw an outline around each group of pixels
        for group in range(1, num_groups):
            if stats[group, cv2.CC_STAT_AREA] < min_area:
                continue # Ignore tiny groups of pixels

            group_mask = (matrix == group).astype(np.uint8) * 255 # Mask of the specific pixel group
            group_mask = cv2.morphologyEx(group_mask, cv2.MORPH_CLOSE, morph_kernel) # Close tiny gaps in the mask
            group_mask = cv2.morphologyEx(group_mask, cv2.MORPH_OPEN, morph_kernel) # Erode then dilate to remove nosie

            contours, _ = cv2.findContours(group_mask, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE )
            simplified_contours = [cv2.approxPolyDP(contour, 3, True) for contour in contours]

            # Draw each contour as a path path
            for contour in simplified_contours:
                if len(contour) < 3:
                    continue # Skip invalid countours
                points = [point[0] for point in contour]
                svg.add_path(points)

    return svg

def generate_template(in_path, out_path, num_colors):
    image, width, height, channels = preprocess_image(in_path)
    new_image, colors = reduce_colors(image, num_colors, 3, 3)
    new_image = new_image.reshape(height, width, channels)
    svg = convert_to_template(new_image, colors)
    svg.write(out_path)