|
import gradio as gr |
|
from transformers import AutoFeatureExtractor, SwinForImageClassification |
|
from PIL import Image |
|
import requests |
|
|
|
feature_extractor = AutoFeatureExtractor.from_pretrained("microsoft/swin-small-patch4-window7-224") |
|
model = SwinForImageClassification.from_pretrained("microsoft/swin-small-patch4-window7-224") |
|
|
|
def classify_image(url): |
|
image = Image.open(requests.get(url, stream=True).raw) |
|
inputs = feature_extractor(images=image, return_tensors="pt") |
|
outputs = model(**inputs) |
|
logits = outputs.logits |
|
predicted_class_idx = logits.argmax(-1).item() |
|
return model.config.id2label[predicted_class_idx] |
|
|
|
examples = [ |
|
["http://images.cocodataset.org/val2017/000000039769.jpg"], |
|
] |
|
|
|
iface = gr.Interface(fn=classify_image, inputs="text", outputs="text", examples=examples) |
|
iface.launch() |
|
|
|
|