brendenc's picture
Update app.py
5d2db50
raw
history blame
1.36 kB
import datasets
import gradio as gr
from transformers import AutoFeatureExtractor, AutoModelForImageClassification
dataset = datasets.load_dataset("beans")
feature_extractor = AutoFeatureExtractor.from_pretrained("saved_model_files")
model = AutoModelForImageClassification.from_pretrained("saved_model_files")
labels = dataset['train'].features['labels'].names
example_imgs = ["example_0.jpg", "example_1.jpg","example_2.jpg"]
def classify(im):
features = feature_extractor(im, return_tensors='pt')
logits = model(features["pixel_values"])[-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)}
return confidences
interface = gr.Interface(fn = classify,
inputs="image",
outputs = "label",
title = "Plant Leaf Disease Classifier",
description = """Below is a simple app to detect Angular Leaf Spot and Bean Rust diseases on leaves.
Data was annotated by experts from the National Crops Resources Research Institute (NaCRRI)
in Uganda and collected by the Makerere AI research lab.""",
examples = example_imgs)
interface.launch(debug=True)