import gradio as gr 
import pandas as pd 

from PIL import Image
from torchkeras import plots 
from torchkeras.data import get_url_img

from pathlib import Path
from ultralytics import YOLO
import ultralytics
from ultralytics.yolo.data import utils 

# Pfad zu Ihrem Modell aktualisieren
model = YOLO('2023-05-12-foodAndDrinks.pt')

# Pfad zur YAML-Datei aktualisieren
yaml_path = '2023-05-12-foodAndDrinks.yaml' 
class_names = utils.yaml_load(yaml_path)['names']

def detect(img):
    if isinstance(img,str):
        img = get_url_img(img) if img.startswith('http') else Image.open(img).convert('RGB')
    result = model.predict(source=img)
    if len(result[0].boxes.boxes)>0:
        vis = plots.plot_detection(img,boxes=result[0].boxes.boxes,
                     class_names=class_names, min_score=0.2)
    else:
        vis = img
    return vis

with gr.Blocks() as demo:

    with gr.Tab("Upload"):
        gr.Markdown("# foodServed, drinkServed, person, V0.0.10")  # Dieser Text wird am Anfang des Tabs angezeigt.
        
        # Pfad zu Ihren Demo-Bildern
        demo_images = ["demoImages/demo01.jpg", "demoImages/demo02.jpg", "demoImages/demo03.jpg", "demoImages/demo04.jpg"]

        input_img = gr.Image(type='pil')
        out_img = gr.Image(type='pil')
        
        gr.Examples(examples=[[img] for img in demo_images], 
                     inputs=[input_img], 
                     outputs=[out_img], 
                     fn=detect)
        
        button = gr.Button("Detect",variant="primary")
        button.click(detect,inputs=input_img, outputs=out_img)
        
        gr.Markdown("## Output")

    with gr.Tab("Url"):
        default_url = 'https://i.postimg.cc/0jdbK03h/food-Image.jpg'
        url = gr.Textbox(value=default_url)
        button = gr.Button("Detect",variant="primary")
        
        gr.Markdown("## Output")
        out_img = gr.Image(type='pil')
        
        button.click(detect,
                     inputs=url, 
                     outputs=out_img)

gr.close_all() 
demo.queue()
demo.launch()