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

learn = load_learner('model.pkl')

labels = learn.dls.vocab

def predict(img):
    img = PILImage.create(img)
    pred, pred_idx, probs = learn.predict(img)
    
    # Check if the predicted label is 'Other'
    if pred == 'Other':
        return {'Message': 'The model can\'t recognize the image, adjust your camera angle and take a close picture of a leaf'}
    
    return {labels[i]: float(probs[i]) for i in range(len(labels))}

title = "Common beans diseases classifier"
description = "An app for Common beans diseases Classification"
article = "<p style='text-align: center'>The app identifies and classifies common beans diseases: Anthracnose and Bean rust.</p>"
examples = ['image.jpeg']
interpretation = 'default'
enable_queue = True

# Modify the output to Textbox to handle string messages as well
gr.Interface(fn=predict,
             inputs=gr.inputs.Image(shape=(512, 512)),
             outputs=gr.outputs.Textbox(),
             title=title,
             description=description,
             article=article,
             examples=examples,
             interpretation=interpretation,
             enable_queue=enable_queue).launch()