app/utils/image_processing.py

"""
Image processing utilities for the FastAPI application.
"""

import base64
import io
from typing import Tuple
import numpy as np
from PIL import Image
import cv2

from app.core.config import settings
from app.core.logging import get_logger

logger = get_logger(__name__)


def decode_base64_image(image_data: str) -> Tuple[np.ndarray, Tuple[int, int]]:
    """
    Decode base64 image data to numpy array.
    
    Args:
        image_data: Base64 encoded image string
        
    Returns:
        Tuple of (image_array, (width, height))
    """
    try:
        # Remove data URL prefix if present
        if image_data.startswith('data:image'):
            image_data = image_data.split(',')[1]
        
        # Decode base64
        image_bytes = base64.b64decode(image_data)
        
        # Open with PIL
        pil_image = Image.open(io.BytesIO(image_bytes))
        
        # Convert to RGB if necessary
        if pil_image.mode != 'RGB':
            pil_image = pil_image.convert('RGB')
        
        # Get original dimensions
        original_dims = pil_image.size  # (width, height)
        
        # Convert to numpy array
        image_array = np.array(pil_image)
        
        logger.debug(f"Decoded image with shape: {image_array.shape}")
        
        return image_array, original_dims
        
    except Exception as e:
        logger.error(f"Failed to decode base64 image: {str(e)}")
        raise ValueError(f"Invalid image data: {str(e)}")


def encode_image_to_base64(image: np.ndarray, format: str = "JPEG", quality: int = 95) -> str:
    """
    Encode numpy array image to base64 string.
    
    Args:
        image: Image as numpy array
        format: Image format (JPEG, PNG, etc.)
        quality: JPEG quality (1-100)
        
    Returns:
        Base64 encoded image string
    """
    try:
        # Convert numpy array to PIL Image
        if image.dtype != np.uint8:
            image = (image * 255).astype(np.uint8)
        
        pil_image = Image.fromarray(image)
        
        # Save to bytes buffer
        buffer = io.BytesIO()
        save_kwargs = {"format": format}
        
        if format.upper() == "JPEG":
            save_kwargs["quality"] = quality
            save_kwargs["optimize"] = True
        
        pil_image.save(buffer, **save_kwargs)
        
        # Encode to base64
        image_bytes = buffer.getvalue()
        base64_string = base64.b64encode(image_bytes).decode('utf-8')
        
        return base64_string
        
    except Exception as e:
        logger.error(f"Failed to encode image to base64: {str(e)}")
        raise ValueError(f"Image encoding failed: {str(e)}")


def validate_image_size(image: np.ndarray) -> bool:
    """
    Validate image dimensions.
    
    Args:
        image: Image as numpy array
        
    Returns:
        True if image size is valid
    """
    height, width = image.shape[:2]
    
    # Check minimum and maximum dimensions
    min_dim = min(width, height)
    max_dim = max(width, height)
    
    if min_dim < 32:  # Too small
        return False
    
    if max_dim > 4096:  # Too large
        return False
    
    return True


def resize_image_if_needed(image: np.ndarray, max_size: int = 1280) -> np.ndarray:
    """
    Resize image if it's too large while maintaining aspect ratio.
    
    Args:
        image: Image as numpy array
        max_size: Maximum dimension size
        
    Returns:
        Resized image
    """
    height, width = image.shape[:2]
    
    if max(height, width) <= max_size:
        return image
    
    # Calculate new dimensions
    if width > height:
        new_width = max_size
        new_height = int(height * (max_size / width))
    else:
        new_height = max_size
        new_width = int(width * (max_size / height))
    
    # Resize using PIL for better quality
    pil_image = Image.fromarray(image)
    resized_pil = pil_image.resize((new_width, new_height), Image.Resampling.LANCZOS)
    
    return np.array(resized_pil)


def validate_image_format(image_bytes: bytes) -> bool:
    """
    Validate if the image format is supported.
    
    Args:
        image_bytes: Raw image bytes
        
    Returns:
        True if format is supported
    """
    try:
        with Image.open(io.BytesIO(image_bytes)) as img:
            # Check if format is in allowed extensions
            format_lower = img.format.lower() if img.format else ""
            allowed_formats = {"jpeg", "jpg", "png", "bmp", "tiff", "webp"}
            return format_lower in allowed_formats
    except Exception:
        return False


def get_image_info(image: np.ndarray) -> dict:
    """
    Get information about an image.
    
    Args:
        image: Image as numpy array
        
    Returns:
        Dictionary with image information
    """
    height, width = image.shape[:2]
    channels = image.shape[2] if len(image.shape) > 2 else 1
    
    return {
        "width": width,
        "height": height,
        "channels": channels,
        "dtype": str(image.dtype),
        "size_mb": image.nbytes / (1024 * 1024)
    }