Update app.py

app.py

@@ -851,6 +851,216 @@ 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:
+            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:
+            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)}")
+
+
 # --------------------- Mount Gradio ---------------------
 
 multi_faceswap_app = build_multi_faceswap_gradio()