srishti-hf1110 commited on
Commit
6cb0ee4
1 Parent(s): 9e49044

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +10 -1
app.py CHANGED
@@ -7,7 +7,16 @@ Original file is located at
7
  https://colab.research.google.com/drive/1XX8pCT291obpzL4fc1vu5L_HTG027lle
8
  """
9
 
10
- import gradio as gr
 
 
 
 
 
 
 
 
 
11
  def classify(im):
12
  features = feature_extractor(im, return_tensors='pt')
13
  logits = model(features["pixel_values"])[-1]
 
7
  https://colab.research.google.com/drive/1XX8pCT291obpzL4fc1vu5L_HTG027lle
8
  """
9
 
10
+ import datasets
11
+ from transformers import AutoFeatureExtractor, AutoModelForImageClassification
12
+
13
+ dataset = load_dataset("beans") # This should be the same as the first line of Python code in this Colab notebook
14
+
15
+ extractor = AutoFeatureExtractor.from_pretrained("saved_model_files")
16
+ model = AutoModelForImageClassification.from_pretrained("saved_model_files")
17
+
18
+ labels = dataset['train'].features['labels'].names
19
+
20
  def classify(im):
21
  features = feature_extractor(im, return_tensors='pt')
22
  logits = model(features["pixel_values"])[-1]