Posterize a PNG and create layers for PCB art

pcb_art_helper.py

#!/usr/bin/env python3

import sys
import os.path

import numpy as np
import skimage.io, skimage.color, skimage.transform
import sklearn.cluster
import matplotlib.pyplot as plt


downscale = 1

n_colors = 5

# Recommended: design with background being light green (or dark purple)
#     so that it blends in with standard fills.
# Add graphics after everything else since this massively slows zone fills
# When exporting:
# Negative threshold 50
# Export as Pcbnew (.kicad_mod file)
# After export, find and replace "Eco1.User" with "F.Cu"


# Optionally remap image color brightness to PCB brightness.
# This is good e.g. if your darkest color isn't the background color.
# shuffle_mapping = [
#     1,
#     0,
#     2,
#     3,
#     4,
# ]
shuffle_mapping = False


# https://blogdotoshparkdotcom.files.wordpress.com/2020/02/palette.png

#       black: 0: solder mask on bare FR4
# dark purple: 1: solder mask on copper
#         tan: 2: bare FR4 (no-mask)
#        gold: 3: copper (no-mask)
#       white: 4: silkscreen on copper (mask or no???)

# Note, the "mask" layer is actually a keepout, so it's
# opposite of what you'd expect.

layer_by_brightness = {
    "mask": [2, 3],
    "silk": [4],
    "copper": [1, 3, 4]
}

false_colors = [
    (48, 50, 65),  # black purple
    (87, 54, 97),  # purple
    (171, 151, 101),  # tan
    (255, 213, 70),  # gold
    (252, 252, 252),  # white
]



# # Generic Green PCB

# #  dark green: 0: solder mask on bare FR4
# #         tan: 1: bare FR4 (no-mask)
# # light green: 2: solder mask on copper
# #      silver: 3: copper (no-mask)
# #       white: 4: silkscreen on copper (with mask)

# false_colors = [
#     (0, 80, 0),  # DarkGreen
#     (210, 180, 140),  # Tan
#     (34, 139, 34),  # ForestGreen
#     (192, 192, 192),  # Silver
#     (255, 255, 255),  # White
# ]

# layer_by_brightness = {
#     "mask": [1, 3],
#     "silk": [4],  # needs inverted???
#     "copper": [2, 3, 4]  # needs inverted???
# }



# n_colors = 3

# false_colors = [
#     (34, 139, 34),  # ForestGreen
#     (192, 192, 192),  # Silver
#     (255, 255, 255),  # White
# ]


# layer_by_brightness = {
#     "mask": [0, 2],
#     "silk": [2],  # needs inverted
# }


# Do shuffling if necessary
if shuffle_mapping:
    layer_by_brightness = {k: [shuffle_mapping[i] for i in v] for k, v in layer_by_brightness.items()}
    false_colors = [false_colors[shuffle_mapping[i]] for i in range(len(false_colors))]




# Image preprocessing
img = skimage.io.imread(sys.argv[-1])

img = skimage.transform.resize(img, (img.shape[0] // downscale, img.shape[1] // downscale), anti_aliasing=True)

if img.shape[2] == 4:
    print("Converting RGBA to RGB")
    img = skimage.color.rgba2rgb(img)
print("Converting colorspace")
lab_img = skimage.color.rgb2lab(img)
y, x, d = lab_img.shape
output_shape = (y, x)
flat = np.reshape(lab_img, (-1, d))



# Kmeans stuff to palettize the image
print("Finding centroids")
kmeans = sklearn.cluster.KMeans(n_clusters=n_colors).fit(flat)
print("Running predictions")
predictions = kmeans.predict(flat)
print("Finding distances")
distances = kmeans.transform(flat)


pixel_labels = predictions.copy()

# If the kmeans prediction has a very small distance to the centroid,
# accept the label as-is.
min_distances = np.amin(distances, axis=-1)
for pred, distance, i in zip(predictions, min_distances, range(flat.shape[0])):
    if distance < 1.5:
        # Close enough to accept the predicted label
        continue
    else:
        # Mark as too far from the prediction
        pixel_labels[i] = -1


# For the pixels that were inbetween the centroids,
# choose the closest centroid among the neighbors.
# This guarantees that color regions are contiguous.
unflat_labels = pixel_labels.reshape(output_shape).copy()
for i in range(flat.shape[0]):
    if pixel_labels[i] == -1:
        y, x = np.unravel_index(i, output_shape)
        neighbors = unflat_labels[y-1:y+2, x-1:x+2]
        neighbors = neighbors.reshape(-1)
        neighbors = neighbors[neighbors != -1]
        local_labels = np.unique(neighbors)
        mask = np.isin(list(range(n_colors)), local_labels)
        mask = (~mask) * 1000 + mask
        local_distances = mask * distances[i]
        best = local_distances.argmin()
        pixel_labels[i] = best


basename = os.path.splitext(sys.argv[-1])[0]


# Map the PCB art layers to centroid labels
cluster_and_label = [(list(c), i) for i, c in enumerate(kmeans.cluster_centers_)]
cluster_and_label.sort()

label_to_brightness_index = {cl[1]:i for i, cl in enumerate(cluster_and_label)}

labeled_image = pixel_labels.reshape(output_shape).copy()
colored_image = [false_colors[label_to_brightness_index[i]] for i in pixel_labels]
colored_image = np.array(colored_image).reshape((*output_shape, 3)).astype("uint8")
colored_fname = f"{basename}_falsecolor.png"
skimage.io.imsave(colored_fname, colored_image)



# # Show the colors of our palette in order of brightness.
# # Useful for debugging layer order.
# display_palette = np.zeros((100, 100 * n_colors, 3))
# for i in range(n_colors):
#     display_palette[:, 100*i:100*i + 100] = cluster_and_label[i][0]
# display_palette_rgb = skimage.color.lab2rgb(display_palette)
# plt.imshow(display_palette_rgb)



pcb_layer_to_labels = {}
for layer in layer_by_brightness.keys():
    labels = [cluster_and_label[i][1] for i in layer_by_brightness[layer]]
    pcb_layer_to_labels[layer] = labels



for layer_name, layer_labels in pcb_layer_to_labels.items():
    outimg = np.isin(pixel_labels, layer_labels)
    outimg = outimg.reshape(output_shape)
    outimg = outimg.astype("uint8") * 255
    fname = f"{basename}_{layer_name}.bmp"
    skimage.io.imsave(fname, outimg)