cupkake14 commited on
Commit
d12f7bd
1 Parent(s): ae69021

create app.py

Browse files
Files changed (1) hide show
  1. app.py +31 -0
app.py ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import datasets
3
+ import transformers
4
+ import torch
5
+
6
+ from transformers import AutoFeatureExtractor, AutoModelForImageClassification
7
+
8
+ dataset = datasets.load_dataset('beans')
9
+
10
+ extractor = AutoFeatureExtractor.from_pretrained("saved_model_files")
11
+ model = AutoModelForImageClassification.from_pretrained("saved_model_files")
12
+
13
+ labels = dataset['train'].features['labels'].names
14
+
15
+ def classify(im):
16
+ features = feature_extractor(im, return_tensors='pt')
17
+ with torch.no_grad():
18
+ logits = model(**features).logits
19
+ probability = torch.nn.functional.softmax(logits, dim=-1)
20
+ probs = probability[0].detach().numpy()
21
+ confidences = {label: float(probs[i]) for i, label in enumerate(labels)}
22
+ return confidences
23
+
24
+
25
+ gr.Interface(fn = classify,
26
+ inputs = "image",
27
+ outputs = "label",
28
+ examples='/examples',
29
+ title='Leaf classification on beans dataset',
30
+ description='Fine-tuning a ViT for bean plant health classification'
31
+ ).launch(debug=True)