|
| 1 | +import sys |
| 2 | +from PIL import Image |
| 3 | +from sklearn.cluster import KMeans |
| 4 | +import numpy as np |
| 5 | +from matplotlib import pyplot as plt |
| 6 | +from matplotlib import colors as mcolors |
| 7 | + |
| 8 | +def extract_colors(image_path, num_colors=4, output_file="color_palette.pdf"): |
| 9 | + try: |
| 10 | + # Load the image |
| 11 | + image = Image.open(image_path) |
| 12 | + |
| 13 | + # Resize and preprocess the image |
| 14 | + small_image = image.resize((100, 100)) # Resize to speed up processing |
| 15 | + image_data = np.array(small_image) |
| 16 | + |
| 17 | + # Handle different image modes |
| 18 | + if len(image_data.shape) == 2: # Grayscale |
| 19 | + image_data = np.stack((image_data,) * 3, axis=-1) |
| 20 | + elif image_data.shape[2] == 4: # RGBA |
| 21 | + image_data = image_data[:, :, :3] # Remove alpha channel |
| 22 | + |
| 23 | + # Flatten the image data to a list of RGB values |
| 24 | + pixels = image_data.reshape((-1, 3)) |
| 25 | + |
| 26 | + # Use KMeans to find dominant colors |
| 27 | + kmeans = KMeans(n_clusters=num_colors, random_state=42) |
| 28 | + kmeans.fit(pixels) |
| 29 | + dominant_colors = kmeans.cluster_centers_.astype(int) |
| 30 | + |
| 31 | + # Convert RGB to HEX |
| 32 | + dominant_colors_hex = [mcolors.rgb2hex(color / 255) for color in dominant_colors] |
| 33 | + |
| 34 | + # Create a visual palette |
| 35 | + fig, ax = plt.subplots(figsize=(8, 2)) |
| 36 | + for i, hex_color in enumerate(dominant_colors_hex): |
| 37 | + ax.add_patch(plt.Rectangle((i, 0), 1, 1, color=hex_color)) |
| 38 | + ax.text(i + 0.5, -0.3, f"Color {i + 1}\n{hex_color}", ha='center', va='center', fontsize=10) |
| 39 | + |
| 40 | + ax.set_xlim(0, len(dominant_colors_hex)) |
| 41 | + ax.set_ylim(0, 1) |
| 42 | + ax.axis('off') |
| 43 | + |
| 44 | + # Save the palette to a PDF file |
| 45 | + plt.savefig(output_file, bbox_inches='tight') |
| 46 | + plt.close() |
| 47 | + |
| 48 | + print(f"Color palette saved to {output_file}") |
| 49 | + except Exception as e: |
| 50 | + print(f"An error occurred: {e}") |
| 51 | + |
| 52 | +if __name__ == "__main__": |
| 53 | + # Request image path from the user |
| 54 | + image_path = input("Please enter the path to the image: ").strip() |
| 55 | + extract_colors(image_path) |
0 commit comments