Update app.py

app.py

@@ -182,73 +182,168 @@ def enhance_image_with_codeformer(rgb_img, temp_dir=None):
     enhanced = cv2.imread(final_path)
     return cv2.cvtColor(enhanced, cv2.COLOR_BGR2RGB)
     
-def multi_face_swap(src_img, tgt_img):
+import numpy as np
+from sklearn.metrics.pairwise import cosine_similarity
+
+SIM_THRESHOLD = 0.35   # strict
+MAX_YAW = 35
+MAX_PITCH = 25
+MIN_FACE_AREA = 10000
+
+
+def face_ok(face):
+    yaw, pitch, roll = face.pose
+    x1, y1, x2, y2 = face.bbox
+    area = (x2 - x1) * (y2 - y1)
+
+    return (
+        abs(yaw) <= MAX_YAW and
+        abs(pitch) <= MAX_PITCH and
+        area >= MIN_FACE_AREA
+    )
+
+
+def quality_score(face):
+    yaw, pitch, roll = face.pose
+    x1, y1, x2, y2 = face.bbox
+    area = (x2 - x1) * (y2 - y1)
+    return area - abs(yaw)*200 - abs(pitch)*150
+
+
+def multi_face_swap_industry(src_img, tgt_img):
+    """
+    Industry-grade 2-face swap
+    Returns swapped RGB image or raises error
+    """
+
     src_bgr = cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR)
     tgt_bgr = cv2.cvtColor(tgt_img, cv2.COLOR_RGB2BGR)
 
     src_faces = face_analysis_app.get(src_bgr)
     tgt_faces = face_analysis_app.get(tgt_bgr)
 
-    if not src_faces or not tgt_faces:
+    if len(src_faces) < 1 or len(tgt_faces) < 1:
         raise ValueError("No faces detected")
 
-    def face_sort_key(face):
-        x1, y1, x2, y2 = face.bbox
-        area = (x2 - x1) * (y2 - y1)
-        cx = (x1 + x2) / 2
-        return (-area, cx)
+    if len(src_faces) > 2 or len(tgt_faces) > 2:
+        raise ValueError("More than 2 faces detected (blocked)")
 
-    # Split by gender
-    src_male = [f for f in src_faces if f.gender == 1]
-    src_female = [f for f in src_faces if f.gender == 0]
+    # ---- Quality filtering ----
+    src_faces = [f for f in src_faces if face_ok(f)]
+    tgt_faces = [f for f in tgt_faces if face_ok(f)]
 
-    tgt_male = [f for f in tgt_faces if f.gender == 1]
-    tgt_female = [f for f in tgt_faces if f.gender == 0]
+    if not src_faces or not tgt_faces:
+        raise ValueError("Faces rejected due to pose/quality")
 
-    # Sort inside gender groups
-    src_male = sorted(src_male, key=face_sort_key)
-    src_female = sorted(src_female, key=face_sort_key)
+    # ---- Sort by quality (best first) ----
+    src_faces = sorted(src_faces, key=quality_score, reverse=True)
+    tgt_faces = sorted(tgt_faces, key=quality_score, reverse=True)
 
-    tgt_male = sorted(tgt_male, key=face_sort_key)
-    tgt_female = sorted(tgt_female, key=face_sort_key)
+    # ---- Extract embeddings ----
+    src_embs = np.array([f.normed_embedding for f in src_faces])
+    tgt_embs = np.array([f.normed_embedding for f in tgt_faces])
 
-    # Build final swap pairs
-    pairs = []
+    sim = cosine_similarity(src_embs, tgt_embs)
 
-    for s, t in zip(src_male, tgt_male):
-        pairs.append((s, t))
+    used_tgts = set()
+    pairs = []
 
-    for s, t in zip(src_female, tgt_female):
-        pairs.append((s, t))
+    for i in range(len(src_faces)):
+        j = np.argmax(sim[i])
+        score = sim[i][j]
 
-    # Fallback if gender mismatch
-    if not pairs:
-        src_faces = sorted(src_faces, key=face_sort_key)
-        tgt_faces = sorted(tgt_faces, key=face_sort_key)
-        pairs = list(zip(src_faces, tgt_faces))
+        if score < SIM_THRESHOLD:
+            raise ValueError(f"Low identity similarity ({score:.2f})")
 
-    result_img = tgt_bgr.copy()
+        if j in used_tgts:
+            continue
 
-    for src_face, _ in pairs:
-        # 🔁 re-detect current target faces
-        current_faces = face_analysis_app.get(result_img)
-        current_faces = sorted(current_faces, key=face_sort_key)
+        pairs.append((src_faces[i], tgt_faces[j], score))
+        used_tgts.add(j)
 
-        # choose best matching gender
-        candidates = [
-            f for f in current_faces if f.gender == src_face.gender
-        ] or current_faces
+    if not pairs:
+        raise ValueError("No valid face pairs")
 
-        target_face = candidates[0]
+    # ---- Swap WITHOUT re-detection ----
+    result = tgt_bgr.copy()
 
-        result_img = swapper.get(
-            result_img,
-            target_face,
+    for src_face, tgt_face, score in pairs:
+        result = swapper.get(
+            result,
+            tgt_face,
             src_face,
             paste_back=True
         )
 
-    return cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB)
+    return cv2.cvtColor(result, cv2.COLOR_BGR2RGB)
+    
+# def multi_face_swap(src_img, tgt_img):
+#     src_bgr = cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR)
+#     tgt_bgr = cv2.cvtColor(tgt_img, cv2.COLOR_RGB2BGR)
+
+#     src_faces = face_analysis_app.get(src_bgr)
+#     tgt_faces = face_analysis_app.get(tgt_bgr)
+
+#     if not src_faces or not tgt_faces:
+#         raise ValueError("No faces detected")
+
+#     def face_sort_key(face):
+#         x1, y1, x2, y2 = face.bbox
+#         area = (x2 - x1) * (y2 - y1)
+#         cx = (x1 + x2) / 2
+#         return (-area, cx)
+
+#     # Split by gender
+#     src_male = [f for f in src_faces if f.gender == 1]
+#     src_female = [f for f in src_faces if f.gender == 0]
+
+#     tgt_male = [f for f in tgt_faces if f.gender == 1]
+#     tgt_female = [f for f in tgt_faces if f.gender == 0]
+
+#     # Sort inside gender groups
+#     src_male = sorted(src_male, key=face_sort_key)
+#     src_female = sorted(src_female, key=face_sort_key)
+
+#     tgt_male = sorted(tgt_male, key=face_sort_key)
+#     tgt_female = sorted(tgt_female, key=face_sort_key)
+
+#     # Build final swap pairs
+#     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_faces = sorted(src_faces, key=face_sort_key)
+#         tgt_faces = sorted(tgt_faces, key=face_sort_key)
+#         pairs = list(zip(src_faces, tgt_faces))
+
+#     result_img = tgt_bgr.copy()
+
+#     for src_face, _ in pairs:
+#         # 🔁 re-detect current target faces
+#         current_faces = face_analysis_app.get(result_img)
+#         current_faces = sorted(current_faces, key=face_sort_key)
+
+#         # choose best matching gender
+#         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
+#         )
+
+#     return cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB)