Update api/main.py

api/main.py

@@ -13,6 +13,9 @@ from PIL import Image
 import cv2
 import logging
 
+from pymongo import MongoClient
+import time
+
 logging.basicConfig(level=logging.INFO)
 log = logging.getLogger("api")
 
@@ -39,6 +42,10 @@ app = FastAPI(title="Photo Object Removal API", version="1.0.0")
 file_store: Dict[str, Dict[str, str]] = {}
 logs: List[Dict[str, str]] = []
 
+MONGO_URI = "mongodb+srv://harilogicgo_db_user:pdnh6UCMsWvuTCoi@kiddoimages.k2a4nuv.mongodb.net/?appName=KiddoImages"
+mongo_client = MongoClient(MONGO_URI)
+mongo_db = mongo_client["object_remover"]
+mongo_logs = mongo_db["api_logs"]
 
 def bearer_auth(authorization: Optional[str] = Header(default=None)) -> None:
     if not ENV_TOKEN:
@@ -182,32 +189,94 @@ def _load_rgba_mask_from_image(img: Image.Image) -> np.ndarray:
     log.info(f"Loaded RGBA mask (alpha-based): {int((mask_bw > 0).sum())} white pixels (to remove)")
     return rgba
 
-
 @app.post("/inpaint")
 def inpaint(req: InpaintRequest, _: None = Depends(bearer_auth)) -> Dict[str, str]:
-    if req.image_id not in file_store or file_store[req.image_id]["type"] != "image":
-        raise HTTPException(status_code=404, detail="image_id not found")
-    if req.mask_id not in file_store or file_store[req.mask_id]["type"] != "mask":
-        raise HTTPException(status_code=404, detail="mask_id not found")
+    start_time = time.time()
+    status = "success"
+    error_msg = None
+    output_name = None
 
-    img_rgba = _load_rgba_image(file_store[req.image_id]["path"])
-    mask_img = Image.open(file_store[req.mask_id]["path"])  # may be RGB/gray/RGBA
-    mask_rgba = _load_rgba_mask_from_image(mask_img)
+    try:
+        if req.image_id not in file_store or file_store[req.image_id]["type"] != "image":
+            raise HTTPException(status_code=404, detail="image_id not found")
 
-    # Debug: check mask before processing
-    white_pixels = int((mask_rgba[:,:,0] > 128).sum())
-    log.info(f"Inpaint request: mask has {white_pixels} white pixels, invert_mask={req.invert_mask}")
+        if req.mask_id not in file_store or file_store[req.mask_id]["type"] != "mask":
+            raise HTTPException(status_code=404, detail="mask_id not found")
 
-    if req.passthrough:
-        result = np.array(img_rgba.convert("RGB"))
-    else:
-        result = process_inpaint(np.array(img_rgba), mask_rgba, invert_mask=req.invert_mask)
-    result_name = f"output_{uuid.uuid4().hex}.png"
-    result_path = os.path.join(OUTPUT_DIR, result_name)
-    Image.fromarray(result).save(result_path, "PNG", optimize=False, compress_level=1)
+        img_rgba = _load_rgba_image(file_store[req.image_id]["path"])
+        mask_img = Image.open(file_store[req.mask_id]["path"])
+        mask_rgba = _load_rgba_mask_from_image(mask_img)
 
-    logs.append({"result": result_name, "timestamp": datetime.utcnow().isoformat()})
-    return {"result": result_name}
+        if req.passthrough:
+            result = np.array(img_rgba.convert("RGB"))
+        else:
+            result = process_inpaint(
+                np.array(img_rgba), 
+                mask_rgba, 
+                invert_mask=req.invert_mask
+            )
+
+        output_name = f"output_{uuid.uuid4().hex}.png"
+        output_path = os.path.join(OUTPUT_DIR, output_name)
+
+        Image.fromarray(result).save(
+            output_path, "PNG", optimize=False, compress_level=1
+        )
+
+        return {"result": output_name}
+
+    except Exception as e:
+        status = "fail"
+        error_msg = str(e)
+        raise
+
+    finally:
+        end_time = time.time()
+        response_time_ms = (end_time - start_time) * 1000
+
+        log_doc = {
+            "input_image_id": req.image_id,
+            "input_mask_id": req.mask_id,
+            "output_id": output_name,
+            "status": status,
+            "timestamp": datetime.utcnow(),
+            "ts": int(time.time()),
+            "response_time_ms": response_time_ms
+        }
+
+        if error_msg:
+            log_doc["error"] = error_msg
+
+        try:
+            mongo_logs.insert_one(log_doc)
+        except Exception as mongo_err:
+            log.error(f"Mongo log insert failed: {mongo_err}")
+
+# @app.post("/inpaint")
+# def inpaint(req: InpaintRequest, _: None = Depends(bearer_auth)) -> Dict[str, str]:
+#     if req.image_id not in file_store or file_store[req.image_id]["type"] != "image":
+#         raise HTTPException(status_code=404, detail="image_id not found")
+#     if req.mask_id not in file_store or file_store[req.mask_id]["type"] != "mask":
+#         raise HTTPException(status_code=404, detail="mask_id not found")
+
+#     img_rgba = _load_rgba_image(file_store[req.image_id]["path"])
+#     mask_img = Image.open(file_store[req.mask_id]["path"])  # may be RGB/gray/RGBA
+#     mask_rgba = _load_rgba_mask_from_image(mask_img)
+
+#     # Debug: check mask before processing
+#     white_pixels = int((mask_rgba[:,:,0] > 128).sum())
+#     log.info(f"Inpaint request: mask has {white_pixels} white pixels, invert_mask={req.invert_mask}")
+
+#     if req.passthrough:
+#         result = np.array(img_rgba.convert("RGB"))
+#     else:
+#         result = process_inpaint(np.array(img_rgba), mask_rgba, invert_mask=req.invert_mask)
+#     result_name = f"output_{uuid.uuid4().hex}.png"
+#     result_path = os.path.join(OUTPUT_DIR, result_name)
+#     Image.fromarray(result).save(result_path, "PNG", optimize=False, compress_level=1)
+
+#     logs.append({"result": result_name, "timestamp": datetime.utcnow().isoformat()})
+#     return {"result": result_name}
 
 
 @app.post("/inpaint-url")