Update app.py

app.py

@@ -10,6 +10,7 @@ import tempfile
 import os
 import time
 import uuid
+import shutil
 
 # Load the pre-trained model
 print("Loading model...")
@@ -33,7 +34,7 @@ def resize_image(image):
 # Background removal process
 def process(image, progress=gr.Progress()):
     if image is None:
-        return None, gr.update(visible=False)
+        return None, None, gr.update(visible=False)
     
     progress(0, desc="Starting processing...")
     
@@ -84,8 +85,23 @@ def process(image, progress=gr.Progress()):
     output_array = np.array(new_im.convert("RGBA"))
     
     progress(1.0, desc="Done!")
-    # Return the processed image and update download component visibility
-    return output_array, gr.update(visible=True, value=filepath)
+    # Return the processed image and filepath and update button visibility
+    return output_array, filepath, gr.update(visible=True)
+
+# Function to handle the download
+def download_image(filepath):
+    if filepath and os.path.exists(filepath):
+        filename = os.path.basename(filepath)
+        # Create a new file in a web-accessible directory
+        download_path = os.path.join("downloads", filename)
+        os.makedirs("downloads", exist_ok=True)
+        
+        # Copy the file to the download directory
+        shutil.copy2(filepath, download_path)
+        
+        # Return the file path for download
+        return download_path
+    return None
 
 # Gradio interface setup
 title = "Background Removal"
@@ -95,27 +111,47 @@ examples = [['./input.jpg']]
 
 # Create the Gradio interface
 with gr.Blocks(css="""
-    .download-btn {
-        margin-top: 10px !important;
+    .download-button {
+        background-color: #4CAF50 !important;
+        color: white !important;
+        border: none !important;
+        padding: 10px 24px !important;
+        text-align: center !important;
+        text-decoration: none !important;
+        display: inline-block !important;
+        font-size: 16px !important;
+        margin: 10px 0 !important;
+        cursor: pointer !important;
+        border-radius: 4px !important;
+        transition: background-color 0.3s !important;
     }
-    .result-container {
-        display: flex;
-        flex-direction: column;
+    .download-button:hover {
+        background-color: #45a049 !important;
     }
 """) as demo:
     gr.Markdown(f"# {title}")
     gr.Markdown(description)
     
+    # Store the filepath
+    filepath = gr.State()
+    
     with gr.Row():
         with gr.Column(scale=1):
             input_image = gr.Image(type="numpy", label="Upload Image")
         
-        with gr.Column(scale=1, elem_classes="result-container"):
+        with gr.Column(scale=1):
             output_image = gr.Image(type="numpy", label="Result")
-            # Use File component as a download button
-            download_file = gr.File(label="Download Processed Image", 
-                                  elem_classes="download-btn",
-                                  visible=False)
+            
+            # Hidden file component for actual download
+            download_file = gr.File(visible=False)
+            
+            # Custom styled download button
+            with gr.Row():
+                download_btn = gr.Button(
+                    "Download Image", 
+                    visible=False,
+                    elem_classes=["download-button"]
+                )
     
     # Set up example images
     gr.Examples(examples, inputs=input_image)
@@ -124,9 +160,35 @@ with gr.Blocks(css="""
     input_image.change(
         fn=process,
         inputs=input_image,
-        outputs=[output_image, download_file],
+        outputs=[output_image, filepath, download_btn],
         show_progress="full"
     )
+    
+    # Handle download button click
+    download_btn.click(
+        fn=download_image,
+        inputs=[filepath],
+        outputs=[download_file]
+    )
+
+    # Add custom HTML for button styling enhancement
+    gr.HTML("""
+    <script>
+    // This script will ensure buttons are properly styled after DOM updates
+    document.addEventListener('DOMContentLoaded', function() {
+        // Process might be needed to re-apply styles after Gradio updates the DOM
+        const observer = new MutationObserver(function() {
+            const buttons = document.querySelectorAll('.download-button');
+            buttons.forEach(button => {
+                button.style.backgroundColor = '#4CAF50';
+                button.style.color = 'white';
+            });
+        });
+        
+        observer.observe(document.body, { childList: true, subtree: true });
+    });
+    </script>
+    """)
 
 if __name__ == "__main__":
     demo.launch(share=False)