Update app.py

app.py

@@ -148,6 +148,13 @@ async def log_faceswap_hit(token: str, status: str = "success"):
 swap_lock = threading.Lock()
 
 def face_swap_and_enhance(src_img, tgt_img, temp_dir="/tmp/faceswap_work"):
+    """
+    Improved face-swap function:
+    - If target is a full-body image (small face), crop an expanded face region,
+      run swap on the crop, then blend back into the full image using seamlessClone.
+    - Uses insightface detection on crops to get reliable landmarks for swapping.
+    - Runs CodeFormer as before on the final blended image.
+    """
     try:
         with swap_lock:
             # Use a temp dir for intermediate files
@@ -155,21 +162,98 @@ def face_swap_and_enhance(src_img, tgt_img, temp_dir="/tmp/faceswap_work"):
                 shutil.rmtree(temp_dir)
             os.makedirs(temp_dir, exist_ok=True)
 
+            # Convert RGB -> BGR for OpenCV / insightface
             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)
 
+            # Detect faces on original images
             src_faces = face_analysis_app.get(src_bgr)
-            tgt_faces = face_analysis_app.get(tgt_bgr)
-            if not src_faces or not tgt_faces:
-                return None, None, "❌ Face not detected in one of the images"
-
-            swapped_path = os.path.join(temp_dir, f"swapped_{uuid.uuid4().hex[:8]}.jpg")
-            swapped_bgr = swapper.get(tgt_bgr, tgt_faces[0], src_faces[0])
-            if swapped_bgr is None:
-                return None, None, "❌ Face swap failed"
-
-            cv2.imwrite(swapped_path, swapped_bgr)
-
+            tgt_faces = face_analysis_app.get(tgt_bgr_full)
+
+            if not src_faces:
+                return None, None, "❌ Face not detected in source image"
+            if not tgt_faces:
+                return None, None, "❌ Face not detected in target image"
+
+            # Helper: expand bbox by factor, keep inside image bounds
+            def expand_bbox(bbox, img_shape, scale=1.6):
+                # bbox expected as [x1, y1, x2, y2]
+                ih, iw = img_shape[:2]
+                x1, y1, x2, y2 = map(int, bbox)
+                w = x2 - x1
+                h = y2 - y1
+                cx = x1 + w // 2
+                cy = y1 + h // 2
+                new_w = int(w * scale)
+                new_h = 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
+
+            # Choose primary faces (first detected)
+            src_face0 = src_faces[0]
+            tgt_face0 = tgt_faces[0]
+
+            # Crop source face to improve alignment (detect again inside crop)
+            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]
+            # Re-run detection on source crop to get accurate face object relative to crop
+            src_crop_faces = face_analysis_app.get(src_crop)
+            if not src_crop_faces:
+                # fallback: continue with original face object
+                src_for_swap = src_bgr
+                src_face_for_swap = src_face0
+            else:
+                src_for_swap = src_crop
+                src_face_for_swap = src_crop_faces[0]
+
+            # For target: expand bbox more aggressively (full-body images)
+            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]
+            # Re-run detection on target crop
+            tgt_crop_faces = face_analysis_app.get(tgt_crop)
+            if not tgt_crop_faces:
+                # If re-detect fails (rare), try original face object on full image
+                # Do swap directly 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"
+                # Save and proceed to enhancement as before
+                swapped_path = os.path.join(temp_dir, f"swapped_{uuid.uuid4().hex[:8]}.jpg")
+                cv2.imwrite(swapped_path, swapped_bgr_full)
+            else:
+                # Swap on the cropped face region (safer and higher-quality)
+                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"
+
+                # blended result: place swapped_crop back into full target image using seamlessClone
+                # Prepare mask for seamlessClone (use non-black pixels)
+                mask = cv2.cvtColor(swapped_crop, cv2.COLOR_BGR2GRAY)
+                _, mask = cv2.threshold(mask, 1, 255, cv2.THRESH_BINARY)
+
+                # Destination center for seamlessClone (center of crop in full image coords)
+                center = ( (t_x1 + t_x2) // 2, (t_y1 + t_y2) // 2 )
+
+                # Use NORMAL_CLONE for better blending
+                try:
+                    # seamlessClone expects same depth and types
+                    blended = cv2.seamlessClone(swapped_crop, tgt_bgr_full, mask, center, cv2.NORMAL_CLONE)
+                except Exception as e:
+                    # If seamlessClone fails (e.g., mismatched sizes), fallback to simple paste
+                    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)
+
+            # Run CodeFormer on the resulting swapped image (same as before)
             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:
@@ -188,6 +272,7 @@ def face_swap_and_enhance(src_img, tgt_img, temp_dir="/tmp/faceswap_work"):
     except Exception as e:
         return None, None, f"❌ Error: {str(e)}"
 
+
 # --------------------- Gradio ---------------------
 with gr.Blocks() as demo:
     gr.Markdown("Face Swap")