Update app.py

app.py

@@ -7,7 +7,7 @@ import onnxruntime
 import os
 
 app = FaceAnalysis(name='buffalo_l', providers=['CPUExecutionProvider'])
-app.prepare(ctx_id=0, det_size=(640, 640))
+app.prepare(ctx_id=0, det_size=(640, 640), det_thresh=0.3)
 
 inswapper_path = "checkpoints/inswapper_128.onnx"
 if not os.path.exists(inswapper_path):
@@ -15,28 +15,36 @@ if not os.path.exists(inswapper_path):
 swapper = insightface.model_zoo.get_model(inswapper_path, providers=['CPUExecutionProvider'])
 
 def preprocess_image(img):
-    # Adjust brightness and contrast to improve quality
-    alpha = 1.2  # Contrast control
-    beta = 10    # Brightness control
+    alpha = 1.3
+    beta = 15
     adjusted = cv2.convertScaleAbs(img, alpha=alpha, beta=beta)
     return adjusted
 
-def sharpen_image(img):
-    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 swap_faces(src_img, dst_img, blur_strength=3, sharpen=False):
+def smooth_blend(img, swapped_img, face_mask):
+    # Dilate the mask for better blending
+    kernel = np.ones((5, 5), np.uint8)
+    mask = cv2.dilate(face_mask, kernel, iterations=2)
+    mask = cv2.GaussianBlur(mask, (15, 15), 0)
+    mask = mask / 255.0
+    mask = np.repeat(mask[:, :, np.newaxis], 3, axis=2)
+    result = (swapped_img * mask + img * (1 - mask)).astype(np.uint8)
+    return result
+
+def create_face_mask(img, keypoints):
+    # Create a convex hull mask around the face using keypoints
+    hull = cv2.convexHull(keypoints.astype(np.int32))
+    mask = np.zeros(img.shape[:2], dtype=np.uint8)
+    cv2.fillConvexPoly(mask, hull, 255)
+    return mask
+
+def swap_faces(src_img, dst_img, blur_strength=5, sharpen=False):
     try:
-        # Preprocess while preserving color
         src = preprocess_image(src_img)
         dst = preprocess_image(dst_img)
 
-        # Convert to RGB for face detection
         src_rgb = cv2.cvtColor(src, cv2.COLOR_BGR2RGB)
         dst_rgb = cv2.cvtColor(dst, cv2.COLOR_BGR2RGB)
 
-        # Detect faces
         src_faces = app.get(src_rgb)
         dst_faces = app.get(dst_rgb)
 
@@ -46,19 +54,26 @@ def swap_faces(src_img, dst_img, blur_strength=3, sharpen=False):
         src_face = src_faces[0]
         dst_face = dst_faces[0]
 
-        # Perform face swap using inswapper
-        swapped_img = swapper.get(dst_rgb, dst_face, src_face, paste_back=True)
+        # Create a mask for smoother blending
+        dst_kps = dst_face.kps
+        face_mask = create_face_mask(dst_rgb, dst_kps)
+
+        # Perform face swap
+        swapped_img = swapper.get(dst_rgb, dst_face, src_face, paste_back=False)
+
+        # Blend the swapped face with the original image
+        blended_img = smooth_blend(dst_rgb, swapped_img, face_mask)
 
-        # Post-processing: Apply slight blur for smoother blending
+        # Post-processing
         if blur_strength > 0:
-            swapped_img = cv2.GaussianBlur(swapped_img, (blur_strength, blur_strength), 0)
+            blended_img = cv2.GaussianBlur(blended_img, (blur_strength, blur_strength), 0)
 
-        # Optional: Apply lightweight sharpening
         if sharpen:
-            swapped_img = sharpen_image(swapped_img)
+            kernel = np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]])
+            blended_img = cv2.filter2D(blended_img, -1, kernel)
+            blended_img = np.clip(blended_img, 0, 255).astype(np.uint8)
 
-        # Convert back to BGR for output
-        result = cv2.cvtColor(swapped_img, cv2.COLOR_RGB2BGR)
+        result = cv2.cvtColor(blended_img, cv2.COLOR_RGB2BGR)
         return result
     except Exception as e:
         print(f"Error: {str(e)}")
@@ -75,7 +90,7 @@ demo = gr.Interface(
     inputs=[
         gr.Image(label="Source Face", type="numpy"),
         gr.Image(label="Target Image", type="numpy"),
-        gr.Slider(label="Blur Strength", minimum=0, maximum=5, step=1, value=3),
+        gr.Slider(label="Blur Strength", minimum=0, maximum=7, step=1, value=5),
         gr.Checkbox(label="Enable Sharpening", value=False)
     ],
     outputs=gr.Image(label="Face Swapped Output"),