Spaces:
Running on Zero
Running on Zero
| import gradio as gr | |
| import numpy as np | |
| from PIL import Image | |
| import time | |
| import os | |
| import sys | |
| # Ensure the pixelfixer package can be imported | |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) | |
| from pixelfixer import detect | |
| from pixelfixer.cli import _overlay, _extract | |
| import spaces | |
| import tempfile | |
| def process_image(input_image_path, color_limit, export_scales, force_step_val): | |
| if input_image_path is None: | |
| return None, None, "Please upload an image.", None, None | |
| # Parse settings | |
| k_colors = 0 | |
| if color_limit != "Unlimited": | |
| k_colors = int(color_limit) | |
| force_step = None | |
| if force_step_val > 0: | |
| force_step = float(force_step_val) | |
| original_name = os.path.basename(input_image_path) | |
| base, _ = os.path.splitext(original_name) | |
| if base.startswith("image") and len(base) > 5: | |
| base = "image" | |
| from pixelfixer.api import process | |
| with open(input_image_path, "rb") as f: | |
| img_bytes = f.read() | |
| try: | |
| result = process( | |
| img_bytes, | |
| mode="full", | |
| two_stage=True, | |
| k_colors=k_colors, | |
| force_step=force_step, | |
| return_png=False | |
| ) | |
| except Exception as e: | |
| return None, None, f"Error processing image: {str(e)}", None, None | |
| tmp_dir = tempfile.mkdtemp() | |
| low_res_array = result["array"] | |
| if k_colors > 0: | |
| from pixelfixer.quantize import kmeans_quantize | |
| low_res_array, _, _ = kmeans_quantize(low_res_array, k=k_colors) | |
| base_img = Image.fromarray(low_res_array) | |
| # 1. Create a Crisp Preview Image (Upscaled heavily for UI viewing) | |
| preview_scale = max(1, 512 // max(result["cols"], result["rows"])) | |
| preview_w = result["cols"] * preview_scale | |
| preview_h = result["rows"] * preview_scale | |
| preview_img = base_img.resize((preview_w, preview_h), Image.NEAREST) | |
| preview_path = os.path.join(tmp_dir, f"{base}-preview.png") | |
| preview_img.save(preview_path, format="PNG") | |
| # 2. Generate Export Files based on user selection | |
| export_paths = [] | |
| if not export_scales: | |
| export_scales = ["1x"] # Fallback if nothing selected | |
| for scale_str in export_scales: | |
| scale_mult = int(scale_str.replace("x", "")) | |
| target_w = result["cols"] * scale_mult | |
| target_h = result["rows"] * scale_mult | |
| export_img = base_img.resize((target_w, target_h), Image.NEAREST) | |
| export_path = os.path.join(tmp_dir, f"{base}-{scale_str}.png") | |
| export_img.save(export_path, format="PNG") | |
| export_paths.append(export_path) | |
| # 3. Generate Overlay | |
| try: | |
| input_img = Image.open(input_image_path).convert("RGBA") | |
| rgba = np.array(input_img) | |
| overlay_img = Image.fromarray(_overlay(rgba, result)) | |
| overlay_path = os.path.join(tmp_dir, f"{base}-overlay.png") | |
| overlay_img.save(overlay_path, format="PNG") | |
| except Exception: | |
| overlay_path = None | |
| # 4. Info String | |
| info = f"Grid detected: {result['cols']}x{result['rows']}\n" | |
| info += f"Step size: {result['step_x']}x{result['step_y']}\n" | |
| info += f"Confidence: {result['confidence'].upper()} ({result['consensus']})\n" | |
| if k_colors > 0: | |
| info += f"Colors: {k_colors}\n" | |
| info += f"Detect: {result['detect_s']}s | Recon: {result['recon_s']}s" | |
| return preview_path, export_paths, info, overlay_path | |
| with gr.Blocks(title="Pixel Art Fixer", css=".pixelated-image { image-rendering: pixelated; }") as app: | |
| gr.Markdown("# Pixel Art Fixer") | |
| gr.Markdown("Convert AI-generated 'fake' pixel art into real, grid-aligned pixel art.") | |
| with gr.Row(): | |
| input_image = gr.Image(type="filepath", label="Input Image", elem_classes="pixelated-image") | |
| preview_image = gr.Image(type="filepath", label="Crisp Preview", elem_classes="pixelated-image") | |
| with gr.Row(): | |
| with gr.Column(): | |
| color_limit = gr.Radio(choices=["Unlimited", "8", "12", "16", "32"], value="Unlimited", label="Color Palette Limit") | |
| export_scales = gr.CheckboxGroup(choices=["1x", "2x", "5x", "10x"], value=["1x", "5x"], label="Export Scales (Select Multiple)") | |
| force_step = gr.Number(label="Force Block Size (px) - Set to 0 for Auto", value=0, minimum=0, precision=1) | |
| process_btn = gr.Button("Process Image", variant="primary") | |
| with gr.Column(): | |
| export_files = gr.File(label="Download Extracted Pixel Art", file_count="multiple") | |
| info_text = gr.Textbox(label="Detection Info", lines=5) | |
| with gr.Column(): | |
| overlay_image = gr.Image(type="filepath", label="Grid Overlay") | |
| process_btn.click( | |
| fn=process_image, | |
| inputs=[input_image, color_limit, export_scales, force_step], | |
| outputs=[preview_image, export_files, info_text, overlay_image] | |
| ) | |
| if __name__ == "__main__": | |
| app.launch() | |