Update app.py

app.py

@@ -9,6 +9,7 @@ from PIL import Image
 import tempfile
 import os
 import time
+import uuid
 
 # Load the pre-trained model
 print("Loading model...")
@@ -18,6 +19,10 @@ net.to(device)
 net.eval()
 print(f"Model loaded on {device}")
 
+# Create output directory if it doesn't exist
+OUTPUT_DIR = "output_images"
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+
 # Resize the input image for model compatibility
 def resize_image(image):
     image = image.convert('RGB')
@@ -28,7 +33,7 @@ def resize_image(image):
 # Background removal process
 def process(image, progress=gr.Progress()):
     if image is None:
-        return None, None
+        return None, None, gr.update(visible=False)
     
     progress(0, desc="Starting processing...")
     
@@ -67,17 +72,26 @@ def process(image, progress=gr.Progress()):
     new_im.putalpha(pil_mask)
     
     progress(0.8, desc="Preparing download...")
-    # Save the processed image to a temporary file
-    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png')
-    new_im.save(temp_file.name, format='PNG')
-    temp_file.close()
+    # Generate a unique filename
+    unique_id = str(uuid.uuid4())[:8]
+    filename = f"background_removed_{unique_id}.png"
+    filepath = os.path.join(OUTPUT_DIR, filename)
+    
+    # Save the processed image
+    new_im.save(filepath, format='PNG')
     
     # Convert to numpy array for display
     output_array = np.array(new_im.convert("RGBA"))
     
     progress(1.0, desc="Done!")
-    # Return both the image for display and the path for download
-    return output_array, temp_file.name
+    # Return image for display and path for download
+    return output_array, filepath, gr.update(visible=True)
+
+# Function to handle the download button click
+def download_image(filepath):
+    if filepath and os.path.exists(filepath):
+        return filepath
+    return None
 
 # Gradio interface setup
 title = "Background Removal"
@@ -86,17 +100,40 @@ description = r"""Background removal model developed by <a href='https://BRIA.AI
 examples = [['./input.jpg']]
 
 # Create the Gradio interface
-with gr.Blocks() as demo:
+with gr.Blocks(css="""
+    .download-btn {
+        background-color: #4CAF50;
+        border: none;
+        color: white;
+        padding: 10px 24px;
+        text-align: center;
+        text-decoration: none;
+        display: inline-block;
+        font-size: 16px;
+        margin: 4px 2px;
+        cursor: pointer;
+        border-radius: 4px;
+    }
+    .download-btn:hover {
+        background-color: #45a049;
+    }
+""") as demo:
     gr.Markdown(f"# {title}")
     gr.Markdown(description)
     
+    # Store the processed image path
+    image_path = gr.State(None)
+    
     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")
-            download_btn = gr.File(label="Download Processed Image", elem_id="download_button")
+            download_btn = gr.Button("Download Image", elem_id="download_button", 
+                                     variant="primary", visible=False, 
+                                     elem_classes="download-btn")
+            download_output = gr.File(visible=False)
     
     # Set up example images
     gr.Examples(examples, inputs=input_image)
@@ -105,18 +142,16 @@ with gr.Blocks() as demo:
     input_image.change(
         fn=process,
         inputs=input_image,
-        outputs=[output_image, download_btn],
+        outputs=[output_image, image_path, download_btn],
         show_progress="full"
     )
-
-    # Style to ensure download button is visible
-    gr.HTML("""
-    <style>
-    #download_button {
-        margin-top: 10px;
-    }
-    </style>
-    """)
+    
+    # Handle download button click
+    download_btn.click(
+        fn=download_image,
+        inputs=[image_path],
+        outputs=[download_output]
+    )
 
 if __name__ == "__main__":
     demo.launch(share=False)