from fastapi import FastAPI, UploadFile, File, HTTPException, Header
from fastapi.responses import JSONResponse
from PIL import Image
import numpy as np
import tensorflow as tf
import io
import os
app = FastAPI()

# Load your model
import os




import keras

from keras.models import load_model


# Read the model file into memory
with open("2.keras", "rb") as f:
    byte_data = f.read()

# Wrap in a BytesIO object
model_file = io.BytesIO(byte_data)

# Load the model, giving a valid .keras filename
model = tf.keras.models.load_model(("2.keras", model_file))




# Now load the model


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

# Define your API key (keep it secret in prod)
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()

        # Process the image
        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)

        # Predict
        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)})