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]}")