from fastapi import FastAPI, File, UploadFile
import numpy as np
import cv2
import pickle
from tensorflow.keras.models import load_model

app = FastAPI()

model = load_model("models/ecg_model.keras")
with open("models/ecg_class_indices.pkl", "rb") as f:
    class_to_index = pickle.load(f)
index_to_class = {v: k for k, v in class_to_index.items()}

IMAGE_SIZE = 224

@app.get("/")
def health():
    return {"status": "ok"}

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    contents = await file.read()
    nparr = np.frombuffer(contents, np.uint8)
    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
    img = cv2.resize(img, (IMAGE_SIZE, IMAGE_SIZE))
    img = img / 255.0
    img = np.expand_dims(img, axis=0)

    preds = model.predict(img)
    pred_idx = int(np.argmax(preds, axis=1)[0])
    confidence = float(np.max(preds))

    return {
        "class": index_to_class[pred_idx],
        "confidence": confidence
    }