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

app = FastAPI(title="Animal Classifier")

# 1) Laad model lokaal vanuit model/
MODEL_PATH = os.path.join(os.path.dirname(__file__), "model")
model = tf.keras.models.load_model(MODEL_PATH)
CLASS_NAMES = ["cat", "dog", "panda"]

# 2) Predict functie
def predict_image(image: Image.Image):
    image = image.convert("RGB")
    resized_image = tf.image.resize(np.array(image), (64, 64))
    images_to_predict = np.expand_dims(np.array(resized_image), axis=0)
    probs = model.predict(images_to_predict)[0]
    return {c: float(p) for c, p in zip(CLASS_NAMES, probs)}

# 3) FastAPI endpoint
@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    contents = await file.read()
    image = Image.open(io.BytesIO(contents))
    result = predict_image(image)
    return JSONResponse(result)