Decoder24/ViT-Classifier

1import os2os.makedirs(os.path.expanduser("~/.streamlit"), exist_ok=True)3 4import streamlit as st5from transformers import ViTFeatureExtractor, ViTForImageClassification6from PIL import Image7import torch8 9st.set_page_config(page_title="Cataract Detection with ViT", layout="wide")10st.title("👁️ Cataract Detection using Vision Transformer (ViT)")11 12uploaded_file = st.file_uploader("Upload an eye image (JPG/PNG)", type=["jpg", "jpeg", "png"])13if uploaded_file:14    image = Image.open(uploaded_file).convert("RGB")15    st.image(image, caption="Uploaded Image", use_column_width=True)16 17    model_name = "Decoder24/Cataract-ViT"18    model = ViTForImageClassification.from_pretrained(model_name)19    extractor = ViTFeatureExtractor.from_pretrained(model_name)20 21    inputs = extractor(images=image, return_tensors="pt")22    outputs = model(**inputs)23    preds = outputs.logits.softmax(dim=-1)24    label = preds.argmax(dim=-1).item()25 26    st.success(f"Predicted class: {model.config.id2label[label]}")27