from fastapi import FastAPI, File, UploadFile
from transformers import ViTForImageClassification, AutoImageProcessor
from transformers import pipeline
import io
from PIL import Image
import torch
import os

# --- Setup ---
app = FastAPI()
MODEL_PATH = "." 

# --- CRITICAL CHANGE: REMOVE THE try/except BLOCK TEMPORARILY ---
# The deployment will crash, but the logs will show the exact reason.

processor = AutoImageProcessor.from_pretrained(MODEL_PATH)
model = ViTForImageClassification.from_pretrained(MODEL_PATH)

device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

# --- Prediction Endpoint ---
@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    """Accepts an image file and returns the probability of a leak."""
    try:
        # 1. Read the uploaded file bytes
        data = await file.read()
        
        # 2. Open the image using PIL
        img = Image.open(io.BytesIO(data)).convert("RGB")
        
        # 3. Preprocess the image (resize, normalize)
        inputs = processor(images=img, return_tensors="pt")
        inputs = {k:v.to(device) for k,v in inputs.items()}
        
        # 4. Run inference
        with torch.no_grad():
            outputs = model(**inputs)
            # Apply softmax to get probabilities
            probs = torch.softmax(outputs.logits, dim=-1).cpu().numpy()[0]
            
            # The 'leak' label has ID 1 (based on your id2label = {0: "no_leak", 1: "leak"})
            prob_leak = float(probs[1]) 
            
        return {
            "prediction": "leak" if prob_leak >= 0.5 else "no_leak",
            "leak_probability": prob_leak
        }
        
    except Exception as e:
        return {"error": str(e), "message": "Prediction failed."}

# --- Root Endpoint (for health check) ---
@app.get("/")
def home():
    return {"status": "Model API is running"}