Production stage 1

model/api.py

@@ -1,7 +1,7 @@
 from slowapi import Limiter
 from slowapi.util import get_remote_address
 from slowapi.errors import RateLimitExceeded
-from slowapi.responses import JSONResponse
+from fastapi.responses import JSONResponse
 from starlette.requests import Request
 from fastapi import FastAPI, UploadFile, File
 from fastapi.middleware.cors import CORSMiddleware
@@ -13,6 +13,12 @@ from model.inference import predict
 
 app = FastAPI()
 
+CURRENT_STATE = {
+    "emotion": None,
+    "confidence": 0.0
+}
+
+
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -21,8 +27,6 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-
-
 @app.get("/")
 def root():
     return {"status": "API running"}
@@ -46,5 +50,50 @@ def rate_limit_handler(request: Request, exc: RateLimitExceeded):
 
 @app.post("/api/predict")
 @limiter.limit("5/minute")
-async def predict_emotion(file: UploadFile = File(...)):
-    ...
+async def predict_emotion(request: Request, file: UploadFile = File(...)):
+    contents = await file.read()
+
+    try:
+        image = Image.open(io.BytesIO(contents)).convert("RGB")
+    except Exception:
+        return {
+            "state": "error",
+            "reason": "invalid_image"
+        }
+
+    result = predict(image)
+
+    if result["confidence"] >= 0.6:
+        RECENT_PREDICTIONS.append(result["emotion"])
+        if len(RECENT_PREDICTIONS) > WINDOW_SIZE:
+            RECENT_PREDICTIONS.pop(0)
+
+
+    if result["confidence"] < 0.6:
+        return {
+            "state": "uncertain",
+            "emotion": CURRENT_STATE["emotion"],
+            "confidence": result["confidence"],
+            "is_confident": False
+        }
+
+    # update memory
+    if RECENT_PREDICTIONS:
+        dominant_emotion = max(
+            set(RECENT_PREDICTIONS),
+            key=RECENT_PREDICTIONS.count
+        )
+
+        CURRENT_STATE["emotion"] = dominant_emotion
+        CURRENT_STATE["confidence"] = result["confidence"]
+
+    RECENT_PREDICTIONS = []
+    WINDOW_SIZE = 5
+
+
+    return {
+        "state": "stable",
+        "emotion": CURRENT_STATE["emotion"],
+        "confidence": CURRENT_STATE["confidence"],
+        "is_confident": True
+    }

model/inference.py

@@ -34,9 +34,16 @@ transform = transforms.Compose([
     )
 ])
 
+
 @torch.no_grad()
 def predict(pil_image: Image.Image):
     x = transform(pil_image).unsqueeze(0).to(device)
     logits = model(x)
-    idx = logits.argmax(dim=1).item()
-    return classes[idx]
+
+    probs = torch.softmax(logits, dim=1)
+    conf, idx = probs.max(dim=1)
+
+    return {
+        "emotion": classes[idx.item()],
+        "confidence": conf.item()
+    }