from fastai.vision.all import *
import gradio as gr

labels = [
    "20km/h",
    "30km/h",
    "50km/h",
    "60km/h",
    "70km/h",
    "80km/h",
    "80km/h aufgehoben",
    "100km/h",
    "120km/h",
    "Überholverbot",
    "Überholverbot (nur LKW)",
    "Vorfahrt an der nächsten Kreuzung",
    "Vorfahrtsstraße",
    "Vorfahrt gewähren",
    "Stop",
    "Einfahrt verboten",
    "Einfahrt verboten (LKW)",
    "Durchfahrt verboten",
    "Gefahr",
    "scharfe Kurve (links)",
    "scharfe Kurve (rechts)",
    "kurvige Strecke",
    "Bodenwellen",
    "Schleudergefahr",
    "verengte Fahrbahn (rechts)",
    "Bauarbeiten",
    "Ampel",
    "Zebrastreifen (Achtung)",
    "spielende Kinder",
    "Fahrradfahrer",
    "Glätte",
    "Wildwechsel",
    "Aufhebung aller Beschränkungen",
    "Nur rechts abbiegen",
    "Nur links abbiegen",
    "Nur geradeaus",
    "Nur geradeaus / rechts",
    "Nur geradeaus / links",
    "Rechts vorbeifahren",
    "Links vorbeifahren",
    "Kreisverkehr",
    "Ende Überholverbot",
    "Ende Überholverbot (LKW)",
]


def get_y(y):
    idx = int(parent_label(y))
    return labels[idx]


learner = load_learner("model.pkl")


lables = learner.dls.vocab


def predict(img):
    img = PILImage.create(img)
    pred, pred_idx, probs = learner.predict(img)
    return {lables[i]: float(probs[i]) for i in range(len(lables))}


gr.Interface(
    fn=predict, inputs=gr.Image(), outputs=gr.Label(num_top_classes=3)
).launch()