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()