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

# 下载模型
model_path = snapshot_download(repo_id='Hongyang-Li/SuperRetinaDemo')
inference_engine = Inference(model_path=model_path)

def main(source_image, target_image):
    # 只有当两个图像都存在时才执行推理
    if source_image is not None and target_image is not None:
        merged, match_show = inference_engine.inference(source_image, target_image)
        return merged, match_show
    else:
        # 如果任何一个图像为空，则返回空结果
        return None, None

with gr.Blocks() as demo:
    gr.Markdown("# SuperRetina CFP Registration Demo")
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("Input Images")
            source_image_input = gr.Image(
                type="numpy", 
                label="Source Image",
                height=300  # 固定高度
            )
            target_image_input = gr.Image(
                type="numpy", 
                label="Target Image", 
                height=300  # 固定高度
            )
            
            # 添加推理按钮
            infer_button = gr.Button("Registration", variant="primary")
            
        with gr.Column():
            gr.Markdown("Result")
            output_image = gr.Image(
                type="numpy", 
                label="Registration Result",
                height=300  # 固定高度
            )
            
            gr.Markdown("Keypoints Matching")
            match_show_image = gr.Image(
                type="numpy", 
                label="Keypoints Matching",
                height=300  # 固定高度
            )
    
    # 只有当按钮被点击时才触发推理
    infer_button.click(
        fn=main,
        inputs=[source_image_input, target_image_input],
        outputs=[output_image, match_show_image]
    )
    
    # 可选：添加清除按钮
    clear_button = gr.Button("Clean all")
    
    def clear_all():
        return None, None, None, None
    
    clear_button.click(
        fn=clear_all,
        inputs=[],
        outputs=[source_image_input, target_image_input, output_image, match_show_image]
    )

demo.launch()