bean-classifier / app.py
Hitesh Karunakara Shetty
Add application file
245c91f
raw history blame
No virus
1.39 kB
import datasets
from transformers import AutoFeatureExtractor, AutoModelForImageClassification
import gradio as gr
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]*100) for i, label in enumerate(labels)}
print(confidences)
return confidences
interface = gr.Interface(fn=classify,
inputs=gr.inputs.Image(type="pil"),
outputs=gr.Label(num_top_classes=3),
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/test/70/image/image.jpg"])
# FILL HERE
interface.launch()