Spaces:
Sleeping
Sleeping
File size: 1,013 Bytes
96229d9 5c4fd9a 96229d9 5c4fd9a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 |
from transformers import ViTImageProcessor, ViTForImageClassification
from PIL import Image
import gradio as gr
# Load the model and processor
processor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224')
model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')
def predict(image):
inputs = processor(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]
def classify_image(image):
image = Image.fromarray(image.astype('uint8'), 'RGB')
label = predict(image)
return label
iface = gr.Interface(
fn=classify_image,
inputs=gr.Image(type="numpy", label="Upload an Image"),
outputs=gr.Textbox(label="Predicted Class"),
title="Image Classification with ViT",
description="Upload an image to classify it using the Vision Transformer (ViT) model."
)
if __name__ == "__main__":
iface.launch()
|