from fastapi import FastAPI, UploadFile, File
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
import tensorflow as tf
import numpy as np
from PIL import Image
from io import BytesIO

app = FastAPI()

app.mount("/", StaticFiles(directory="frontend", html=True), name="static")

# Load TFLite model
interpreter = tf.lite.Interpreter(model_path="model.tflite")
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

CLASS_NAMES = ["cat", "dog", "..."]  # Your classes here

def preprocess(data):
    image = Image.open(BytesIO(data)).convert("RGB")
    image = image.resize((256, 256))
    img = np.array(image) / 255.0
    return np.expand_dims(img, axis=0).astype(np.float32)

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    image = preprocess(await file.read())
    interpreter.set_tensor(input_details[0]['index'], image)
    interpreter.invoke()
    output = interpreter.get_tensor(output_details[0]['index'])
    pred = CLASS_NAMES[np.argmax(output[0])]
    conf = float(np.max(output[0]))
    return {"class": pred, "confidence": conf}