|
import streamlit as st |
|
from transformers import ViTImageProcessor, ViTForImageClassification |
|
from PIL import Image as img |
|
|
|
x = st.file_uploader("Upload Images", type=["png","jpg","jpeg"]) |
|
if x is not None: |
|
st.image(img.open(x),width=255) |
|
i = img.open(x) |
|
processor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224') |
|
model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224') |
|
inputs = processor(images=i, return_tensors="pt") |
|
outputs = model(**inputs) |
|
logits = outputs.logits |
|
predicted_class_idx = logits.argmax(-1).item() |
|
st.text("Our Model Predicts : ") |
|
st.write(model.config.id2label[predicted_class_idx]) |