app.py

import gradio as gr
from rembg import remove
from PIL import Image
import io
import os
import zipfile
from datetime import datetime

def remove_background(image):
    output = remove(image)
    return output

def process_images(images):
    outputs = []
    for img in images:
        img_pil = Image.open(img.name)
        output = remove_background(img_pil)
        outputs.append(output)
    return outputs

def save_images(images):
    if not images:
        return None, "No images to save."
    
    # Create a directory to store processed images
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    save_dir = f"processed_images_{timestamp}"
    os.makedirs(save_dir, exist_ok=True)
    
    # Save individual images
    for i, img in enumerate(images):
        img.save(os.path.join(save_dir, f"processed_image_{i+1}.png"))
    
    # Create a zip file
    zip_filename = f"processed_images_{timestamp}.zip"
    with zipfile.ZipFile(zip_filename, 'w') as zipf:
        for root, dirs, files in os.walk(save_dir):
            for file in files:
                zipf.write(os.path.join(root, file), file)
    
    return zip_filename, f"Images saved and zipped as {zip_filename}"

def process_and_save(images):
    processed_images = process_images(images)
    zip_file, save_message = save_images(processed_images)
    return processed_images, save_message, zip_file

iface = gr.Interface(
    fn=process_and_save,
    inputs=gr.File(file_count="multiple", label="Upload Images"),
    outputs=[
        gr.Gallery(label="Processed Images"),
        gr.Textbox(label="Save Status"),
        gr.File(label="Download Processed Images")
    ],
    title="Batch Background Remover",
    description="Upload multiple images to remove their backgrounds in batch. All processed images will be saved and available for download as a zip file."
)

if __name__ == "__main__":
    iface.launch()