import gradio as gr
from inference import Inference
import os
from huggingface_hub import snapshot_download

#MODEL_ID = os.getenv("MODEL_ID", "your_username/your_model_name")  # 替换为你的模型ID

model_path = snapshot_download(repo_id='AIMClab-RUC/UNet_DCP_1024')

TEXT_OPTIONS = ["CFP", "UWF", "FFA", "SLO", "OCTA"]

inference_engine = Inference(model_path=model_path)

def main(image):
    out = inference_engine.inference(image, "CFP")
    return out

with gr.Blocks() as demo:
    gr.Markdown("# [ICASSP 2025] Broad domain retinal vessel segmentation")
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="numpy", label="Input Image")
        with gr.Column():
            image_output = gr.Image(type="numpy", label="Output")
    
    # 当图像输入发生变化时自动触发推理
    image_input.change(
        fn=main,
        inputs=image_input,
        outputs=image_output
    )

demo.launch()