AIGame/Tools/model_tools/png_compress.py

"""Lossless PNG size reduction for bake outputs.

Blender's default PNG compression is low (~15/100). Final _base/_mix (and
bake maps) are rewritten with zlib level 9. Pillow `optimize=True` is used
when available (system Python); Blender's interpreter falls back to IDAT
recompress via stdlib only. Pixels are unchanged.
"""
from __future__ import annotations

import os
import struct
import zlib
from typing import Dict, Iterable, List, Optional

PNG_SIG = b"\x89PNG\r\n\x1a\n"


def prepare_blender_png(compression: int = 100) -> None:
    """Blender Image.save() uses scene render PNG compression (0-100)."""
    import bpy

    settings = bpy.context.scene.render.image_settings
    settings.file_format = "PNG"
    settings.color_depth = "8"
    settings.compression = int(compression)


def compress_png(path: str, level: int = 9) -> Dict[str, int]:
    """Rewrite `path` in place. Returns {before, after} byte sizes."""
    path = os.path.abspath(path)
    before = os.path.getsize(path)
    after = before
    tmp = path + ".pngcompress.tmp"
    try:
        if _compress_with_pillow(path, tmp, level) or _compress_idat(path, tmp, level):
            new_size = os.path.getsize(tmp)
            if new_size < before:
                os.replace(tmp, path)
                after = new_size
            elif os.path.isfile(tmp):
                os.remove(tmp)
    finally:
        if os.path.isfile(tmp):
            try:
                os.remove(tmp)
            except OSError:
                pass
    return {"before": before, "after": after}


def compress_pngs(paths: Iterable[str], level: int = 9) -> List[Dict[str, object]]:
    reports = []
    for path in paths:
        if not path or not os.path.isfile(path):
            continue
        if os.path.splitext(path)[1].lower() != ".png":
            continue
        sizes = compress_png(path, level=level)
        reports.append({"path": os.path.abspath(path), **sizes})
    return reports


def compress_dir_pngs(directory: str, suffixes: Optional[Iterable[str]] = None) -> List[Dict[str, object]]:
    """Compress PNGs in a directory. `suffixes` like ('_base.png', '_mix.png')."""
    directory = os.path.abspath(directory)
    names = sorted(os.listdir(directory)) if os.path.isdir(directory) else []
    suffix_tuple = tuple(suffixes) if suffixes else None
    paths = []
    for name in names:
        if not name.lower().endswith(".png"):
            continue
        if suffix_tuple and not name.endswith(suffix_tuple):
            continue
        paths.append(os.path.join(directory, name))
    return compress_pngs(paths)


def _compress_with_pillow(src: str, dst: str, level: int) -> bool:
    try:
        from PIL import Image
    except ImportError:
        return False
    try:
        with Image.open(src) as im:
            im.load()
            im.save(dst, format="PNG", optimize=True, compress_level=level)
    except Exception:
        return False
    return os.path.isfile(dst)


def _write_chunk(buf: bytearray, ctype: bytes, data: bytes) -> None:
    buf.extend(struct.pack(">I", len(data)))
    buf.extend(ctype)
    buf.extend(data)
    buf.extend(struct.pack(">I", zlib.crc32(ctype + data) & 0xFFFFFFFF))


def _compress_idat(src: str, dst: str, level: int) -> bool:
    with open(src, "rb") as f:
        data = f.read()
    if len(data) < 8 or data[:8] != PNG_SIG:
        return False
    pos = 8
    before_idat: List[tuple] = []
    after_idat: List[tuple] = []
    idat_parts: List[bytes] = []
    seen_idat = False
    while pos + 12 <= len(data):
        length = struct.unpack(">I", data[pos : pos + 4])[0]
        ctype = data[pos + 4 : pos + 8]
        end = pos + 8 + length
        if end + 4 > len(data):
            return False
        cdata = data[pos + 8 : end]
        pos = end + 4
        if ctype == b"IDAT":
            idat_parts.append(cdata)
            seen_idat = True
            continue
        if seen_idat:
            after_idat.append((ctype, cdata))
        else:
            before_idat.append((ctype, cdata))
        if ctype == b"IEND":
            break
    if not idat_parts:
        return False
    raw = zlib.decompress(b"".join(idat_parts))
    packed = zlib.compress(raw, level)
    out = bytearray(PNG_SIG)
    for ctype, cdata in before_idat:
        _write_chunk(out, ctype, cdata)
    _write_chunk(out, b"IDAT", packed)
    for ctype, cdata in after_idat:
        _write_chunk(out, ctype, cdata)
    with open(dst, "wb") as f:
        f.write(out)
    return True