ScubaSpotter / app.py
patrickerichsen's picture
init
ab30774
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)