app.py

from fastapi import FastAPI, File, UploadFile
from fastapi.middleware.cors import CORSMiddleware
import tensorflow as tf
from keras.layers import TFSMLayer
import numpy as np
from PIL import Image
import io

# 🔥 preprocess khusus DenseNet
from tensorflow.keras.applications.densenet import preprocess_input

app = FastAPI()

# 🔥 CORS (biar bisa diakses dari frontend)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 🔥 label mapping (WAJIB sesuai training)
labels = ['butterfly', 'cow', 'dog', 'elephant', 'goat', 'hen', 'horse', 'spyder']

# 🔥 load model (Keras 3 fix)
model = None

def get_model():
    global model
    if model is None:
        print("🔥 Loading model...")
        model = TFSMLayer("saved_model", call_endpoint="serving_default")
    return model
    
def preprocess(image):
    image = image.resize((224, 224))
    image = np.array(image).astype("float32")

    # 🔥 preprocessing ImageNet (DenseNet)
    image = preprocess_input(image)

    return np.expand_dims(image, axis=0)

@app.get("/")
def root():
    return {"message": "Model is running!"}

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    try:
        image = Image.open(io.BytesIO(await file.read())).convert("RGB")
        input_data = preprocess(image)

        # 🔥 inference
        pred = get_model()(input_data)

        # handle dict output
        if isinstance(pred, dict):
            pred = list(pred.values())[0]

        pred = pred.numpy()

        # 🔥 ambil hasil utama
        class_idx = int(np.argmax(pred))
        confidence = float(np.max(pred))
        class_name = labels[class_idx]

        # 🔥 top-3 prediction
        top_indices = pred[0].argsort()[-3:][::-1]
        top_predictions = [
            {
                "class": labels[i],
                "confidence": float(pred[0][i])
            }
            for i in top_indices
        ]

        return {
            "class": class_name,
            "confidence": confidence,
            "top_predictions": top_predictions
        }

    except Exception as e:
        return {"error": str(e)}