import gradio as gr
from tongue_model import TongueModelWrapper 

wrapper = TongueModelWrapper(model_path="best_model.pt")

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 中醫舌象自動診斷系統")
    
    with gr.Row():
        with gr.Column():
            input_img = gr.Image(label="上傳舌象照片")
            btn = gr.Button("開始分析", variant="primary")
        
        with gr.Column():
            output_label = gr.Label(label="預測結果機率")

    btn.click(fn=wrapper.predict, inputs=input_img, outputs=output_label)

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