From 7d1577a6e6f621d5c44d5f7abeb11db15e4cb5fb Mon Sep 17 00:00:00 2001 From: Azalea <22280294+hykilpikonna@users.noreply.github.com> Date: Sat, 21 Dec 2024 00:50:34 -0500 Subject: [PATCH] [+] k-means clustering and PCA --- server/colors.py | 131 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 131 insertions(+) create mode 100644 server/colors.py diff --git a/server/colors.py b/server/colors.py new file mode 100644 index 0000000..7e7947d --- /dev/null +++ b/server/colors.py @@ -0,0 +1,131 @@ +import numpy as np +from sklearn.cluster import KMeans +from sklearn.decomposition import PCA +from PIL import Image +import matplotlib.pyplot as plt +import json + + +def generate_hv_colors(cluster_labels, sorted_indices): + """ + Generate hColors and vColors from the clustered image. + + Parameters: + cluster_labels (ndarray): The label of each pixel in the image corresponding to its color cluster. + sorted_indices (ndarray): The index mapping of the sorted cluster centers. + + Returns: + hColors (ndarray): (n+1) x (n+1) array with the color indices of horizontal lines. + vColors (ndarray): (n+1) x (n+1) array with the color indices of vertical lines. + """ + h, w = cluster_labels.shape + hColors = np.zeros((h + 1, w + 1), dtype=np.uint8) + vColors = np.zeros((h + 1, w + 1), dtype=np.uint8) + + # Fill hColors and vColors + for y in range(h): + for x in range(w): + color_idx = cluster_labels[y, x] + color_idx = np.where(sorted_indices == color_idx)[0][0] + hColors[y, x] = color_idx + # if x >= w - 1: + # hColors[y, x + 1] = color_idx + + vColors[y, x] = color_idx + # if y >= h - 1: + # vColors[y + 1, x] = color_idx + + # Handle the bottom-right corner + # hColors[h, w] = cluster_labels[h - 1, w - 1] + # vColors[h, w] = cluster_labels[h - 1, w - 1] + + return hColors, vColors + + +def convert_colors_to_hex(cluster_centers): + """ + Convert RGB colors to hex color strings. + + Parameters: + cluster_centers (ndarray): The array of RGB colors (k, 3). + + Returns: + list: List of hex color strings. + """ + return [f"#{r:02x}{g:02x}{b:02x}" for r, g, b in cluster_centers] + + +if __name__ == '__main__': + # Load image and convert to RGB + image_path = r"C:\Users\Azalea\Downloads\Untitled_Artwork.png" + image = Image.open(image_path).convert('RGB') + image_array = np.array(image) + + # Reshape the image to a 2D array (num_pixels, 3) + height, width, _ = image_array.shape + pixels = image_array.reshape((-1, 3)) + + # Cluster the pixels using KMeans + n_colors = 32 + kmeans = KMeans(n_clusters=n_colors, random_state=42) + kmeans.fit(pixels) + + # Get the mean color of each cluster (centroids) + cluster_centers = kmeans.cluster_centers_.astype(int) + + # Replace each pixel with the closest cluster center + labels = kmeans.labels_.reshape(height, width) + new_image_array = cluster_centers[labels].reshape(image_array.shape) + + # Use PCA to sort the colors + pca = PCA(n_components=1) + color_positions = pca.fit_transform(cluster_centers) + sorted_indices = np.argsort(color_positions[:, 0]) + sorted_cluster_centers = cluster_centers[sorted_indices] + + # Create a color palette from the sorted cluster centers + palette_height = 20 * n_colors # Each color swatch height + palette_width = 100 # Width of each color swatch + palette = np.zeros((palette_height, palette_width, 3), dtype=np.uint8) + + for i, color in enumerate(sorted_cluster_centers): + start_y = i * 20 + end_y = (i + 1) * 20 + palette[start_y:end_y, :] = color + + # Display the original, reduced color image, and palette + fig, axes = plt.subplots(1, 3, figsize=(15, 6)) + + axes[0].set_title('Original Image') + axes[0].imshow(image) + axes[0].axis('off') + + axes[1].set_title(f'Image with {n_colors} Colors') + axes[1].imshow(new_image_array) + axes[1].axis('off') + + axes[2].set_title('Sorted Color Palette (PCA)') + axes[2].imshow(palette) + axes[2].axis('off') + + plt.tight_layout() + plt.show() + + # Generate hColors and vColors + hColors, vColors = generate_hv_colors(labels, sorted_indices) + + # Prepare output JSON structure + output_data = { + "hColors": {"__typedArray__": True, "type": "Int8Array", + "data": hColors.flatten().tolist()}, + "vColors": {"__typedArray__": True, "type": "Int8Array", + "data": vColors.flatten().tolist()}, + "colors": convert_colors_to_hex(sorted_cluster_centers) + } + + # Save to JSON file + output_path = r"C:\Users\Azalea\Downloads\output.json" + with open(output_path, 'w') as f: + json.dump(output_data, f) + + print(f"Saved output to {output_path}")