app.py

import cv2
import numpy as np
import json
import math
import random
import gradio as gr

# Set random seed for reproducibility
random.seed(42)
np.random.seed(42)

# --- Utility Functions ---

def generate_id(prefix="el"):
    """Generate a random ID for Excalidraw elements."""
    return f"{prefix}_{''.join(random.choices('abcdefghijklmnopqrstuvwxyz0123456789', k=9))}"

def bgr_to_hex(bgr_color):
    """Convert BGR color tuple to a hex string."""
    b, g, r = [int(c) for c in bgr_color]
    return f"#{r:02x}{g:02x}{b:02x}"

def simplify_contour(contour, epsilon_factor=0.002):
    """
    Simplifies a contour to create smooth, crisp lines using a small epsilon factor.
    
    Args:
        contour: OpenCV contour (numpy array of points).
        epsilon_factor: Factor to determine the approximation accuracy. A smaller
                        value results in more points and a line closer to the original shape.
    
    Returns:
        List of [x, y] coordinate pairs representing the simplified line.
    """
    if len(contour) < 2:
        return []
    
    # Calculate epsilon based on the contour's perimeter for responsive simplification
    epsilon = epsilon_factor * cv2.arcLength(contour, True)
    
    # Apply the Douglas-Peucker algorithm to simplify the contour
    simplified = cv2.approxPolyDP(contour, epsilon, True)
    
    # Convert from OpenCV's format (n, 1, 2) to a simple list of [x, y] points
    return [[float(point[0][0]), float(point[0][1])] for point in simplified]

def image_to_excalidraw_json(image_np, min_shape_size):
    """
    Main function to convert an image with a black background into Excalidraw JSON.
    
    Args:
        image_np: The input image as a NumPy array (from Gradio).
        min_shape_size: The minimum area for a contour to be considered a shape.
    
    Returns:
        A string containing the Excalidraw JSON, or an error message.
    """
    if image_np is None:
        return "Please upload an image to begin."
    
    # Convert the input image from RGB (Gradio) to BGR (OpenCV)
    img_bgr = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
    
    # --- Shape Detection ---
    # Convert the image to grayscale to easily find non-black areas
    img_gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
    
    # Threshold the image to create a binary mask. All pixels that are not black (value > 5)
    # will be turned white (255). This isolates the shapes from the background.
    _, binary_mask = cv2.threshold(img_gray, 5, 255, cv2.THRESH_BINARY)
    
    # Find contours (outlines of the shapes) in the binary mask
    # RETR_EXTERNAL gets only the outer contours, which is what we want.
    # CHAIN_APPROX_SIMPLE compresses horizontal, vertical, and diagonal segments, saving memory.
    contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    excalidraw_elements = []
    
    # Process each detected contour
    for contour in contours:
        # Filter out very small contours that are likely just noise
        if cv2.contourArea(contour) < min_shape_size:
            continue
            
        # --- Color Extraction ---
        # Create a mask for the current contour to isolate it
        mask = np.zeros_like(img_gray)
        cv2.drawContours(mask, [contour], -1, (255), thickness=cv2.FILLED)
        
        # Calculate the average color of the shape from the original color image
        # The mask ensures we only consider pixels belonging to the current shape.
        mean_color_bgr = cv2.mean(img_bgr, mask=mask)[:3]
        stroke_color_hex = bgr_to_hex(mean_color_bgr)
        
        # --- Point Simplification ---
        # Simplify the contour to get smooth points for the Excalidraw line
        points = simplify_contour(contour)
        
        if len(points) < 2:
            continue
            
        # --- Element Creation ---
        # The first point is the anchor (x, y), subsequent points are relative to it.
        start_x, start_y = points[0]
        relative_points = [[p[0] - start_x, p[1] - start_y] for p in points]
        
        # Define the style for the Excalidraw element
        element_style = {
            "strokeColor": stroke_color_hex,
            "backgroundColor": "transparent",
            "fillStyle": "hachure",
            "strokeWidth": 2,
            "strokeStyle": "solid",
            "roughness": 0,
            "opacity": 100,
            "roundness": {"type": 2}, # Creates smooth, rounded corners
            "seed": random.randint(1_000_000, 9_999_999),
            "version": 1,
            "versionNonce": random.randint(1_000_000, 9_999_999),
        }
        
        # Create the Excalidraw line element dictionary
        line_element = {
            "id": generate_id("line"),
            "type": "line",
            "x": float(start_x),
            "y": float(start_y),
            "angle": 0,
            "points": relative_points,
            **element_style,
        }
        excalidraw_elements.append(line_element)
        
    # Create the final Excalidraw clipboard structure
    final_excalidraw_structure = {
        "type": "excalidraw/clipboard",
        "elements": excalidraw_elements,
        "files": {}
    }
    
    # Convert the Python dictionary to a nicely formatted JSON string
    return json.dumps(final_excalidraw_structure, indent=2)

# --- Gradio User Interface ---
def create_interface():
    """Creates and configures the Gradio web interface."""
    with gr.Blocks(title="Image to Excalidraw", theme=gr.themes.Soft()) as iface:
        gr.Markdown("# ✏️ Image to Excalidraw Converter")
        gr.Markdown("Upload an image with colored shapes on a **black background**. This tool will trace the shapes and generate Excalidraw JSON for you to copy and paste.")
        
        with gr.Row(variant="panel"):
            with gr.Column(scale=1):
                image_input = gr.Image(type="numpy", label="Upload Your Image")
                
                min_shape_size = gr.Slider(
                    minimum=10,
                    maximum=1000,
                    value=50,
                    step=10,
                    label="Minimum Shape Size",
                    info="Filters out small noise. Increase if you see unwanted tiny specks."
                )
                
                process_btn = gr.Button("✅ Generate Excalidraw Code", variant="primary")
            
            with gr.Column(scale=2):
                output_json = gr.Textbox(
                    label="Excalidraw JSON Output",
                    info="Click to copy the JSON, then paste it directly into Excalidraw (Ctrl+V or Cmd+V).",
                    lines=20,
                    max_lines=30,
                    show_copy_button=True
                )
        
        # Connect the button click event to the processing function
        process_btn.click(
            fn=image_to_excalidraw_json,
            inputs=[image_input, min_shape_size],
            outputs=output_json
        )
        
        # For better user experience, also trigger processing when the image or slider changes
        image_input.change(
            fn=image_to_excalidraw_json,
            inputs=[image_input, min_shape_size],
            outputs=output_json
        )
        min_shape_size.release(
            fn=image_to_excalidraw_json,
            inputs=[image_input, min_shape_size],
            outputs=output_json
        )
        
    return iface

# --- Launch the App ---
if __name__ == "__main__":
    app = create_interface()
    app.launch()