Update app.py

app.py

@@ -10,7 +10,6 @@ import tempfile
 import os
 import time
 import uuid
-import shutil
 
 # Load the pre-trained model
 print("Loading model...")
@@ -34,7 +33,7 @@ def resize_image(image):
 # Background removal process
 def process(image, progress=gr.Progress()):
     if image is None:
-        return None, None, gr.update(visible=False)
+        return None, None
     
     progress(0, desc="Starting processing...")
     
@@ -85,23 +84,8 @@ def process(image, progress=gr.Progress()):
     output_array = np.array(new_im.convert("RGBA"))
     
     progress(1.0, desc="Done!")
-    # 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
+    # Return both the image for display and the filepath for download
+    return output_array, filepath
 
 # Gradio interface setup
 title = "Background Removal"
@@ -109,86 +93,76 @@ description = r"""Background removal model developed by <a href='https://BRIA.AI
 
 examples = [['./input.jpg']]
 
-# Create the Gradio interface
-with gr.Blocks(css="""
-    .download-button {
+# Create HTML with custom CSS
+custom_css = """
+<style>
+.container {
+    display: flex;
+    flex-direction: row;
+    width: 100%;
+}
+.column {
+    flex: 1;
+    padding: 10px;
+}
+.result-area {
+    position: relative;
+}
+.download-btn {
+    background-color: #4CAF50;
+    color: white;
+    border: none;
+    padding: 10px 20px;
+    text-align: center;
+    text-decoration: none;
+    display: inline-block;
+    font-size: 16px;
+    margin-top: 10px;
+    border-radius: 4px;
+    cursor: pointer;
+    transition: background-color 0.3s;
+}
+.download-btn:hover {
+    background-color: #45a049;
+}
+</style>
+"""
+
+# Create the Gradio interface with a simpler approach
+demo = gr.Interface(
+    fn=process,
+    inputs=gr.Image(type="numpy", label="Upload Image"),
+    outputs=[
+        gr.Image(type="numpy", label="Result"),
+        gr.File(label="Download Processed Image")
+    ],
+    examples=examples,
+    title=title,
+    description=description + custom_css,
+    allow_flagging="never",
+    theme=gr.themes.Default(
+        primary_hue="green",
+        secondary_hue="green",
+    ),
+    css="""
+    .gradio-container .footer {
+        display: none !important;
+    }
+    /* Style the download button to look like our green button */
+    .file-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;
+        padding: 10px 20px !important;
         font-size: 16px !important;
-        margin: 10px 0 !important;
-        cursor: pointer !important;
         border-radius: 4px !important;
-        transition: background-color 0.3s !important;
+        margin-top: 10px !important;
     }
-    .download-button:hover {
+    .file-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):
-            output_image = gr.Image(type="numpy", label="Result")
-            
-            # 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)
-    
-    # Process automatically when image is uploaded
-    input_image.change(
-        fn=process,
-        inputs=input_image,
-        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)