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

app = FastAPI()

# Load the model once when server starts
model = tf.keras.models.load_model("orange_disease_model.h5")

# Define your classes (Make sure these match your labels.txt!)
CLASS_NAMES = [
    'Citrus canker', 'Citrus greening', 'Citrus mealybugs', 'Die back', 
    'Foliage damaged', 'Healthy leaf', 'Powdery mildew', 'Shot hole', 
    'Spiny whitefly', 'Yellow dragon', 'Yellow leaves'
]

@app.get("/")
def home():
    return {"message": "Orange Disease Detection API is Running!"}

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    # 1. Read the image uploaded by the user
    image_data = await file.read()
    image = Image.open(io.BytesIO(image_data))
    
    # 2. Preprocess (Resize to 224x224 and Normalize)
    image = image.resize((224, 224))
    img_array = tf.keras.preprocessing.image.img_to_array(image)
    img_array = tf.expand_dims(img_array, 0) # Create a batch
    img_array = img_array / 255.0

    # 3. Predict
    predictions = model.predict(img_array)
    predicted_class = CLASS_NAMES[np.argmax(predictions[0])]
    confidence = float(np.max(predictions[0]))

    # 4. Return JSON
    return {
        "class": predicted_class,
        "confidence": confidence
    }