Update app.py

app.py

@@ -5,45 +5,62 @@ from PIL import Image
 import insightface
 from insightface.app import FaceAnalysis
 
-# Load face detector and swapper
+# Initialize face analysis and swapper model
 app = FaceAnalysis(name='buffalo_l')
 app.prepare(ctx_id=0, det_size=(640, 640))
+
 swapper = insightface.model_zoo.get_model('inswapper_128.onnx')
 
 def to_numpy(img):
+    """Convert PIL image to NumPy array in RGB."""
     if isinstance(img, Image.Image):
         return np.array(img.convert("RGB"))
     return img
 
-def swap_faces(source_img, dest_img, source_index=1, dest_index=1):
+def enhance_image_colors(image_np):
+    """Apply slight contrast and brightness enhancements."""
+    lab = cv2.cvtColor(image_np, cv2.COLOR_RGB2LAB)
+    l, a, b = cv2.split(lab)
+    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
+    l_clahe = clahe.apply(l)
+    lab_enhanced = cv2.merge((l_clahe, a, b))
+    enhanced = cv2.cvtColor(lab_enhanced, cv2.COLOR_LAB2RGB)
+    return enhanced
+
+def swap_faces(source_img, target_img, source_index=1, target_index=1):
+    # Convert images to NumPy
     source_np = to_numpy(source_img)
-    dest_np = to_numpy(dest_img)
+    target_np = to_numpy(target_img)
 
     # Detect faces
     source_faces = sorted(app.get(source_np), key=lambda x: x.bbox[0])
-    dest_faces = sorted(app.get(dest_np), key=lambda x: x.bbox[0])
+    target_faces = sorted(app.get(target_np), key=lambda x: x.bbox[0])
 
-    if len(source_faces) < source_index or len(dest_faces) < dest_index:
+    if len(source_faces) < source_index or len(target_faces) < target_index:
         raise gr.Error("Face index out of range!")
 
     source_face = source_faces[source_index - 1]
-    dest_face = dest_faces[dest_index - 1]
+    target_face = target_faces[target_index - 1]
+
+    # Perform face swap
+    result_bgr = swapper.get(target_np, target_face, source_face, paste_back=True)
+    result_rgb = cv2.cvtColor(result_bgr, cv2.COLOR_BGR2RGB)
+
+    # Enhance image for better visual quality
+    enhanced_result = enhance_image_colors(result_rgb)
 
-    # Swap
-    result = swapper.get(dest_np, dest_face, source_face, paste_back=True)
-    result_rgb = cv2.cvtColor(result, cv2.COLOR_BGR2RGB)
-    return Image.fromarray(result_rgb)
+    return Image.fromarray(enhanced_result)
 
-# Gradio app
+# Gradio Interface
 gr.Interface(
     fn=swap_faces,
     inputs=[
         gr.Image(label="Source Image"),
         gr.Image(label="Target Image"),
         gr.Number(value=1, label="Source Face Index"),
-        gr.Number(value=1, label="Target Face Index"),
+        gr.Number(value=1, label="Target Face Index")
     ],
     outputs=gr.Image(label="Swapped Output"),
-    title="Face Swap with InsightFace",
-    description="Swap faces using InsightFace and ONNX model (inswapper_128.onnx)."
+    title="✨ Face Swap with InsightFace",
+    description="Upload two images and swap faces using InsightFace (inswapper_128.onnx). Adjust face indices if needed."
 ).launch()