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

def predict(input_image):
    """
    Predict model output
    """
    price = "0"
    return [input_image, 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")]
    # )
    gr.Interface(inputs=["image"],outputs=["image"],fn=lambda img:model(img).render()[0]).launch()


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