kunalsharma/ComputerVisionDemo

1import gradio as gr2 3from transformers import ViTImageProcessor, ViTForImageClassification4from PIL import Image5import requests6 7def predict(url):8    image = Image.open(requests.get(url, stream=True).raw)9 10    processor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224')11    model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')12 13    inputs = processor(images=image, return_tensors="pt")14    outputs = model(**inputs)15    logits = outputs.logits16    # model predicts one of the 1000 ImageNet classes17    predicted_class_idx = logits.argmax(-1).item()18    result = ("Predicted class:", [model.config.id2label[predicted_class_idx]])19    return result20 21iface = gr.Interface(fn=predict, inputs="text", outputs="text")22iface.launch()