import tensorflow as tf
import numpy as np
import json
from fastapi import FastAPI, File, UploadFile
from PIL import Image
import io

# ✅ Load .keras model directly
model = tf.keras.models.load_model(
    "leaf_disease_efficientnetB0_fixed.keras"
)

# Load class names
with open("class_names.json") as f:
    class_names = json.load(f)

IMG_SIZE = 224

app = FastAPI()


def preprocess_image(image):
    image = image.resize((IMG_SIZE, IMG_SIZE))

    image = np.array(image).astype(np.float32)

    image = tf.keras.applications.efficientnet.preprocess_input(image)

    image = np.expand_dims(image, axis=0)

    return image


@app.get("/")
def home():
    return {"message": "Leaf Disease Model API is running"}


@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    contents = await file.read()

    image = Image.open(io.BytesIO(contents)).convert("RGB")

    image = preprocess_image(image)

    predictions = model.predict(image)

    predicted_index = int(np.argmax(predictions))
    confidence = float(np.max(predictions))

    predicted_class = class_names[predicted_index]

    return {
        "class": predicted_class,
        "confidence": confidence
    }