ST_Segmentation / app.py
Schrodingers's picture
Update app.py
e54b7cb
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()