embed / app.py
rohith2157's picture
apppppi
a816c37 verified
Raw History Blame Contribute Delete
1.07 kB
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