from fastapi import FastAPI, UploadFile, File, HTTPException, Header
from fastapi.responses import JSONResponse
from PIL import Image
import numpy as np
import io
from keras.models import load_model
import os

model_path = os.path.join(os.path.dirname(__file__), "converted_model.keras")
model = load_model(model_path)


app = FastAPI()

CLASS_NAMES = ['Fungi', 'Healthy', 'Nematode', 'Pest', 'Phytopthora', 'Virus']

API_KEY = "mysecretkey"

@app.post("/predict")
async def predict(file: UploadFile = File(...), x_api_key: str = Header(None)):
    if x_api_key != API_KEY:
        raise HTTPException(status_code=401, detail="Invalid or missing API Key")

    try:
        contents = await file.read()
        image = Image.open(io.BytesIO(contents)).convert("RGB")
        image = image.resize((224, 224))
        img_array = np.array(image).astype("float32")
        img_array = np.expand_dims(img_array, axis=0)

        prediction = model.predict(img_array)
        predicted_class = int(np.argmax(prediction[0]))
        predicted_label = CLASS_NAMES[predicted_class]

        return {
            "prediction": predicted_label,
            "probabilities": {
                CLASS_NAMES[i]: float(round(prediction[0][i], 4)) for i in range(6)
            }
        }

    except Exception as e:
        return JSONResponse(status_code=500, content={"error": str(e)})