Jbrcoleman/ComputerVisionDemo

1import gradio as gr2from transformers import ViTImageProcessor, ViTForImageClassification3from PIL import Image4import requests5 6 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= f"Predicted class: {model.config.id2label[predicted_class_idx]}"19    return result20 21iface = gr.Interface(fn=predict, inputs="text", outputs="text")22iface.launch()