from fastai.vision import *
import gradio as gr
from pathlib import Path
from fastbook import load_learner

model_path = Path('models')
image_path = Path('images')

learn = load_learner(model_path/'cloudmodel.pk1')

categories = (
  'Cirrus', 
  'Cirrostratus',
  'Cirrocumulus',
  'Altostratus',
  'Altocumulus',
  'Stratus',
  'Stratocumulus',
  'Nimbostratus',
  'Cumulus',
  'Cumulonimbus',
  'Lenticular'
)

image = gr.inputs.Image(shape=(192,192))
label = gr.outputs.Label()

learn = load_learner(model_path/'cloudmodel.pkl')
examples = [image_path / f"{c}.jpg" for c in categories]

def classify_image(img):
   pred,idx,probs = learn.predict(img)
   return dict(zip(categories, map(float, probs)))


iface = gr.Interface(fn=classify_image, inputs=image, outputs=label, 
   examples=examples)

iface.launch()