from fastapi import FastAPI, UploadFile, File
import onnxruntime as ort
import numpy as np
from PIL import Image
import io

app = FastAPI()

MODEL_PATH = "/app/w600k_r50.onnx"

session = ort.InferenceSession(
    MODEL_PATH,
    providers=["CPUExecutionProvider"]
)
input_name = session.get_inputs()[0].name

def preprocess(img):
    img = img.resize((112, 112))
    img = np.array(img).astype("float32")
    img = (img - 127.5) / 128.0
    img = np.transpose(img, (2, 0, 1))
    return np.expand_dims(img, axis=0)

@app.post("/embed")
async def embed_face(file: UploadFile = File(...)):
    img = Image.open(io.BytesIO(await file.read())).convert("RGB")
    inp = preprocess(img)
    emb = session.run(None, {input_name: inp})[0][0]
    return {"embedding": emb.tolist()}
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(
        "app:app",
        host="0.0.0.0",
        port=7860,
        reload=False
    )
det_sess = ort.InferenceSession(
    "/app/det_10g.onnx",
    providers=["CPUExecutionProvider"]
)
det_input = det_sess.get_inputs()[0].name