diff --git a/gui/sixel.py b/gui/sixel.py index ae0c008..9809262 100644 --- a/gui/sixel.py +++ b/gui/sixel.py @@ -115,7 +115,7 @@ def register_image_colors(colors: List[Tuple[int, int, int]]) -> Dict[Tuple[int, Returns: Dictionary mapping RGB tuples to their palette indices """ - global _NEXT_IMAGE_COLOR_INDEX, _ITERM2_COLOR_TABLE + global _NEXT_IMAGE_COLOR_INDEX, _ITERM2_COLOR_TABLE, _ITERM2_PALETTE_BYTES color_map = {} new_colors_added = False @@ -131,6 +131,7 @@ def register_image_colors(colors: List[Tuple[int, int, int]]) -> Dict[Tuple[int, # Invalidate iTerm2 color cache if new colors were added if new_colors_added: _ITERM2_COLOR_TABLE = None + _ITERM2_PALETTE_BYTES = None return color_map @@ -766,6 +767,7 @@ def pixels_to_png( # Cache for iTerm2 color lookup table (built once, reused) _ITERM2_COLOR_TABLE: Optional[List[Tuple[int, int, int]]] = None +_ITERM2_PALETTE_BYTES: Optional[bytes] = None def _get_iterm2_color_table() -> List[Tuple[int, int, int]]: @@ -782,10 +784,26 @@ def _get_iterm2_color_table() -> List[Tuple[int, int, int]]: return _ITERM2_COLOR_TABLE +def _get_iterm2_palette_bytes() -> bytes: + """Get or build a 256-color palette byte string for paletted PNGs.""" + global _ITERM2_PALETTE_BYTES + if _ITERM2_PALETTE_BYTES is None: + color_table = _get_iterm2_color_table() + palette = bytearray(256 * 3) + for idx, (r, g, b) in enumerate(color_table[:256]): + base = idx * 3 + palette[base] = r + palette[base + 1] = g + palette[base + 2] = b + _ITERM2_PALETTE_BYTES = bytes(palette) + return _ITERM2_PALETTE_BYTES + + def invalidate_iterm2_color_cache() -> None: """Invalidate the iTerm2 color cache (call after registering new image colors).""" - global _ITERM2_COLOR_TABLE + global _ITERM2_COLOR_TABLE, _ITERM2_PALETTE_BYTES _ITERM2_COLOR_TABLE = None + _ITERM2_PALETTE_BYTES = None def pixels_to_iterm2(pixels: List[List[int]], width: int, height: int) -> str: @@ -794,7 +812,7 @@ def pixels_to_iterm2(pixels: List[List[int]], width: int, height: int) -> str: Optimized for speed with: - Pre-built color lookup table for O(1) access - - JPEG encoding (faster than PNG, hardware accelerated on macOS) + - Paletted PNG encoding when possible (1 byte per pixel, lossless) - Direct byte array construction Args: @@ -809,31 +827,47 @@ def pixels_to_iterm2(pixels: List[List[int]], width: int, height: int) -> str: # Fall back to sixel if PIL not available return pixels_to_sixel(pixels, width, height) - # Get cached color table for fast lookup + # Prefer paletted PNG when possible (faster + smaller, no quality loss) color_table = _get_iterm2_color_table() - table_len = len(color_table) - - # Build raw RGB bytes directly - rgb_bytes = bytearray(width * height * 3) - idx = 0 - for row in pixels: - for color_idx in row: - if color_idx < table_len: - r, g, b = color_table[color_idx] - else: - r, g, b = 0, 0, 0 - rgb_bytes[idx] = r - rgb_bytes[idx + 1] = g - rgb_bytes[idx + 2] = b - idx += 3 - - # Create image from raw bytes - img = Image.frombytes("RGB", (width, height), bytes(rgb_bytes)) - - # Encode to JPEG - faster than PNG, especially on macOS - buffer = io.BytesIO() - img.save(buffer, format='JPEG', quality=90) - image_data = buffer.getvalue() + max_color_idx = len(color_table) - 1 + use_palette = max_color_idx < 256 + + if use_palette: + # Build palette-indexed bytes (1 byte per pixel) + index_bytes = bytearray(width * height) + idx = 0 + for row in pixels: + for color_idx in row: + index_bytes[idx] = color_idx + idx += 1 + + img = Image.frombytes("P", (width, height), bytes(index_bytes)) + img.putpalette(_get_iterm2_palette_bytes()) + + buffer = io.BytesIO() + img.save(buffer, format='PNG', optimize=False, compress_level=1) + image_data = buffer.getvalue() + else: + # Fallback to JPEG for large color palettes + table_len = len(color_table) + rgb_bytes = bytearray(width * height * 3) + idx = 0 + for row in pixels: + for color_idx in row: + if color_idx < table_len: + r, g, b = color_table[color_idx] + else: + r, g, b = 0, 0, 0 + rgb_bytes[idx] = r + rgb_bytes[idx + 1] = g + rgb_bytes[idx + 2] = b + idx += 3 + + img = Image.frombytes("RGB", (width, height), bytes(rgb_bytes)) + + buffer = io.BytesIO() + img.save(buffer, format='JPEG', quality=90) + image_data = buffer.getvalue() # Base64 encode b64_data = base64.b64encode(image_data).decode('ascii')