danieljanata-ComfyUI-Prompting-System/database/image_handler.py

"""Handle image compression, base64 encoding, and thumbnail management"""

import io
import base64
from pathlib import Path
from typing import Optional, List, Dict
from PIL import Image
from datetime import datetime


class ImageHandler:
    """Manages image thumbnails with lock/unlock mechanism"""

    def __init__(self, thumbnail_size: tuple = (150, 150), quality: int = 70):
        """
        Initialize ImageHandler.

        Args:
            thumbnail_size: Tuple of (width, height) for thumbnails
            quality: JPEG quality (0-100)
        """
        self.thumbnail_size = thumbnail_size
        self.quality = quality

    def compress_and_encode(self, image_path: str) -> Optional[str]:
        """
        Compress image and encode to base64 data URI.

        Args:
            image_path: Path to source image

        Returns:
            Base64 data URI string or None on error
        """
        try:
            # Open and convert to RGB (handles RGBA, grayscale, etc.)
            img = Image.open(image_path)

            # Convert RGBA to RGB if needed
            if img.mode in ('RGBA', 'LA', 'P'):
                background = Image.new('RGB', img.size, (255, 255, 255))
                if img.mode == 'P':
                    img = img.convert('RGBA')
                background.paste(img, mask=img.split()[-1] if img.mode in ('RGBA', 'LA') else None)
                img = background
            elif img.mode != 'RGB':
                img = img.convert('RGB')

            # Create thumbnail (maintains aspect ratio)
            img.thumbnail(self.thumbnail_size, Image.Resampling.LANCZOS)

            # Encode to JPEG in memory
            buffer = io.BytesIO()
            img.save(buffer, format='JPEG', quality=self.quality, optimize=True)
            buffer.seek(0)

            # Encode to base64
            img_base64 = base64.b64encode(buffer.read()).decode('utf-8')

            # Return as data URI
            return f"data:image/jpeg;base64,{img_base64}"

        except Exception as e:
            print(f"[PromptSystem] Error compressing image {image_path}: {e}")
            return None

    def decode_and_save(self, base64_data: str, output_path: str) -> bool:
        """
        Decode base64 data URI and save to file.

        Args:
            base64_data: Base64 data URI string
            output_path: Path to save the image

        Returns:
            True if successful, False otherwise
        """
        try:
            # Remove data URI prefix if present
            if base64_data.startswith('data:image'):
                base64_data = base64_data.split(',', 1)[1]

            # Decode base64
            img_data = base64.b64decode(base64_data)

            # Save to file
            with open(output_path, 'wb') as f:
                f.write(img_data)

            return True

        except Exception as e:
            print(f"[PromptSystem] Error decoding and saving image: {e}")
            return False

    def manage_thumbnails(
        self,
        current_thumbnails: List[Dict],
        new_image_path: str,
        max_thumbnails: int = 3
    ) -> List[Dict]:
        """
        Manage thumbnail slots with lock/unlock mechanism.

        Logic:
        - Unlocked slots are always replaced with the newest image
        - Locked slots are preserved
        - Maximum of max_thumbnails slots
        - When all locked and adding new, remove oldest locked

        Args:
            current_thumbnails: List of thumbnail dicts with 'data', 'locked', 'timestamp'
            new_image_path: Path to new image to add
            max_thumbnails: Maximum number of thumbnails to keep

        Returns:
            Updated list of thumbnails
        """
        # Encode new image
        new_thumbnail_data = self.compress_and_encode(new_image_path)
        if not new_thumbnail_data:
            return current_thumbnails

        new_thumbnail = {
            'data': new_thumbnail_data,
            'locked': False,
            'timestamp': datetime.now().isoformat(),
            'source_image': str(Path(new_image_path).name)
        }

        # If no thumbnails yet, just add it
        if not current_thumbnails:
            return [new_thumbnail]

        # Find unlocked slot
        unlocked_index = None
        for i, thumb in enumerate(current_thumbnails):
            if not thumb.get('locked', False):
                unlocked_index = i
                break

        # Replace unlocked slot
        if unlocked_index is not None:
            current_thumbnails[unlocked_index] = new_thumbnail
            return current_thumbnails

        # All locked - add new slot if under max
        if len(current_thumbnails) < max_thumbnails:
            current_thumbnails.append(new_thumbnail)
            return current_thumbnails

        # All locked and at max - remove oldest locked
        # Sort by timestamp, keep newest (max_thumbnails - 1), then add new
        current_thumbnails.sort(key=lambda x: x.get('timestamp', ''), reverse=True)
        return current_thumbnails[:max_thumbnails - 1] + [new_thumbnail]

    def lock_thumbnail(self, thumbnails: List[Dict], index: int) -> List[Dict]:
        """
        Lock a specific thumbnail by index.

        Args:
            thumbnails: List of thumbnail dicts
            index: Index of thumbnail to lock

        Returns:
            Updated thumbnails list
        """
        if 0 <= index < len(thumbnails):
            thumbnails[index]['locked'] = True
        return thumbnails

    def unlock_thumbnail(self, thumbnails: List[Dict], index: int) -> List[Dict]:
        """
        Unlock a specific thumbnail by index.

        Args:
            thumbnails: List of thumbnail dicts
            index: Index of thumbnail to unlock

        Returns:
            Updated thumbnails list
        """
        if 0 <= index < len(thumbnails):
            thumbnails[index]['locked'] = False
        return thumbnails

    def merge_thumbnails(
        self,
        thumbnails_a: List[Dict],
        thumbnails_b: List[Dict],
        max_thumbnails: int = 3
    ) -> List[Dict]:
        """
        Merge two thumbnail lists, keeping locked ones and newest.

        Args:
            thumbnails_a: First thumbnail list
            thumbnails_b: Second thumbnail list
            max_thumbnails: Maximum thumbnails to keep

        Returns:
            Merged thumbnail list
        """
        # Collect all locked thumbnails
        locked = [t for t in thumbnails_a if t.get('locked', False)]
        locked.extend([t for t in thumbnails_b if t.get('locked', False)])

        # Remove duplicates based on data (same image)
        seen_data = set()
        unique_locked = []
        for thumb in locked:
            data_hash = hash(thumb.get('data', ''))
            if data_hash not in seen_data:
                seen_data.add(data_hash)
                unique_locked.append(thumb)

        # If more than max, keep newest locked
        if len(unique_locked) > max_thumbnails:
            unique_locked.sort(key=lambda x: x.get('timestamp', ''), reverse=True)
            return unique_locked[:max_thumbnails]

        # If under max, add newest unlocked to fill
        if len(unique_locked) < max_thumbnails:
            unlocked = [t for t in thumbnails_a if not t.get('locked', False)]
            unlocked.extend([t for t in thumbnails_b if not t.get('locked', False)])

            # Sort by timestamp, newest first
            unlocked.sort(key=lambda x: x.get('timestamp', ''), reverse=True)

            # Add unlocked until we reach max
            for thumb in unlocked:
                if len(unique_locked) >= max_thumbnails:
                    break
                data_hash = hash(thumb.get('data', ''))
                if data_hash not in seen_data:
                    seen_data.add(data_hash)
                    unique_locked.append(thumb)

        return unique_locked