brendenc's picture
Updated from colab
c8189b0
raw
history blame
1.28 kB
import datasets
import gradio as gr
from transformers import AutoFeatureExtractor, AutoModelForImageClassification
dataset = datasets.load_dataset("beans")
extractor = AutoFeatureExtractor.from_pretrained("saved_model_files")
model = AutoModelForImageClassification.from_pretrained("saved_model_files")
labels = dataset['train'].features['labels'].names
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)