feat: add color palette limit selection

app.py

@@ -14,10 +14,15 @@ import spaces
 import tempfile
 
 @spaces.GPU
-def process_image(input_image_path):
+def process_image(input_image_path, color_limit):
     if input_image_path is None:
         return None, None, "Please upload an image."
         
+    # Parse color limit
+    k_colors = 0
+    if color_limit != "Unlimited":
+        k_colors = int(color_limit)
+        
     # Get original filename
     original_name = os.path.basename(input_image_path)
     base, _ = os.path.splitext(original_name)
@@ -47,7 +52,7 @@ def process_image(input_image_path):
     
     # Generate and save extraction with accurate colors
     from pixelfixer.reconstruct import two_stage_pack
-    low = two_stage_pack(rgba, int(r["cols"]), int(r["rows"]))
+    low = two_stage_pack(rgba, int(r["cols"]), int(r["rows"]), k_colors=k_colors)
     extract_img = Image.fromarray(low)
     
     # Upscale using nearest neighbor to prevent blurriness
@@ -61,6 +66,8 @@ def process_image(input_image_path):
     # Create info string
     info = f"Grid detected: {r['cols']}x{r['rows']}\n"
     info += f"Step size: {r['step_x']}x{r['step_y']}\n"
+    if k_colors > 0:
+        info += f"Colors: {k_colors}\n"
     info += f"Processing time: {dt}s"
     
     return overlay_path, extract_path, info
@@ -72,16 +79,17 @@ with gr.Blocks(title="Pixel Art Fixer") as app:
     with gr.Row():
         with gr.Column():
             input_image = gr.Image(type="filepath", label="Input Image")
+            color_limit = gr.Radio(choices=["Unlimited", "8", "12", "16", "32"], value="Unlimited", label="Color Palette Limit")
             process_btn = gr.Button("Process Image", variant="primary")
             
         with gr.Column():
-            info_text = gr.Textbox(label="Detection Info", lines=3)
+            info_text = gr.Textbox(label="Detection Info", lines=4)
             overlay_image = gr.Image(type="filepath", label="Grid Overlay")
             extract_image = gr.Image(type="filepath", label="Extracted Pixel Art")
             
     process_btn.click(
         fn=process_image,
-        inputs=[input_image],
+        inputs=[input_image, color_limit],
         outputs=[overlay_image, extract_image, info_text]
     )