Update app.py

app.py

@@ -144,129 +144,7 @@ async def log_faceswap_hit(token: str, status: str = "success"):
         "timestamp": datetime.utcnow()
     })
 
-# --------------------- Enhanced Face Swap Utilities ---------------------
-
-def match_color_histogram(src_face, tgt_face):
-    """Match color histogram of source to target for better color consistency"""
-    src_lab = cv2.cvtColor(src_face, cv2.COLOR_BGR2LAB)
-    tgt_lab = cv2.cvtColor(tgt_face, cv2.COLOR_BGR2LAB)
-    
-    matched_channels = []
-    for i in range(3):
-        src_channel = src_lab[:, :, i]
-        tgt_channel = tgt_lab[:, :, i]
-        
-        src_mean, src_std = cv2.meanStdDev(src_channel)
-        tgt_mean, tgt_std = cv2.meanStdDev(tgt_channel)
-        
-        matched_channel = ((src_channel - src_mean) * (tgt_std / (src_std + 1e-6))) + tgt_mean
-        matched_channel = np.clip(matched_channel, 0, 255).astype(np.uint8)
-        matched_channels.append(matched_channel)
-    
-    matched_lab = cv2.merge(matched_channels)
-    return cv2.cvtColor(matched_lab, cv2.COLOR_LAB2BGR)
-
-def create_seamless_mask(face_bbox, img_shape, feather_amount=15):
-    """Create a smooth feathered mask for seamless blending"""
-    mask = np.zeros(img_shape[:2], dtype=np.uint8)
-    x1, y1, x2, y2 = map(int, face_bbox)
-    
-    # Create elliptical mask
-    center = ((x1 + x2) // 2, (y1 + y2) // 2)
-    axes = ((x2 - x1) // 2, (y2 - y1) // 2)
-    cv2.ellipse(mask, center, axes, 0, 0, 360, 255, -1)
-    
-    # Apply Gaussian blur for feathering
-    mask = cv2.GaussianBlur(mask, (feather_amount*2+1, feather_amount*2+1), 0)
-    
-    return mask
-
-def poisson_blend(src, tgt, mask, center):
-    """Enhanced Poisson blending with error handling"""
-    try:
-        # Ensure mask is proper format
-        if len(mask.shape) == 3:
-            mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
-        
-        # Ensure mask is binary
-        _, mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)
-        
-        # Try mixed clone first (better for face swaps)
-        result = cv2.seamlessClone(src, tgt, mask, center, cv2.MIXED_CLONE)
-        return result
-    except Exception as e:
-        logger.warning(f"Poisson blend failed: {e}, using alpha blending")
-        # Fallback to alpha blending
-        mask_3ch = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR) / 255.0
-        return (src * mask_3ch + tgt * (1 - mask_3ch)).astype(np.uint8)
-
-def multi_band_blending(src, tgt, mask, levels=4):
-    """Multi-band blending for smoother transitions"""
-    # Build Gaussian pyramids
-    src_pyr = [src.copy()]
-    tgt_pyr = [tgt.copy()]
-    mask_pyr = [mask.copy()]
-    
-    for i in range(levels):
-        src_pyr.append(cv2.pyrDown(src_pyr[i]))
-        tgt_pyr.append(cv2.pyrDown(tgt_pyr[i]))
-        mask_pyr.append(cv2.pyrDown(mask_pyr[i]))
-    
-    # Build Laplacian pyramids
-    src_lap = [src_pyr[levels]]
-    tgt_lap = [tgt_pyr[levels]]
-    
-    for i in range(levels, 0, -1):
-        size = (src_pyr[i-1].shape[1], src_pyr[i-1].shape[0])
-        src_lap.append(cv2.subtract(src_pyr[i-1], cv2.pyrUp(src_pyr[i], dstsize=size)))
-        tgt_lap.append(cv2.subtract(tgt_pyr[i-1], cv2.pyrUp(tgt_pyr[i], dstsize=size)))
-    
-    # Blend each level
-    blended = []
-    for i in range(len(src_lap)):
-        if len(mask_pyr) > i:
-            mask_norm = mask_pyr[len(mask_pyr)-1-i].astype(float) / 255.0
-            if len(mask_norm.shape) == 2:
-                mask_norm = np.stack([mask_norm]*3, axis=2)
-            ls = src_lap[i] * mask_norm + tgt_lap[i] * (1 - mask_norm)
-            blended.append(ls.astype(src.dtype))
-    
-    # Reconstruct
-    result = blended[0]
-    for i in range(1, len(blended)):
-        size = (blended[i].shape[1], blended[i].shape[0])
-        result = cv2.add(cv2.pyrUp(result, dstsize=size), blended[i])
-    
-    return result
-
-def enhance_face_details(img):
-    """Apply subtle sharpening to enhance facial details"""
-    kernel = np.array([[-1,-1,-1],
-                       [-1, 9,-1],
-                       [-1,-1,-1]]) / 9
-    sharpened = cv2.filter2D(img, -1, kernel)
-    # Blend 70% original with 30% sharpened for subtle enhancement
-    return cv2.addWeighted(img, 0.7, sharpened, 0.3, 0)
-
-def adjust_face_lighting(swapped_face, target_face):
-    """Match lighting conditions between swapped and target face"""
-    # Convert to LAB color space
-    swapped_lab = cv2.cvtColor(swapped_face, cv2.COLOR_BGR2LAB).astype(float)
-    target_lab = cv2.cvtColor(target_face, cv2.COLOR_BGR2LAB).astype(float)
-    
-    # Match L channel (luminance)
-    swapped_l = swapped_lab[:, :, 0]
-    target_l = target_lab[:, :, 0]
-    
-    swapped_mean = swapped_l.mean()
-    target_mean = target_l.mean()
-    
-    # Adjust luminance
-    swapped_lab[:, :, 0] = np.clip(swapped_l + (target_mean - swapped_mean) * 0.5, 0, 255)
-    
-    return cv2.cvtColor(swapped_lab.astype(np.uint8), cv2.COLOR_LAB2BGR)
-
-# --------------------- Enhanced Face Swap Pipeline ---------------------
+# --------------------- Face Swap Pipeline ---------------------
 swap_lock = threading.Lock()
 
 def face_swap_and_enhance(src_img, tgt_img, temp_dir="/tmp/faceswap_work"):
@@ -280,43 +158,30 @@ def face_swap_and_enhance(src_img, tgt_img, temp_dir="/tmp/faceswap_work"):
             src_bgr = cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR)
             tgt_bgr_full = cv2.cvtColor(tgt_img, cv2.COLOR_RGB2BGR)
 
-            # Detect faces with higher confidence
             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"
 
-            # Sort by face size (use largest faces)
-            src_faces.sort(key=lambda x: (x.bbox[2]-x.bbox[0])*(x.bbox[3]-x.bbox[1]), reverse=True)
-            tgt_faces.sort(key=lambda x: (x.bbox[2]-x.bbox[0])*(x.bbox[3]-x.bbox[1]), reverse=True)
-
-            def expand_bbox(bbox, img_shape, scale=1.5):
-                """Smart bbox expansion with padding"""
+            def expand_bbox(bbox, img_shape, scale=1.6):
                 ih, iw = img_shape[:2]
                 x1, y1, x2, y2 = map(int, bbox)
                 w, h = x2 - x1, y2 - y1
                 cx, cy = x1 + w // 2, y1 + h // 2
-                
-                # Use larger dimension for square crop
-                size = max(w, h)
-                new_size = int(size * scale)
-                
-                nx1 = max(0, cx - new_size // 2)
-                ny1 = max(0, cy - new_size // 2)
-                nx2 = min(iw, cx + new_size // 2)
-                ny2 = min(ih, cy + new_size // 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
 
             src_face0 = src_faces[0]
             tgt_face0 = tgt_faces[0]
 
-            # Extract and align source face with optimal crop
-            s_x1, s_y1, s_x2, s_y2 = expand_bbox(src_face0.bbox, src_bgr.shape, scale=1.3)
+            # 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-detect face in crop for better precision
             src_crop_faces = face_analysis_app.get(src_crop)
             if src_crop_faces:
                 src_for_swap = src_crop
@@ -325,93 +190,43 @@ 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
 
-            # Extract target face with larger context for better blending
-            t_x1, t_y1, t_x2, t_y2 = expand_bbox(tgt_face0.bbox, tgt_bgr_full.shape, scale=1.7)
+            # 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-detect in target crop
             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]
 
-                # Perform face swap on crop
                 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"
 
-                # Extract face region from swapped crop for color matching
-                f_x1, f_y1, f_x2, f_y2 = map(int, tgt_face_for_swap.bbox)
-                swapped_face_region = swapped_crop[f_y1:f_y2, f_x1:f_x2]
-                target_face_region = tgt_for_swap[f_y1:f_y2, f_x1:f_x2]
-                
-                # Match colors for better consistency
-                if swapped_face_region.size > 0 and target_face_region.size > 0:
-                    color_matched = match_color_histogram(swapped_face_region, target_face_region)
-                    swapped_crop[f_y1:f_y2, f_x1:f_x2] = color_matched
-                
-                # Adjust lighting
-                swapped_crop = adjust_face_lighting(swapped_crop, tgt_for_swap)
-                
-                # Create advanced seamless mask
-                mask = create_seamless_mask(tgt_face_for_swap.bbox, tgt_for_swap.shape, feather_amount=20)
-                
-                # Multi-band blending for natural transition
+                # 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:
-                    # Resize mask to match crop size
-                    mask_resized = cv2.resize(mask, (swapped_crop.shape[1], swapped_crop.shape[0]))
-                    blended_crop = multi_band_blending(swapped_crop, tgt_for_swap, mask_resized, levels=4)
-                except Exception as e:
-                    logger.warning(f"Multi-band blending failed: {e}, using Poisson")
-                    # Fallback to Poisson blending
-                    center = ((t_x1 + t_x2) // 2 - t_x1, (t_y1 + t_y2) // 2 - t_y1)
-                    mask_3ch = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR)
-                    blended_crop = poisson_blend(swapped_crop, tgt_for_swap, mask_3ch, center)
-                
-                # Place blended crop back into full image
-                blended_full = tgt_bgr_full.copy()
-                h, w = blended_crop.shape[:2]
-                
-                # Create soft edge mask for final compositing
-                edge_mask = np.ones((h, w), dtype=np.float32)
-                edge_feather = 30
-                edge_mask = cv2.GaussianBlur(edge_mask, (edge_feather*2+1, edge_feather*2+1), 0)
-                edge_mask = np.stack([edge_mask]*3, axis=2)
-                
-                # Composite with soft edges
-                crop_region = blended_full[t_y1:t_y1+h, t_x1:t_x1+w]
-                blended_full[t_y1:t_y1+h, t_x1:t_x1+w] = (
-                    blended_crop * edge_mask + crop_region * (1 - edge_mask)
-                ).astype(np.uint8)
-                
-                # Enhance facial details
-                face_region = blended_full[t_y1:t_y2, t_x1:t_x2]
-                enhanced_face = enhance_face_details(face_region)
-                blended_full[t_y1:t_y2, t_x1:t_x2] = enhanced_face
-                
+                    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_full, [cv2.IMWRITE_JPEG_QUALITY, 95])
+                cv2.imwrite(swapped_path, blended)
 
             else:
-                # Fallback: swap on full image with enhancements
+                # 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"
-                
-                # Apply color matching and lighting adjustment
-                t_x1, t_y1, t_x2, t_y2 = map(int, tgt_face0.bbox)
-                swapped_face = swapped_bgr_full[t_y1:t_y2, t_x1:t_x2]
-                target_face = tgt_bgr_full[t_y1:t_y2, t_x1:t_x2]
-                
-                if swapped_face.size > 0 and target_face.size > 0:
-                    color_matched = match_color_histogram(swapped_face, target_face)
-                    swapped_bgr_full[t_y1:t_y2, t_x1:t_x2] = color_matched
-                
-                swapped_bgr_full = adjust_face_lighting(swapped_bgr_full, tgt_bgr_full)
-                
                 swapped_path = os.path.join(temp_dir, f"swapped_{uuid.uuid4().hex[:8]}.jpg")
-                cv2.imwrite(swapped_path, swapped_bgr_full, [cv2.IMWRITE_JPEG_QUALITY, 95])
+                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"
@@ -430,18 +245,19 @@ def face_swap_and_enhance(src_img, tgt_img, temp_dir="/tmp/faceswap_work"):
             return final_img, final_path, ""
 
     except Exception as e:
-        logger.exception("Face swap error")
         return None, None, f"❌ Error: {str(e)}"
 
+
+
 # --------------------- Gradio ---------------------
 with gr.Blocks() as demo:
-    gr.Markdown("# 🎭 Enhanced Face Swap\n### With improved color matching, blending, and detail preservation")
+    gr.Markdown("Face Swap")
 
     with gr.Row():
         src_input = gr.Image(type="numpy", label="Upload Your Face")
         tgt_input = gr.Image(type="numpy", label="Upload Target Image")
 
-    btn = gr.Button("✨ Swap Face", variant="primary")
+    btn = gr.Button("Swap Face")
     output_img = gr.Image(type="numpy", label="Enhanced Output")
     download = gr.File(label="⬇️ Download Enhanced Image")
     error_box = gr.Textbox(label="Logs / Errors", interactive=False)