Update app.py

app.py

@@ -1,37 +1,59 @@
-# -*- coding: UTF-8 -*-
 #!/usr/bin/env python
+# -*- coding: UTF-8 -*-
 
 import os
 import json
 import shutil
+import logging
+import tempfile
 from datetime import datetime
-from dotenv import load_dotenv
+from typing import Tuple, Optional
+
 import numpy as np
 import cv2
 from PIL import Image
 import gradio as gr
+from dotenv import load_dotenv
 from huggingface_hub import HfApi, login
-from roop.globals import (
+from insightface.app import FaceAnalysis
+
+import roop.globals
+from roop.core import (
     start,
     decode_execution_providers,
     suggest_max_memory,
     suggest_execution_threads,
 )
-from roop.core import normalize_output_path
 from roop.processors.frame.core import get_frame_processors_modules
-from insightface.app import FaceAnalysis
+from roop.utilities import normalize_output_path
+
+# Configure logging
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 
 # Load environment variables
 load_dotenv()
 
-# Cosine similarity function
-def cosine_similarity(a, b):
+
+def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
+    """
+    Calculate the cosine similarity between two vectors.
+    
+    Parameters:
+        a (np.ndarray): First vector.
+        b (np.ndarray): Second vector.
+        
+    Returns:
+        float: Cosine similarity.
+    """
     return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-6)
 
-# Dataset handler class
+
 class FaceIntegrDataset:
-    def __init__(self, repo_id="Arrcttacsrks/face_integrData"):
-        self.token = os.getenv("hf_token")
+    """
+    Handler for face integration dataset upload to Hugging Face.
+    """
+    def __init__(self, repo_id: str = "Arrcttacsrks/face_integrData") -> None:
+        self.token = os.getenv('hf_token')
         if not self.token:
             raise ValueError("HF_TOKEN environment variable is not set")
         self.repo_id = repo_id
@@ -40,208 +62,287 @@ class FaceIntegrDataset:
         self.temp_dir = "temp_dataset"
         os.makedirs(self.temp_dir, exist_ok=True)
 
-    def create_date_folder(self):
+    def create_date_folder(self) -> Tuple[str, str]:
+        """
+        Create a folder based on the current date inside the temporary directory.
+        
+        Returns:
+            Tuple[str, str]: The folder path and the current date string.
+        """
         current_date = datetime.now().strftime("%Y-%m-%d")
         folder_path = os.path.join(self.temp_dir, current_date)
         os.makedirs(folder_path, exist_ok=True)
         return folder_path, current_date
 
-    def save_metadata(self, source_path, target_path, output_path, timestamp):
+    def save_metadata(self, source_path: str, target_path: str, output_path: str, timestamp: str) -> dict:
+        """
+        Create metadata dictionary for the face swap process.
+        
+        Parameters:
+            source_path (str): Filename of the source image.
+            target_path (str): Filename of the target image.
+            output_path (str): Filename of the output image.
+            timestamp (str): Timestamp string.
+            
+        Returns:
+            dict: Metadata information.
+        """
         metadata = {
             "timestamp": timestamp,
             "source_image": source_path,
             "target_image": target_path,
             "output_image": output_path,
-            "date_created": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+            "date_created": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
         }
         return metadata
 
-    def upload_to_hf(self, local_folder, date_folder):
+    def upload_to_hf(self, local_folder: str, date_folder: str) -> bool:
+        """
+        Upload a local folder to the Hugging Face dataset repository.
+        
+        Parameters:
+            local_folder (str): The local folder path.
+            date_folder (str): The subfolder in the repository.
+            
+        Returns:
+            bool: True if upload is successful, False otherwise.
+        """
         try:
             self.api.upload_folder(
                 folder_path=local_folder,
                 repo_id=self.repo_id,
                 repo_type="dataset",
-                path_in_repo=date_folder,
+                path_in_repo=date_folder
             )
+            logging.info("Successfully uploaded files to Hugging Face repository.")
             return True
         except Exception as e:
-            print(f"Error uploading to Hugging Face: {str(e)}")
+            logging.error(f"Error uploading to Hugging Face: {str(e)}")
             return False
 
-# Image face swap function
-def swap_face(source_file, target_file, doFaceEnhancer):
-    dataset_handler = FaceIntegrDataset()
-    folder_path, date_folder = dataset_handler.create_date_folder()
-    timestamp = datetime.now().strftime("%S-%M-%H-%d-%m-%Y")
 
+def configure_roop_globals(source_path: str, target_path: str, output_path: str, do_face_enhancer: bool) -> None:
+    """
+    Configure global variables required for the face swap process.
+    
+    Parameters:
+        source_path (str): Path to the source image.
+        target_path (str): Path to the target image.
+        output_path (str): Path to save the output image.
+        do_face_enhancer (bool): Flag to determine if face enhancer should be used.
+    """
+    roop.globals.source_path = source_path
+    roop.globals.target_path = target_path
+    roop.globals.output_path = normalize_output_path(source_path, target_path, output_path)
+    roop.globals.frame_processors = ["face_swapper", "face_enhancer"] if do_face_enhancer else ["face_swapper"]
+    roop.globals.headless = True
+    roop.globals.keep_fps = True
+    roop.globals.keep_audio = True
+    roop.globals.keep_frames = False
+    roop.globals.many_faces = False
+    roop.globals.video_encoder = "libx264"
+    roop.globals.video_quality = 18
+    roop.globals.max_memory = suggest_max_memory()
+    roop.globals.execution_providers = decode_execution_providers(["cuda"])
+    roop.globals.execution_threads = suggest_execution_threads()
+
+
+def swap_face(source_file: np.ndarray, target_file: np.ndarray, doFaceEnhancer: bool) -> Optional[np.ndarray]:
+    """
+    Perform face swapping on static images.
+    
+    Parameters:
+        source_file (np.ndarray): Source image array.
+        target_file (np.ndarray): Target image array.
+        doFaceEnhancer (bool): Flag to apply face enhancer.
+        
+    Returns:
+        Optional[np.ndarray]: The output image array if successful, otherwise None.
+    """
+    folder_path = None
     try:
-        # Save source and target images
+        dataset_handler = FaceIntegrDataset()
+        folder_path, date_folder = dataset_handler.create_date_folder()
+        timestamp = datetime.now().strftime("%S-%M-%H-%d-%m-%Y")
         source_path = os.path.join(folder_path, f"source_{timestamp}.jpg")
         target_path = os.path.join(folder_path, f"target_{timestamp}.jpg")
-        output_path = os.path.join(folder_path, f"OutputImage_{timestamp}.jpg")
+        output_path = os.path.join(folder_path, f"OutputImage{timestamp}.jpg")
 
         if source_file is None or target_file is None:
             raise ValueError("Source and target images are required")
-
+            
         Image.fromarray(source_file).save(source_path)
         Image.fromarray(target_file).save(target_path)
-
-        # Configure Roop globals
-        roop.globals.source_path = source_path
-        roop.globals.target_path = target_path
-        roop.globals.output_path = normalize_output_path(
-            roop.globals.source_path, roop.globals.target_path, output_path
-        )
-
-        roop.globals.frame_processors = (
-            ["face_swapper", "face_enhancer"] if doFaceEnhancer else ["face_swapper"]
-        )
-        roop.globals.headless = True
-        roop.globals.keep_fps = True
-        roop.globals.keep_audio = True
-        roop.globals.keep_frames = False
-        roop.globals.many_faces = False
-        roop.globals.video_encoder = "libx264"
-        roop.globals.video_quality = 18
-        roop.globals.max_memory = suggest_max_memory()
-        roop.globals.execution_providers = decode_execution_providers(["cuda"])
-        roop.globals.execution_threads = suggest_execution_threads()
-
+        
+        logging.info(f"Source image saved at: {source_path}")
+        logging.info(f"Target image saved at: {target_path}")
+        
+        # Configure global parameters for roop
+        configure_roop_globals(source_path, target_path, output_path, doFaceEnhancer)
+        
         # Pre-check frame processors
         for frame_processor in get_frame_processors_modules(roop.globals.frame_processors):
             if not frame_processor.pre_check():
-                raise RuntimeError("Frame processor pre-check failed")
-
-        # Start face swap process
+                logging.error("Pre-check failed for frame processor.")
+                return None
+        
+        logging.info("Starting face swap process...")
         start()
-
-        # Save metadata
+        
         metadata = dataset_handler.save_metadata(
-            f"source_{timestamp}.jpg",
-            f"target_{timestamp}.jpg",
-            f"OutputImage_{timestamp}.jpg",
-            timestamp,
+            os.path.basename(source_path),
+            os.path.basename(target_path),
+            os.path.basename(output_path),
+            timestamp
         )
+        
         metadata_path = os.path.join(folder_path, f"metadata_{timestamp}.json")
-        with open(metadata_path, "w") as f:
+        with open(metadata_path, 'w') as f:
             json.dump(metadata, f, indent=4)
-
-        # Upload to Hugging Face
+        
         upload_success = dataset_handler.upload_to_hf(folder_path, date_folder)
-        if not upload_success:
-            print("Failed to upload files to Hugging Face dataset")
-
-        # Return output image
-        if os.path.exists(output_path):
-            output_image = Image.open(output_path)
-            return np.array(output_image)
+        if upload_success:
+            logging.info(f"Successfully uploaded files to dataset {dataset_handler.repo_id}")
         else:
-            raise FileNotFoundError("Output image not found")
-
+            logging.error("Failed to upload files to Hugging Face dataset")
+        
+        if os.path.exists(roop.globals.output_path):
+            output_image = Image.open(roop.globals.output_path)
+            output_array = np.array(output_image)
+            shutil.rmtree(folder_path, ignore_errors=True)
+            return output_array
+        else:
+            logging.error("Output image not found")
+            shutil.rmtree(folder_path, ignore_errors=True)
+            return None
+        
     except Exception as e:
-        print(f"Error in face swap process: {str(e)}")
-        raise gr.Error(f"Face swap failed: {str(e)}")
-    finally:
+        logging.exception(f"Error in face swap process: {str(e)}")
         if folder_path and os.path.exists(folder_path):
-            shutil.rmtree(folder_path)
+            shutil.rmtree(folder_path, ignore_errors=True)
+        raise gr.Error(f"Face swap failed: {str(e)}")
+
 
-# Video face swap helper function
-def swap_face_frame(frame_bgr, replacement_face_rgb, doFaceEnhancer):
+def swap_face_frame(frame_bgr: np.ndarray, replacement_face_rgb: np.ndarray, doFaceEnhancer: bool) -> np.ndarray:
+    """
+    Swap face in a single video frame.
+    
+    Parameters:
+        frame_bgr (np.ndarray): Video frame in BGR format.
+        replacement_face_rgb (np.ndarray): Replacement face image in RGB format.
+        doFaceEnhancer (bool): Flag to apply face enhancer.
+        
+    Returns:
+        np.ndarray: Processed frame with face swapped (in RGB format).
+    """
+    # Convert BGR to RGB for processing
     frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
-    temp_dir = "temp_faceswap_frame"
-    os.makedirs(temp_dir, exist_ok=True)
+    temp_dir = tempfile.mkdtemp(prefix="temp_faceswap_frame_")
     timestamp = datetime.now().strftime("%S-%M-%H-%d-%m-%Y")
-
+    source_path = os.path.join(temp_dir, f"source_{timestamp}.jpg")
+    target_path = os.path.join(temp_dir, f"target_{timestamp}.jpg")
+    output_path = os.path.join(temp_dir, f"OutputImage_{timestamp}.jpg")
+    
     try:
-        source_path = os.path.join(temp_dir, f"source_{timestamp}.jpg")
-        target_path = os.path.join(temp_dir, f"target_{timestamp}.jpg")
-        output_path = os.path.join(temp_dir, f"OutputImage_{timestamp}.jpg")
-
         Image.fromarray(frame_rgb).save(source_path)
         Image.fromarray(replacement_face_rgb).save(target_path)
-
-        # Configure Roop globals
-        roop.globals.source_path = source_path
-        roop.globals.target_path = target_path
-        roop.globals.output_path = normalize_output_path(
-            source_path, target_path, output_path
-        )
-        roop.globals.frame_processors = (
-            ["face_swapper", "face_enhancer"] if doFaceEnhancer else ["face_swapper"]
-        )
-        roop.globals.headless = True
-        roop.globals.keep_fps = True
-        roop.globals.keep_audio = True
-        roop.globals.keep_frames = False
-        roop.globals.many_faces = False
-        roop.globals.video_encoder = "libx264"
-        roop.globals.video_quality = 18
-        roop.globals.max_memory = suggest_max_memory()
-        roop.globals.execution_providers = decode_execution_providers(["cuda"])
-        roop.globals.execution_threads = suggest_execution_threads()
-
+        
+        configure_roop_globals(source_path, target_path, output_path, doFaceEnhancer)
+        
         start()
-
-        # Return swapped frame
-        if os.path.exists(output_path):
-            return np.array(Image.open(output_path))
+        
+        if os.path.exists(roop.globals.output_path):
+            swapped_img = np.array(Image.open(roop.globals.output_path))
         else:
-            return frame_rgb
+            logging.warning("Output image not found after face swap; returning original frame.")
+            swapped_img = frame_rgb
+    except Exception as e:
+        logging.exception(f"Error in processing frame for face swap: {str(e)}")
+        swapped_img = frame_rgb
     finally:
-        shutil.rmtree(temp_dir)
-
-# Video face swap function
-def swap_face_video(reference_face, replacement_face, video_input, similarity_threshold, doFaceEnhancer):
-    fa = FaceAnalysis()
-    fa.prepare(ctx_id=0)
+        shutil.rmtree(temp_dir, ignore_errors=True)
+    
+    return swapped_img
 
-    ref_detections = fa.get(reference_face)
-    if not ref_detections:
-        raise gr.Error("No face detected in the reference image!")
-    ref_embedding = ref_detections[0].embedding
 
-    cap = cv2.VideoCapture(video_input)
-    if not cap.isOpened():
-        raise gr.Error("Failed to open input video!")
-
-    fps = cap.get(cv2.CAP_PROP_FPS)
-    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-
-    output_video_path = "temp_faceswap_video.mp4"
-    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
-    out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))
-
-    frame_index = 0
-    while True:
-        ret, frame = cap.read()
-        if not ret:
-            break
-
-        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        detections = fa.get(frame_rgb)
-        swap_this_frame = any(
-            cosine_similarity(det.embedding, ref_embedding) >= similarity_threshold
-            for det in detections
-        )
-
-        if swap_this_frame:
-            swapped_frame_rgb = swap_face_frame(frame, replacement_face, doFaceEnhancer)
-            swapped_frame = cv2.cvtColor(swapped_frame_rgb, cv2.COLOR_RGB2BGR)
-        else:
-            swapped_frame = frame
-
-        out.write(swapped_frame)
-        frame_index += 1
-        print(f"Processed frame {frame_index}")
+def swap_face_video(reference_face: np.ndarray, replacement_face: np.ndarray, video_input: str,
+                    similarity_threshold: float, doFaceEnhancer: bool) -> str:
+    """
+    Perform face swapping on a video frame-by-frame.
+    
+    Parameters:
+        reference_face (np.ndarray): Reference face image (RGB) for face locking.
+        replacement_face (np.ndarray): Replacement face image (RGB).
+        video_input (str): Path to the input video file.
+        similarity_threshold (float): Threshold for face similarity (0.0 - 1.0).
+        doFaceEnhancer (bool): Flag to apply face enhancer.
+        
+    Returns:
+        str: Path to the output video file.
+        
+    Raises:
+        gr.Error: If face detection fails or video cannot be processed.
+    """
+    try:
+        # Initialize insightface face analysis
+        fa = FaceAnalysis()
+        fa.prepare(ctx_id=0)
+        
+        # Get embedding for the reference face
+        ref_detections = fa.get(reference_face)
+        if not ref_detections:
+            raise gr.Error("No face detected in the reference image!")
+        ref_embedding = ref_detections[0].embedding
+        
+        # Open video input
+        cap = cv2.VideoCapture(video_input)
+        if not cap.isOpened():
+            raise gr.Error("Cannot open the input video!")
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        
+        output_video_path = "temp_faceswap_video.mp4"
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+        out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))
+        
+        frame_index = 0
+        while True:
+            ret, frame = cap.read()
+            if not ret:
+                break
+            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            detections = fa.get(frame_rgb)
+            swap_this_frame = any(
+                cosine_similarity(det.embedding, ref_embedding) >= similarity_threshold
+                for det in detections
+            )
+            
+            if swap_this_frame:
+                swapped_frame_rgb = swap_face_frame(frame, replacement_face, doFaceEnhancer)
+                swapped_frame = cv2.cvtColor(swapped_frame_rgb, cv2.COLOR_RGB2BGR)
+            else:
+                swapped_frame = frame
+            
+            out.write(swapped_frame)
+            frame_index += 1
+            logging.info(f"Processed frame {frame_index}")
+        
+        cap.release()
+        out.release()
+        return output_video_path
+    except Exception as e:
+        logging.exception(f"Error processing video: {str(e)}")
+        raise gr.Error(f"Face swap video failed: {str(e)}")
 
-    cap.release()
-    out.release()
-    return output_video_path
 
-# Gradio interface
-def create_interface():
+def create_interface() -> gr.Blocks:
+    """
+    Create and return the Gradio interface for face swapping.
+    
+    Returns:
+        gr.Blocks: The Gradio interface.
+    """
     custom_css = """
     .container {
         max-width: 1200px;
@@ -262,55 +363,102 @@ def create_interface():
         <p>This tool performs face swapping with optional enhancement.</p>
     </div>
     """
-
     with gr.Blocks(title=title, css=custom_css) as app:
         gr.Markdown(f"<h1 style='text-align: center;'>{title}</h1>")
         gr.Markdown(description)
-
         with gr.Tabs():
             with gr.TabItem("FaceSwap Image"):
                 with gr.Row():
-                    source_image = gr.Image(label="Source Image", type="numpy", sources=["upload"])
-                    target_image = gr.Image(label="Target Image", type="numpy", sources=["upload"])
-                    output_image = gr.Image(label="Output Image", type="numpy", interactive=False, elem_classes="output-image")
-
-                enhance_checkbox = gr.Checkbox(label="Apply Enhancement?", info="Improve image quality", value=False)
-                process_btn = gr.Button("Process Face Swap", variant="primary", size="lg")
-
+                    with gr.Column(scale=1):
+                        source_image = gr.Image(
+                            label="Source Image",
+                            type="numpy",
+                            sources=["upload"]
+                        )
+                    with gr.Column(scale=1):
+                        target_image = gr.Image(
+                            label="Target Image",
+                            type="numpy",
+                            sources=["upload"]
+                        )
+                    with gr.Column(scale=1):
+                        output_image = gr.Image(
+                            label="Output Image",
+                            type="numpy",
+                            interactive=False,
+                            elem_classes="output-image"
+                        )
+                with gr.Row():
+                    enhance_checkbox = gr.Checkbox(
+                        label="Apply Face Enhancer",
+                        info="Improve image quality",
+                        value=False
+                    )
+                with gr.Row():
+                    process_btn = gr.Button(
+                        "Process Face Swap",
+                        variant="primary",
+                        size="lg"
+                    )
                 process_btn.click(
                     fn=swap_face,
                     inputs=[source_image, target_image, enhance_checkbox],
                     outputs=output_image,
-                    api_name="swap_face",
+                    api_name="swap_face"
                 )
-
             with gr.TabItem("FaceSwap Video"):
                 gr.Markdown("<h2 style='text-align:center;'>FaceSwap Video</h2>")
                 with gr.Row():
-                    ref_image = gr.Image(label="Reference Face", type="numpy", sources=["upload"])
-                    swap_image = gr.Image(label="Replacement Face", type="numpy", sources=["upload"])
-                    video_input = gr.Video(label="Input Video")
-                    similarity_threshold = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.7, label="Similarity Threshold")
-                    enhance_checkbox_video = gr.Checkbox(label="Apply Enhancement?", info="Improve image quality", value=False)
-                    video_output = gr.Video(label="Output Video")
-
-                process_video_btn = gr.Button("Process FaceSwap Video", variant="primary", size="lg")
-
+                    ref_image = gr.Image(
+                        label="Reference Face Image (Lock Face)",
+                        type="numpy",
+                        sources=["upload"]
+                    )
+                    swap_image = gr.Image(
+                        label="Replacement Face Image",
+                        type="numpy",
+                        sources=["upload"]
+                    )
+                video_input = gr.Video(
+                    label="Input Video"
+                )
+                similarity_threshold = gr.Slider(
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.01,
+                    value=0.7,
+                    label="Similarity Threshold"
+                )
+                enhance_checkbox_video = gr.Checkbox(
+                    label="Apply Face Enhancer",
+                    info="Optional quality enhancement",
+                    value=False
+                )
+                process_video_btn = gr.Button(
+                    "Process FaceSwap Video",
+                    variant="primary",
+                    size="lg"
+                )
+                video_output = gr.Video(
+                    label="Output Video"
+                )
                 process_video_btn.click(
                     fn=swap_face_video,
                     inputs=[ref_image, swap_image, video_input, similarity_threshold, enhance_checkbox_video],
                     outputs=video_output,
-                    api_name="swap_face_video",
+                    api_name="swap_face_video"
                 )
-
         gr.Markdown(article)
-
     return app
 
-# Main function
-def main():
+
+def main() -> None:
+    """
+    Launch the Gradio interface.
+    """
     app = create_interface()
     app.launch(share=False)
 
+
 if __name__ == "__main__":
-    main()
+    main()