Update app.py

app.py

@@ -144,7 +144,129 @@ async def log_faceswap_hit(token: str, status: str = "success"):
         "timestamp": datetime.utcnow()
     })
 
-# --------------------- Face Swap Pipeline ---------------------
+# --------------------- 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 ---------------------
 swap_lock = threading.Lock()
 
 def face_swap_and_enhance(src_img, tgt_img, temp_dir="/tmp/faceswap_work"):
@@ -158,30 +280,43 @@ 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"
 
-            def expand_bbox(bbox, img_shape, scale=1.6):
+            # 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"""
                 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
-                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)
+                
+                # 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)
+                
                 return nx1, ny1, nx2, ny2
 
             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)
+            # 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)
             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
@@ -190,43 +325,93 @@ 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
 
-            # 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)
+            # 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)
             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"
 
-                # 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)
-
+                # 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
                 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
-
+                    # 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
+                
                 swapped_path = os.path.join(temp_dir, f"swapped_{uuid.uuid4().hex[:8]}.jpg")
-                cv2.imwrite(swapped_path, blended)
+                cv2.imwrite(swapped_path, blended_full, [cv2.IMWRITE_JPEG_QUALITY, 95])
 
             else:
-                # Fallback: swap on full image if crop detection fails
+                # Fallback: swap on full image with enhancements
                 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(swapped_path, swapped_bgr_full, [cv2.IMWRITE_JPEG_QUALITY, 95])
 
             # 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"
@@ -245,19 +430,18 @@ 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("Face Swap")
+    gr.Markdown("# 🎭 Enhanced Face Swap\n### With improved color matching, blending, and detail preservation")
 
     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")
+    btn = gr.Button("✨ Swap Face", variant="primary")
     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)
@@ -383,4 +567,4 @@ async def preview_result(result_key: str):
 fastapi_app = mount_gradio_app(fastapi_app, demo, path="/gradio")
 
 if __name__ == "__main__":
-    uvicorn.run(fastapi_app, host="0.0.0.0", port=7860)
+    uvicorn.run(fastapi_app, host="0.0.0.0", port=7860)