src/utils/image_utils.py

import base64
from io import BytesIO
from PIL import Image, ImageChops
from PIL import ImageDraw
import math

class ImageUtils:
    def __init__(self):
        pass

    @staticmethod
    def crop_base64(base64_string, output_format='PNG') -> str:
        """
        Takes a base64 encoded image, crops it by removing uniform background,
        and returns the cropped image as base64.
        
        Args:
            base64_string (str or bytes): Base64 encoded image string or raw bytes
            output_format (str): Output image format ('PNG', 'JPEG', etc.)
        
        Returns:
            str: Base64 encoded cropped image, or empty string if cropping fails
        """
        try:
            # Handle both base64 strings and raw bytes
            if isinstance(base64_string, bytes):
                # If it's raw bytes, treat it as image data directly
                image_data = base64_string
            else:
                # If it's a string, decode base64 to image
                image_data = base64.b64decode(base64_string)
            
            im = Image.open(BytesIO(image_data))
            
            # Apply the original trim logic
            bg = Image.new(im.mode, im.size, im.getpixel((0,0)))
            diff = ImageChops.difference(im, bg)
            diff = ImageChops.add(diff, diff, 2.0, -100)
            bbox = diff.getbbox()
            
            if bbox:
                cropped_im = im.crop(bbox)
            else:
                cropped_im = im  # Return original if no cropping needed
            
            # Convert back to base64
            buffer = BytesIO()
            cropped_im.save(buffer, format=output_format)
            cropped_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
            
            return cropped_base64
            
        except Exception as e:
            print(f"Error processing image: {e}")
            return ""

    @staticmethod
    def crop_image(im: Image.Image) -> Image.Image:
        """
        Original trim function for PIL Image objects
        """
        try:
            bg = Image.new(im.mode, im.size, im.getpixel((0,0)))
            diff = ImageChops.difference(im, bg)
            diff = ImageChops.add(diff, diff, 2.0, -100)
            bbox = diff.getbbox()
            if bbox:
                return im.crop(bbox)
            return im
        except Exception as e:
            print(f"Error cropping image: {e}")
            return im

    @staticmethod
    def draw_bounding_boxes(pil_image: Image.Image, boxes: list[tuple[int, int, int, int]], color: str = "red", width: int = 2) -> Image.Image:
        """
        Draw bounding boxes on a PIL image.
        
        Args:
            pil_image: A PIL.Image instance.
            boxes: A list of boxes, each specified as (x1, y1, x2, y2).
            color: The color for the bounding box outline.
            width: The width of the bounding box line.
            
        Returns:
            The PIL.Image with drawn bounding boxes.
        """
        try:
            draw = ImageDraw.Draw(pil_image)
            for box in boxes:
                draw.rectangle(box, outline=color, width=width)
            return pil_image
        except Exception as e:
            print(f"Error drawing bounding boxes: {e}")
            return pil_image

    @staticmethod
    def standardize_image_size(image: Image.Image, target_size: tuple = (1200, 1600), maintain_aspect_ratio: bool = True) -> Image.Image:
        """
        Resize image to target size while optionally maintaining aspect ratio.
        
        Args:
            image: PIL Image to resize
            target_size: Target (width, height) in pixels
            maintain_aspect_ratio: If True, fit within target size while maintaining aspect ratio
            
        Returns:
            Resized PIL Image
        """
        if maintain_aspect_ratio:
            # Calculate aspect ratios
            img_ratio = image.width / image.height
            target_ratio = target_size[0] / target_size[1]
            
            if img_ratio > target_ratio:
                # Image is wider than target, fit to width
                new_width = target_size[0]
                new_height = int(target_size[0] / img_ratio)
            else:
                # Image is taller than target, fit to height
                new_height = target_size[1]
                new_width = int(target_size[1] * img_ratio)
            
            # Resize image
            resized_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
            
            # Create new image with target size and white background
            final_image = Image.new('RGB', target_size, 'white')
            
            # Calculate position to center the resized image
            x_offset = (target_size[0] - new_width) // 2
            y_offset = (target_size[1] - new_height) // 2
            
            # Paste the resized image onto the white background
            final_image.paste(resized_image, (x_offset, y_offset))
            
            return final_image
        else:
            # Direct resize to target size
            return image.resize(target_size, Image.Resampling.LANCZOS)

    @staticmethod
    def optimize_image_quality(image: Image.Image, max_size_bytes: int = 1024 * 1024, initial_quality: int = 95) -> tuple[Image.Image, int]:
        """
        Optimize image quality to fit within specified file size limit.
        
        Args:
            image: PIL Image to optimize
            max_size_bytes: Maximum file size in bytes (default 1MB)
            initial_quality: Starting quality (1-100) - not used for PNG but kept for compatibility
            
        Returns:
            Tuple of (optimized_image, final_quality)
        """
        # For PNG, we'll use compression levels instead of quality
        # PNG compression levels range from 0 (no compression) to 9 (maximum compression)
        compression_levels = [0, 1, 3, 5, 7, 9]  # Try different compression levels
        
        for compression in compression_levels:
            # Save image to buffer with current compression
            buffer = BytesIO()
            image.save(buffer, format='PNG', optimize=True, compress_level=compression)
            current_size = buffer.tell()
            
            # If size is within limit, return the image
            if current_size <= max_size_bytes:
                # Reset buffer position and load the optimized image
                buffer.seek(0)
                optimized_image = Image.open(buffer)
                return optimized_image, 95  # Return a default quality value for compatibility
        
        # If we can't get under the size limit, return the most compressed version
        buffer = BytesIO()
        image.save(buffer, format='PNG', optimize=True, compress_level=9)
        buffer.seek(0)
        optimized_image = Image.open(buffer)
        return optimized_image, 50  # Return a lower quality value for compatibility

    @staticmethod
    def process_image_for_comparison(image: Image.Image, target_size: tuple = (1200, 1600), max_size_bytes: int = 1024 * 1024) -> tuple[Image.Image, int, int]:
        """
        Process image for comparison: standardize size and optimize quality.
        
        Args:
            image: PIL Image to process
            target_size: Target size in pixels (width, height)
            max_size_bytes: Maximum file size in bytes (default 1MB)
            
        Returns:
            Tuple of (processed_image, final_quality, file_size_bytes)
        """
        # First, standardize the size
        sized_image = ImageUtils.standardize_image_size(image, target_size, maintain_aspect_ratio=True)
        
        # Then optimize quality to fit within size limit
        optimized_image, quality = ImageUtils.optimize_image_quality(sized_image, max_size_bytes)
        
        # Get final file size (using PNG format for consistency)
        buffer = BytesIO()
        optimized_image.save(buffer, format='PNG', optimize=True)
        file_size = buffer.tell()
        
        return optimized_image, quality, file_size

    @staticmethod
    def image_to_base64_optimized(image: Image.Image, target_size: tuple = (1200, 1600), max_size_bytes: int = 1024 * 1024) -> str:
        """
        Convert image to base64 with size and quality optimization.
        
        Args:
            image: PIL Image to convert
            target_size: Target size in pixels (width, height)
            max_size_bytes: Maximum file size in bytes (default 1MB)
            
        Returns:
            Base64 encoded string of the optimized image
        """
        processed_image, quality, file_size = ImageUtils.process_image_for_comparison(
            image, target_size, max_size_bytes
        )
        
        # Convert to base64 as PNG format
        buffer = BytesIO()
        processed_image.save(buffer, format='PNG', optimize=True)
        image_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
        
        return image_base64