stickerbot/image_processor.py

"""Image processing module for converting images to Telegram sticker format."""
import io
import logging
from typing import Optional, Tuple
from PIL import Image

logger = logging.getLogger(__name__)

# Telegram sticker requirements
STICKER_SIZE = 512  # One side must be exactly 512px
MAX_FILE_SIZE_MB = 3  # Maximum file size in MB
MAX_FILE_SIZE_BYTES = MAX_FILE_SIZE_MB * 1024 * 1024


async def process_image_for_sticker(file_data: bytes, filename: Optional[str] = None) -> Optional[bytes]:
    """
    Process an image file to meet Telegram sticker requirements.
    
    Args:
        file_data: Raw image bytes from Telegram
        filename: Original filename (optional, for format detection)
        
    Returns:
        Processed image bytes in WebP format, or None if processing fails
    """
    try:
        # Open image from bytes
        image = Image.open(io.BytesIO(file_data))
        
        # Convert to RGBA if necessary (Telegram requires transparency support)
        if image.mode != 'RGBA':
            if image.mode == 'RGB':
                # Add alpha channel
                image = image.convert('RGBA')
            elif image.mode in ('L', 'LA', 'P'):
                # Grayscale or palette - convert to RGBA
                image = image.convert('RGBA')
            else:
                # Any other mode - convert to RGBA
                image = image.convert('RGBA')
        
        # Resize to Telegram sticker format (512px on one side)
        image = resize_for_sticker(image)
        
        # Convert to WebP format with optimization
        output = io.BytesIO()
        
        # Try different quality levels to meet file size requirements
        quality = 95
        while quality >= 60:
            output.seek(0)
            output.truncate(0)
            
            image.save(
                output,
                format='WEBP',
                quality=quality,
                method=6,  # Best compression method
                exact=True  # Preserve transparency
            )
            
            file_size = output.tell()
            
            if file_size <= MAX_FILE_SIZE_BYTES:
                logger.info(f"Image processed: {image.size}, quality={quality}, size={file_size/1024:.1f}KB")
                output.seek(0)
                return output.getvalue()
            
            # Reduce quality and try again
            quality -= 5
            logger.debug(f"File too large ({file_size/1024:.1f}KB), reducing quality to {quality}")
        
        # If we get here, even lowest quality is too large
        logger.error(f"Could not compress image below {MAX_FILE_SIZE_MB}MB")
        return None
        
    except Exception as e:
        logger.error(f"Error processing image: {e}")
        return None


def resize_for_sticker(image: Image.Image) -> Image.Image:
    """
    Resize image to meet Telegram sticker requirements.
    One side must be exactly 512px, other side <= 512px
    
    Args:
        image: PIL Image object
        
    Returns:
        Resized image
    """
    width, height = image.size
    
    # Determine which side is longer
    if width >= height:
        # Landscape or square - set width to 512
        new_width = STICKER_SIZE
        new_height = int((height / width) * STICKER_SIZE)
    else:
        # Portrait - set height to 512
        new_height = STICKER_SIZE
        new_width = int((width / height) * STICKER_SIZE)
    
    # Ensure neither dimension exceeds 512
    new_width = min(new_width, STICKER_SIZE)
    new_height = min(new_height, STICKER_SIZE)
    
    # Resize using LANCZOS for best quality
    resized = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
    
    logger.debug(f"Resized image from {width}x{height} to {new_width}x{new_height}")
    return resized


async def validate_image(file_data: bytes) -> Tuple[bool, str]:
    """
    Validate if file data is a valid image.
    
    Args:
        file_data: Raw file bytes
        
    Returns:
        Tuple of (is_valid, error_message)
    """
    try:
        image = Image.open(io.BytesIO(file_data))
        
        # Check format
        if image.format not in ('PNG', 'JPEG', 'JPG', 'WEBP', 'GIF', 'BMP', 'TIFF'):
            return False, f"Unsupported image format: {image.format}"
        
        # Check if image can be processed
        width, height = image.size
        if width < 64 or height < 64:
            return False, "Image too small (minimum 64x64 pixels)"
        
        # Check file size
        if len(file_data) > 10 * 1024 * 1024:  # 10MB limit for input
            return False, "Image file too large (max 10MB for upload)"
        
        return True, ""
        
    except Exception as e:
        return False, f"Invalid image file: {str(e)}"


async def get_image_info(file_data: bytes) -> dict:
    """
    Get information about an image.
    
    Args:
        file_data: Raw image bytes
        
    Returns:
        Dictionary with image info
    """
    try:
        image = Image.open(io.BytesIO(file_data))
        return {
            "format": image.format,
            "mode": image.mode,
            "size": image.size,
            "width": image.width,
            "height": image.height,
        }
    except Exception as e:
        logger.error(f"Error getting image info: {e}")
        return {}