Update app.py

app.py

@@ -1,41 +1,50 @@
-import os
-import cv2
-import numpy as np
 import insightface
-from insightface.app import FaceAnalysis
+import numpy as np
+import cv2
 import gradio as gr
+from insightface.app import FaceAnalysis
+import os
 
-# Initialize InsightFace for face detection
+# Initialize InsightFace
 app = FaceAnalysis(name='buffalo_l', providers=['CPUExecutionProvider'])
 app.prepare(ctx_id=0, det_size=(640, 640), det_thresh=0.3)
 
-# Load the InSwapper model (lightweight face swapper)
+# Load InSwapper model
 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'])
 
-# Preprocessing function to adjust contrast and brightness
-def preprocess_image(img):
-    alpha = 1.3  # Contrast control
-    beta = 15    # Brightness control
-    adjusted = cv2.convertScaleAbs(img, alpha=alpha, beta=beta)
-    return adjusted
+# Denoising to remove noise from the images
+def denoise_image(img):
+    return cv2.fastNlMeansDenoisingColored(img, None, 10, 10, 7, 21)
+
+# Color transfer to match the source image to the target's color tone
+def color_transfer(src, dst):
+    src_mean, src_std = cv2.meanStdDev(src)
+    dst_mean, dst_std = cv2.meanStdDev(dst)
+    for c in range(3):
+        src[:, :, c] = ((src[:, :, c] - src_mean[c]) * (dst_std[c] / (src_std[c] + 1e-5)) + dst_mean[c])
+    return np.clip(src, 0, 255).astype(np.uint8)
 
-# Function to sharpen the image to enhance details
+# Sharpening to enhance the details on the swapped face
 def sharpen_image(img):
-    kernel = np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]])  # Sharpening kernel
-    sharpened = cv2.filter2D(img, -1, kernel)
-    return np.clip(sharpened, 0, 255).astype(np.uint8)
+    kernel = np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]])
+    return cv2.filter2D(img, -1, kernel)
+
+# Seamless Cloning for better blending
+def seamless_cloning(src, dst, mask, center):
+    return cv2.seamlessClone(src, dst, mask, center, cv2.NORMAL_CLONE)
 
-# Function to perform face swapping
+# Swap faces and improve blending
 def swap_faces(src_img, dst_img, blur_strength=5, sharpen=False):
     try:
-        # Preprocess both images
-        src = preprocess_image(src_img)
-        dst = preprocess_image(dst_img)
+        # Preprocessing (denoise and color match)
+        src = denoise_image(src_img)
+        dst = denoise_image(dst_img)
+        src = color_transfer(src, dst)
 
-        # Convert images to RGB
+        # Convert to RGB
         src_rgb = cv2.cvtColor(src, cv2.COLOR_BGR2RGB)
         dst_rgb = cv2.cvtColor(dst, cv2.COLOR_BGR2RGB)
 
@@ -46,33 +55,42 @@ def swap_faces(src_img, dst_img, blur_strength=5, sharpen=False):
         if not src_faces or not dst_faces:
             raise ValueError("No faces detected in one of the images.")
 
-        # Use the first detected faces
+        # Select the first detected faces
         src_face = src_faces[0]
         dst_face = dst_faces[0]
 
-        # Perform face swapping using the inswapper model
+        # Get the swapped image
         swapped_img = swapper.get(dst_rgb, dst_face, src_face, paste_back=True)
 
-        # Apply Gaussian blur if required
+        # Apply blur if necessary
         if blur_strength > 0:
             swapped_img = cv2.GaussianBlur(swapped_img, (blur_strength, blur_strength), 0)
 
-        # Apply sharpening if enabled
+        # Apply sharpening if requested
         if sharpen:
             swapped_img = sharpen_image(swapped_img)
 
-        # Convert the final image back to BGR
-        result = cv2.cvtColor(swapped_img, cv2.COLOR_RGB2BGR)
-        return result
+        # Convert back to BGR for final output
+        swapped_bgr = cv2.cvtColor(swapped_img, cv2.COLOR_RGB2BGR)
+
+        # Create mask for seamless cloning
+        mask = np.zeros(swapped_bgr.shape, dtype=np.uint8)
+        mask[dst_face.ymin:dst_face.ymax, dst_face.xmin:dst_face.xmax] = 255
+        center = (dst_face.center[0], dst_face.center[1])
+
+        # Perform seamless cloning
+        final_result = seamless_cloning(swapped_bgr, dst_img, mask, center)
+
+        return final_result
     except Exception as e:
         print(f"Error: {str(e)}")
-        return np.zeros((640, 640, 3), dtype=np.uint8)  # Return a blank image on error
+        return np.zeros((640, 640, 3), dtype=np.uint8)
 
 # Gradio interface setup
-title = "🧠 Futuristic Face Swapper with inswapper_128"
+title = "🧠 Futuristic Face Swapper with Enhanced Blending"
 description = (
     "Upload a source face and a target image. The AI swaps the face using inswapper_128.onnx "
-    "for clean, smooth results. Adjust blur strength or enable sharpening for enhanced output."
+    "with enhancements such as color matching, sharpening, and seamless blending for smoother results."
 )
 
 demo = gr.Interface(