Update app.py

app.py

@@ -50,7 +50,11 @@ DO_SPACES_ENDPOINT = os.getenv("DO_SPACES_ENDPOINT", f"https://{DO_SPACES_REGION
 DO_SPACES_KEY = os.getenv("DO_SPACES_KEY")  # Your Access Key
 DO_SPACES_SECRET = os.getenv("DO_SPACES_SECRET")  # Your Secret Key
 DO_SPACES_BUCKET = os.getenv("DO_SPACES_BUCKET")  # Bucket Name
-
+ADMIN_MONGO_URL = os.getenv("ADMIN_MONGO_URL") #MONGODB Storage for usage count
+admin_client = AsyncIOMotorClient(ADMIN_MONGO_URL)
+admin_db = admin_client.adminPanel
+subcategories_col = admin_db.subcategories
+media_clicks_col = admin_db.media_clicks
 # --------------------- Download Models ---------------------
 def download_models():
     logger.info("Downloading models from private HF repo...")
@@ -239,134 +243,179 @@ def root():
 @fastapi_app.get("/health")
 async def health():
     return {"status": "healthy"}
-# @fastapi_app.get("/health")
-# async def health():
-#     mongo_ok = False
-#     spaces_ok = False
-#     mongo_error = None
-#     spaces_error = None
-
-#     # ---- MongoDB Check (ASYNC SAFE) ----
-#     try:
-#         if database is None:
-#             raise Exception("MongoDB client not initialized")
-
-#         await database.command("ping")
-#         mongo_ok = True
-
-#     except Exception as e:
-#         mongo_error = str(e)
-
-#     # ---- DigitalOcean Spaces Check (RUN IN THREADPOOL) ----
-#     try:
-#         def check_spaces():
-#             client = get_spaces_client()
-#             client.list_objects_v2(
-#                 Bucket=DO_SPACES_BUCKET,
-#                 MaxKeys=1
-#             )
-
-#         await run_in_threadpool(check_spaces)
-#         spaces_ok = True
-
-#     except Exception as e:
-#         spaces_error = str(e)
-
-#     # ---- FINAL RESPONSE ----
-#     if mongo_ok and spaces_ok:
-#         return {
-#             "status": "healthy",
-#             "mongo": "connected",
-#             "spaces": "connected",
-#             "timestamp": datetime.utcnow().isoformat()
-#         }
-
-#     raise HTTPException(
-#         status_code=503,
-#         detail={
-#             "status": "unhealthy",
-#             "mongo": "connected" if mongo_ok else "failed",
-#             "spaces": "connected" if spaces_ok else "failed",
-#             "mongo_error": mongo_error,
-#             "spaces_error": spaces_error,
-#             "timestamp": datetime.utcnow().isoformat()
-#         }
-#     )
 
 from fastapi import Form
 import requests
-
 @fastapi_app.post("/face-swap", dependencies=[Depends(verify_token)])
 async def face_swap_api(
     source: UploadFile = File(...),
-    target_category_id: str = Form(...),
+    target_category_id: str = Form(None),
+    new_category_id: str = Form(None),
+    user_id: Optional[str] = Form(None),
     credentials: HTTPAuthorizationCredentials = Security(security)
 ):
-    start_time = datetime.utcnow()  # start timer
+    start_time = datetime.utcnow()
+
     try:
-        # Read source image
+        # ---------- VALIDATION ----------
+        if target_category_id and new_category_id:
+            raise HTTPException(
+                status_code=400,
+                detail="Provide only one of new_category_id or target_category_id, not both."
+            )
+
+        if not target_category_id and not new_category_id:
+            raise HTTPException(
+                status_code=400,
+                detail="Either new_category_id or target_category_id is required."
+            )
+
+        # ---------- READ SOURCE ----------
         src_bytes = await source.read()
 
         # Save source to Spaces
         src_key = f"faceswap/source/{uuid.uuid4().hex}_{source.filename}"
         upload_to_spaces(src_bytes, src_key, content_type=source.content_type)
 
-        # Find target image URL from categoryId
-        client = get_spaces_client()
-        base_prefix = "faceswap/target/"
-        resp = client.list_objects_v2(Bucket=DO_SPACES_BUCKET, Prefix=base_prefix, Delimiter="/")
-        categories = [prefix["Prefix"].split("/")[2] for prefix in resp.get("CommonPrefixes", [])]
-        target_url = None
-        for category in categories:
-            original_prefix = f"faceswap/target/{category}/original/"
-            thumb_prefix = f"faceswap/target/{category}/thumb/"
-            original_objects = client.list_objects_v2(Bucket=DO_SPACES_BUCKET, Prefix=original_prefix)
-            thumb_objects = client.list_objects_v2(Bucket=DO_SPACES_BUCKET, Prefix=thumb_prefix)
-            original_files = [obj["Key"].split("/")[-1] for obj in original_objects.get("Contents", []) if obj["Key"].endswith(".png")]
-            thumb_files = [obj["Key"].split("/")[-1] for obj in thumb_objects.get("Contents", []) if obj["Key"].endswith(".png")]
-            for idx, filename in enumerate(sorted(original_files), start=1):
-                cid = f"{category.lower()}image_{idx}"
-                if filename in thumb_files and cid == target_category_id:
-                    target_url = f"https://{DO_SPACES_BUCKET}.blr1.digitaloceanspaces.com/{original_prefix}{filename}"
+        # ==========================================================
+        # CASE 1: NEW_CATEGORY_ID → Fetch from MongoDB
+        # ==========================================================
+        if new_category_id:
+            from bson import ObjectId
+
+            doc = await subcategories_col.find_one({
+                "asset_images._id": ObjectId(new_category_id)
+            })
+
+            if not doc:
+                raise HTTPException(status_code=404, detail="Asset image not found in database")
+
+            # Find the asset image inside array
+            asset = None
+            for img in doc["asset_images"]:
+                if str(img["_id"]) == str(new_category_id):
+                    asset = img
+                    break
+
+            if not asset:
+                raise HTTPException(status_code=404, detail="Asset image URL not found")
+
+            target_url = asset["url"]
+            category_id = str(doc["categoryId"])
+
+            # ---------- media_clicks update only if user_id exists ----------
+            if user_id:
+                user_oid = ObjectId(user_id)
+                cat_oid = ObjectId(category_id)
+
+                now = datetime.utcnow()
+
+                # Try to update existing category entry
+                update_result = await media_clicks_col.update_one(
+                    {
+                        "userId": user_oid,
+                        "categories.categoryId": cat_oid
+                    },
+                    {
+                        "$set": {
+                            "updatedAt": now,
+                            "categories.$.lastClickedAt": now
+                        },
+                        "$inc": {
+                            "categories.$.click_count": 1
+                        }
+                    }
+                )
+
+                # If category did not exist, push new category object
+                if update_result.matched_count == 0:
+                    await media_clicks_col.update_one(
+                        {"userId": user_oid},
+                        {
+                            "$setOnInsert": {"createdAt": now},
+                            "$set": {"updatedAt": now},
+                            "$push": {
+                                "categories": {
+                                    "categoryId": cat_oid,
+                                    "click_count": 1,
+                                    "lastClickedAt": now
+                                }
+                            }
+                        },
+                        upsert=True
+                    )
+
+        # ==========================================================
+        # CASE 2: EXISTING TARGET_CATEGORY_ID (DigitalOcean)
+        # ==========================================================
+        if target_category_id:
+            client = get_spaces_client()
+            base_prefix = "faceswap/target/"
+            resp = client.list_objects_v2(Bucket=DO_SPACES_BUCKET, Prefix=base_prefix, Delimiter="/")
+            categories = [prefix["Prefix"].split("/")[2] for prefix in resp.get("CommonPrefixes", [])]
+
+            target_url = None
+            for category in categories:
+                original_prefix = f"faceswap/target/{category}/original/"
+                thumb_prefix = f"faceswap/target/{category}/thumb/"
+
+                original_objects = client.list_objects_v2(Bucket=DO_SPACES_BUCKET, Prefix=original_prefix)
+                thumb_objects = client.list_objects_v2(Bucket=DO_SPACES_BUCKET, Prefix=thumb_prefix)
+
+                original_files = [obj["Key"].split("/")[-1] for obj in original_objects.get("Contents", []) if obj["Key"].endswith(".png")]
+                thumb_files = [obj["Key"].split("/")[-1] for obj in thumb_objects.get("Contents", []) if obj["Key"].endswith(".png")]
+
+                for idx, filename in enumerate(sorted(original_files), start=1):
+                    cid = f"{category.lower()}image_{idx}"
+                    if filename in thumb_files and cid == target_category_id:
+                        target_url = f"https://{DO_SPACES_BUCKET}.blr1.digitaloceanspaces.com/{original_prefix}{filename}"
+                        break
+                if target_url:
                     break
-            if target_url:
-                break
-        if not target_url:
-            raise HTTPException(status_code=404, detail="Target categoryId not found")
 
-        # Download target image from Spaces
+            if not target_url:
+                raise HTTPException(status_code=404, detail="Target categoryId not found")
+
+        # ==========================================================
+        # DOWNLOAD TARGET IMAGE
+        # ==========================================================
         resp = requests.get(target_url)
         if resp.status_code != 200:
-            raise HTTPException(status_code=404, detail="Target image not found in Spaces")
+            raise HTTPException(status_code=404, detail="Target image could not be fetched")
+
         tgt_bytes = resp.content
 
-        # Decode for processing
+        # Decode images
         src_array = np.frombuffer(src_bytes, np.uint8)
         tgt_array = np.frombuffer(tgt_bytes, np.uint8)
+
         src_bgr = cv2.imdecode(src_array, cv2.IMREAD_COLOR)
         tgt_bgr = cv2.imdecode(tgt_array, cv2.IMREAD_COLOR)
+
         if src_bgr is None or tgt_bgr is None:
             raise HTTPException(status_code=400, detail="Invalid image data")
 
         src_rgb = cv2.cvtColor(src_bgr, cv2.COLOR_BGR2RGB)
         tgt_rgb = cv2.cvtColor(tgt_bgr, cv2.COLOR_BGR2RGB)
 
-        # Run face swap pipeline
+        # ==========================================================
+        # FACE SWAP
+        # ==========================================================
         final_img, final_path, err = face_swap_and_enhance(src_rgb, tgt_rgb)
         if err:
             raise HTTPException(status_code=500, detail=err)
 
-        # Upload result to Spaces
+        # Upload output to Spaces
         with open(final_path, "rb") as f:
             result_bytes = f.read()
+
         result_key = f"faceswap/result/{uuid.uuid4().hex}_enhanced.png"
         result_url = upload_to_spaces(result_bytes, result_key, content_type="image/png")
 
-        # Calculate response time
+        # Log success
         end_time = datetime.utcnow()
         response_time_ms = (end_time - start_time).total_seconds() * 1000
 
-        # Log response time only
         if database is not None:
             await database.api_logs.insert_one({
                 "endpoint": "/face-swap",
@@ -378,7 +427,7 @@ async def face_swap_api(
         return {"result_key": result_key, "result_url": result_url}
 
     except Exception as e:
-        # Log response time even on error
+        # Log failure
         end_time = datetime.utcnow()
         response_time_ms = (end_time - start_time).total_seconds() * 1000
         if database is not None:
@@ -391,6 +440,101 @@ async def face_swap_api(
             })
         raise HTTPException(status_code=500, detail=f"Face swap failed: {str(e)}")
 
+# @fastapi_app.post("/face-swap", dependencies=[Depends(verify_token)])
+# async def face_swap_api(
+#     source: UploadFile = File(...),
+#     target_category_id: str = Form(...),
+#     credentials: HTTPAuthorizationCredentials = Security(security)
+# ):
+#     start_time = datetime.utcnow()  # start timer
+#     try:
+#         # Read source image
+#         src_bytes = await source.read()
+
+#         # Save source to Spaces
+#         src_key = f"faceswap/source/{uuid.uuid4().hex}_{source.filename}"
+#         upload_to_spaces(src_bytes, src_key, content_type=source.content_type)
+
+#         # Find target image URL from categoryId
+#         client = get_spaces_client()
+#         base_prefix = "faceswap/target/"
+#         resp = client.list_objects_v2(Bucket=DO_SPACES_BUCKET, Prefix=base_prefix, Delimiter="/")
+#         categories = [prefix["Prefix"].split("/")[2] for prefix in resp.get("CommonPrefixes", [])]
+#         target_url = None
+#         for category in categories:
+#             original_prefix = f"faceswap/target/{category}/original/"
+#             thumb_prefix = f"faceswap/target/{category}/thumb/"
+#             original_objects = client.list_objects_v2(Bucket=DO_SPACES_BUCKET, Prefix=original_prefix)
+#             thumb_objects = client.list_objects_v2(Bucket=DO_SPACES_BUCKET, Prefix=thumb_prefix)
+#             original_files = [obj["Key"].split("/")[-1] for obj in original_objects.get("Contents", []) if obj["Key"].endswith(".png")]
+#             thumb_files = [obj["Key"].split("/")[-1] for obj in thumb_objects.get("Contents", []) if obj["Key"].endswith(".png")]
+#             for idx, filename in enumerate(sorted(original_files), start=1):
+#                 cid = f"{category.lower()}image_{idx}"
+#                 if filename in thumb_files and cid == target_category_id:
+#                     target_url = f"https://{DO_SPACES_BUCKET}.blr1.digitaloceanspaces.com/{original_prefix}{filename}"
+#                     break
+#             if target_url:
+#                 break
+#         if not target_url:
+#             raise HTTPException(status_code=404, detail="Target categoryId not found")
+
+#         # Download target image from Spaces
+#         resp = requests.get(target_url)
+#         if resp.status_code != 200:
+#             raise HTTPException(status_code=404, detail="Target image not found in Spaces")
+#         tgt_bytes = resp.content
+
+#         # Decode for processing
+#         src_array = np.frombuffer(src_bytes, np.uint8)
+#         tgt_array = np.frombuffer(tgt_bytes, np.uint8)
+#         src_bgr = cv2.imdecode(src_array, cv2.IMREAD_COLOR)
+#         tgt_bgr = cv2.imdecode(tgt_array, cv2.IMREAD_COLOR)
+#         if src_bgr is None or tgt_bgr is None:
+#             raise HTTPException(status_code=400, detail="Invalid image data")
+
+#         src_rgb = cv2.cvtColor(src_bgr, cv2.COLOR_BGR2RGB)
+#         tgt_rgb = cv2.cvtColor(tgt_bgr, cv2.COLOR_BGR2RGB)
+
+#         # Run face swap pipeline
+#         final_img, final_path, err = face_swap_and_enhance(src_rgb, tgt_rgb)
+#         if err:
+#             raise HTTPException(status_code=500, detail=err)
+
+#         # Upload result to Spaces
+#         with open(final_path, "rb") as f:
+#             result_bytes = f.read()
+#         result_key = f"faceswap/result/{uuid.uuid4().hex}_enhanced.png"
+#         result_url = upload_to_spaces(result_bytes, result_key, content_type="image/png")
+
+#         # Calculate response time
+#         end_time = datetime.utcnow()
+#         response_time_ms = (end_time - start_time).total_seconds() * 1000
+
+#         # Log response time only
+#         if database is not None:
+#             await database.api_logs.insert_one({
+#                 "endpoint": "/face-swap",
+#                 "status": "success",
+#                 "response_time_ms": response_time_ms,
+#                 "timestamp": end_time
+#             })
+
+#         return {"result_key": result_key, "result_url": result_url}
+
+#     except Exception as e:
+#         # Log response time even on error
+#         end_time = datetime.utcnow()
+#         response_time_ms = (end_time - start_time).total_seconds() * 1000
+#         if database is not None:
+#             await database.api_logs.insert_one({
+#                 "endpoint": "/face-swap",
+#                 "status": "fail",
+#                 "response_time_ms": response_time_ms,
+#                 "timestamp": end_time,
+#                 "error": str(e)
+#             })
+#         raise HTTPException(status_code=500, detail=f"Face swap failed: {str(e)}")
+
 
 
 @fastapi_app.get("/preview/{result_key:path}")