stickerbot/background_remover.py

"""Background removal module supporting local (rembg) and API (remove.bg) methods."""
import io
import logging
import time
from typing import Optional, Tuple, Union
import requests
from PIL import Image

logger = logging.getLogger(__name__)

# Remove.bg API endpoint
REMOVE_BG_API_URL = "https://api.remove.bg/v1.0/removebg"
MAX_API_CALLS_PER_MONTH = 50

# Retry configuration
MAX_RETRIES = 3
RETRY_DELAYS = [2, 5, 10]  # seconds


async def remove_background_local(image_bytes: bytes) -> Optional[bytes]:
    """
    Remove background using local rembg library (U²Net model).
    
    Args:
        image_bytes: Raw image bytes
        
    Returns:
        PNG image bytes with transparent background, or None if failed
    """
    try:
        from rembg import remove
        
        # Open image
        input_image = Image.open(io.BytesIO(image_bytes))
        
        # Remove background
        output_image = remove(input_image)
        
        # Save to bytes
        output_buffer = io.BytesIO()
        output_image.save(output_buffer, format='PNG')
        output_buffer.seek(0)
        
        logger.info("Background removed successfully (local)")
        return output_buffer.getvalue()
        
    except Exception as e:
        logger.error(f"Error removing background locally: {e}")
        return None


async def remove_background_api(
    image_bytes: bytes, 
    api_key: str
) -> Tuple[Union[bytes, str], str]:
    """
    Remove background using remove.bg API with retry logic.
    
    Args:
        image_bytes: Raw image bytes
        api_key: Remove.bg API key
        
    Returns:
        Tuple of (result, message)
        - Success: (image_bytes, "success")
        - Rate limit: ("LIMIT_REACHED", error_message)
        - Other error: ("API_ERROR", error_message)
    """
    headers = {
        "X-Api-Key": api_key
    }
    
    last_error = None
    
    # Retry loop
    for attempt in range(MAX_RETRIES):
        try:
            # Prepare multipart form data
            files = {
                'image_file': ('image.png', io.BytesIO(image_bytes), 'image/png')
            }
            
            response = requests.post(
                REMOVE_BG_API_URL,
                headers=headers,
                files=files,
                timeout=30
            )
            
            # Check response
            if response.status_code == 200:
                # Success - return image
                logger.info("Background removed successfully (API)")
                return response.content, "success"
            
            elif response.status_code in (402, 429):
                # Rate limit / Payment required (limit reached)
                error_data = response.json()
                error_msg = error_data.get('errors', [{}])[0].get('title', 'Rate limit exceeded')
                logger.warning(f"Remove.bg API limit reached: {error_msg}")
                return "LIMIT_REACHED", error_msg
            
            else:
                # Other API error
                error_data = response.json() if response.text else {}
                error_msg = error_data.get('errors', [{}])[0].get('title', f'HTTP {response.status_code}')
                logger.error(f"Remove.bg API error: {error_msg}")
                last_error = error_msg
                
        except requests.exceptions.Timeout:
            last_error = "Timeout - API antwortet nicht"
            logger.warning(f"API timeout (attempt {attempt + 1}/{MAX_RETRIES})")
            
        except requests.exceptions.ConnectionError:
            last_error = "Verbindungsfehler - API nicht erreichbar"
            logger.warning(f"API connection error (attempt {attempt + 1}/{MAX_RETRIES})")
            
        except Exception as e:
            last_error = str(e)
            logger.error(f"API call error (attempt {attempt + 1}/{MAX_RETRIES}): {e}")
        
        # Wait before retry (except for last attempt)
        if attempt < MAX_RETRIES - 1:
            delay = RETRY_DELAYS[attempt]
            logger.info(f"Retrying in {delay} seconds...")
            time.sleep(delay)
    
    # All retries failed
    logger.error(f"All API retries failed: {last_error}")
    return "API_ERROR", f"API nicht erreichbar nach {MAX_RETRIES} Versuchen: {last_error}"


async def should_show_api_limit_warning(api_key: str, current_usage: int) -> bool:
    """
    Check if we should show warning on the API button.
    
    Args:
        api_key: API key (if None, no warning)
        current_usage: Current month usage count
        
    Returns:
        True if limit reached or very close
    """
    if not api_key:
        return False
    
    # Show warning when at or above limit
    return current_usage >= MAX_API_CALLS_PER_MONTH


async def validate_image_for_bg_removal(image_bytes: bytes) -> Tuple[bool, str]:
    """
    Validate if image is suitable for background removal.
    
    Args:
        image_bytes: Raw image bytes
        
    Returns:
        Tuple of (is_valid, error_message)
    """
    try:
        image = Image.open(io.BytesIO(image_bytes))
        
        # Check format
        if image.format not in ('PNG', 'JPEG', 'JPG', 'WEBP', 'BMP'):
            return False, f"Format '{image.format}' nicht unterstützt. Nutze PNG, JPG oder WebP."
        
        # Check size
        width, height = image.size
        if width < 64 or height < 64:
            return False, "Bild zu klein (min. 64x64 Pixel)"
        
        if width > 4096 or height > 4096:
            return False, "Bild zu groß (max. 4096x4096 Pixel)"
        
        # Check file size
        if len(image_bytes) > 10 * 1024 * 1024:  # 10MB
            return False, "Datei zu groß (max. 10MB)"
        
        return True, ""
        
    except Exception as e:
        return False, f"Ungültiges Bild: {str(e)}"