andresgtn commited on
Commit
0eae57d
β€’
1 Parent(s): e503868

create app.py

Browse files
Files changed (1) hide show
  1. app.py +29 -0
app.py ADDED
@@ -0,0 +1,29 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import datasets
2
+ from transformers import AutoFeatureExtractor, AutoModelForImageClassification
3
+ import numpy as np
4
+ import gradio as gr
5
+
6
+ dataset = load_dataset("beans") # This should be the same as the first line of Python code in this Colab notebook
7
+
8
+ extractor = AutoFeatureExtractor.from_pretrained("saved_model_files")
9
+ model = AutoModelForImageClassification.from_pretrained("saved_model_files")
10
+
11
+ # add to cuda?
12
+ model.eval()
13
+ model.to(device)
14
+
15
+ labels = dataset['train'].features['labels'].names
16
+
17
+ def classify(im):
18
+ features = extractor(im, return_tensors='pt')
19
+ features.to(device) # move to gpu as model, if available
20
+ with torch.no_grad():
21
+ logits = model(**features).logits
22
+ probability = torch.nn.functional.softmax(logits, dim=-1)
23
+ probs = probability[0].to('cpu').detach().numpy()
24
+ confidences = {label: float(probs[i]) for i, label in enumerate(labels)}
25
+ return confidences
26
+
27
+ interface = gr.Interface(classify, gr.Image(shape=(200, 200)), 'text')
28
+ #demo.launch()
29
+ interface.launch(debug=False)