rahult commited on
Commit
de48d8d
1 Parent(s): ccbbf66

Add app.py and requirements.txt

Browse files
Files changed (2) hide show
  1. app.py +25 -0
  2. requirements.txt +4 -0
app.py ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import datasets
2
+ import gradio as gr
3
+ from transformers import AutoFeatureExtractor, AutoModelForImageClassification
4
+ import torch
5
+
6
+ dataset = datasets.load_dataset('beans')
7
+
8
+ extractor = AutoFeatureExtractor.from_pretrained("rahult/bean_classification")
9
+ model = AutoModelForImageClassification.from_pretrained("rahult/bean_classification")
10
+ model.eval()
11
+
12
+ labels = dataset['train'].features['labels'].names
13
+
14
+ def classify(im):
15
+ features = extractor(im, return_tensors='pt')
16
+ with torch.no_grad():
17
+ logits = model(**features).logits
18
+ probability = torch.nn.functional.softmax(logits, dim=-1)
19
+ probs = probability[0].detach().numpy()
20
+ confidences = {label: float(probs[i]) for i, label in enumerate(labels)}
21
+ return confidences
22
+
23
+
24
+ interface = gr.Interface(fn=classify, inputs="image", outputs="label", allow_flagging='manual')
25
+ interface.launch()
requirements.txt ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ datasets
2
+ transformers
3
+ gradio
4
+ torch