Update app.py

app.py

@@ -4,6 +4,7 @@ import numpy as np
 import matplotlib.pyplot as plt
 from sklearn.cluster import KMeans
 from collections import Counter
+import random
 
 # Function to extract dominant colors
 def extract_colors(image, num_colors=5):
@@ -18,42 +19,63 @@ def extract_colors(image, num_colors=5):
     kmeans.fit(pixels)
     colors = kmeans.cluster_centers_.astype(int)
     
-    # Count occurrences of each color
-    counts = Counter(kmeans.labels_)
-    sorted_colors = [colors[i] for i in np.argsort(list(counts.values()))[::-1]]
-    
-    return sorted_colors
+    return colors
 
 # Convert RGB to Hex
 def rgb_to_hex(rgb):
     return "#" + "".join(f"{int(c):02x}" for c in rgb)
 
-# Function to generate color palette
-def generate_palette(image, num_colors=5):
-    colors = extract_colors(image, num_colors)
-    hex_colors = [rgb_to_hex(color) for color in colors]
+# Generate AI-based harmonious color palettes
+def generate_ai_palette():
+    return [rgb_to_hex([random.randint(0, 255) for _ in range(3)]) for _ in range(5)]
+
+# Function to update palette with user selection
+def update_palette(image, selected_x, selected_y, locked_colors):
+    if image is not None:
+        selected_color = image[int(selected_y), int(selected_x)]
+        selected_hex = rgb_to_hex(selected_color)
+        
+        colors = extract_colors(image, 5)
+        hex_colors = [rgb_to_hex(color) for color in colors]
+        
+        # Replace unlocked colors with extracted ones
+        for i in range(5):
+            if locked_colors[i] == "":
+                locked_colors[i] = selected_hex if i == 0 else hex_colors[i]
+    else:
+        # If no image, generate AI colors
+        locked_colors = [color if color != "" else rgb_to_hex([random.randint(0, 255) for _ in range(3)]) for color in locked_colors]
     
     # Create color swatches
     fig, ax = plt.subplots(figsize=(8, 2))
-    ax.imshow([colors], aspect="auto")
+    rgb_colors = [[int(h[i:i+2], 16) for i in (1, 3, 5)] for h in locked_colors]
+    ax.imshow([rgb_colors], aspect="auto")
     ax.set_xticks([])
     ax.set_yticks([])
     
-    return fig, " | ".join(hex_colors)
+    return fig, locked_colors
 
 # Gradio UI
 with gr.Blocks(theme=gr.themes.Soft()) as iface:
     gr.Markdown("# 🎨 AI-Powered Image Color Palette Generator")
-    gr.Markdown("Upload an image to extract its dominant colors and generate a color palette.")
+    gr.Markdown("Upload an image to extract its dominant colors, or generate AI-based color palettes.")
+    
+    image_input = gr.Image(type="numpy", label="Upload an Image (Optional)")
     
-    image_input = gr.Image(type="numpy", label="Upload an Image")
-    num_colors = gr.Slider(3, 10, value=5, step=1, label="Number of Colors")
-    generate_button = gr.Button("Generate Palette 🎨")
+    with gr.Row():
+        generate_ai_button = gr.Button("Generate AI Palette 🎨")
+        regenerate_button = gr.Button("Re-run Unlocked Colors 🔄")
     
     palette_output = gr.Plot(label="Color Palette")
-    hex_output = gr.Textbox(label="Hex Codes", interactive=False)
+    hex_output = gr.Textbox(label="Hex Codes", interactive=True)
+    
+    with gr.Row():
+        locked_colors = [gr.Textbox(label=f"Color {i+1}", interactive=True) for i in range(5)]
+    
+    generate_ai_button.click(generate_ai_palette, inputs=[], outputs=locked_colors)
+    regenerate_button.click(update_palette, inputs=[image_input, None, None, locked_colors], outputs=[palette_output, locked_colors])
     
-    generate_button.click(generate_palette, inputs=[image_input, num_colors], outputs=[palette_output, hex_output])
+    image_input.select(update_palette, inputs=[image_input, image_input.select, image_input.select, locked_colors], outputs=[palette_output, locked_colors])
 
 # Launch the app
 if __name__ == "__main__":