Update app.py

app.py

@@ -1,3 +1,4 @@
+import os
 import cv2
 import numpy as np
 import gradio as gr
@@ -5,169 +6,284 @@ from insightface.app import FaceAnalysis
 from insightface.model_zoo import get_model
 from PIL import Image
 import tempfile
-import os
+import logging
+
+# --- Configuration & Setup ---
+# Configure logging for better debugging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+
+# Constants for model paths and settings
+SWAPPER_MODEL_PATH = "models/inswapper_128.onnx" # Consider making this a script argument or env variable
+FACE_ANALYZER_NAME = 'buffalo_l'
+DETECTION_SIZE = (640, 640)
+CPU_PROVIDERS = ['CPUExecutionProvider'] # Or ['CUDAExecutionProvider', 'CPUExecutionProvider'] if GPU is available
 
-# Initialize models
-face_analyzer = FaceAnalysis(
-    name='buffalo_l',
-    providers=['CPUExecutionProvider'],
-    root=os.environ.get('MODEL_DIR', '.')
-)
-face_analyzer.prepare(ctx_id=0, det_size=(640, 640))
-
-# Load SimSwap model (256x256 version)
-simswap_model = get_model(
-    'models/simswap_256.onnx',
-    download=False,
-    download_zip=False
-)
-
-def get_faces(img):
-    """Detect faces in image with enhanced error handling"""
+# --- Global Variables (Lazy Loaded) ---
+face_analyzer = None
+swapper = None
+
+# --- Initialization Functions ---
+def initialize_models():
+    """Initialize the face analyzer and swapper models."""
+    global face_analyzer, swapper
     try:
-        return face_analyzer.get(img)
+        if face_analyzer is None:
+            logging.info(f"Initializing FaceAnalysis model: {FACE_ANALYZER_NAME}")
+            face_analyzer = FaceAnalysis(name=FACE_ANALYZER_NAME, providers=CPU_PROVIDERS)
+            face_analyzer.prepare(ctx_id=0, det_size=DETECTION_SIZE)
+            logging.info("FaceAnalysis model initialized.")
+
+        if swapper is None:
+            if not os.path.exists(SWAPPER_MODEL_PATH):
+                logging.error(f"Swapper model not found at {SWAPPER_MODEL_PATH}. Please download it.")
+                # You might want to raise an exception here or try to download it
+                # For now, let's assume get_model can handle download if 'download=True' was intended
+                # but the original code had 'download=False'
+                raise FileNotFoundError(f"Swapper model not found: {SWAPPER_MODEL_PATH}")
+            logging.info(f"Loading swapper model from: {SWAPPER_MODEL_PATH}")
+            swapper = get_model(SWAPPER_MODEL_PATH, download=False) # Set download=True if you want to auto-download
+            logging.info("Swapper model loaded.")
     except Exception as e:
-        raise gr.Error(f"Face detection failed: {str(e)}") from e
+        logging.error(f"Error during model initialization: {e}")
+        raise # Re-raise the exception to stop the app if models can't load
+
+# Call initialization at the start
+initialize_models()
 
-def draw_faces(img, faces):
-    """Draw face bounding boxes with improved visualization"""
-    img_with_boxes = img.copy()
+# --- Core Functions ---
+def get_faces_from_image(img_np: np.ndarray):
+    """
+    Detects faces in a NumPy image array (BGR format).
+
+    Args:
+        img_np: NumPy array representing the image (BGR).
+
+    Returns:
+        A list of face objects detected by FaceAnalysis.
+    """
+    if face_analyzer is None:
+        raise gr.Error("Face analyzer not initialized. Please check logs.")
+    faces = face_analyzer.get(img_np)
+    return faces
+
+def draw_detected_faces(img_np: np.ndarray, faces: list):
+    """
+    Draws bounding boxes and labels on faces in an image.
+
+    Args:
+        img_np: NumPy array representing the image (BGR).
+        faces: A list of face objects.
+
+    Returns:
+        NumPy array with faces drawn (BGR).
+    """
+    img_with_boxes = img_np.copy()
     for i, face in enumerate(faces):
         box = face.bbox.astype(int)
-        color = (0, 255, 0)  # Green
-        thickness = max(1, int(min(box[2]-box[0], box[3]-box[1]) / 150))
-        
-        # Draw bounding box
-        cv2.rectangle(img_with_boxes, 
-                     (box[0], box[1]), 
-                     (box[2], box[3]), 
-                     color, thickness)
-        
-        # Draw label
-        label = f"Face {i}"
-        font_scale = max(0.5, min(box[2]-box[0], box[3]-box[1]) / 1000)
-        cv2.putText(img_with_boxes, label, 
-                   (box[0], box[1] - 10), 
-                   cv2.FONT_HERSHEY_SIMPLEX,
-                   font_scale, color, thickness)
+        # Ensure coordinates are within image bounds
+        x1, y1, x2, y2 = box[0], box[1], box[2], box[3]
+        cv2.rectangle(img_with_boxes, (x1, y1), (x2, y2), (0, 255, 0), 2)
+        label_position = (x1, max(0, y1 - 10)) # Ensure label is not drawn outside top
+        cv2.putText(img_with_boxes, f"Face {i}", label_position,
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (36, 255, 12), 2) # Changed color for visibility
     return img_with_boxes
 
-def process_image(img):
-    """Convert PIL/numpy to BGR numpy array"""
-    if isinstance(img, np.ndarray):
-        return img[:, :, ::-1] if img.shape[2] == 3 else img
-    return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
-
-def face_swap(source_img, target_img, source_index, target_index):
-    """Perform face swap using SimSwap 256 model"""
-    # Convert inputs
-    source_np = process_image(source_img)
-    target_np = process_image(target_img)
-    
-    # Detect faces
-    source_faces = get_faces(source_np)
-    target_faces = get_faces(target_np)
-    
-    # Validate selections
+def convert_pil_to_cv2(pil_image: Image.Image) -> np.ndarray:
+    """Converts a PIL Image to an OpenCV NumPy array (BGR)."""
+    return cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
+
+def convert_cv2_to_pil(cv2_image: np.ndarray) -> Image.Image:
+    """Converts an OpenCV NumPy array (BGR) to a PIL Image."""
+    return Image.fromarray(cv2.cvtColor(cv2_image, cv2.COLOR_BGR2RGB))
+
+# --- Gradio Interface Functions ---
+def process_face_swap(source_pil_img: Image.Image, target_pil_img: Image.Image, target_face_index: int):
+    """
+    Performs the face swap operation.
+
+    Args:
+        source_pil_img: PIL Image of the source face.
+        target_pil_img: PIL Image of the target scene.
+        target_face_index: Index of the face in the target image to be swapped.
+
+    Returns:
+        Tuple: (PIL Image of the swapped result, path to temporary file for download)
+               or raises gr.Error on failure.
+    """
+    if swapper is None:
+        raise gr.Error("Swapper model not initialized. Please check logs.")
+    if source_pil_img is None:
+        raise gr.Error("Source image not provided.")
+    if target_pil_img is None:
+        raise gr.Error("Target image not provided.")
+
+    source_np = convert_pil_to_cv2(source_pil_img)
+    target_np = convert_pil_to_cv2(target_pil_img)
+
+    # Get face from source image
+    source_faces = get_faces_from_image(source_np)
     if not source_faces:
-        raise gr.Error("No faces found in source image!")
+        raise gr.Error("No face found in the source image. Please use a clear image of a face.")
+    source_face = source_faces[0] # Assuming the first detected face is the one to use
+
+    # Get faces from target image
+    target_faces = get_faces_from_image(target_np)
     if not target_faces:
-        raise gr.Error("No faces found in target image!")
-    if source_index >= len(source_faces):
-        raise gr.Error(f"Source face index {source_index} out of range (max {len(source_faces)-1})")
-    if target_index >= len(target_faces):
-        raise gr.Error(f"Target face index {target_index} out of range (max {len(target_faces)-1})")
-
-    # Perform swap with SimSwap
-    result = simswap_model.get(
-        target_np,
-        target_faces[target_index],
-        source_faces[source_index],
-        paste_back=True
-    )
-    
-    # Convert and save
-    result_rgb = cv2.cvtColor(result, cv2.COLOR_BGR2RGB)
-    with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp:
-        Image.fromarray(result_rgb).save(tmp.name, quality=95)
-        return result_rgb, tmp.name
-
-def update_faces_preview(img, image_type):
-    """Update face preview with selection options"""
-    if img is None:
-        return None, 0, gr.Slider(visible=False)
-    
-    img_np = process_image(img)
-    faces = get_faces(img_np)
-    
-    if not faces:
-        return img, 0, gr.Slider(visible=False)
-    
-    preview_img = draw_faces(img_np, faces)
-    preview_rgb = cv2.cvtColor(preview_img, cv2.COLOR_BGR2RGB)
-    
-    return (
-        preview_rgb,
-        len(faces),
-        gr.Slider(
-            maximum=max(0, len(faces)-1),
-            visible=len(faces) > 0,
-            label=f"Select {image_type} Face Index"
-        )
+        raise gr.Error("No faces found in the target image.")
+    if not (0 <= target_face_index < len(target_faces)):
+        # This case should ideally be prevented by the slider's dynamic range
+        raise gr.Error(f"Selected face index ({target_face_index}) is out of range. "
+                       f"Detected {len(target_faces)} faces (indices 0 to {len(target_faces)-1}).")
+
+    try:
+        logging.info(f"Swapping face from source to target face index {target_face_index}.")
+        # Ensure target_face_index is an integer for indexing
+        target_face_to_swap = target_faces[int(target_face_index)]
+        swapped_bgr_img = swapper.get(target_np, target_face_to_swap, source_face, paste_back=True)
+    except Exception as e:
+        logging.error(f"Error during face swapping: {e}")
+        raise gr.Error(f"An error occurred during the swap process: {str(e)}")
+
+    swapped_pil_img = convert_cv2_to_pil(swapped_bgr_img)
+
+    # Save to a temporary file for download
+    try:
+        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp_file:
+            swapped_pil_img.save(tmp_file.name)
+            temp_file_path = tmp_file.name
+        logging.info(f"Swapped image saved to temporary file: {temp_file_path}")
+    except Exception as e:
+        logging.error(f"Error saving to temporary file: {e}")
+        raise gr.Error("Could not save the swapped image for download.")
+
+    return swapped_pil_img, temp_file_path
+
+def preview_target_faces(target_pil_img: Image.Image):
+    """
+    Updates the preview of detected faces in the target image and adjusts the slider.
+
+    Args:
+        target_pil_img: PIL Image of the target scene.
+
+    Returns:
+        Tuple: (PIL Image with detected faces, Gradio Slider update)
+    """
+    if target_pil_img is None:
+        # Return a blank image and default slider if no image is provided
+        blank_image_pil = Image.new('RGB', (DETECTION_SIZE[0], DETECTION_SIZE[1]), color = 'lightgray')
+        return blank_image_pil, gr.Slider(minimum=0, maximum=0, value=0, step=1, interactive=False)
+
+    target_np = convert_pil_to_cv2(target_pil_img)
+    faces = get_faces_from_image(target_np)
+
+    preview_np_img = draw_detected_faces(target_np, faces)
+    preview_pil_img = convert_cv2_to_pil(preview_np_img)
+
+    num_faces = len(faces)
+    if num_faces > 0:
+        # Update slider: max index is num_faces - 1
+        slider_update = gr.Slider(minimum=0, maximum=num_faces - 1, value=0, step=1, interactive=True)
+    else:
+        # No faces, disable slider
+        slider_update = gr.Slider(minimum=0, maximum=0, value=0, step=1, interactive=False)
+
+    return preview_pil_img, slider_update
+
+# --- Gradio UI Definition ---
+with gr.Blocks(title="Face Swap Pro 🔄", theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        """
+        # 🎭 Face Swap Pro 🚀
+        Upload a source image with the face you want to use, and a target image where you want to swap a face.
+        Use the 'Preview Detected Faces' button to see faces in the target image and select which one to replace.
+        """
     )
 
-with gr.Blocks(title="SimSwap Pro", css=".gradio-container {background-color: #f0f2f6}") as demo:
-    gr.Markdown("""
-    # 🔄 SimSwap Face Swapper (256x256)
-    *Higher quality face swapping using SimSwap 256 model*
-    """)
-    
     with gr.Row():
-        with gr.Column():
-            gr.Markdown("## Source Face")
-            source_img = gr.Image(label="Source Image", type="pil")
-            source_preview_btn = gr.Button("Preview Source Faces")
-            source_preview = gr.Image(label="Detected Faces", interactive=False)
-            source_index = gr.Slider(visible=False, step=1)
-        
-        with gr.Column():
-            gr.Markdown("## Target Image")
-            target_img = gr.Image(label="Target Image", type="pil")
-            target_preview_btn = gr.Button("Preview Target Faces")
-            target_preview = gr.Image(label="Detected Faces", interactive=False)
-            target_index = gr.Slider(visible=False, step=1)
-    
-    swap_btn = gr.Button("✨ Swap Faces", variant="primary")
-    result_title = gr.Markdown("## 🔮 Swapped Result", visible=False)
-    
+        with gr.Column(scale=1):
+            source_image_input = gr.Image(label="👤 Source Face Image", type="pil", sources=["upload", "clipboard"])
+        with gr.Column(scale=1):
+            target_image_input = gr.Image(label="🖼️ Target Scene Image", type="pil", sources=["upload", "clipboard"])
+
+    with gr.Row():
+        preview_button = gr.Button("🔍 Preview Detected Faces in Target", variant="secondary")
+        face_index_slider = gr.Slider(
+            label="🎯 Select Target Face Index (0-indexed)",
+            minimum=0,
+            maximum=0, # Will be updated dynamically
+            step=1,
+            value=0,
+            interactive=False # Initially not interactive
+        )
+
+    target_faces_preview_output = gr.Image(label="👀 Detected Faces in Target", interactive=False)
+
+    gr.HTML("<hr>") # Visual separator
+
+    swap_button = gr.Button("🔁 SWAP FACES NOW!", variant="primary")
+
     with gr.Row():
-        output_img = gr.Image(label="Output", interactive=False)
-        download_btn = gr.File(label="Download Result", visible=False)
-    
-    # Event handlers
-    source_preview_btn.click(
-        fn=lambda x: update_faces_preview(x, "Source"),
-        inputs=source_img,
-        outputs=[source_preview, source_index, source_index]
+        swapped_image_output = gr.Image(label="✨ Swapped Result", interactive=False)
+        download_output_file = gr.File(label="⬇️ Download Swapped Image")
+
+    # --- Event Handlers ---
+    def on_target_image_change(target_img):
+        """Called when the target image is uploaded or cleared."""
+        if target_img is None:
+            # Reset preview and slider if target image is cleared
+            blank_image_pil = Image.new('RGB', (DETECTION_SIZE[0], DETECTION_SIZE[1]), color = 'lightgray')
+            return blank_image_pil, gr.Slider(minimum=0, maximum=0, value=0, step=1, interactive=False)
+        # If an image is uploaded, the preview button click will handle the update.
+        # This function primarily handles the reset case.
+        # Or, you could auto-trigger preview here:
+        # return preview_target_faces(target_img)
+        # For now, let's keep it explicit with the button.
+        return target_faces_preview_output.value, face_index_slider.value # Return current values
+
+
+    # Connect target image change to potentially reset/update UI elements
+    # (e.g., if user clears the target image)
+    target_image_input.change(
+        fn=on_target_image_change,
+        inputs=[target_image_input],
+        outputs=[target_faces_preview_output, face_index_slider]
+    )
+
+    preview_button.click(
+        fn=preview_target_faces,
+        inputs=[target_image_input],
+        outputs=[target_faces_preview_output, face_index_slider]
     )
-    
-    target_preview_btn.click(
-        fn=lambda x: update_faces_preview(x, "Target"),
-        inputs=target_img,
-        outputs=[target_preview, target_index, target_index]
+
+    swap_button.click(
+        fn=process_face_swap,
+        inputs=[source_image_input, target_image_input, face_index_slider],
+        outputs=[swapped_image_output, download_output_file]
     )
-    
-    swap_btn.click(
-        fn=face_swap,
-        inputs=[source_img, target_img, source_index, target_index],
-        outputs=[output_img, download_btn]
-    ).then(
-        fn=lambda: [
-            gr.Markdown(visible=True),
-            gr.File(visible=True)
+
+    # --- Examples ---
+    gr.Examples(
+        examples=[
+            ["examples/source_face.jpg", "examples/target_group.jpg", 0], # Create these example files
+            ["examples/source_actor.png", "examples/target_scene.png", 1]
         ],
-        outputs=[result_title, download_btn]
+        inputs=[source_image_input, target_image_input, face_index_slider],
+        outputs=[swapped_image_output, download_output_file],
+        fn=process_face_swap,
+        cache_examples=False, # Set to True if your examples are static and processing is slow
+        label="Example Face Swaps"
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    # Ensure 'models' and 'examples' directories exist or handle their absence
+    os.makedirs("models", exist_ok=True)
+    os.makedirs("examples", exist_ok=True)
+    # You'd typically place your 'inswapper_128.onnx' in the 'models' directory
+    # and example images in the 'examples' directory.
+
+    # Check if the swapper model exists before launching
+    if not os.path.exists(SWAPPER_MODEL_PATH):
+        print(f"ERROR: Swapper model not found at {SWAPPER_MODEL_PATH}")
+        print("Please download the 'inswapper_128.onnx' model and place it in the 'models' directory.")
+    else:
+        demo.launch()