Update app.py

app.py

@@ -3,59 +3,96 @@ import numpy as np
 import cv2
 import gradio as gr
 from insightface.app import FaceAnalysis
-import onnxruntime
 import os
 
+# Suppress NumPy rcond warning (optional)
+import warnings
+warnings.filterwarnings("ignore", category=FutureWarning)
+
+# Initialize InsightFace
 app = FaceAnalysis(name='buffalo_l', providers=['CPUExecutionProvider'])
 app.prepare(ctx_id=0, det_size=(640, 640), det_thresh=0.3)
 
+# Load inswapper model for face swapping
 inswapper_path = "checkpoints/inswapper_128.onnx"
 if not os.path.exists(inswapper_path):
     raise FileNotFoundError(f"Model not found at {inswapper_path}")
 swapper = insightface.model_zoo.get_model(inswapper_path, providers=['CPUExecutionProvider'])
 
 def preprocess_image(img):
-    alpha = 1.3
-    beta = 15
+    """Preprocess image to enhance it for better swap results."""
+    # Adjust contrast and brightness for better facial features extraction
+    alpha = 1.3  # contrast control
+    beta = 15  # brightness control
     adjusted = cv2.convertScaleAbs(img, alpha=alpha, beta=beta)
-    return adjusted
+    
+    # Optional: histogram equalization to enhance details
+    adjusted = cv2.cvtColor(adjusted, cv2.COLOR_BGR2GRAY)
+    adjusted = cv2.equalizeHist(adjusted)
+    return cv2.cvtColor(adjusted, cv2.COLOR_GRAY2BGR)
 
 def sharpen_image(img):
-    kernel = np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]])
+    """Enhance the sharpness of the face swapped image."""
+    kernel = np.array([[0, -1, 0], [-1, 5,-1], [0, -1, 0]])
     sharpened = cv2.filter2D(img, -1, kernel)
     return np.clip(sharpened, 0, 255).astype(np.uint8)
 
+def color_correction(src_img, dst_img):
+    """Color correction to ensure the source face blends seamlessly with the target face."""
+    src_img = cv2.cvtColor(src_img, cv2.COLOR_BGR2RGB)
+    dst_img = cv2.cvtColor(dst_img, cv2.COLOR_BGR2RGB)
+    
+    src_mean, src_std = cv2.meanStdDev(src_img)
+    dst_mean, dst_std = cv2.meanStdDev(dst_img)
+    
+    for c in range(3):
+        src_img[:, :, c] = ((src_img[:, :, c] - src_mean[c]) * (dst_std[c] / (src_std[c] + 1e-5)) + dst_mean[c])
+    
+    return cv2.cvtColor(np.clip(src_img, 0, 255).astype(np.uint8), cv2.COLOR_RGB2BGR)
+
 def swap_faces(src_img, dst_img, blur_strength=5, sharpen=False):
     try:
+        # Preprocess images for better face swapping
         src = preprocess_image(src_img)
         dst = preprocess_image(dst_img)
 
+        # Convert images to RGB for InsightFace processing
         src_rgb = cv2.cvtColor(src, cv2.COLOR_BGR2RGB)
         dst_rgb = cv2.cvtColor(dst, cv2.COLOR_BGR2RGB)
 
+        # Detect faces in source and destination images
         src_faces = app.get(src_rgb)
         dst_faces = app.get(dst_rgb)
 
         if not src_faces or not dst_faces:
             raise ValueError("No faces detected in one of the images.")
 
+        # Extract the first face from each image
         src_face = src_faces[0]
         dst_face = dst_faces[0]
 
+        # Perform face swapping using the inswapper model
         swapped_img = swapper.get(dst_rgb, dst_face, src_face, paste_back=True)
 
+        # Apply Gaussian blur if specified
         if blur_strength > 0:
             swapped_img = cv2.GaussianBlur(swapped_img, (blur_strength, blur_strength), 0)
 
+        # Apply sharpening if enabled
         if sharpen:
             swapped_img = sharpen_image(swapped_img)
 
+        # Color correction to match the destination face
+        swapped_img = color_correction(swapped_img, dst)
+
+        # Return the final swapped face image
         result = cv2.cvtColor(swapped_img, cv2.COLOR_RGB2BGR)
         return result
     except Exception as e:
         print(f"Error: {str(e)}")
-        return np.zeros((640, 640, 3), dtype=np.uint8)
+        return np.zeros((640, 640, 3), dtype=np.uint8)  # Fallback image
 
+# Gradio UI setup
 title = "🧠 Futuristic Face Swapper with inswapper_128"
 description = (
     "Upload a source face and a target image. The AI swaps the face using inswapper_128.onnx "
@@ -77,4 +114,4 @@ demo = gr.Interface(
 )
 
 if __name__ == "__main__":
-    demo.queue().launch()
+    demo.queue().launch()