from transformers import AutoImageProcessor, AutoModelForImageClassification
from fastapi import FastAPI, UploadFile, File
from PIL import Image
import torch
import io

processor = AutoImageProcessor.from_pretrained("NeuronZero/EyeDiseaseClassifier")
model = AutoModelForImageClassification.from_pretrained("NeuronZero/EyeDiseaseClassifier")

app = FastAPI()

@app.get("/")
def home():
    return {"message": "Habeeb"}

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    # Read file into memory
    contents = await file.read()
    
    # Load as PIL image
    image = Image.open(io.BytesIO(contents)).convert("RGB")
    
    # Preprocess
    inputs = processor(images=image, return_tensors="pt")
    
    # Run inference
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        predicted_class_idx = logits.argmax(-1).item()
        confidence = torch.softmax(logits, dim=-1)[0][predicted_class_idx].item()
    
    result = model.config.id2label[predicted_class_idx]
    
    return {
        "filename": file.filename,
        "classification": result,
        "confidence": round(confidence, 4)
    }