from fastapi import FastAPI
from pydantic import BaseModel
import tensorflow as tf
import base64
import numpy as np

app = FastAPI(title="TF Serving Base64 API")

class ImagePayload(BaseModel):
    image_base64: str

@app.get("/")
def home():
    return {"status": "ok", "message": "API is running! Use POST /predict"}

@app.post("/predict")
def predict(payload: ImagePayload):
    img_bytes = base64.b64decode(payload.image_base64)
    img = tf.image.decode_image(img_bytes, channels=3)
    img = tf.image.resize(img, (224, 224))
    img = img / 255.0
    img = tf.expand_dims(img, 0)

    model = tf.keras.models.load_model("model.keras")

    pred = model.predict(img).tolist()
    return {"prediction": pred}