Update app.py

app.py

@@ -1,5 +1,5 @@
-# -* coding:UTF-8 -*
-# !/usr/bin/env python
+# -*- coding:UTF-8 -*-
+#!/usr/bin/env python
 import numpy as np
 import gradio as gr
 import roop.globals
@@ -13,72 +13,183 @@ 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()
 
-def swap_face(source_file, target_file,doFaceEnhancer):
+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)
 
-    source_path = "input.jpg"
-    target_path = "target.jpg"
+    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
 
-    source_image = Image.fromarray(source_file)
-    source_image.save(source_path)
-    target_image = Image.fromarray(target_file)
-    target_image.save(target_path)
+    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
 
-    print("source_path: ", source_path)
-    print("target_path: ", target_path)
+    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
 
-    roop.globals.source_path = source_path
-    roop.globals.target_path = target_path
-    output_path = "output.jpg"
-    roop.globals.output_path = normalize_output_path(
-        roop.globals.source_path, roop.globals.target_path, output_path
-    )
-    if doFaceEnhancer == True:
-        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()
-
-    print(
-        "start process",
-        roop.globals.source_path,
-        roop.globals.target_path,
-        roop.globals.output_path,
-    )
-
-    for frame_processor in get_frame_processors_modules(
-        roop.globals.frame_processors
-    ):
-        if not frame_processor.pre_check():
-            return
-
-    start()
-    return output_path
+def swap_face(source_file, target_file, doFaceEnhancer):
+    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"Image{timestamp}.jpg")
+        
+        # Save the input images
+        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
+        
+        # Process the face swap
+        start()
+        
+        # Save metadata
+        metadata = dataset_handler.save_metadata(
+            f"source_{timestamp}.jpg",
+            f"target_{timestamp}.jpg",
+            f"Image{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")
+        
+        # Clean up temp folder after upload
+        shutil.rmtree(folder_path)
+        
+        return output_path
+        
+    except Exception as e:
+        print(f"Error in face swap process: {str(e)}")
+        raise gr.Error("Face swap failed. Please check your Hugging Face token and try again.")
 
+# Gradio interface setup
 title = "Face - Интегратор"
-
 description = r"""
+The application will save the image history to Hugging Face dataset using the environment variable token.
 """
 article = r"""
 <br><br><br><br><br>
 """
 
-
 app = gr.Interface(
     fn=swap_face, 
     title=title,
     description=description,
     article=article,
-    inputs=[gr.Image(), gr.Image(),gr.Checkbox(label="Применить алгоритм?", info="Улучшение качества изображения")], outputs="image"
+    inputs=[
+        gr.Image(), 
+        gr.Image(),
+        gr.Checkbox(label="Применить алгоритм?", info="Улучшение качества изображения")
+    ], 
+    outputs="image"
 )
-app.launch()
+
+app.launch()