francescafr commited on
Commit
5d972b5
·
1 Parent(s): 2724d5b

added inital app code

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