Update app.py

app.py

@@ -1,38 +1,37 @@
-# -*- coding:UTF-8 -*-
+# -*- coding: UTF-8 -*-
 #!/usr/bin/env python
+
+import os
+import json
+import shutil
+from datetime import datetime
+from dotenv import load_dotenv
 import numpy as np
+import cv2
+from PIL import Image
 import gradio as gr
-import roop.globals
-from roop.core import (
+from huggingface_hub import HfApi, login
+from roop.globals 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 roop.utilities import normalize_output_path
-import os
-from PIL import Image
-from datetime import datetime
-from huggingface_hub import HfApi, login
-from datasets import load_dataset, Dataset
-import json
-import shutil
-from dotenv import load_dotenv
-import cv2
 from insightface.app import FaceAnalysis
 
 # Load environment variables
 load_dotenv()
 
-# Hàm tính cosine similarity để mày so sánh "điểm tương đồng" của khuôn mặt
+# Cosine similarity function
 def cosine_similarity(a, b):
     return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-6)
 
-# Class FaceIntegrDataset nguyên bản (cho image swap, không cần "xịn" cho video)
+# Dataset handler class
 class FaceIntegrDataset:
     def __init__(self, repo_id="Arrcttacsrks/face_integrData"):
-        self.token = os.getenv('hf_token')
+        self.token = os.getenv("hf_token")
         if not self.token:
             raise ValueError("HF_TOKEN environment variable is not set")
         self.repo_id = repo_id
@@ -53,7 +52,7 @@ class FaceIntegrDataset:
             "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
 
@@ -63,46 +62,41 @@ class FaceIntegrDataset:
                 folder_path=local_folder,
                 repo_id=self.repo_id,
                 repo_type="dataset",
-                path_in_repo=date_folder
+                path_in_repo=date_folder,
             )
             return True
         except Exception as e:
             print(f"Error uploading to Hugging Face: {str(e)}")
             return False
 
-# Hàm swap_face nguyên bản dành cho ghép ảnh tĩnh
+# Image face swap function
 def swap_face(source_file, target_file, doFaceEnhancer):
-    folder_path = None
+    dataset_handler = FaceIntegrDataset()
+    folder_path, date_folder = dataset_handler.create_date_folder()
+    timestamp = datetime.now().strftime("%S-%M-%H-%d-%m-%Y")
+
     try:
-        dataset_handler = FaceIntegrDataset()
-        folder_path, date_folder = dataset_handler.create_date_folder()
-        timestamp = datetime.now().strftime("%S-%M-%H-%d-%m-%Y")
+        # Save source and target images
         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)
-        
-        print("source_path: ", source_path)
-        print("target_path: ", 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.source_path, roop.globals.target_path, output_path
+        )
+
+        roop.globals.frame_processors = (
+            ["face_swapper", "face_enhancer"] if doFaceEnhancer else ["face_swapper"]
         )
-        
-        if doFaceEnhancer:
-            roop.globals.frame_processors = ["face_swapper", "face_enhancer"]
-        else:
-            roop.globals.frame_processors = ["face_swapper"]
-        
         roop.globals.headless = True
         roop.globals.keep_fps = True
         roop.globals.keep_audio = True
@@ -113,162 +107,140 @@ def swap_face(source_file, target_file, doFaceEnhancer):
         roop.globals.max_memory = suggest_max_memory()
         roop.globals.execution_providers = decode_execution_providers(["cuda"])
         roop.globals.execution_threads = suggest_execution_threads()
-        
-        print(
-            "start process",
-            roop.globals.source_path,
-            roop.globals.target_path,
-            roop.globals.output_path,
-        )
-        
+
+        # Pre-check frame processors
         for frame_processor in get_frame_processors_modules(roop.globals.frame_processors):
             if not frame_processor.pre_check():
-                return None
-        
+                raise RuntimeError("Frame processor pre-check failed")
+
+        # Start face swap process
         start()
-        
+
+        # Save metadata
         metadata = dataset_handler.save_metadata(
             f"source_{timestamp}.jpg",
             f"target_{timestamp}.jpg",
-            f"OutputImage{timestamp}.jpg",
-            timestamp
+            f"OutputImage_{timestamp}.jpg",
+            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 upload_success:
-            print(f"Successfully uploaded files to dataset {dataset_handler.repo_id}")
-        else:
+        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)
-            output_array = np.array(output_image)
-            shutil.rmtree(folder_path)
-            return output_array
+            return np.array(output_image)
         else:
-            print("Output image not found")
-            if folder_path and os.path.exists(folder_path):
-                shutil.rmtree(folder_path)
-            return None
-        
+            raise FileNotFoundError("Output image not found")
+
     except Exception as e:
         print(f"Error in face swap process: {str(e)}")
+        raise gr.Error(f"Face swap failed: {str(e)}")
+    finally:
         if folder_path and os.path.exists(folder_path):
             shutil.rmtree(folder_path)
-        raise gr.Error(f"Face swap failed: {str(e)}")
 
-# Hàm xử lý ghép mặt cho 1 frame video bằng cách "mượn" thuật toán của roop
+# Video face swap helper function
 def swap_face_frame(frame_bgr, replacement_face_rgb, doFaceEnhancer):
-    # Chuyển frame từ BGR sang RGB vì PIL làm việc với RGB
     frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
     temp_dir = "temp_faceswap_frame"
     os.makedirs(temp_dir, exist_ok=True)
     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")
-    Image.fromarray(frame_rgb).save(source_path)
-    Image.fromarray(replacement_face_rgb).save(target_path)
-    
-    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)
-    
-    if doFaceEnhancer:
-        roop.globals.frame_processors = ["face_swapper", "face_enhancer"]
-    else:
-        roop.globals.frame_processors = ["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()
-    
-    start()
-    
-    if os.path.exists(output_path):
-        swapped_img = np.array(Image.open(output_path))
-    else:
-        swapped_img = frame_rgb
-    shutil.rmtree(temp_dir)
-    return swapped_img
-
-# Hàm xử lý ghép mặt cho video frame-by-frame với insightface để so sánh khuôn mặt
-def swap_face_video(reference_face, replacement_face, video_input, similarity_threshold, doFaceEnhancer):
-    """
-    reference_face: Ảnh tham chiếu (RGB) để khóa khuôn mặt
-    replacement_face: Ảnh ghép (RGB)
-    video_input: Đường dẫn file video đầu vào
-    similarity_threshold: Ngưỡng (0.0 - 1.0) cho tỉ lệ tương đồng
-    doFaceEnhancer: Boolean, có áp dụng cải thiện chất lượng hay không
-    """
+
     try:
-        # Chuẩn bị insightface
-        fa = FaceAnalysis()
-        # Loại bỏ nms=0.4 vì hàm prepare() không hỗ trợ argument này
-        fa.prepare(ctx_id=0)
-        
-        # Lấy embedding của khuôn mặt tham chiếu
-        ref_detections = fa.get(reference_face)
-        if not ref_detections:
-            raise gr.Error("Không phát hiện khuôn mặt trong ảnh tham chiếu!")
-        ref_embedding = ref_detections[0].embedding
-        
-        # Mở video đầu vào
-        cap = cv2.VideoCapture(video_input)
-        if not cap.isOpened():
-            raise gr.Error("Không mở được video đầu vào!")
-        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
-            # Chuyển frame sang RGB để insightface xử lý
-            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-            detections = fa.get(frame_rgb)
-            swap_this_frame = False
-            for det in detections:
-                sim = cosine_similarity(det.embedding, ref_embedding)
-                if sim >= similarity_threshold:
-                    swap_this_frame = True
-                    break
-            if swap_this_frame:
-                # Ghép mặt từ replacement_face vào frame
-                swapped_frame_rgb = swap_face_frame(frame, replacement_face, doFaceEnhancer)
-                # Chuyển ngược lại sang BGR để ghi video
-                swapped_frame = cv2.cvtColor(swapped_frame_rgb, cv2.COLOR_RGB2BGR)
-            else:
-                swapped_frame = frame
-            out.write(swapped_frame)
-            frame_index += 1
-            print(f"Đã xử lý frame {frame_index}")
-        cap.release()
-        out.release()
-        return output_video_path
-    except Exception as e:
-        print(f"Lỗi khi xử lý video: {str(e)}")
-        raise gr.Error(f"Face swap video failed: {str(e)}")
+        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()
+
+        start()
+
+        # Return swapped frame
+        if os.path.exists(output_path):
+            return np.array(Image.open(output_path))
+        else:
+            return 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)
+
+    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
 
-# Giao diện Gradio được xây dựng với hai tab: Image và Video
+        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}")
+
+    cap.release()
+    out.release()
+    return output_video_path
+
+# Gradio interface
 def create_interface():
     custom_css = """
     .container {
@@ -284,107 +256,61 @@ def create_interface():
     }
     """
     title = "Face - Integrator"
-    description = r"""
-    Upload source and target images to perform face swap.
-    """
-    article = r"""
+    description = "Upload source and target images to perform face swap."
+    article = """
     <div style="text-align: center; max-width: 650px; margin: 40px auto;">
-        <p>
-            This tool performs face swapping with optional enhancement.
-        </p>
+        <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():
-                    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 the algorithm?",
-                        info="Image Quality Improvement",
-                        value=False
-                    )
-                with gr.Row():
-                    process_btn = gr.Button(
-                        "Process Face Swap",
-                        variant="primary",
-                        size="lg"
-                    )
+                    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")
+
                 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="Ảnh mặt tham chiếu (khóa khuôn mặt)",
-                        type="numpy",
-                        sources=["upload"]
-                    )
-                    swap_image = gr.Image(
-                        label="Ảnh mặt ghép",
-                        type="numpy",
-                        sources=["upload"]
-                    )
-                video_input = gr.Video(
-                    label="Video đầu vào"
-                )
-                similarity_threshold = gr.Slider(
-                    minimum=0.0,
-                    maximum=1.0,
-                    step=0.01,
-                    value=0.7,
-                    label="Tỉ lệ tương đồng"
-                )
-                enhance_checkbox_video = gr.Checkbox(
-                    label="Áp dụng cải thiện chất lượng ảnh",
-                    info="Tùy chọn cải thiện",
-                    value=False
-                )
-                process_video_btn = gr.Button(
-                    "Xử lý FaceSwap Video",
-                    variant="primary",
-                    size="lg"
-                )
-                video_output = gr.Video(
-                    label="Video kết quả"
-                )
+                    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")
+
                 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():
     app = create_interface()
     app.launch(share=False)
 
 if __name__ == "__main__":
-    main()
+    main()