Deploy ViT classifier to Hugging Face

app.py

@@ -0,0 +1,26 @@
+import os
+os.makedirs(os.path.expanduser("~/.streamlit"), exist_ok=True)
+
+import streamlit as st
+from transformers import ViTFeatureExtractor, ViTForImageClassification
+from PIL import Image
+import torch
+
+st.set_page_config(page_title="Cataract Detection with ViT", layout="wide")
+st.title("👁️ Cataract Detection using Vision Transformer (ViT)")
+
+uploaded_file = st.file_uploader("Upload an eye image (JPG/PNG)", type=["jpg", "jpeg", "png"])
+if uploaded_file:
+    image = Image.open(uploaded_file).convert("RGB")
+    st.image(image, caption="Uploaded Image", use_column_width=True)
+
+    model_name = "Decoder24/Cataract-ViT"
+    model = ViTForImageClassification.from_pretrained(model_name)
+    extractor = ViTFeatureExtractor.from_pretrained(model_name)
+
+    inputs = extractor(images=image, return_tensors="pt")
+    outputs = model(**inputs)
+    preds = outputs.logits.softmax(dim=-1)
+    label = preds.argmax(dim=-1).item()
+
+    st.success(f"Predicted class: {model.config.id2label[label]}")

requirements.txt

@@ -1,3 +1,13 @@
-altair
-pandas
-streamlit
+torch
+torchvision
+torchaudio
+timm
+scikit-learn
+opencv-python
+matplotlib
+seaborn
+albumentations
+wandb
+streamlit
+transformers
+pillow