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

title = "ScubaSpotter 🤿"
description = "An image classifier for underwater marine life, including scuba divers themselves. Trained on the [Sea Animals Image Dataset](https://www.kaggle.com/datasets/vencerlanz09/sea-animals-image-dataste)."
examples = ['./examples/clam.jpg', './examples/scuba_diver.jpg', './examples/turtle.jpg']

learn = load_learner('export.pkl')

labels = learn.dls.vocab

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

def most_confident_prediction(img):
    img = PILImage.create(img)
    pred, pred_idx, probs = learn.predict(img)
    predictions = {labels[i]: float(probs[i]) for i in range(len(labels))}
    most_confident_prediction = max(predictions.items(), key=lambda x: x[1])
    return {most_confident_prediction[0]: float(most_confident_prediction[1])}

gr.Interface(fn=most_confident_prediction, inputs="image", outputs="label", title=title,description=description,examples=examples).launch(share=True)