Update app.py

app.py

@@ -148,83 +148,54 @@ 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
+            # Prepare temporary directory
             if os.path.exists(temp_dir):
                 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_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_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"
+            if not src_faces or not tgt_faces:
+                return None, None, "❌ Face not detected in source or 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)
+                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
 
-            # 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)
+            # 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]
-            # 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:
+            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
 
-            # For target: expand bbox more aggressively (full-body images)
+            # 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]
-            # 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)
+
+            if tgt_crop_faces:
                 tgt_for_swap = tgt_crop
                 tgt_face_for_swap = tgt_crop_faces[0]
 
@@ -232,20 +203,16 @@ def face_swap_and_enhance(src_img, tgt_img, temp_dir="/tmp/faceswap_work"):
                 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)
+                # Create mask with threshold for seamlessClone
                 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 )
+                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
+                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
@@ -253,7 +220,15 @@ def face_swap_and_enhance(src_img, tgt_img, temp_dir="/tmp/faceswap_work"):
                 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)
+            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:
@@ -273,6 +248,7 @@ def face_swap_and_enhance(src_img, tgt_img, temp_dir="/tmp/faceswap_work"):
         return None, None, f"❌ Error: {str(e)}"
 
 
+
 # --------------------- Gradio ---------------------
 with gr.Blocks() as demo:
     gr.Markdown("Face Swap")