abidlabs HF Staff commited on
Commit
8b7f1e2
·
1 Parent(s): 76a7e13

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -0
app.py CHANGED
@@ -1,6 +1,7 @@
1
  from transformers import ViTFeatureExtractor, ViTForImageClassification
2
  from PIL import Image
3
  import torch.nn.functional as F
 
4
 
5
 
6
  feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224')
@@ -13,6 +14,7 @@ def predict(image):
13
  predicted_class_prob = F.softmax(logits, dim=-1).detach().numpy().max()
14
  predicted_class_idx = logits.argmax(-1).item()
15
  label = model.config.id2label[predicted_class_idx].split(",")[0]
 
16
  return {label: float(predicted_class_prob)}
17
 
18
  import gradio as gr
 
1
  from transformers import ViTFeatureExtractor, ViTForImageClassification
2
  from PIL import Image
3
  import torch.nn.functional as F
4
+ import time
5
 
6
 
7
  feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224')
 
14
  predicted_class_prob = F.softmax(logits, dim=-1).detach().numpy().max()
15
  predicted_class_idx = logits.argmax(-1).item()
16
  label = model.config.id2label[predicted_class_idx].split(",")[0]
17
+ time.sleep(2)
18
  return {label: float(predicted_class_prob)}
19
 
20
  import gradio as gr