xferdie commited on
Commit
2f789a6
1 Parent(s): eeab7a4

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +16 -9
app.py CHANGED
@@ -1,16 +1,23 @@
 
1
  from transformers import AutoFeatureExtractor, SwinForImageClassification
2
  from PIL import Image
3
  import requests
4
 
5
- url = "http://images.cocodataset.org/val2017/000000039769.jpg"
6
- image = Image.open(requests.get(url, stream=True).raw)
7
-
8
  feature_extractor = AutoFeatureExtractor.from_pretrained("microsoft/swin-small-patch4-window7-224")
9
  model = SwinForImageClassification.from_pretrained("microsoft/swin-small-patch4-window7-224")
10
 
11
- inputs = feature_extractor(images=image, return_tensors="pt")
12
- outputs = model(**inputs)
13
- logits = outputs.logits
14
- # model predicts one of the 1000 ImageNet classes
15
- predicted_class_idx = logits.argmax(-1).item()
16
- print("Predicted class:", model.config.id2label[predicted_class_idx])
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
  from transformers import AutoFeatureExtractor, SwinForImageClassification
3
  from PIL import Image
4
  import requests
5
 
 
 
 
6
  feature_extractor = AutoFeatureExtractor.from_pretrained("microsoft/swin-small-patch4-window7-224")
7
  model = SwinForImageClassification.from_pretrained("microsoft/swin-small-patch4-window7-224")
8
 
9
+ def classify_image(url):
10
+ image = Image.open(requests.get(url, stream=True).raw)
11
+ inputs = feature_extractor(images=image, return_tensors="pt")
12
+ outputs = model(**inputs)
13
+ logits = outputs.logits
14
+ predicted_class_idx = logits.argmax(-1).item()
15
+ return model.config.id2label[predicted_class_idx]
16
+
17
+ examples = [
18
+ ["http://images.cocodataset.org/val2017/000000039769.jpg"],
19
+ ]
20
+
21
+ iface = gr.Interface(fn=classify_image, inputs="text", outputs="text", examples=examples)
22
+ iface.launch()
23
+