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

app = FastAPI()

# Load the TensorFlow model
model = tf.keras.models.load_model('./plant_disease_detection.h5')

# Load categories
with open('./categories.json') as f:
    categories = json.load(f)

def preprocess_image(image_bytes):
    # Convert the image to a NumPy array
    image = Image.open(io.BytesIO(image_bytes))
    image = image.resize((224, 224))  # Adjust size as needed
    image_array = np.array(image) / 255.0  # Normalize to [0, 1]
    image_array = np.expand_dims(image_array, axis=0)  # Add batch dimension
    return image_array

@app.post('/predict')
async def predict(file: UploadFile = File(...)):
    if file.content_type.startswith('image/') is False:
        raise HTTPException(status_code=400, detail='Invalid file type')

    image_bytes = await file.read()
    image_array = preprocess_image(image_bytes)

    # Make prediction
    predictions = model.predict(image_array)
    predicted_class = np.argmax(predictions, axis=1)[0]

    # Map to category names
    predicted_label = categories.get(str(predicted_class), 'Unknown')

    return JSONResponse(content={
        'class': predicted_label,
        'confidence': float(predictions[0][predicted_class])
    })


if __name__ == '__main__':
    import uvicorn
    uvicorn.run(app, host='0.0.0.0', port=8080)