background_remover.py

#!/usr/bin/env python3
"""
Image Background Removal Tool

A Python program that removes backgrounds from images using AI-powered rembg library.
Supports various image formats and provides both single file and batch processing.
"""

import os
import sys
import argparse
from pathlib import Path
from typing import List, Optional

import numpy as np
from PIL import Image
from rembg import remove
import cv2


class BackgroundRemover:
    """Main class for handling image background removal operations."""
    
    def __init__(self):
        self.supported_formats = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp'}
        
    def is_supported_format(self, file_path: str) -> bool:
        """Check if the file format is supported."""
        return Path(file_path).suffix.lower() in self.supported_formats
    
    def remove_background(self, input_path: str, output_path: Optional[str] = None) -> str:
        """
        Remove background from a single image.
        
        Args:
            input_path: Path to the input image
            output_path: Path for the output image (optional)
            
        Returns:
            Path to the output image
        """
        if not os.path.exists(input_path):
            raise FileNotFoundError(f"Input file not found: {input_path}")
            
        if not self.is_supported_format(input_path):
            raise ValueError(f"Unsupported file format: {Path(input_path).suffix}")
        
        # Generate output path if not provided
        if output_path is None:
            input_file = Path(input_path)
            output_path = str(input_file.parent / f"{input_file.stem}_no_bg.png")
        
        # Ensure output directory exists
        output_file = Path(output_path)
        try:
            output_file.parent.mkdir(parents=True, exist_ok=True)
        except Exception as e:
            raise RuntimeError(f"Failed to create output directory {output_file.parent}: {str(e)}")
        
        try:
            # Load the image
            with open(input_path, 'rb') as input_file:
                input_data = input_file.read()
            
            # Remove background using rembg
            output_data = remove(input_data)
            
            # Save the result - rembg returns bytes
            with open(output_path, 'wb') as output_file:
                # rembg.remove() returns bytes in practice
                output_file.write(output_data)  # type: ignore
            
            print(f"✓ Background removed: {input_path} -> {output_path}")
            return output_path
            
        except Exception as e:
            raise RuntimeError(f"Failed to process {input_path}: {str(e)}")
    
    def batch_process(self, input_dir: str, output_dir: Optional[str] = None) -> List[str]:
        """
        Process all supported images in a directory.
        
        Args:
            input_dir: Directory containing input images
            output_dir: Directory for output images (optional)
            
        Returns:
            List of output file paths
        """
        input_path = Path(input_dir)
        
        if not input_path.exists() or not input_path.is_dir():
            raise ValueError(f"Input directory does not exist: {input_dir}")
        
        # Set up output directory
        if output_dir is None:
            output_path = input_path / "no_background"
        else:
            output_path = Path(output_dir)
        
        try:
            output_path.mkdir(parents=True, exist_ok=True)
        except Exception as e:
            raise RuntimeError(f"Failed to create output directory {output_path}: {str(e)}")
        
        # Find all supported image files
        image_files = []
        for ext in self.supported_formats:
            image_files.extend(input_path.glob(f"*{ext}"))
            image_files.extend(input_path.glob(f"*{ext.upper()}"))
        
        if not image_files:
            print(f"No supported image files found in: {input_dir}")
            return []
        
        print(f"Found {len(image_files)} image(s) to process...")
        
        processed_files = []
        for img_file in image_files:
            try:
                output_file = output_path / f"{img_file.stem}_no_bg.png"
                self.remove_background(str(img_file), str(output_file))
                processed_files.append(str(output_file))
            except Exception as e:
                print(f"✗ Error processing {img_file}: {str(e)}")
        
        return processed_files
    
    def get_image_info(self, image_path: str) -> dict:
        """Get basic information about an image."""
        if not os.path.exists(image_path):
            raise FileNotFoundError(f"Image not found: {image_path}")
        
        try:
            with Image.open(image_path) as img:
                return {
                    'path': image_path,
                    'format': img.format,
                    'mode': img.mode,
                    'size': img.size,
                    'has_transparency': img.mode in ('RGBA', 'LA') or 'transparency' in img.info
                }
        except Exception as e:
            raise RuntimeError(f"Failed to read image info: {str(e)}")


def main():
    """Main function to handle command-line interface."""
    parser = argparse.ArgumentParser(
        description="Remove backgrounds from images using AI",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python background_remover.py image.jpg
  python background_remover.py image.jpg -o output.png
  python background_remover.py -d /path/to/images
  python background_remover.py -d /path/to/images -o /path/to/output
  python background_remover.py -i image.jpg
        """
    )
    
    parser.add_argument(
        'input',
        nargs='?',
        help='Input image file path'
    )
    
    parser.add_argument(
        '-o', '--output',
        help='Output file or directory path'
    )
    
    parser.add_argument(
        '-d', '--directory',
        help='Process all images in a directory'
    )
    
    parser.add_argument(
        '-i', '--info',
        help='Show information about an image file'
    )
    
    args = parser.parse_args()
    
    # Create background remover instance
    remover = BackgroundRemover()
    
    try:
        # Show image info
        if args.info:
            info = remover.get_image_info(args.info)
            print("\nImage Information:")
            print(f"Path: {info['path']}")
            print(f"Format: {info['format']}")
            print(f"Mode: {info['mode']}")
            print(f"Size: {info['size'][0]} x {info['size'][1]} pixels")
            print(f"Has transparency: {info['has_transparency']}")
            return
        
        # Batch processing
        if args.directory:
            print(f"Processing images in directory: {args.directory}")
            processed = remover.batch_process(args.directory, args.output)
            print(f"\n✓ Successfully processed {len(processed)} image(s)")
            return
        
        # Single file processing
        if args.input:
            if not os.path.exists(args.input):
                print(f"Error: File not found: {args.input}")
                sys.exit(1)
            
            output_path = remover.remove_background(args.input, args.output)
            
            # Show before/after info
            original_info = remover.get_image_info(args.input)
            result_info = remover.get_image_info(output_path)
            
            print(f"\nProcessing complete!")
            print(f"Original: {original_info['size'][0]}x{original_info['size'][1]} ({original_info['format']})")
            print(f"Result: {result_info['size'][0]}x{result_info['size'][1]} ({result_info['format']})")
            return
        
        # If no arguments provided, show help
        parser.print_help()
        
    except Exception as e:
        print(f"Error: {str(e)}")
        sys.exit(1)


if __name__ == "__main__":
    main()