108 lines
3.2 KiB
Python
108 lines
3.2 KiB
Python
|
|
# app.py
|
|
from fastapi import FastAPI, UploadFile, File, Body
|
|
from fastapi.responses import JSONResponse
|
|
import numpy as np
|
|
from PIL import Image
|
|
import io
|
|
import tensorflow as tf
|
|
|
|
app = FastAPI()
|
|
|
|
class Service:
|
|
def __init__(self):
|
|
self.model = None
|
|
self.models={}
|
|
self.active_model = "default"
|
|
|
|
def activate_model(self, name):
|
|
|
|
self.model = self.models[name]
|
|
self.active_model = name
|
|
return 0
|
|
|
|
def load_model(self, path, name):
|
|
self.models[name] = tf.keras.models.load_model(path)
|
|
return 0
|
|
|
|
def get_accept_size(self):
|
|
print(self.model.layers[0].input)
|
|
|
|
return list(self.model.input_shape)
|
|
|
|
def get_output_size(self):
|
|
return list(self.model.layers[-1].output[0].shape[1:4])
|
|
|
|
|
|
def predict(self, data: bytes):
|
|
target_size = self.get_accept_size()
|
|
image = np.reshape(
|
|
np.frombuffer(data, dtype=np.uint8),
|
|
(target_size[1], target_size[2], target_size[3])
|
|
) / 255.0
|
|
|
|
predictions = self.model.predict(np.array([image]))
|
|
|
|
if predictions[0].ndim == 2:
|
|
classed = predictions[0]
|
|
else:
|
|
classed = tf.argmax(predictions[0], axis=2)
|
|
|
|
return JSONResponse({"result": classed.numpy().reshape(-1).tolist()})
|
|
|
|
def predict_raw(self, data):
|
|
target_size = self.get_accept_size()
|
|
image = np.reshape(np.frombuffer(data, dtype=np.uint8), (target_size[1], target_size[2], target_size[3])) / 255.
|
|
|
|
predictions = self.model(np.array([image])).numpy()
|
|
|
|
if predictions[0].ndim==2:
|
|
classed = predictions[0][:,:,np.newaxis].clip(0,255)
|
|
return (classed).astype(np.uint8).reshape(-1).tolist()
|
|
else:
|
|
classed = predictions[0][:,:,:]
|
|
|
|
return (classed*255).astype(np.uint8).reshape(-1).tolist()
|
|
|
|
svc = Service()
|
|
|
|
def _prep_raw_bytes_for_predict(buf: bytes) -> bytes:
|
|
"""Decode JPEG/PNG -> RGB, resize to model size, return raw uint8 bytes HxWxC."""
|
|
h, w, c = svc.get_accept_size()[1:4]
|
|
img = Image.open(io.BytesIO(buf)).convert("RGB")
|
|
if img.size != (w, h):
|
|
img = img.resize((w, h), Image.BILINEAR)
|
|
arr = np.asarray(img, dtype=np.uint8) # HxWx3 uint8
|
|
return arr.tobytes()
|
|
|
|
@app.post("/predict")
|
|
async def predict(file: UploadFile = File(...)):
|
|
buf = await file.read()
|
|
raw = _prep_raw_bytes_for_predict(buf)
|
|
return svc.predict(raw)
|
|
|
|
@app.post("/predict_raw")
|
|
async def predict_raw(file: UploadFile = File(...)):
|
|
buf = await file.read()
|
|
raw = _prep_raw_bytes_for_predict(buf)
|
|
return svc.predict_raw(raw)
|
|
|
|
@app.post("/load_model")
|
|
async def load_model(path: str = Body(...), name: str = Body(...)):
|
|
try:
|
|
svc.load_model(path, name)
|
|
return JSONResponse({"status": "success", "message": f"Model {name} loaded successfully."})
|
|
except Exception as e:
|
|
return JSONResponse({"status": "error", "message": str(e)}, status_code=500)
|
|
|
|
@app.post("/activate_model")
|
|
async def activate_model(name: str = Body(...)):
|
|
try:
|
|
svc.activate_model(name)
|
|
return JSONResponse({"status": "success", "message": f"Model {name} activated successfully."})
|
|
except Exception as e:
|
|
return JSONResponse({"status": "error", "message": str(e)}, status_code=500)
|
|
|
|
|
|
|