Update app.py

app.py

@@ -19,36 +19,35 @@ 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
+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)
 class FaceIntegrDataset:
     def __init__(self, repo_id="Arrcttacsrks/face_integrData"):
-        # Get token from environment variable
         self.token = os.getenv('hf_token')
         if not self.token:
             raise ValueError("HF_TOKEN environment variable is not set")
-            
         self.repo_id = repo_id
         self.api = HfApi()
-        
-        # Login to Hugging Face
         login(self.token)
-        
-        # Create local temp directory for organizing files
         self.temp_dir = "temp_dataset"
         os.makedirs(self.temp_dir, exist_ok=True)
 
     def create_date_folder(self):
-        """Create folder structure based on current date"""
         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):
-        """Save metadata for the face swap operation"""
         metadata = {
             "timestamp": timestamp,
             "source_image": source_path,
@@ -59,9 +58,7 @@ class FaceIntegrDataset:
         return metadata
 
     def upload_to_hf(self, local_folder, date_folder):
-        """Upload files to Hugging Face dataset"""
         try:
-            # Upload the files
             self.api.upload_folder(
                 folder_path=local_folder,
                 repo_id=self.repo_id,
@@ -73,36 +70,26 @@ class FaceIntegrDataset:
             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
 def swap_face(source_file, target_file, doFaceEnhancer):
     folder_path = None
     try:
-        # Initialize dataset handler
         dataset_handler = FaceIntegrDataset()
-        
-        # Create date-based folder
         folder_path, date_folder = dataset_handler.create_date_folder()
-        
-        # Generate timestamp for unique identification
         timestamp = datetime.now().strftime("%S-%M-%H-%d-%m-%Y")
-        
-        # Save input images with timestamp in folder
         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")
-        
-        # Save the input images
+
         if source_file is None or target_file is None:
             raise ValueError("Source and target images are required")
             
-        source_image = Image.fromarray(source_file)
-        source_image.save(source_path)
-        target_image = Image.fromarray(target_file)
-        target_image.save(target_path)
+        Image.fromarray(source_file).save(source_path)
+        Image.fromarray(target_file).save(target_path)
         
         print("source_path: ", source_path)
         print("target_path: ", target_path)
         
-        # Set global paths
         roop.globals.source_path = source_path
         roop.globals.target_path = target_path
         roop.globals.output_path = normalize_output_path(
@@ -111,13 +98,11 @@ def swap_face(source_file, target_file, doFaceEnhancer):
             output_path
         )
         
-        # Configure face processing options
         if doFaceEnhancer:
             roop.globals.frame_processors = ["face_swapper", "face_enhancer"]
         else:
             roop.globals.frame_processors = ["face_swapper"]
         
-        # Set global parameters
         roop.globals.headless = True
         roop.globals.keep_fps = True
         roop.globals.keep_audio = True
@@ -136,15 +121,12 @@ def swap_face(source_file, target_file, doFaceEnhancer):
             roop.globals.output_path,
         )
         
-        # Check frame processors
         for frame_processor in get_frame_processors_modules(roop.globals.frame_processors):
             if not frame_processor.pre_check():
                 return None
         
-        # Process the face swap
         start()
         
-        # Save metadata
         metadata = dataset_handler.save_metadata(
             f"source_{timestamp}.jpg",
             f"target_{timestamp}.jpg",
@@ -152,12 +134,10 @@ def swap_face(source_file, target_file, doFaceEnhancer):
             timestamp
         )
         
-        # Save metadata to JSON file in the same folder
         metadata_path = os.path.join(folder_path, f"metadata_{timestamp}.json")
         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:
@@ -165,11 +145,9 @@ def swap_face(source_file, target_file, doFaceEnhancer):
         else:
             print("Failed to upload files to Hugging Face dataset")
         
-        # Read the output image before cleaning up
         if os.path.exists(output_path):
             output_image = Image.open(output_path)
             output_array = np.array(output_image)
-            # Clean up temp folder after reading the image
             shutil.rmtree(folder_path)
             return output_array
         else:
@@ -184,8 +162,113 @@ def swap_face(source_file, target_file, doFaceEnhancer):
             shutil.rmtree(folder_path)
         raise gr.Error(f"Face swap failed: {str(e)}")
 
+# HÀM MỚI: Xử lý ghép mặt cho 1 frame video bằng cách "mượn" thuật toán của roop
+def swap_face_frame(frame_bgr, replacement_face_rgb, doFaceEnhancer):
+    # Tao convert frame từ BGR sang RGB vì PIL làm việc với RGB – không cho mày chê!
+    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 MỚI: 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()
+        fa.prepare(ctx_id=0, nms=0.4)
+        
+        # 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
+            # Convert 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)
+                # Convert 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)}")
+
+# Giao diện Gradio được "đấm" thêm mục chuyển đổi giữa image và video
 def create_interface():
-    # Create custom style
     custom_css = """
     .container {
         max-width: 1200px;
@@ -199,11 +282,9 @@ def create_interface():
         padding: 10px;
     }
     """
-
-    # Gradio interface setup
     title = "Face - Integrator"
     description = r"""
-    Please upload source and target images to begin the face swap process.
+    Upload source and target images to perform face swap.
     """
     article = r"""
     <div style="text-align: center; max-width: 650px; margin: 40px auto;">
@@ -212,65 +293,99 @@ def create_interface():
         </p>
     </div>
     """
-
-    # Create Gradio interface with improved layout
     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.Row():
-            with gr.Column(scale=1):
-                source_image = gr.Image(
-                    label="Source Image",
-                    type="numpy",
-                    sources=["upload"]
+        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"
+                    )
+                process_btn.click(
+                    fn=swap_face,
+                    inputs=[source_image, target_image, enhance_checkbox],
+                    outputs=output_image,
+                    api_name="swap_face"
                 )
-                
-            with gr.Column(scale=1):
-                target_image = gr.Image(
-                    label="Target Image",
-                    type="numpy",
-                    sources=["upload"]
+            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",
+                    type="filepath"
                 )
-                
-            with gr.Column(scale=1):
-                output_image = gr.Image(
-                    label="Output Image",
-                    type="numpy",
-                    interactive=False,
-                    elem_classes="output-image"
+                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ả",
+                    type="filepath"
+                )
+                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"
                 )
-        
-        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"
-            )
-        
-        # Set up the processing event
-        process_btn.click(
-            fn=swap_face,
-            inputs=[source_image, target_image, enhance_checkbox],
-            outputs=output_image,
-            api_name="swap_face"
-        )
-        
         gr.Markdown(article)
-    
     return app
 
 def main():
-    # Create and launch the interface
     app = create_interface()
     app.launch(share=False)
 
 if __name__ == "__main__":
-    main()
+    main()