Update app.py

app.py

@@ -147,123 +147,116 @@ 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:
-            # Cleanup temp and prepare dirs
+            # Clean workspace
             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)
+            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_full)
+            tgt_faces = face_analysis_app.get(tgt_bgr)
+
             if not src_faces or not tgt_faces:
                 return None, None, "❌ Face not detected in source or target image"
 
-            def expand_bbox(bbox, img_shape, scale=1.6):
-                ih, iw = img_shape[:2]
+            src_face = src_faces[0]
+            tgt_face = tgt_faces[0]
+
+            # --- Adaptive expansion based on face size ---
+            def adaptive_bbox(bbox, shape, factor=1.7):
+                h, w = shape[:2]
                 x1, y1, x2, y2 = map(int, bbox)
-                w, h = x2 - x1, y2 - y1
-                cx, cy = x1 + w // 2, y1 + h // 2
-                new_w, new_h = int(w * scale), int(h * scale)
-                nx1 = max(0, cx - new_w // 2)
-                ny1 = max(0, cy - new_h // 2)
-                nx2 = min(iw, cx + new_w // 2)
-                ny2 = min(ih, cy + new_h // 2)
+                bw, bh = x2 - x1, y2 - y1
+                cx, cy = x1 + bw // 2, y1 + bh // 2
+                new_w, new_h = int(bw * factor), int(bh * factor)
+                nx1, ny1 = max(0, cx - new_w // 2), max(0, cy - new_h // 2)
+                nx2, ny2 = min(w, cx + new_w // 2), min(h, cy + new_h // 2)
                 return nx1, ny1, nx2, ny2
 
-            src_face0 = src_faces[0]
-            tgt_face0 = tgt_faces[0]
-
-            # 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)
-            if src_crop_faces:
-                src_for_swap = src_crop
-                src_face_for_swap = src_crop_faces[0]
-            else:
-                src_for_swap = src_bgr
-                src_face_for_swap = src_face0
-
-            # 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"
-                # 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)[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, (t_x1 + center[0], t_y1 + center[1]), cv2.NORMAL_CLONE)
-                except Exception:
-                    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
-                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)
-            # CodeFormer enhancement post-processing as last step
+            sx1, sy1, sx2, sy2 = adaptive_bbox(src_face.bbox, src_bgr.shape)
+            tx1, ty1, tx2, ty2 = adaptive_bbox(tgt_face.bbox, tgt_bgr.shape)
+
+            src_crop = src_bgr[sy1:sy2, sx1:sx2]
+            tgt_crop = tgt_bgr[ty1:ty2, tx1:tx2]
+
+            src_face_crop = face_analysis_app.get(src_crop)
+            tgt_face_crop = face_analysis_app.get(tgt_crop)
+            if not src_face_crop or not tgt_face_crop:
+                return None, None, "❌ Re-detection failed in cropped region"
+
+            src_face_refined = src_face_crop[0]
+            tgt_face_refined = tgt_face_crop[0]
+
+            # --- Alignment correction using landmarks (similarity transform) ---
+            def align_face(img, src_pts, tgt_pts):
+                M, _ = cv2.estimateAffinePartial2D(src_pts, tgt_pts)
+                if M is not None:
+                    aligned = cv2.warpAffine(img, M, (img.shape[1], img.shape[0]), flags=cv2.INTER_LINEAR)
+                    return aligned
+                return img
+
+            src_pts = src_face_refined.landmark_2d_106[:68].astype(np.float32)
+            tgt_pts = tgt_face_refined.landmark_2d_106[:68].astype(np.float32)
+            src_aligned = align_face(src_crop, src_pts, tgt_pts)
+
+            # --- Perform swap on aligned faces ---
+            swapped_crop = swapper.get(tgt_crop, tgt_face_refined, src_face_refined)
+            if swapped_crop is None:
+                return None, None, "❌ Swap failed"
+
+            # --- Feathered mask for smooth blending ---
+            gray = cv2.cvtColor(swapped_crop, cv2.COLOR_BGR2GRAY)
+            mask = np.where(gray > 5, 255, 0).astype(np.uint8)
+            mask = cv2.GaussianBlur(mask, (41, 41), 15)
+
+            # --- Tone normalization between faces ---
+            def match_tone(src, ref):
+                src_mean, src_std = cv2.meanStdDev(src)
+                ref_mean, ref_std = cv2.meanStdDev(ref)
+                adj = (src - src_mean) * (ref_std / (src_std + 1e-6)) + ref_mean
+                return np.clip(adj, 0, 255).astype(np.uint8)
+
+            swapped_tone = match_tone(swapped_crop, tgt_crop)
+
+            # --- Seamless composite ---
+            center = (tx1 + (tx2 - tx1)//2, ty1 + (ty2 - ty1)//2)
+            try:
+                blended = cv2.seamlessClone(swapped_tone, tgt_bgr, mask, center, cv2.NORMAL_CLONE)
+            except Exception:
+                blended = tgt_bgr.copy()
+                h, w = swapped_tone.shape[:2]
+                blended[ty1:ty1+h, tx1:tx1+w] = swapped_tone
+
+            swapped_path = os.path.join(temp_dir, f"swapped_{uuid.uuid4().hex}.jpg")
+            cv2.imwrite(swapped_path, blended)
+
+            # --- CodeFormer enhancement ---
             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")]
+
+            final_dir = os.path.join(temp_dir, "final_results")
+            final_files = [f for f in os.listdir(final_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])
+                return None, None, "❌ No enhanced output found"
+            final_path = os.path.join(final_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")