app.py

import gradio as gr
import rembg
from PIL import Image
import io
import os
import tempfile
import uuid

def remove_background(input_image):
    """Remove background from image and return both image and download data"""
    if input_image is None:
        return None, None
    
    try:
        # Remove background using AI
        output_image = rembg.remove(input_image)
        
        # Create a temporary file with a unique name
        temp_dir = tempfile.gettempdir()
        filename = f"background_removed_{uuid.uuid4().hex[:8]}.png"
        file_path = os.path.join(temp_dir, filename)
        
        # Save the image to the temporary file
        output_image.save(file_path, format='PNG')
        
        return output_image, file_path
        
    except Exception as e:
        print(f"Error in remove_background: {e}")
        return None, None

def clear_all():
    """Clear all images and reset the interface"""
    return None, None, None

# Create the Gradio interface using Blocks for more control
with gr.Blocks(title="AI Background Remover", theme="soft") as demo:
    gr.Markdown("# 🎨 AI Background Remover")
    gr.Markdown("Upload any image to automatically remove the background. Perfect for e-commerce, marketing, and content creation!")
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="📷 Upload Image")
            submit_btn = gr.Button("Remove Background 🚀", variant="primary")
        
        with gr.Column():
            output_image = gr.Image(type="pil", label="🎨 Background Removed", interactive=False)
            download_btn = gr.DownloadButton("📥 Download Image", visible=True, label="Download your processed image")
            clear_btn = gr.Button("🗑️ Clear All", variant="secondary")
    
    # Handle examples if they exist
    example_files = []
    for i in range(1, 6):
        example_path = f"example_{i:02d}.jpg"
        if os.path.exists(example_path):
            example_files.append([example_path])
    
    if example_files:
        gr.Examples(
            examples=example_files,
            inputs=input_image,
            label="Try these examples:"
        )
    
    # Set up the interaction
    submit_btn.click(
        fn=remove_background,
        inputs=input_image,
        outputs=[output_image, download_btn]
    )

    # Clear Button
    clear_btn.click(fn=clear_all, inputs=None, outputs=[input_image, output_image, download_btn])

if __name__ == "__main__":
    # For Docker deployment
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )