from  hubconf import custom
model = custom(path_or_model='best.pt')  # custom example
model.eval()
# model = create(name='yolov7', pretrained=True, channels=3, classes=80, autoshape=True)  # pretrained example

# Verify inference
import numpy as np
import torch
from PIL import Image
import gradio as gr

# imgs = [np.zeros((640, 480, 3))]
# imgs = 'inference/images/meal.jpg'

# results = model(imgs)  # batched inference
# results.print()
# results.save()
def total_price(predicted):
    price = 0
    for name, confidence in predicted:
        if name == "side dish":
            price += 10
        elif name == "purple rice" or name == "white rice" or name == "brown rice":
            price += 20
        elif name == "40dollars meal":
            price += 40
        elif name == "30dollars meal":
            price += 30
        elif name == "25dollars meal":
            price += 25

    return price

def predict(input_image):
    """
    Predict model output
    """
    # Disable gradient computation
    with torch.no_grad():
        results = model(input_image)

    results_pd = results.pandas().xyxy[0]
    name, confidence = results_pd["name"], results_pd["confidence"]
    predicted = list(zip(name, confidence))
    output_image = results.render()[0]
    price = total_price(predicted)

    # Return the output image and price
    return [output_image, price]

with gr.Blocks() as demo:
    # Title
    gr.HTML(
    """
    <h1 align="center">Group 29 - AI Cafeteria Price Evaluator</h1>
    """)
    examples = ["./examples/img_1.jpg", "./examples/img_2.jpg", "./examples/img_3.jpg"]
    # gr.Interface(inputs=["image"],outputs=["image"],fn=lambda img:model(img).render()[0]).launch()
    gr.Interface(inputs=["image"], outputs=["image", "text"], fn=predict, examples=examples)

    gr.HTML(
    """
    <style>
        table.center {
            margin-left: auto; 
            margin-right: auto;
        }
        .center th{
            text-align: center;
        }
    </style>
    <h1 align="center">Price List</h1>
    <h3>
    <table class="center">
        <tr>
            <th style="text-align: center">Menu</th>
            <th style="text-align: center">Price</th>
        </tr>
        <tr>
            <td>White/Brown/Purple Rice</td>
            <td>20</td>
        </tr>
        <tr>
            <td>Side Dish</td>
            <td>10</td>
        </tr>
        <tr>
            <td>25 Dollar Meal</td>
            <td>25</td>
        </tr>
        <tr>
            <td>30 Dollar Meal</td>
            <td>30</td>
        </tr>
        <tr>
            <td>40 Dollar Meal</td>
            <td>40</td>
        </tr>
    </table>
    </h3>
    """)
    
if __name__ == "__main__":
    demo.launch()