import gradio as gr
import torch
# model = YOLO('best.pt')
# path = 'best.pt'
# model = model = torch.hub.load('WongKinYiu/yolov7', 'custom', 'best.pt',
#                         force_reload=True, trust_repo=True)

model = torch.hub.load('./', 'custom', 'best.pt',force_reload=True, source='local',trust_repo=True)

def predict(input_image):
    """
    Predict model output
    """
    # 使用模型進行預測
    output = model(input_image)

    # 將模型輸出轉為可讀的文字，這部分需要依據你的模型輸出的實際格式進行處理
    price = str(output)

    # 返回預測結果
    return [output, price]

with gr.Blocks() as demo:
    # Title
    gr.Markdown(
    """
    <h1 align="center">AI Cafeteria Price Evaluator</h1>
    """)

    # Model Evaluation
    gr.Interface(
        fn=predict,
        inputs=gr.Image(type="pil"),
        outputs=[gr.Image(type="pil", label="Image Prediction"),
                 gr.Textbox(type="text", label="Price Prediction")]
    )

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