Update app.py

app.py

@@ -31,6 +31,7 @@ from typing import Optional
 import requests
 import json
 from bson import ObjectId  
+
 # --------------------- Logging ---------------------
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -148,12 +149,16 @@ client = None
 database = None
 
 # --------------------- Admin Panel DB (categories + media_clicks) ---------------------
+# --------------------- Admin Panel DB (categories + subcategories + media_clicks) ---------------------
 ADMIN_MONGO_URL = os.getenv("ADMIN_MONGO_URL")
 admin_client = AsyncIOMotorClient(ADMIN_MONGO_URL)
 admin_db = admin_client.adminPanel
 
+# Collections
 categories_col = admin_db.categories
+subcategories_col = admin_db.subcategories
 media_clicks_col = admin_db.media_clicks
+users_col = admin_db.users  # optional, only if needed
 
 # --------------------- FastAPI ---------------------
 fastapi_app = FastAPI()
@@ -289,185 +294,217 @@ def root():
 async def health():
     return {"status": "healthy"}
 
+@fastapi_app.post("/face-swap")
+async def face_swap_api(
+    source: UploadFile = File(...),
+    target_category_id: str = Form(...),  # REQUIRED (old behavior preserved)
+    category_id: Optional[str] = Form(None),
+    user_id: Optional[str] = Form(None),
+    new_subcategory_id: Optional[str] = Form(None),
+    user_email: str = Depends(verify_firebase_token)
+):
+    start_time = datetime.now(timezone.utc)
+
+    try:
+        # ---------------------------------------------------------
+        # NORMALIZE EMPTY STRINGS (Android older versions)
+        # ---------------------------------------------------------
+        if target_category_id == "":
+            target_category_id = None
+
+        if new_subcategory_id == "":
+            new_subcategory_id = None
+
+        if category_id == "":
+            category_id = None
+
+        if user_id == "":
+            user_id = None
+
+        # ---------------------------------------------------------
+        # STRICT XOR VALIDATION
+        # ---------------------------------------------------------
+        if target_category_id and new_subcategory_id:
+            raise HTTPException(
+                status_code=400,
+                detail="Provide ONLY ONE of: target_category_id OR new_subcategory_id"
+            )
+
+        if not target_category_id and not new_subcategory_id:
+            raise HTTPException(
+                status_code=400,
+                detail="Either target_category_id OR new_subcategory_id is required"
+            )
+
+        # ---------------------------------------------------------
+        # READ SOURCE IMAGE
+        # ---------------------------------------------------------
+        src_bytes = await source.read()
+        src_key = f"bikini-theme/source/{uuid.uuid4().hex}_{source.filename}"
+        upload_to_spaces(src_bytes, src_key, content_type=source.content_type)
+
+        # ---------------------------------------------------------
+        # CASE 1 — Old behavior (use DO Spaces target image)
+        # ---------------------------------------------------------
+        if target_category_id:
+            target_filename = f"{target_category_id}.png"
+            target_url = (
+                f"https://{DO_SPACES_BUCKET}.{DO_SPACES_REGION}."
+                f"digitaloceanspaces.com/bikini-theme/target/{target_filename}"
+            )
+
+            resp = requests.get(target_url)
+            if resp.status_code != 200:
+                await log_faceswap_hit(user_email, "error", start_time, datetime.now(timezone.utc))
+                raise HTTPException(status_code=404, detail=f"Target image not found: {target_url}")
+
+            tgt_bytes = resp.content
+
+        # ---------------------------------------------------------
+        # CASE 2 — New behavior (use subcategory asset image)
+        # ---------------------------------------------------------
+        else:
+            # Find subcategory asset by asset_images._id
+            asset = await admin_db.subcategories.find_one(
+                {"asset_images._id": ObjectId(new_subcategory_id)},
+                {"asset_images.$": 1}
+            )
+
+            if not asset or "asset_images" not in asset:
+                await log_faceswap_hit(user_email, "error", start_time, datetime.now(timezone.utc))
+                raise HTTPException(
+                    status_code=404,
+                    detail="Subcategory asset image not found"
+                )
+
+            # Extract the single matching image URL
+            asset_url = asset["asset_images"][0]["url"]
+
+            resp = requests.get(asset_url)
+            if resp.status_code != 200:
+                raise HTTPException(
+                    status_code=404,
+                    detail=f"Failed to download asset image: {asset_url}"
+                )
+
+            tgt_bytes = resp.content
+
+        # ---------------------------------------------------------
+        # DECODE BOTH 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:
+            await log_faceswap_hit(user_email, "error", start_time, datetime.now(timezone.utc))
+            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)
+
+        # ---------------------------------------------------------
+        # FACE SWAP & ENHANCE
+        # ---------------------------------------------------------
+        final_img, final_path, err = face_swap_and_enhance(src_rgb, tgt_rgb)
+        if err:
+            await log_faceswap_hit(user_email, "error", start_time, datetime.now(timezone.utc))
+            raise HTTPException(status_code=500, detail=err)
+
+        # Save final output to DO Spaces
+        with open(final_path, "rb") as f:
+            result_bytes = f.read()
+
+        result_key = f"bikini-theme/result/{uuid.uuid4().hex}_enhanced.png"
+        result_url = upload_to_spaces(result_bytes, result_key, "image/png")
+
+        await log_faceswap_hit(user_email, "success", start_time, datetime.now(timezone.utc))
+
+        # ---------------------------------------------------------
+        # SUCCESS RESPONSE
+        # ---------------------------------------------------------
+        return {
+            "result_url": result_url,
+            "category_id": category_id,
+            "user_id": user_id,
+            "new_subcategory_id": new_subcategory_id
+        }
+
+    except Exception as e:
+        await log_faceswap_hit(user_email, "error", start_time, datetime.now(timezone.utc))
+        raise HTTPException(status_code=500, detail=f"Face swap failed: {str(e)}")
+
+
+####------------------------------------OLD CODE------------------------------------####
 # @fastapi_app.post("/face-swap")
 # async def face_swap_api(
 #     source: UploadFile = File(...),
 #     target_category_id: str = Form(...),
-#     user_id: Optional[str] = Form(None),
-#     category_id: Optional[str] = Form(None),
 #     user_email: str = Depends(verify_firebase_token)
 # ):
+#     # start_time = datetime.utcnow()
 #     start_time = datetime.now(timezone.utc)
-
 #     try:
-#         # READ SOURCE
 #         src_bytes = await source.read()
 #         src_key = f"bikini-theme/source/{uuid.uuid4().hex}_{source.filename}"
 #         upload_to_spaces(src_bytes, src_key, content_type=source.content_type)
 
-#         # ------------------------------------------------------------------
-#         # OPTIONAL CATEGORY CLICK LOGGING
-#         # ------------------------------------------------------------------
-#         if user_id and category_id:
-#             try:
-#                 user_oid = ObjectId(user_id)
-#                 category_oid = ObjectId(category_id)
-#                 now = datetime.utcnow()
-
-#                 # check category exists
-#                 category_doc = await categories_col.find_one({"_id": category_oid})
-#                 if category_doc:
-#                     update_res = await media_clicks_col.update_one(
-#                         {"userId": user_oid, "categories.categoryId": category_oid},
-#                         {
-#                             "$set": {
-#                                 "updatedAt": now,
-#                                 "categories.$.lastClickedAt": now
-#                             },
-#                             "$inc": {"categories.$.click_count": 1}
-#                         }
-#                     )
-
-#                     if update_res.matched_count == 0:
-#                         await media_clicks_col.update_one(
-#                             {"userId": user_oid},
-#                             {
-#                                 "$setOnInsert": {"createdAt": now},
-#                                 "$set": {"updatedAt": now},
-#                                 "$push": {
-#                                     "categories": {
-#                                         "categoryId": category_oid,
-#                                         "click_count": 1,
-#                                         "lastClickedAt": now
-#                                     }
-#                                 }
-#                             },
-#                             upsert=True
-#                         )
-#                 else:
-#                     logger.error(f"Invalid category_id (not in categories collection): {category_id}")
-
-#             except Exception as log_err:
-#                 logger.error(f"MEDIA_CLICK ERROR: {log_err}")
-
-#         # ------------------------------------------------------------------
-#         # TARGET IMAGE DOWNLOAD
-#         # ------------------------------------------------------------------
 #         target_filename = f"{target_category_id}.png"
 #         target_url = f"https://{DO_SPACES_BUCKET}.{DO_SPACES_REGION}.digitaloceanspaces.com/bikini-theme/target/{target_filename}"
 
 #         resp = requests.get(target_url)
 #         if resp.status_code != 200:
+#             # end_time = datetime.utcnow()
 #             end_time = datetime.now(timezone.utc)
-#             await log_faceswap_hit(user_email, "error", start_time, end_time)
-#             raise HTTPException(404, f"Target not found at {target_url}")
+#             await log_faceswap_hit(user_email, status="error", start_time=start_time, end_time=end_time)
+#             raise HTTPException(status_code=404, detail=f"Target image not found at {target_url}")
 
 #         tgt_bytes = resp.content
 
-#         # FACE SWAP
-#         src_bgr = cv2.imdecode(np.frombuffer(src_bytes, np.uint8), cv2.IMREAD_COLOR)
-#         tgt_bgr = cv2.imdecode(np.frombuffer(tgt_bytes, np.uint8), cv2.IMREAD_COLOR)
+#         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:
+#             #end_time = datetime.utcnow()
 #             end_time = datetime.now(timezone.utc)
-#             await log_faceswap_hit(user_email, "error", start_time, end_time)
-#             raise HTTPException(400, "Invalid image data")
+
+#             await log_faceswap_hit(user_email, status="error", start_time=start_time, end_time=end_time)
+#             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)
 
 #         final_img, final_path, err = face_swap_and_enhance(src_rgb, tgt_rgb)
 #         if err:
-#             end_time = datetime.now(timezone.utc)
-#             await log_faceswap_hit(user_email, "error", start_time, end_time)
-#             raise HTTPException(500, err)
+#             end_time = datetime.utcnow()
+#             await log_faceswap_hit(user_email, status="error", start_time=start_time, end_time=end_time)
+#             raise HTTPException(status_code=500, detail=err)
 
 #         with open(final_path, "rb") as f:
 #             result_bytes = f.read()
-
 #         result_key = f"bikini-theme/result/{uuid.uuid4().hex}_enhanced.png"
-#         result_url = upload_to_spaces(result_bytes, result_key, "image/png")
+#         result_url = upload_to_spaces(result_bytes, result_key, content_type="image/png")
 
+#         #end_time = datetime.utcnow()
 #         end_time = datetime.now(timezone.utc)
-#         await log_faceswap_hit(user_email, "success", start_time, end_time)
+#         await log_faceswap_hit(user_email, status="success", start_time=start_time, end_time=end_time)
 
 #         return {"result_url": result_url}
 
 #     except Exception as e:
+#         #end_time = datetime.utcnow()
 #         end_time = datetime.now(timezone.utc)
-#         try:
-#             await log_faceswap_hit(user_email, "error", start_time, end_time)
-#         except:
-#             pass
-#         raise HTTPException(500, f"Face swap failed: {e}")
-
-@fastapi_app.post("/face-swap")
-async def face_swap_api(
-    source: UploadFile = File(...),
-    target_category_id: str = Form(...),
-    user_email: str = Depends(verify_firebase_token)
-):
-    # start_time = datetime.utcnow()
-    start_time = datetime.now(timezone.utc)
-    try:
-        src_bytes = await source.read()
-        src_key = f"bikini-theme/source/{uuid.uuid4().hex}_{source.filename}"
-        upload_to_spaces(src_bytes, src_key, content_type=source.content_type)
-
-        target_filename = f"{target_category_id}.png"
-        target_url = f"https://{DO_SPACES_BUCKET}.{DO_SPACES_REGION}.digitaloceanspaces.com/bikini-theme/target/{target_filename}"
-
-        resp = requests.get(target_url)
-        if resp.status_code != 200:
-            # end_time = datetime.utcnow()
-            end_time = datetime.now(timezone.utc)
-            await log_faceswap_hit(user_email, status="error", start_time=start_time, end_time=end_time)
-            raise HTTPException(status_code=404, detail=f"Target image not found at {target_url}")
-
-        tgt_bytes = resp.content
-
-        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:
-            #end_time = datetime.utcnow()
-            end_time = datetime.now(timezone.utc)
 
-            await log_faceswap_hit(user_email, status="error", start_time=start_time, end_time=end_time)
-            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)
-
-        final_img, final_path, err = face_swap_and_enhance(src_rgb, tgt_rgb)
-        if err:
-            end_time = datetime.utcnow()
-            await log_faceswap_hit(user_email, status="error", start_time=start_time, end_time=end_time)
-            raise HTTPException(status_code=500, detail=err)
-
-        with open(final_path, "rb") as f:
-            result_bytes = f.read()
-        result_key = f"bikini-theme/result/{uuid.uuid4().hex}_enhanced.png"
-        result_url = upload_to_spaces(result_bytes, result_key, content_type="image/png")
-
-        #end_time = datetime.utcnow()
-        end_time = datetime.now(timezone.utc)
-        await log_faceswap_hit(user_email, status="success", start_time=start_time, end_time=end_time)
-
-        return {"result_url": result_url}
-
-    except Exception as e:
-        #end_time = datetime.utcnow()
-        end_time = datetime.now(timezone.utc)
-
-        # Ensure we log the error with timestamps before raising
-        try:
-            await log_faceswap_hit(user_email, status="error", start_time=start_time, end_time=end_time)
-        except Exception as log_exc:
-            logger.error("Failed to write log_faceswap_hit: %s", log_exc)
-        raise HTTPException(status_code=500, detail=f"Face swap failed: {str(e)}")
+#         # Ensure we log the error with timestamps before raising
+#         try:
+#             await log_faceswap_hit(user_email, status="error", start_time=start_time, end_time=end_time)
+#         except Exception as log_exc:
+#             logger.error("Failed to write log_faceswap_hit: %s", log_exc)
+#         raise HTTPException(status_code=500, detail=f"Face swap failed: {str(e)}")
 
 
 @fastapi_app.get("/preview/{result_key:path}")