Upload folder using huggingface_hub

app.py

@@ -0,0 +1,251 @@
+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)

requirements.txt

@@ -1,4 +1,4 @@
-gradio>=4.0.0
+gradio[mcp]>=4.0.0
 rembg>=2.0.0
 pillow>=9.0.0
 numpy>=1.21.0