Spaces:
Sleeping
Sleeping
| # app.py | |
| # A user-friendly version with a custom download button, inspired by the watermark remover app. | |
| import gradio as gr | |
| from rembg import remove | |
| import traceback | |
| from PIL import Image | |
| import tempfile # Used to create temporary files for downloading | |
| # --- Core Function --- | |
| # This function processes the image. It remains the same. | |
| def process_image(input_image): | |
| """ | |
| Removes the background from an image. | |
| Handles errors gracefully. | |
| """ | |
| if input_image is None: | |
| raise gr.Error("Please upload an image or select one of the examples.") | |
| try: | |
| # Use the default rembg model for stability | |
| return remove(input_image) | |
| except Exception as e: | |
| print(traceback.format_exc()) | |
| raise gr.Error(f"Failed to process image: {e}") | |
| # --- Download Function --- | |
| # This function correctly prepares the image for download. | |
| def prepare_for_download(image): | |
| """ | |
| Saves the processed PIL image to a temporary file and returns the path. | |
| This is the correct way to trigger a download in Gradio. | |
| """ | |
| if image is None: | |
| # Prevent an error if the download button is clicked before an image is processed. | |
| return None | |
| try: | |
| # Create a temporary file with a '.png' extension. | |
| # 'delete=False' is important so the file is not deleted before Gradio serves it. | |
| with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp: | |
| image.save(tmp.name) | |
| # Return the path to the temporary file | |
| return tmp.name | |
| except Exception as e: | |
| print(f"Error creating download file: {e}") | |
| return None | |
| # --- Gradio App with Modern UI --- | |
| # We inject CSS to hide the default utility buttons on the output image. | |
| with gr.Blocks( | |
| theme=gr.themes.Soft(primary_hue="violet"), | |
| css="#output_image .meta-buttons { display: none !important; }" | |
| ) as demo: | |
| gr.Markdown( | |
| """ | |
| Upload an image, use your webcam, or try one of the examples below to instantly remove the background. | |
| """ | |
| ) | |
| # Main two-column layout | |
| with gr.Row(variant="panel"): | |
| # --- Input Column --- | |
| with gr.Column(scale=1, min_width=300): | |
| image_input = gr.Image( | |
| type="pil", | |
| label="Upload or Drag an Image", | |
| sources=["upload", "webcam"], | |
| height=400 | |
| ) | |
| with gr.Row(): | |
| clear_btn = gr.ClearButton( | |
| value="Clear", | |
| components=[image_input], | |
| variant="secondary" | |
| ) | |
| submit_btn = gr.Button("Remove Background", variant="primary") | |
| # --- Output Column --- | |
| with gr.Column(scale=1, min_width=300): | |
| # We give the output image an ID so we can target it with CSS | |
| image_output = gr.Image( | |
| elem_id="output_image", | |
| type="pil", | |
| label="Result", | |
| height=400, | |
| interactive=False # This disables the upload overlay on the result | |
| ) | |
| # ✅ NEW: Use a gr.File component for the download, which is more reliable. | |
| # This will appear as a download link/button when populated. | |
| download_output = gr.File(label="Download Result", visible=False) | |
| # --- Event Listeners --- | |
| def on_submit(input_img): | |
| # This function runs when the user clicks submit. | |
| # It processes the image and then prepares it for download. | |
| processed_img = process_image(input_img) | |
| download_path = prepare_for_download(processed_img) | |
| # It returns the processed image to the preview box, | |
| # and the downloadable file to the (now visible) download component. | |
| return processed_img, gr.update(value=download_path, visible=True) | |
| submit_btn.click( | |
| fn=on_submit, | |
| inputs=image_input, | |
| outputs=[image_output, download_output] | |
| ) | |
| # When the user clicks "Clear", also hide the download button. | |
| clear_btn.click(lambda: gr.update(visible=False), outputs=[download_output]) | |
| # --- Launch the App --- | |
| if __name__ == "__main__": | |
| demo.launch() | |