Update app.py

app.py

@@ -150,104 +150,98 @@ swap_lock = threading.Lock()
 def face_swap_and_enhance(src_img, tgt_img, temp_dir="/tmp/faceswap_work"):
     try:
         with swap_lock:
-            # Clean workspace
+            # Prepare temporary directory
             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 = cv2.cvtColor(tgt_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)
+            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"
 
-            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]
+            def expand_bbox(bbox, img_shape, scale=1.6):
+                ih, iw = img_shape[:2]
                 x1, y1, x2, y2 = map(int, bbox)
-                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)
+                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)
                 return nx1, ny1, nx2, ny2
 
-            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 ---
+            src_face0 = src_faces[0]
+            tgt_face0 = tgt_faces[0]
+
+            # More accurate source face crop with slight expansion
+            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
+
+            # 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)
+            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]
+
+                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
+                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)
+
+                try:
+                    blended = cv2.seamlessClone(swapped_crop, tgt_bgr_full, mask, center, 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
+                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
             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_dir = os.path.join(temp_dir, "final_results")
-            final_files = [f for f in os.listdir(final_dir) if f.endswith(".png")]
+            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 output found"
-            final_path = os.path.join(final_dir, final_files[0])
+                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:
@@ -255,8 +249,6 @@ def face_swap_and_enhance(src_img, tgt_img, temp_dir="/tmp/faceswap_work"):
 
 
 
-
-
 # --------------------- Gradio ---------------------
 with gr.Blocks() as demo:
     gr.Markdown("Face Swap")