app.py

import gradio as gr
import os
from rembg import remove
from PIL import Image
import io
import zipfile
from typing import List, Tuple, Union, IO, Optional
import spaces

@spaces.GPU
def remove_background_single(image_file: Union[str, IO[bytes]]) -> Optional[Image.Image]:
    """Removes the background from a single image.

    This function takes an image file (either as a path or a file-like object),
    processes it to remove the background, and returns the result as a 
    Pillow (PIL) Image object with a transparent background.

    Args:
        image_file: The input image. Can be a path to a
            file on disk or a file-like object (e.g., from a file upload).

    Returns:
        A Pillow Image object with the background removed, or None if 
        the input was None.
    """
    if image_file is None:
        return None
    
    # Read the image data
    if hasattr(image_file, 'read'):
        input_data = image_file.read()
    else:
        with open(image_file, 'rb') as f:
            input_data = f.read()
    
    # Remove background
    output_data = remove(input_data)
    
    # Convert to PIL Image
    output_image = Image.open(io.BytesIO(output_data))
    return output_image

@spaces.GPU
def remove_background_multiple(image_files: List[Union[str, IO[bytes]]]) -> Tuple[Optional[str], List[Image.Image]]:
    """Removes backgrounds from multiple images and provides a zip archive.

    This function processes a list of image files. For each image, it removes
    the background, then bundles all the processed images (in PNG format)
    into a single zip file for download. It also returns a small number of
    processed images as previews for display in the UI.

    Args:
        image_files: A list of input images. Each element can be a 
            path to a file or a file-like object.

    Returns:
        A tuple containing:
            - A string with the file path to the generated zip archive. This 
              is `None` if the input list was empty.
            - A list of up to 5 processed Pillow Image objects to be used 
              as a preview.
    """
    if not image_files:
        return None, []
    
    processed_images = []
    preview_images = []
    
    # Create a zip file in memory
    zip_buffer = io.BytesIO()
    
    with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
        for i, image_file in enumerate(image_files):
            try:
                # Get original filename or create one
                if hasattr(image_file, 'name') and image_file.name:
                    original_name = os.path.basename(image_file.name)
                    name_without_ext = os.path.splitext(original_name)[0]
                else:
                    name_without_ext = f"image_{i+1}"
                
                # Remove background
                processed_image = remove_background_single(image_file)
                
                if processed_image:
                    # Save to zip file
                    img_buffer = io.BytesIO()
                    processed_image.save(img_buffer, format='PNG')
                    img_buffer.seek(0)
                    
                    zip_file.writestr(f"{name_without_ext}_no_bg.png", img_buffer.getvalue())
                    
                    # Add to preview (limit to first 5 images for display)
                    if len(preview_images) < 5:
                        preview_images.append(processed_image)
                    
            except Exception as e:
                print(f"Error processing image {i+1}: {str(e)}")
                continue
    
    zip_buffer.seek(0)
    
    # Save zip file temporarily
    zip_path = "processed_images.zip"
    with open(zip_path, 'wb') as f:
        f.write(zip_buffer.getvalue())
    
    return zip_path, preview_images

def process_images(single_image, multiple_images, processing_mode):
    """Main processing function based on selected mode"""
    if processing_mode == "Single Image":
        if single_image is None:
            return None, None, None, "Please upload an image first."
        
        try:
            result = remove_background_single(single_image)
            return result, None, None, "Background removed successfully!"
        except Exception as e:
            return None, None, None, f"Error processing image: {str(e)}"
    
    else:  # Multiple Images
        if not multiple_images:
            return None, None, None, "Please upload at least one image."
        
        try:
            zip_file, preview_images = remove_background_multiple(multiple_images)
            if zip_file and preview_images:
                return None, preview_images, zip_file, f"Processed {len(multiple_images)} images successfully! Download the zip file to get all results."
            else:
                return None, None, None, "No images were processed successfully."
        except Exception as e:
            return None, None, None, f"Error processing images: {str(e)}"

# Create Gradio interface
with gr.Blocks(title="Background Removal Tool", theme=gr.themes.Default()) as app:
    gr.Markdown(
        """
        # 🖼️ Background Removal Tool
        
        Upload your images and get them back with transparent backgrounds!
        
        **Choose your mode:**
        - **Single Image**: Upload one image and preview the result
        - **Multiple Images**: Upload multiple images and download them as a zip file
        """
    )
    
    with gr.Row():
        processing_mode = gr.Radio(
            choices=["Single Image", "Multiple Images"],
            value="Single Image",
            label="Processing Mode"
        )
    
    with gr.Row():
        with gr.Column():
            # Single image input (visible by default)
            single_image_input = gr.File(
                label="Upload Single Image",
                file_types=["image"],
                visible=True
            )
            
            # Multiple images input (hidden by default)
            multiple_images_input = gr.File(
                label="Upload Multiple Images",
                file_count="multiple",
                file_types=["image"],
                visible=False
            )
            
            process_btn = gr.Button("Remove Background", variant="primary", size="lg")
        
        with gr.Column():
            # Output for single image
            single_output = gr.Image(
                label="Result",
                visible=True
            )
            
            # Output for multiple images
            multiple_output_gallery = gr.Gallery(
                label="Preview (first 5 images)",
                visible=False,
                columns=3,
                rows=2,
                height="auto"
            )
            
            download_file = gr.File(
                label="Download All Processed Images",
                visible=False
            )
    
    status_message = gr.Textbox(label="Status", interactive=False)
    
    # Function to toggle visibility based on mode
    def toggle_inputs(mode):
        if mode == "Single Image":
            return (
                gr.update(visible=True),   # single_image_input
                gr.update(visible=False),  # multiple_images_input
                gr.update(visible=True),   # single_output
                gr.update(visible=False),  # multiple_output_gallery
                gr.update(visible=False)   # download_file
            )
        else:
            return (
                gr.update(visible=False),  # single_image_input
                gr.update(visible=True),   # multiple_images_input
                gr.update(visible=False),  # single_output
                gr.update(visible=True),   # multiple_output_gallery
                gr.update(visible=True)    # download_file
            )
    
    # Toggle inputs when mode changes
    processing_mode.change(
        fn=toggle_inputs,
        inputs=[processing_mode],
        outputs=[single_image_input, multiple_images_input, single_output, multiple_output_gallery, download_file]
    )
    
    # Process images when button is clicked
    process_btn.click(
        fn=process_images,
        inputs=[single_image_input, multiple_images_input, processing_mode],
        outputs=[single_output, multiple_output_gallery, download_file, status_message]
    )
    
    # Remove the separate gallery update function since we're handling it directly now

    gr.Markdown(
        """
        ### 📝 Instructions:
        1. Choose your processing mode (Single or Multiple images)
        2. Upload your image(s) - supports PNG, JPG, JPEG formats
        3. Click "Remove Background" to process
        4. For single images: preview the result directly
        5. For multiple images: download the zip file containing all processed images
        
        ### ⚡ Features:
        - Automatic background removal using AI
        - Support for batch processing
        - Transparent PNG output
        - Easy download of results
        """
    )

if __name__ == "__main__":
    app.launch(mcp_server=True)