Update app.py

app.py

@@ -1,8 +1,9 @@
 import io
 import os
 import threading
+import time
 from fastapi import FastAPI, File, UploadFile
-from fastapi.responses import StreamingResponse
+from fastapi.responses import StreamingResponse, JSONResponse
 from PIL import Image
 import torch
 import torchvision.transforms as transforms
@@ -30,6 +31,13 @@ print(f"ONNX model loaded with providers: {providers}")
 # Lock for thread-safe ONNX inference
 onnx_lock = threading.Lock()
 
+# Stats
+stats = {
+    "completed": 0,
+    "total_time": 0.0,
+    "avg_time": 0.0,
+}
+
 # Helper functions
 def run_model_onnx(input_tensor: torch.Tensor) -> torch.Tensor:
     with onnx_lock:  # ensure thread safety
@@ -56,11 +64,25 @@ app = FastAPI(title="Background Removal API")
 
 @app.post("/remove-background")
 async def remove_background(file: UploadFile = File(...)):
+    start_time = time.time()
+
     image = Image.open(file.file).convert("RGB")
     result_image = process_image(image)
 
+    elapsed = time.time() - start_time
+    stats["completed"] += 1
+    stats["total_time"] += elapsed
+    stats["avg_time"] = stats["total_time"] / stats["completed"]
+
     buf = io.BytesIO()
     result_image.save(buf, format="PNG")
     buf.seek(0)
 
     return StreamingResponse(buf, media_type="image/png")
+
+@app.get("/stats")
+async def get_stats():
+    return JSONResponse({
+        "completed": stats["completed"],
+        "avg_time": round(stats["avg_time"], 3),
+    })