Update app.py

app.py

@@ -1,70 +1,105 @@
-import gradio as gr
 import os
 import cv2
-import torch
 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 tempfile
 
-from gfpgan import GFPGANer
-from insightface.app import FaceAnalysis
-from insightface.model_zoo import model_zoo
-
-# Paths
-MODEL_DIR = "models"
-GFPGAN_MODEL_PATH = os.path.join(MODEL_DIR, "GFPGANv1.4.pth")
-INSWAPPER_PATH = os.path.join(MODEL_DIR, "inswapper_128.onnx")
-
-# Initialize GFPGAN
-gfpganer = GFPGANer(
-    model_path=GFPGAN_MODEL_PATH,
-    upscale=2,
-    arch='clean',
-    channel_multiplier=2,
-    bg_upsampler=None
-)
-
-# Initialize InsightFace (for face detection + swapping)
-faceapp = FaceAnalysis(name='buffalo_l', providers=['CPUExecutionProvider'])
-faceapp.prepare(ctx_id=0)
-
-swapper = model_zoo.get_model(INSWAPPER_PATH, providers=['CPUExecutionProvider'])
-
-# Function: Face Restoration + Swapping
-def restore_and_swap(input_img: Image.Image, target_img: Image.Image):
-    input_np = np.array(input_img.convert("RGB"))
-    target_np = np.array(target_img.convert("RGB"))
-
-    # Step 1: Face restore
-    cropped_faces, restored_img = gfpganer.enhance(input_np, has_aligned=False, only_center_face=False, paste_back=True)
-
-    # Step 2: Detect face in target
-    target_faces = faceapp.get(target_np)
-    if len(target_faces) == 0:
-        return Image.fromarray(restored_img), "No face found in target image."
+# Load InsightFace components
+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)
 
-    # Step 3: Detect face in restored image
-    source_faces = faceapp.get(restored_img)
+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):
+    if source_img is None or target_img is None:
+        raise gr.Error("Please upload both source and target images.")
+
+    source_np = cv2.cvtColor(np.array(source_img), cv2.COLOR_RGB2BGR)
+    target_np = cv2.cvtColor(np.array(target_img), cv2.COLOR_RGB2BGR)
+
+    source_faces = face_analyzer.get(source_np)
     if len(source_faces) == 0:
-        return Image.fromarray(restored_img), "No face found in restored image."
-
-    # Step 4: Face swap using first face
-    swapped = swapper.get(restored_img, source_faces[0], target_faces[0], paste_back=True)
-
-    return Image.fromarray(swapped), "Success! Face restored and swapped."
-
-# Gradio Interface
-iface = gr.Interface(
-    fn=restore_and_swap,
-    inputs=[
-        gr.Image(type="pil", label="Input Image (to restore + swap)"),
-        gr.Image(type="pil", label="Target Image (face to apply)")
-    ],
-    outputs=[
-        gr.Image(type="pil", label="Output Image"),
-        gr.Textbox(label="Status")
-    ],
-    title="GFPGAN + InsightFace Face Swapper",
-    description="This app restores old faces using GFPGAN and swaps them with another face using InsightFace InSwapper."
-)
-
-iface.launch()
+        raise gr.Error("No face found in source image.")
+
+    target_faces = face_analyzer.get(target_np)
+    if len(target_faces) == 0:
+        raise gr.Error("No faces found in target image.")
+
+    face_index = int(face_index_str)
+    if face_index >= len(target_faces):
+        raise gr.Error("Face index out of range.")
+
+    # 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
+        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!"
+
+    except Exception as e:
+        raise gr.Error(f"Face swap failed: {e}")
+
+# Gradio UI
+with gr.Blocks(title="Enhanced Face Swap") as demo:
+    gr.Markdown("## 🔁 Enhanced Face Swapper (CPU-only)")
+
+    with gr.Row():
+        source_image = gr.Image(label="Source Face", type="pil")
+        target_image = 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")
+
+    with gr.Row():
+        output_image = gr.Image(label="Swapped Result")
+        download_file = gr.File(label="Download Swapped Image")
+
+    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])
+
+if __name__ == "__main__":
+    demo.launch()