Update app.py

app.py

@@ -851,6 +851,215 @@ async def multi_face_swap_api(
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
+# @fastapi_app.post("/face-swap-couple", dependencies=[Depends(verify_token)])
+# async def face_swap_api(
+#     image1: UploadFile = File(...),
+#     image2: Optional[UploadFile] = File(None),
+#     target_category_id: str = Form(None),
+#     new_category_id: str = Form(None),
+#     user_id: Optional[str] = Form(None),
+#     credentials: HTTPAuthorizationCredentials = Security(security)
+# ):
+#     """
+#     Production-ready face swap endpoint supporting:
+#     - Multiple source images (image1 + optional image2)
+#     - Gender-based pairing
+#     - Merged faces from multiple sources
+#     - Mandatory CodeFormer enhancement
+#     """
+#     start_time = datetime.utcnow()
+    
+#     try:
+#         # -----------------------------
+#         # Validate input
+#         # -----------------------------
+#         if target_category_id == "":
+#             target_category_id = None
+#         if new_category_id == "":
+#             new_category_id = None
+#         if user_id == "":
+#             user_id = None
+
+#         if target_category_id and new_category_id:
+#             raise HTTPException(400, "Provide only one of new_category_id or target_category_id.")
+#         if not target_category_id and not new_category_id:
+#             raise HTTPException(400, "Either new_category_id or target_category_id is required.")
+
+#         logger.info(f"[FaceSwap] Incoming request → target_category_id={target_category_id}, new_category_id={new_category_id}, user_id={user_id}")
+
+#         # -----------------------------
+#         # Read source images
+#         # -----------------------------
+#         src_images = []
+#         img1_bytes = await image1.read()
+#         src1 = cv2.imdecode(np.frombuffer(img1_bytes, np.uint8), cv2.IMREAD_COLOR)
+#         if src1 is None:
+#             raise HTTPException(400, "Invalid image1 data")
+#         src_images.append(cv2.cvtColor(src1, cv2.COLOR_BGR2RGB))
+
+#         if image2:
+#             img2_bytes = await image2.read()
+#             src2 = cv2.imdecode(np.frombuffer(img2_bytes, np.uint8), cv2.IMREAD_COLOR)
+#             if src2 is not None:
+#                 src_images.append(cv2.cvtColor(src2, cv2.COLOR_BGR2RGB))
+
+#         # -----------------------------
+#         # Determine target image
+#         # -----------------------------
+#         target_bytes = None
+
+#         if new_category_id:
+#             doc = await subcategories_col.find_one({"asset_images._id": ObjectId(new_category_id)})
+#             if not doc:
+#                 raise HTTPException(404, "Asset image not found")
+#             asset = next((img for img in doc["asset_images"] if str(img["_id"]) == new_category_id), None)
+#             if not asset:
+#                 raise HTTPException(404, "Asset image URL not found")
+#             target_url = asset["url"]
+
+#         elif 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 = [p["Prefix"].split("/")[2] for p in resp.get("CommonPrefixes", [])]
+#             target_url = None
+#             for category in categories:
+#                 original_prefix = f"faceswap/target/{category}/original/"
+#                 objects = client.list_objects_v2(Bucket=DO_SPACES_BUCKET, Prefix=original_prefix).get("Contents", [])
+#                 original_filenames = sorted([obj["Key"].split("/")[-1] for obj in objects if obj["Key"].endswith(".png")])
+#                 for idx, filename in enumerate(original_filenames, start=1):
+#                     cid = f"{category.lower()}image_{idx}"
+#                     if cid == target_category_id:
+#                         target_url = f"{DO_SPACES_ENDPOINT}/{DO_SPACES_BUCKET}/{original_prefix}{filename}"
+#                         break
+#                 if target_url:
+#                     break
+#             if not target_url:
+#                 raise HTTPException(404, "Target categoryId not found")
+
+#         # -----------------------------
+#         # Download target image
+#         # -----------------------------
+#         async with httpx.AsyncClient(timeout=30.0) as client:
+#             resp = await client.get(target_url)
+#             resp.raise_for_status()
+#             target_bytes = resp.content
+
+#         tgt_bgr = cv2.imdecode(np.frombuffer(target_bytes, np.uint8), cv2.IMREAD_COLOR)
+#         if tgt_bgr is None:
+#             raise HTTPException(400, "Invalid target image data")
+#         target_rgb = cv2.cvtColor(tgt_bgr, cv2.COLOR_BGR2RGB)
+
+#         # -----------------------------
+#         # Merge all source faces
+#         # -----------------------------
+#         all_src_faces = []
+#         for img in src_images:
+#             faces = face_analysis_app.get(cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
+#             all_src_faces.extend(faces)
+
+#         if not all_src_faces:
+#             raise HTTPException(400, "No faces detected in source images")
+
+#         tgt_faces = face_analysis_app.get(tgt_bgr)
+#         if not tgt_faces:
+#             raise HTTPException(400, "No faces detected in target image")
+
+#         # -----------------------------
+#         # Gender-based pairing
+#         # -----------------------------
+#         def face_sort_key(face):
+#             x1, y1, x2, y2 = face.bbox
+#             area = (x2 - x1) * (y2 - y1)
+#             cx = (x1 + x2) / 2
+#             return (-area, cx)
+
+#         # Separate by gender
+#         src_male = sorted([f for f in all_src_faces if f.gender == 1], key=face_sort_key)
+#         src_female = sorted([f for f in all_src_faces if f.gender == 0], key=face_sort_key)
+#         tgt_male = sorted([f for f in tgt_faces if f.gender == 1], key=face_sort_key)
+#         tgt_female = sorted([f for f in tgt_faces if f.gender == 0], key=face_sort_key)
+
+#         pairs = []
+#         for s, t in zip(src_male, tgt_male):
+#             pairs.append((s, t))
+#         for s, t in zip(src_female, tgt_female):
+#             pairs.append((s, t))
+
+#         # fallback if gender mismatch
+#         if not pairs:
+#             src_all = sorted(all_src_faces, key=face_sort_key)
+#             tgt_all = sorted(tgt_faces, key=face_sort_key)
+#             pairs = list(zip(src_all, tgt_all))
+
+#         # -----------------------------
+#         # Perform face swap
+#         # -----------------------------
+#         with swap_lock:
+#             result_img = tgt_bgr.copy()
+#             for src_face, _ in pairs:
+#                 current_faces = sorted(face_analysis_app.get(result_img), key=face_sort_key)
+#                 candidates = [f for f in current_faces if f.gender == src_face.gender] or current_faces
+#                 target_face = candidates[0]
+#                 result_img = swapper.get(result_img, target_face, src_face, paste_back=True)
+
+#         result_rgb = cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB)
+
+#         # -----------------------------
+#         # Mandatory enhancement
+#         # -----------------------------
+#         enhanced_rgb = mandatory_enhancement(result_rgb)
+#         enhanced_bgr = cv2.cvtColor(enhanced_rgb, cv2.COLOR_RGB2BGR)
+
+#         # -----------------------------
+#         # Save, upload, compress
+#         # -----------------------------
+#         temp_dir = tempfile.mkdtemp(prefix="faceswap_")
+#         final_path = os.path.join(temp_dir, "result.png")
+#         cv2.imwrite(final_path, enhanced_bgr)
+
+#         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)
+
+#         compressed_bytes = compress_image(result_bytes, max_size=(1280, 1280), quality=72)
+#         compressed_key = f"faceswap/result/{uuid.uuid4().hex}_enhanced_compressed.jpg"
+#         compressed_url = upload_to_spaces(compressed_bytes, compressed_key, content_type="image/jpeg")
+
+#         # -----------------------------
+#         # Log API usage
+#         # -----------------------------
+#         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": "success",
+#                 "response_time_ms": response_time_ms,
+#                 "timestamp": end_time
+#             })
+
+#         return {
+#             "result_key": result_key,
+#             "result_url": result_url,
+#             "compressed_url": compressed_url
+#         }
+
+#     except Exception as e:
+#         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(500, f"Face swap failed: {str(e)}")
+
 @fastapi_app.post("/face-swap-couple", dependencies=[Depends(verify_token)])
 async def face_swap_api(
     image1: UploadFile = File(...),
@@ -904,51 +1113,208 @@ async def face_swap_api(
                 src_images.append(cv2.cvtColor(src2, cv2.COLOR_BGR2RGB))
 
         # -----------------------------
-        # Determine target image
+        # Resolve target image
         # -----------------------------
-        target_bytes = None
-
+        target_url = None
         if new_category_id:
-            doc = await subcategories_col.find_one({"asset_images._id": ObjectId(new_category_id)})
+            doc = await subcategories_col.find_one({
+                "asset_images._id": ObjectId(new_category_id)
+            })
+
             if not doc:
-                raise HTTPException(404, "Asset image not found")
-            asset = next((img for img in doc["asset_images"] if str(img["_id"]) == new_category_id), None)
+                raise HTTPException(404, "Asset image not found in database")
+
+            asset = next(
+                (img for img in doc["asset_images"] if str(img["_id"]) == new_category_id),
+                None
+            )
+
             if not asset:
                 raise HTTPException(404, "Asset image URL not found")
+
             target_url = asset["url"]
+            subcategory_oid = doc["_id"]
 
-        elif target_category_id:
+            if user_id:
+                try:
+                    user_id_clean = user_id.strip()
+                    if not user_id_clean:
+                        raise ValueError("user_id cannot be empty")
+                    try:
+                        user_oid = ObjectId(user_id_clean)
+                    except (InvalidId, ValueError):
+                        logger.error(f"Invalid user_id format: {user_id_clean}")
+                        raise ValueError(f"Invalid user_id format: {user_id_clean}")
+
+                    now = datetime.utcnow()
+
+                    # Step 1: ensure root document exists
+                    await media_clicks_col.update_one(
+                        {"userId": user_oid},
+                        {
+                            "$setOnInsert": {
+                                "userId": user_oid,
+                                "createdAt": now,
+                                "ai_edit_complete": 0,
+                                "ai_edit_daily_count": []
+                            }
+                        },
+                        upsert=True
+                    )
+
+                    # Step 2: handle daily usage (binary, no duplicates)
+                    doc = await media_clicks_col.find_one(
+                        {"userId": user_oid},
+                        {"ai_edit_daily_count": 1}
+                    )
+
+                    daily_entries = doc.get("ai_edit_daily_count", []) if doc else []
+
+                    today_date = datetime(now.year, now.month, now.day)
+
+                    daily_map = {}
+                    for entry in daily_entries:
+                        d = entry["date"]
+                        if isinstance(d, datetime):
+                            d = datetime(d.year, d.month, d.day)
+                        daily_map[d] = entry["count"]
+
+                    last_date = max(daily_map.keys()) if daily_map else None
+
+                    if last_date != today_date:
+                        daily_map[today_date] = 1
+
+                    final_daily_entries = [
+                        {"date": d, "count": daily_map[d]}
+                        for d in sorted(daily_map.keys())
+                    ]
+
+                    final_daily_entries = final_daily_entries[-32:]
+
+                    await media_clicks_col.update_one(
+                        {"userId": user_oid},
+                        {
+                            "$set": {
+                                "ai_edit_daily_count": final_daily_entries,
+                                "updatedAt": now
+                            }
+                        }
+                    )
+
+                    # Step 3: try updating existing subCategory
+                    update_result = await media_clicks_col.update_one(
+                        {
+                            "userId": user_oid,
+                            "subCategories.subCategoryId": subcategory_oid
+                        },
+                        {
+                            "$inc": {
+                                "subCategories.$.click_count": 1,
+                                "ai_edit_complete": 1
+                            },
+                            "$set": {
+                                "subCategories.$.lastClickedAt": now,
+                                "ai_edit_last_date": now,
+                                "updatedAt": now
+                            }
+                        }
+                    )
+
+                    # Step 4: push subCategory if missing
+                    if update_result.matched_count == 0:
+                        await media_clicks_col.update_one(
+                            {"userId": user_oid},
+                            {
+                                "$inc": {
+                                    "ai_edit_complete": 1
+                                },
+                                "$set": {
+                                    "ai_edit_last_date": now,
+                                    "updatedAt": now
+                                },
+                                "$push": {
+                                    "subCategories": {
+                                        "subCategoryId": subcategory_oid,
+                                        "click_count": 1,
+                                        "lastClickedAt": now
+                                    }
+                                }
+                            }
+                        )
+
+                    # Step 5: sort subCategories by lastClickedAt (ascending)
+                    user_doc = await media_clicks_col.find_one({"userId": user_oid})
+                    if user_doc and "subCategories" in user_doc:
+                        subcategories = user_doc["subCategories"]
+                        subcategories_sorted = sorted(
+                            subcategories,
+                            key=lambda x: x.get("lastClickedAt") if x.get("lastClickedAt") is not None else datetime.min
+                        )
+                        await media_clicks_col.update_one(
+                            {"userId": user_oid},
+                            {
+                                "$set": {
+                                    "subCategories": subcategories_sorted,
+                                    "updatedAt": now
+                                }
+                            }
+                        )
+
+                    logger.info(
+                        "[MEDIA_CLICK] user=%s subCategory=%s ai_edit_complete++ daily_tracked",
+                        user_id,
+                        str(subcategory_oid)
+                    )
+
+                except Exception as media_err:
+                    logger.error(f"MEDIA_CLICK ERROR: {media_err}")
+
+        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="/")
+            resp = client.list_objects_v2(
+                Bucket=DO_SPACES_BUCKET, Prefix=base_prefix, Delimiter="/"
+            )
+
             categories = [p["Prefix"].split("/")[2] for p in resp.get("CommonPrefixes", [])]
-            target_url = None
+
             for category in categories:
                 original_prefix = f"faceswap/target/{category}/original/"
-                objects = client.list_objects_v2(Bucket=DO_SPACES_BUCKET, Prefix=original_prefix).get("Contents", [])
-                original_filenames = sorted([obj["Key"].split("/")[-1] for obj in objects if obj["Key"].endswith(".png")])
+                thumb_prefix = f"faceswap/target/{category}/thumb/"
+
+                original_objects = client.list_objects_v2(
+                    Bucket=DO_SPACES_BUCKET, Prefix=original_prefix
+                ).get("Contents", [])
+
+                thumb_objects = client.list_objects_v2(
+                    Bucket=DO_SPACES_BUCKET, Prefix=thumb_prefix
+                ).get("Contents", [])
+
+                original_filenames = sorted([
+                    obj["Key"].split("/")[-1] for obj in original_objects
+                    if obj["Key"].split("/")[-1].endswith(".png")
+                ])
+
                 for idx, filename in enumerate(original_filenames, start=1):
                     cid = f"{category.lower()}image_{idx}"
                     if cid == target_category_id:
                         target_url = f"{DO_SPACES_ENDPOINT}/{DO_SPACES_BUCKET}/{original_prefix}{filename}"
                         break
+
                 if target_url:
                     break
+
             if not target_url:
                 raise HTTPException(404, "Target categoryId not found")
 
-        # -----------------------------
-        # Download target image
-        # -----------------------------
         async with httpx.AsyncClient(timeout=30.0) as client:
-            resp = await client.get(target_url)
-            resp.raise_for_status()
-            target_bytes = resp.content
+            response = await client.get(target_url)
+            response.raise_for_status()
+            tgt_bytes = response.content
 
-        tgt_bgr = cv2.imdecode(np.frombuffer(target_bytes, np.uint8), cv2.IMREAD_COLOR)
+        tgt_bgr = cv2.imdecode(np.frombuffer(tgt_bytes, np.uint8), cv2.IMREAD_COLOR)
         if tgt_bgr is None:
             raise HTTPException(400, "Invalid target image data")
-        target_rgb = cv2.cvtColor(tgt_bgr, cv2.COLOR_BGR2RGB)
 
         # -----------------------------
         # Merge all source faces
@@ -1062,6 +1428,7 @@ async def face_swap_api(
 
 
 
+
 # --------------------- Mount Gradio ---------------------
 
 multi_faceswap_app = build_multi_faceswap_gradio()