Spaces:
Runtime error
Runtime error
from transformers import pipeline | |
from transformers import ViTFeatureExtractor, ViTForImageClassification | |
from PIL import Image as img | |
import numpy as np | |
import gradio as gr | |
featureextractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224') | |
model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224') | |
def classify(input_img): | |
filename = input_img | |
imagearray = input_img | |
inputs = featureextractor(images = imagearray, return_tensors="pt") | |
outputs = model(**inputs) | |
logits = outputs.logits | |
predicted_class_idx = logits.argmax(-1).item() | |
return model.config.id2label[predicted_class_idx] | |
demo = gr.Interface(fn=classify, inputs="image", outputs="text") | |
demo.launch() |