app.py

import cv2 as cv
import numpy as np
import gradio as gr
from huggingface_hub import hf_hub_download
from efficientSAM import EfficientSAM

# Download EfficientSAM ONNX model
model_path = hf_hub_download(
    repo_id="opencv/image_segmentation_efficientsam",
    filename="image_segmentation_efficientsam_ti_2025april.onnx"
)

# Initialize the EfficientSAM model
model = EfficientSAM(modelPath=model_path)

# Global state
state = {
    "points": [],
    "labels": [],
    "image": None,             # BGR image for inference
    "original_image_rgb": None # Clean RGB image to display/edit
}

MAX_POINTS = 6

def segment_image_with_prompts(image, evt: gr.SelectData):
    if state["image"] is None:
        bgr = cv.cvtColor(image, cv.COLOR_RGB2BGR)
        state["image"] = bgr
        state["original_image_rgb"] = image.copy()

    x, y = int(evt.index[0]), int(evt.index[1])

    if len(state["points"]) < MAX_POINTS:
        state["points"].append([x, y])
        state["labels"].append(1)  # foreground
    else:
        return state["original_image_rgb"]

    marked_img = state["original_image_rgb"].copy()
    for (px, py), lbl in zip(state["points"], state["labels"]):
        color = (0, 255, 0) if lbl == 1 else (0, 0, 255)
        cv.circle(marked_img, (px, py), 5, color, -1)

    return marked_img

def run_segmentation():
    if state["image"] is None or len(state["points"]) == 0:
        return None

    result = model.infer(image=state["image"], points=state["points"], labels=state["labels"])
    if len(result) == 0:
        return None

    vis_result = visualize(state["image"], result)
    return cv.cvtColor(vis_result, cv.COLOR_BGR2RGB)

def visualize(image, result):
    vis_result = np.copy(image)
    mask = np.copy(result)
    _, binary = cv.threshold(mask, 127, 255, cv.THRESH_BINARY)
    red_channel = vis_result[:, :, 2]
    red_channel = np.where(binary == 255, np.minimum(red_channel * 1.8, 255), red_channel)
    vis_result[:, :, 2] = red_channel
    contours, _ = cv.findContours(binary, cv.RETR_LIST, cv.CHAIN_APPROX_TC89_L1)
    cv.drawContours(vis_result, contours, -1, (255, 255, 255), 2)
    return vis_result

def clear_all():
    state["points"].clear()
    state["labels"].clear()
    state["image"] = None
    state["original_image_rgb"] = None
    return None, None

def clear_points():
    state["points"].clear()
    state["labels"].clear()
    return state["original_image_rgb"], None

def update_image_state(new_image):
    """
    Called only when the user actually loads or picks a new image.
    Reset everything here.
    """
    if new_image is None:
        return None, None

    bgr = cv.cvtColor(new_image, cv.COLOR_RGB2BGR)
    state["image"] = bgr
    state["original_image_rgb"] = new_image.copy()

    state["points"].clear()
    state["labels"].clear()

    return None, None

# Gradio UI
with gr.Blocks(css='''.example * {
    font-style: italic;
    font-size: 18px !important;
    color: #0ea5e9 !important;
    }''') as demo:

    gr.Markdown("### EfficientSAM: Interactive Image Segmentation")
    gr.Markdown("Click to add up to 6 **foreground points**, then click **Segment**. Use **Clear All** to reset, or **Clear Points** to remove only clicks.")

    with gr.Row():
        input_image = gr.Image(type="numpy", label="Upload Image", interactive=True)
        prompt_overlay = gr.Image(type="numpy", label="Your Clicks", interactive=False)
        output_image = gr.Image(type="numpy", label="Segmentation Result")

    input_image.change(fn=update_image_state, inputs=input_image, outputs=[output_image, prompt_overlay])

    with gr.Row():
        segment_btn = gr.Button("Segment", variant="primary")
        clear_points_btn = gr.Button("Clear Points")
        clear_all_btn = gr.Button("Clear All")

    segment_btn.click(fn=run_segmentation, outputs=output_image)
    clear_points_btn.click(fn=clear_points, outputs=[input_image, output_image])
    clear_all_btn.click(fn=clear_all, outputs=[input_image, output_image])
    input_image.select(fn=segment_image_with_prompts, inputs=input_image, outputs=prompt_overlay)

    gr.Markdown("Click on any example to try it.", elem_classes=["example"])

    gr.Examples(
        examples=[
            ["examples/dog416.png"],
            ["examples/messi5.jpg"],
            ["examples/selfie.jpg"],
            ["examples/space_shuttle.jpg"]
        ],
        inputs=input_image
    )

if __name__ == "__main__":
    demo.launch()