import gradio as gr
import cv2
import numpy as np

# Function to perform image segmentation using OpenCV Watershed Algorithm
def watershed_segmentation(input_image, scribble_image):
    # Load the input image and scribble image
    image = cv2.cvtColor(input_image.astype('uint8'), cv2.COLOR_RGBA2BGR)
    scribble = cv2.cvtColor(scribble_image.astype('uint8'), cv2.COLOR_RGBA2GRAY)

    # Convert scribble to markers (0 for background, 1 for unknown, 2 for foreground)
    markers = np.zeros_like(scribble)
    markers[scribble == 0] = 0
    markers[scribble == 255] = 1
    markers[scribble == 128] = 2

    # Apply watershed algorithm
    cv2.watershed(input_image, markers)

    # Create a segmented mask
    segmented_mask = np.zeros_like(input_image, dtype=np.uint8)
    segmented_mask[markers == 2] = [0, 0, 255]  # Red color for segmented regions

    return segmented_mask

# Gradio interface
input_image = gr.inputs.Image(type='pil', label="Upload an image")
scribble_image = gr.inputs.Image(type='pil', label="Scribble on the image")
output_image = gr.outputs.Image(type='pil', label="Segmented Image")

gr.Interface(
    fn=watershed_segmentation,
    inputs=[input_image, scribble_image],
    outputs=output_image,
    title="Swallowtail Image Segmentation Algorithm",
    description="Upload an image and scribble on it to perform segmentation.",
).launch()