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

path = Path()
path.ls(file_exts='.pkl')
learn = load_learner(path/'model.pkl')

labels = {
        "cirrus": "Cirrus (Federwolke)",
        "cirrocumulus": "Cirrocumulus (kleine Schäfchenwolke)",
        "cirrostratus": "Cirrostratus (hohe Schleierwolke)",
        "altocumulus": "Altocumulus (große Schäfchenwolke)",
        "altostratus": "Altostratus (mittelhohe Schichtwolke)",
        "stratocumulus": "Stratocumulus (Haufenschichtwolke)",
        "stratus": "Stratus (tiefe Schichtwolke)",
        "cumulus": "Cumulus (Haufenwolke)",
        "nimbostratus": "Nimbostratus (Regenwolke)",
        "cumulonimbus": "Cumulonimbus (Gewitterwolke)"
        }

def predict_cloud_class(img):
  cloudtype,_,probs = learn.predict(PILImage.create(img))
  return {labels[learn.dls.vocab[i]]: float(probs[i]) for i in range(len(learn.dls.vocab))}

image = gr.inputs.Image()
label = gr.outputs.Label(num_top_classes=5)
gr.Interface(fn=predict_cloud_class, inputs=image, outputs=label,examples=["test_cirrus.jpg","test_cumulus.jpg","test_altocumulus.jpg"]).launch(server_name="0.0.0.0")