WSL + API support

PythonModelAPI/PythonModelAPI.py

@@ -0,0 +1,107 @@
+
+# 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)
+
+
+