from fastapi import FastAPI, File, UploadFile
import numpy as np
import tensorflow as tf
from tensorflow import keras
from PIL import Image
import os
from huggingface_hub import hf_hub_download

app = FastAPI(title="🐾 Animal Classifier API")

# Download the model (same as your Gradio version)
repo_id = "Juli-Kath/animal-classification-azure"
os.makedirs("./model/unpacked_keras/variables", exist_ok=True)
hf_hub_download(repo_id, filename="unpacked_keras/saved_model.pb", repo_type="model", local_dir="./model")
hf_hub_download(repo_id, filename="unpacked_keras/variables/variables.index", repo_type="model", local_dir="./model")
hf_hub_download(repo_id, filename="unpacked_keras/variables/variables.data-00000-of-00001", repo_type="model", local_dir="./model")

# Load TensorFlow SavedModel
model_layer = keras.layers.TFSMLayer("./model/unpacked_keras", call_endpoint="serving_default")
inp = tf.keras.Input(shape=(64, 64, 3))
out = model_layer(inp)
model = tf.keras.Model(inp, out)

CLASSES = ["cat", "dog", "panda"]

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    image = Image.open(file.file).resize((64, 64))
    img = np.expand_dims(np.array(image) / 255.0, axis=0)
    outputs = model(img)
    preds = outputs["output_0"].numpy().flatten()
    return {c: float(p) for c, p in zip(CLASSES, preds)}


@app.get("/")
def home():
    return {"message": "🐾 Animal Classifier API is running! Go to /docs to test."}