demo_generate/rembg_method.py

from PIL import Image
import io
import logging
from datetime import datetime

logger = logging.getLogger(__name__)

def remove_background(image_data):
    start_time = datetime.now()
    logger.info(f"Removing background - {start_time}")
    
    try:
        # Try to import rembg
        try:
            from rembg import remove
            logger.info("rembg module imported successfully")
        except ImportError:
            logger.error("rembg module not found!")
            raise ImportError('rembg library not installed')
            
        input_image = Image.open(io.BytesIO(image_data))
        logger.info(f"Image opened. Mode: {input_image.mode}, Size: {input_image.size}")
        
        if input_image.mode != 'RGB':
            logger.info("Converting image to RGB...")
            input_image = input_image.convert('RGB')
            
        logger.info("Removing background...")
        output_image = remove(input_image)
        
        img_io = io.BytesIO()
        output_image.save(img_io, 'PNG')
        img_io.seek(0)
        
        duration = (datetime.now() - start_time).total_seconds()
        logger.info(f"Successfully processed. Duration: {duration} seconds")
        return img_io
    except Exception as e:
        logger.error(f"Background removal error: {str(e)}", exc_info=True)
        raise