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

learn = load_learner('nationalflagidentifier.pkl')

def predict(image):
  img = PILImage.create(image)
  country, idx, probs = learn.predict(img)
  return f'This is the national flag of {country}------------ Probability: {probs[idx]:.4f}'

title = 'National Flag Classifier'
description = """Upload a flag photo and get to know to which country or dependency it belongs to. 
                                   267 countries or dependencies supported."""
interpretation = 'default'

infer = gr.Interface(
    fn = predict,
    inputs = gr.inputs.Image(shape = (256, 256)),
    outputs = gr.outputs.Label(num_top_classes = len(learn.dls.vocab)),
    title = title,
    description = description,
    interpretation = interpretation
)

infer.launch()