Update api/main.py

api/main.py

@@ -11,7 +11,20 @@ from pydantic import BaseModel
 from PIL import Image
 
 from src.core import process_inpaint
+from pymongo import MongoClient
 
+# Connect to MongoDB (URI must be in HF env variables)
+MONGO_URI = os.environ.get("MONGO_URI")
+mongo_client = None
+mongo_collection = None
+
+if MONGO_URI:
+    try:
+        mongo_client = MongoClient(MONGO_URI)
+        mongo_db = mongo_client["object-remove-logs"]
+        mongo_collection = mongo_db["api_logs"]
+    except Exception as e:
+        print("Mongo connection failed:", e)
 # Directories (use writable space on HF Spaces)
 BASE_DIR = os.environ.get("DATA_DIR", "/data")
 if not os.path.isdir(BASE_DIR):
@@ -137,7 +150,7 @@ def inpaint(req: InpaintRequest, _: None = Depends(bearer_auth)) -> Dict[str, st
         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_img = Image.open(file_store[req.mask_id]["path"])
     mask_rgba = _load_rgba_mask_from_image(mask_img)
 
     result = process_inpaint(np.array(img_rgba), mask_rgba)
@@ -146,6 +159,15 @@ def inpaint(req: InpaintRequest, _: None = Depends(bearer_auth)) -> Dict[str, st
     Image.fromarray(result).save(result_path)
 
     logs.append({"result": result_name, "timestamp": datetime.utcnow().isoformat()})
+
+    # Mongo logging
+    if mongo_collection:
+        mongo_collection.insert_one({
+            "endpoint": "/inpaint",
+            "result": result_name,
+            "timestamp": datetime.utcnow().isoformat()
+        })
+
     return {"result": result_name}
 
 
@@ -178,7 +200,6 @@ def inpaint_multipart(
     request: Request = None,
     _: None = Depends(bearer_auth),
 ) -> Dict[str, str]:
-    # Load in-memory
     img = Image.open(image.file).convert("RGBA")
     m = Image.open(mask.file)
     mask_rgba = _load_rgba_mask_from_image(m)
@@ -199,6 +220,16 @@ def inpaint_multipart(
     if url:
         entry["url"] = url
     logs.append(entry)
+
+    # Mongo logging
+    if mongo_collection:
+        mongo_collection.insert_one({
+            "endpoint": "/inpaint-multipart",
+            "result": result_name,
+            "url": url,
+            "timestamp": datetime.utcnow().isoformat()
+        })
+
     resp: Dict[str, str] = {"result": result_name}
     if url:
         resp["url"] = url