import gradio as gr
import opencv
import numpy as np
import watershed  # Your Python bindings module

def segment_image(input_image):
    # Convert the image received from Gradio into a format suitable for our C++ function
    is_success, im_buf_arr = cv2.imencode(".jpg", input_image)
    byte_im = im_buf_arr.tobytes()

    # You can save this to a file or directly pass it if your C++ function can handle byte streams
    with open("temp_input.jpg", "wb") as f:
        f.write(byte_im)

    # Call the C++ watershed_segmentation function
    watershed.watershed_segmentation("temp_input.jpg", "output.jpg")

    # Read the output image to display in Gradio
    output_img = cv2.imread("output.jpg")

    # Convert from BGR (OpenCV format) to RGB (for displaying in Gradio)
    output_img = cv2.cvtColor(output_img, cv2.COLOR_BGR2RGB)
    return output_img

# Define Gradio interface
interface = gr.Interface(
    fn=segment_image, 
    inputs=gr.inputs.Image(type="numpy", label="Input Image"), 
    outputs=gr.outputs.Image(type="numpy", label="Segmented Image"),
    live=True,
    title="Watershed Image Segmentation",
    description="Upload an image for watershed segmentation."
)

interface.launch()