bean-classifier / app.py
hshetty's picture
Update app.py
275d7bf
import datasets
from transformers import AutoFeatureExtractor, AutoModelForImageClassification
import gradio as gr
import torch
dataset = datasets.load_dataset('beans','full_size') # This should be the same as the first line of Python code in this Colab notebook
extractor = AutoFeatureExtractor.from_pretrained("saved_model_files")
model = AutoModelForImageClassification.from_pretrained("saved_model_files")
labels = dataset['train'].features['labels'].names
def classify(im):
features = extractor(im, return_tensors='pt')
with torch.no_grad():
logits = model(features["pixel_values"])[-1]
#logits = torch.nn.functional.softmax(logits, dim=-1)
probability = torch.nn.functional.softmax(logits, dim=-1)
probs = probability[0].detach().numpy()
confidences = {label: float(probs[i]) for i, label in enumerate(labels)}
print(confidences)
return confidences
interface = gr.Interface(fn=classify,
inputs="image",
outputs=gr.Label(num_top_classes=3),
title="Bean leaf health classifier",
description="This is a bean leaf health classifier app created as part of CoRise End to End Vision application project",
examples=["https://datasets-server.huggingface.co/assets/beans/--/default/validation/3/image/image.jpg",
"https://datasets-server.huggingface.co/assets/beans/--/default/test/20/image/image.jpg",
"https://datasets-server.huggingface.co/assets/beans/--/default/validation/97/image/image.jpg"])
# FILL HERE
interface.launch()