import numpy as np
import cv2
from scipy.ndimage import label, binary_dilation
import gradio as gr

def create_bridged_stencil(image, min_bridges=2, bridge_width=2):
    # Convert the image to grayscale if it's not already
    if len(image.shape) == 3:
        img = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    else:
        img = image
    
    # Threshold the image to create a binary image
    _, binary = cv2.threshold(img, 127, 255, cv2.THRESH_BINARY_INV)
    
    # Find connected components (islands)
    labeled, num_features = label(binary)
    
    # Create a copy of the binary image for adding bridges
    bridged = binary.copy()
    
    # Function to add a bridge between two points
    def add_bridge(start, end):
        cv2.line(bridged, start, end, 255, bridge_width)
    
    # For each island
    for i in range(1, num_features + 1):
        island = (labeled == i)
        
        # Find the contour of the island
        contours, _ = cv2.findContours((island * 255).astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        contour = contours[0]
        
        # Get the convex hull of the contour
        hull = cv2.convexHull(contour)
        
        # Find the center of the island
        M = cv2.moments(contour)
        if M["m00"] != 0:
            cx = int(M["m10"] / M["m00"])
            cy = int(M["m01"] / M["m00"])
        else:
            cx, cy = np.mean(contour, axis=0)[0]
        
        # Add bridges from the center to the convex hull points
        num_bridges = 0
        for point in hull:
            if num_bridges >= min_bridges:
                break
            add_bridge((cx, cy), tuple(point[0]))
            num_bridges += 1
    
    # Dilate the bridged image to make bridges more prominent
    bridged = binary_dilation(bridged, iterations=1)
    
    # Convert back to RGB for Gradio output
    return cv2.cvtColor((bridged * 255).astype(np.uint8), cv2.COLOR_GRAY2RGB)

def process_image(input_image, min_bridges, bridge_width):
    result = create_bridged_stencil(input_image, min_bridges, bridge_width)
    return result

# Create the Gradio interface
iface = gr.Interface(
    fn=process_image,
    inputs=[
        gr.Image(type="numpy", label="Input Image"),
        gr.Slider(minimum=1, maximum=5, step=1, value=2, label="Minimum Bridges"),
        gr.Slider(minimum=1, maximum=5, step=1, value=2, label="Bridge Width")
    ],
    outputs=gr.Image(type="numpy", label="Bridged Stencil"),
    title="Bridged Stencil Generator",
    description="Upload an image to create a bridged stencil. Adjust the minimum number of bridges and bridge width as needed."
)

# Launch the interface
iface.launch()