Update app.py

app.py

@@ -1,42 +1,47 @@
-import os
+# app.py
+
+import gradio as gr
 import cv2
 import numpy as np
-import gradio as gr
-from insightface.app import FaceAnalysis
-from insightface.model_zoo import get_model
 from PIL import Image
+import os
 import tempfile
-
-# Load InsightFace components
+from insightface.app import FaceAnalysis
+from insightface.model_zoo.inswapper import INSwapper
+from skimage.exposure import match_histograms
+import torch
+from basicsr.archs.rrdbnet_arch import RRDBNet
+from realesrgan import RealESRGANer
+from gfpgan import GFPGANer
+
+# Setup face analyzer
 face_analyzer = FaceAnalysis(name='buffalo_l', providers=['CPUExecutionProvider'])
 face_analyzer.prepare(ctx_id=0, det_size=(640, 640))
 
-swapper_model_path = "models/inswapper_128.onnx"
-swapper = get_model(swapper_model_path, download=False)
-
-def draw_faces(image_np, faces):
-    """Draw bounding boxes and labels on detected faces."""
-    image_draw = image_np.copy()
-    for i, face in enumerate(faces):
-        box = face.bbox.astype(int)
-        cv2.rectangle(image_draw, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 2)
-        cv2.putText(image_draw, f"Face {i}", (box[0], box[1] - 10),
-                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
-    return image_draw
-
-def preview_faces(target_image):
-    """Preview all detected faces in target image."""
-    if target_image is None:
-        raise gr.Error("Please upload a target image.")
-    target_np = cv2.cvtColor(np.array(target_image), cv2.COLOR_RGB2BGR)
-    faces = face_analyzer.get(target_np)
-    if len(faces) == 0:
-        raise gr.Error("No faces found in target image.")
-    preview_img = draw_faces(target_np, faces)
-    preview_img_rgb = cv2.cvtColor(preview_img, cv2.COLOR_BGR2RGB)
-    return Image.fromarray(preview_img_rgb), [str(i) for i in range(len(faces))]
-
-def swap_faces(source_img, target_img, face_index_str, blend_alpha):
+# Load face swapper
+swapper = INSwapper("models/inswapper_128.onnx")
+
+# Load GFPGAN and Real-ESRGAN
+bg_upsampler = RealESRGANer(
+    scale=2,
+    model_path='models/RealESRGAN_x2plus.pth',
+    model=RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64,
+                  num_block=23, num_grow_ch=32, scale=2),
+    tile=400, tile_pad=10, pre_pad=0, half=False
+)
+
+face_enhancer = GFPGANer(
+    model_path='models/GFPGANv1.3.pth',
+    upscale=2,
+    arch='clean',
+    channel_multiplier=2,
+    bg_upsampler=bg_upsampler
+)
+
+def match_histogram(source_face, target_face):
+    return match_histograms(source_face, target_face, channel_axis=-1)
+
+def swap_faces(source_img, target_img, face_index_str, blend_alpha, enhance_face):
     if source_img is None or target_img is None:
         raise gr.Error("Please upload both source and target images.")
 
@@ -44,62 +49,81 @@ def swap_faces(source_img, target_img, face_index_str, blend_alpha):
     target_np = cv2.cvtColor(np.array(target_img), cv2.COLOR_RGB2BGR)
 
     source_faces = face_analyzer.get(source_np)
-    if len(source_faces) == 0:
-        raise gr.Error("No face found in source image.")
-
     target_faces = face_analyzer.get(target_np)
+
+    if len(source_faces) == 0:
+        raise gr.Error("No face detected in source image.")
     if len(target_faces) == 0:
-        raise gr.Error("No faces found in target image.")
+        raise gr.Error("No faces detected in target image.")
 
     face_index = int(face_index_str)
     if face_index >= len(target_faces):
-        raise gr.Error("Face index out of range.")
+        raise gr.Error(f"Face index {face_index} out of range (found {len(target_faces)} faces).")
 
-    # Perform face swap
     try:
         swapped = swapper.get(target_np, target_faces[face_index], source_faces[0], paste_back=True)
-        # Blend original + swapped if alpha < 1.0
-        if 0 <= blend_alpha < 1.0:
-            blended = cv2.addWeighted(target_np, 1 - blend_alpha, swapped, blend_alpha, 0)
-        else:
-            blended = swapped
+
+        box = target_faces[face_index].bbox.astype(int)
+        x1, y1, x2, y2 = box
+        original_face = target_np[y1:y2, x1:x2]
+        swapped_face = swapped[y1:y2, x1:x2]
+        matched = match_histogram(swapped_face, original_face)
+        swapped[y1:y2, x1:x2] = matched.astype(np.uint8)
+
+        mask = np.zeros_like(swapped[:, :, 0])
+        cv2.rectangle(mask, (x1, y1), (x2, y2), 255, -1)
+        mask = cv2.GaussianBlur(mask, (15, 15), 0)
+        blended = np.where(mask[..., None] > 0,
+                           cv2.addWeighted(swapped, blend_alpha, target_np, 1 - blend_alpha, 0),
+                           target_np)
+
+        if enhance_face:
+            _, _, enhanced = face_enhancer.enhance(blended, has_aligned=False, only_center_face=False, paste_back=True)
+            blended = enhanced
+
         final_rgb = cv2.cvtColor(blended, cv2.COLOR_BGR2RGB)
         output_img = Image.fromarray(final_rgb)
 
-        # Save temp file for download
         tmp_file = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
         output_img.save(tmp_file.name)
-        return output_img, tmp_file.name, "✅ Face swap successful!"
-
+        return output_img, tmp_file.name, "✅ Swap & Enhancement Completed!"
     except Exception as e:
-        raise gr.Error(f"Face swap failed: {e}")
+        raise gr.Error(f"Swap failed: {str(e)}")
+
+def count_faces(image):
+    if image is None:
+        return "Please upload a target image."
+    img_np = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+    faces = face_analyzer.get(img_np)
+    return f"{len(faces)} face(s) detected in the target image."
 
 # Gradio UI
-with gr.Blocks(title="Enhanced Face Swap") as demo:
-    gr.Markdown("## 🔁 Enhanced Face Swapper (CPU-only)")
+with gr.Blocks(title="💫 Face Swap & Enhance App") as demo:
+    gr.Markdown("""
+    # 🔄 Face Swapper + 👩‍🔧 GFPGAN Enhancer + 🧙 Real-ESRGAN
+    Upload a face and swap it into another photo. Optional face enhancement with GFPGAN!
+    """)
 
     with gr.Row():
-        source_image = gr.Image(label="Source Face", type="pil")
-        target_image = gr.Image(label="Target Image", type="pil")
+        source_input = gr.Image(label="🧑 Source Face", type="pil")
+        target_input = gr.Image(label="🧑 Target Image", type="pil")
 
-    preview_btn = gr.Button("👁️ Preview Detected Faces")
-    preview_output = gr.Image(label="Target Faces Preview")
-    face_dropdown = gr.Dropdown(choices=[], label="Select Face Index")
-    blend_slider = gr.Slider(label="Swap Strength", minimum=0.0, maximum=1.0, step=0.1, value=1.0)
-    swap_button = gr.Button("🔄 Swap Face")
+    detect_button = gr.Button("🔍 Detect Faces in Target")
+    face_count = gr.Textbox(label="Detected Faces", interactive=False)
+    detect_button.click(count_faces, inputs=target_input, outputs=face_count)
 
     with gr.Row():
-        output_image = gr.Image(label="Swapped Result")
-        download_file = gr.File(label="Download Swapped Image")
+        face_index = gr.Textbox(label="🎯 Face Index", value="0")
+        blend_slider = gr.Slider(label="💫 Blend Smoothness", minimum=0, maximum=1, value=0.8, step=0.05)
+        enhance_face = gr.Checkbox(label="✨ Enhance Face with GFPGAN", value=True)
 
+    swap_button = gr.Button("🔁 Run Face Swap")
+    output_img = gr.Image(label="🖼️ Output")
+    download_path = gr.File(label="📥 Download")
     status = gr.Textbox(label="Status", interactive=False)
 
-    preview_btn.click(preview_faces, inputs=target_image,
-                      outputs=[preview_output, face_dropdown])
-
-    swap_button.click(swap_faces,
-                      inputs=[source_image, target_image, face_dropdown, blend_slider],
-                      outputs=[output_image, download_file, status])
+    swap_button.click(swap_faces, inputs=[source_input, target_input, face_index, blend_slider, enhance_face],
+                      outputs=[output_img, download_path, status])
 
-if __name__ == "__main__":
-    demo.launch()
+if __name__ == '__main__':
+    demo.launch()