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, None
+        return None, gr.update(visible=False)
     
     progress(0, desc="Starting processing...")
     
@@ -84,8 +85,9 @@ def process(image, progress=gr.Progress()):
     output_array = np.array(new_im.convert("RGBA"))
     
     progress(1.0, desc="Done!")
-    # Return both the image for display and the filepath for download
-    return output_array, filepath
+    
+    # Return the processed image and the file for download
+    return output_array, gr.update(visible=True, value=filepath, interactive=True)
 
 # Gradio interface setup
 title = "Background Removal"
@@ -93,76 +95,83 @@ description = r"""Background removal model developed by <a href='https://BRIA.AI
 
 examples = [['./input.jpg']]
 
-# 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;
+# Create the Gradio interface
+with gr.Blocks(css="""
+    .download-btn {
+        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;
     }
-    /* Style the download button to look like our green button */
-    .file-button {
+    .download-btn:hover {
+        background-color: #45a049 !important;
+    }
+    /* Style the file download component to look like a button */
+    .download-container [data-testid="file"] {
+        display: block !important;
+    }
+    .download-container [data-testid="file"] > div:first-child {
+        display: none !important; /* Hide the file icon */
+    }
+    .download-container [data-testid="file"] > div:last-child {
+        margin-left: 0 !important;
+    }
+    .download-container [data-testid="file"] button {
         background-color: #4CAF50 !important;
         color: white !important;
         border: none !important;
-        padding: 10px 20px !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;
-        margin-top: 10px !important;
+        transition: background-color 0.3s !important;
+        width: 100% !important;
     }
-    .file-button:hover {
+    .download-container [data-testid="file"] button:hover {
         background-color: #45a049 !important;
     }
-    """
-)
+""") as demo:
+    gr.Markdown(f"# {title}")
+    gr.Markdown(description)
+    
+    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")
+            
+            # Direct download file component styled as a button
+            with gr.Row(elem_classes="download-container"):
+                download_file = gr.File(
+                    label="", 
+                    file_count="single",
+                    interactive=True,
+                    visible=False
+                )
+    
+    # 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, download_file],
+        show_progress="full"
+    )
 
 if __name__ == "__main__":
     demo.launch(share=False)