Update app.py

app.py

@@ -147,20 +147,39 @@ async def log_faceswap_hit(token: str, status: str = "success"):
 # --------------------- Face Swap Pipeline ---------------------
 swap_lock = threading.Lock()
 
+def match_histogram(src_face, tgt_face):
+    """Match the histogram of src_face to that of tgt_face for color harmony."""
+    matched = cv2.cvtColor(src_face, cv2.COLOR_BGR2LAB)
+    target_lab = cv2.cvtColor(tgt_face, cv2.COLOR_BGR2LAB)
+    for i in range(3):
+        matched[..., i] = cv2.equalizeHist(matched[..., i])
+        matched[..., i] = cv2.matchTemplate(matched[..., i], target_lab[..., i], cv2.TM_CCOEFF_NORMED)
+    return cv2.cvtColor(matched, cv2.COLOR_LAB2BGR)
+
+def find_mask_center(mask):
+    """Compute center for seamlessClone using mask moments."""
+    moments = cv2.moments(mask)
+    if moments["m00"] != 0:
+        center_x = int(moments["m10"] / moments["m00"])
+        center_y = int(moments["m01"] / moments["m00"])
+        return (center_x, center_y)
+    else:
+        # Fallback to geometric center
+        h, w = mask.shape
+        return (w // 2, h // 2)
+
 def face_swap_and_enhance(src_img, tgt_img, temp_dir="/tmp/faceswap_work"):
     try:
         with swap_lock:
-            # Prepare temporary directory
+            # Cleanup temp and prepare dirs
             if os.path.exists(temp_dir):
                 shutil.rmtree(temp_dir)
             os.makedirs(temp_dir, exist_ok=True)
-
             src_bgr = cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR)
             tgt_bgr_full = cv2.cvtColor(tgt_img, cv2.COLOR_RGB2BGR)
 
             src_faces = face_analysis_app.get(src_bgr)
             tgt_faces = face_analysis_app.get(tgt_bgr_full)
-
             if not src_faces or not tgt_faces:
                 return None, None, "❌ Face not detected in source or target image"
 
@@ -179,7 +198,7 @@ def face_swap_and_enhance(src_img, tgt_img, temp_dir="/tmp/faceswap_work"):
             src_face0 = src_faces[0]
             tgt_face0 = tgt_faces[0]
 
-            # More accurate source face crop with slight expansion
+            # Source crop for precise landmark
             s_x1, s_y1, s_x2, s_y2 = expand_bbox(src_face0.bbox, src_bgr.shape, scale=1.4)
             src_crop = src_bgr[s_y1:s_y2, s_x1:s_x2]
             src_crop_faces = face_analysis_app.get(src_crop)
@@ -190,65 +209,61 @@ def face_swap_and_enhance(src_img, tgt_img, temp_dir="/tmp/faceswap_work"):
                 src_for_swap = src_bgr
                 src_face_for_swap = src_face0
 
-            # More aggressive target crop for precise landmark detection
-            t_x1, t_y1, t_x2, t_y2 = expand_bbox(tgt_face0.bbox, tgt_bgr_full.shape, scale=1.6)
+            # Target crop for even more precision
+            t_x1, t_y1, t_x2, t_y2 = expand_bbox(tgt_face0.bbox, tgt_bgr_full.shape, scale=1.7)
             tgt_crop = tgt_bgr_full[t_y1:t_y2, t_x1:t_x2]
             tgt_crop_faces = face_analysis_app.get(tgt_crop)
-
             if tgt_crop_faces:
                 tgt_for_swap = tgt_crop
                 tgt_face_for_swap = tgt_crop_faces[0]
-
+                # Face swap
                 swapped_crop = swapper.get(tgt_for_swap, tgt_face_for_swap, src_face_for_swap)
                 if swapped_crop is None:
                     return None, None, "❌ Face swap failed on crop"
-
-                # Create mask with threshold for seamlessClone
+                # Histogram matching for color realism
+                swapped_crop = match_histogram(swapped_crop, tgt_for_swap)
+                # Mask for seamless blending
                 mask = cv2.cvtColor(swapped_crop, cv2.COLOR_BGR2GRAY)
-                _, mask = cv2.threshold(mask, 1, 255, cv2.THRESH_BINARY)
-
-                center = ((t_x1 + t_x2) // 2, (t_y1 + t_y2) // 2)
-
+                mask = cv2.threshold(mask, 1, 255, cv2.THRESH_BINARY)[1]
+                mask = cv2.GaussianBlur(mask, (15, 15), 8)
+                # Get the center using moments
+                center = find_mask_center(mask)
+                # Smoothing step to reduce artifacts
+                swapped_crop = cv2.bilateralFilter(swapped_crop, 9, 75, 75)
+                # Seamless clone for perfect embedding
                 try:
-                    blended = cv2.seamlessClone(swapped_crop, tgt_bgr_full, mask, center, cv2.NORMAL_CLONE)
+                    blended = cv2.seamlessClone(swapped_crop, tgt_bgr_full, mask, (t_x1 + center[0], t_y1 + center[1]), cv2.NORMAL_CLONE)
                 except Exception:
-                    # Fallback to direct paste if seamlessClone fails
                     blended = tgt_bgr_full.copy()
                     h, w = swapped_crop.shape[:2]
                     blended[t_y1:t_y1+h, t_x1:t_x1+w] = swapped_crop
-
                 swapped_path = os.path.join(temp_dir, f"swapped_{uuid.uuid4().hex[:8]}.jpg")
                 cv2.imwrite(swapped_path, blended)
-
             else:
-                # Fallback: swap on full image if crop detection fails
+                # Fallback: swap on full image
                 swapped_bgr_full = swapper.get(tgt_bgr_full, tgt_face0, src_face0)
                 if swapped_bgr_full is None:
                     return None, None, "❌ Face swap failed on full image"
                 swapped_path = os.path.join(temp_dir, f"swapped_{uuid.uuid4().hex[:8]}.jpg")
                 cv2.imwrite(swapped_path, swapped_bgr_full)
-
-            # Run CodeFormer enhancement on the swapped image
+            # CodeFormer enhancement post-processing as last step
             cmd = f"python {CODEFORMER_PATH} -w 0.7 --input_path {swapped_path} --output_path {temp_dir} --bg_upsampler realesrgan --face_upsample"
             result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
             if result.returncode != 0:
                 return None, None, f"❌ CodeFormer failed:\n{result.stderr}"
-
             final_results_dir = os.path.join(temp_dir, "final_results")
             final_files = [f for f in os.listdir(final_results_dir) if f.endswith(".png")]
             if not final_files:
                 return None, None, "❌ No enhanced image found"
-
             final_path = os.path.join(final_results_dir, final_files[0])
             final_img = cv2.cvtColor(cv2.imread(final_path), cv2.COLOR_BGR2RGB)
-
             return final_img, final_path, ""
-
     except Exception as e:
         return None, None, f"❌ Error: {str(e)}"
 
 
 
+
 # --------------------- Gradio ---------------------
 with gr.Blocks() as demo:
     gr.Markdown("Face Swap")