Update app.py

app.py

@@ -24,10 +24,57 @@ from dotenv import load_dotenv
 load_dotenv()
 
 class FaceIntegrDataset:
-    # [Previous FaceIntegrDataset class implementation remains the same]
-    pass
+    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()
@@ -44,6 +91,9 @@ def swap_face(source_file, target_file, doFaceEnhancer):
         output_path = os.path.join(folder_path, f"Image{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)
@@ -117,50 +167,103 @@ def swap_face(source_file, target_file, doFaceEnhancer):
         
         # Read the output image before cleaning up
         if os.path.exists(output_path):
-            output_image = np.array(Image.open(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_image
+            return output_array
         else:
             print("Output image not found")
-            shutil.rmtree(folder_path)
+            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 os.path.exists(folder_path):
+        if folder_path and os.path.exists(folder_path):
             shutil.rmtree(folder_path)
-        raise gr.Error("Face swap failed. Please check your Hugging Face token and try again.")
+        raise gr.Error(f"Face swap failed: {str(e)}")
+
+# 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 - Интегратор"
 description = r"""
 The application will save the image history to Hugging Face dataset using the environment variable token.
+Please upload source and target images to begin the face swap process.
 """
 article = r"""
-<br><br><br><br><br>
+<div style="text-align: center; max-width: 650px; margin: 40px auto;">
+    <p>
+        This tool performs face swapping with optional enhancement.
+        The processed images are automatically saved to the Hugging Face dataset.
+    </p>
+</div>
 """
 
-with gr.Blocks(title=title) as app:
+# 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():
-            source_image = gr.Image(label="Source Image")
-            target_image = gr.Image(label="Target Image")
-            enhance_checkbox = gr.Checkbox(
-                label="Применить алгоритм?", 
-                info="Улучшение качества изображения"
+        with gr.Column(scale=1):
+            source_image = gr.Image(
+                label="Source Image",
+                type="numpy",
+                tool="upload"
+            )
+            
+        with gr.Column(scale=1):
+            target_image = gr.Image(
+                label="Target Image",
+                type="numpy",
+                tool="upload"
+            )
+            
+        with gr.Column(scale=1):
+            output_image = gr.Image(
+                label="Output Image",
+                type="numpy",
+                elem_classes="output-image"
             )
-        with gr.Column():
-            output_image = gr.Image(label="Output Image")
     
-    process_btn = gr.Button("Process")
+    with gr.Row():
+        enhance_checkbox = gr.Checkbox(
+            label="Применить алгоритм?",
+            info="Улучшение качества изображения",
+            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
+        outputs=output_image,
+        api_name="swap_face"
     )
     
     gr.Markdown(article)
 
-app.launch()
+# Launch the application
+app.launch(share=False)