Python 3 script to optimize textures in simple(!) glTF 2.0 GLB files.

optimize_simple_glb.py

import json
from io import BytesIO
from pathlib import Path
from struct import Struct
from typing import TYPE_CHECKING

from PIL import Image
from PIL.PngImagePlugin import PngImageFile

if TYPE_CHECKING:
    from io import BufferedIOBase


class SimpleGLB:
    HEADER_STRUCT = Struct("<4sII")
    CHUNK_STRUCT = Struct("<I4s")

    def __init__(self, file_path: str | Path) -> None:
        if not isinstance(file_path, Path):
            assert isinstance(file_path, str)
            file_path = Path(file_path)

        assert file_path.is_file()

        self.file_path = file_path

        with open(file_path, "rb") as file:
            magic, version, length = SimpleGLB.HEADER_STRUCT.unpack(file.read(12))

            assert magic == b"glTF"
            assert version == 2

            self._original_size: int = length
            self._read_json_chunk(file)
            self._read_bin_chunk(file)

    def _read_json_chunk(self, file: "BufferedIOBase"):
        chunk_length, chunk_type = SimpleGLB.CHUNK_STRUCT.unpack(file.read(8))

        assert chunk_type == b"JSON"

        self.json: dict = json.loads(file.read(chunk_length))

        assert self.json["asset"]["version"] == "2.0"
        assert len(self.json["buffers"]) == 1
        assert len(self.json["buffers"][0].keys()) == 1

    def _read_bin_chunk(self, file: "BufferedIOBase"):
        chunk_length, chunk_type = SimpleGLB.CHUNK_STRUCT.unpack(file.read(8))

        assert chunk_type == b"BIN\0"

        self.bin_chunk_offset = file.tell()

        self.buffers: list[bytes] = [None] * len(self.json["bufferViews"])
        for i, buffer in enumerate(self.json["bufferViews"]):
            file.seek(self.bin_chunk_offset + buffer["byteOffset"])
            self.buffers[i] = bytearray(file.read(buffer["byteLength"]))

    def _update_file(self):
        with open(self.file_path, "wb") as file:
            json_data = json.dumps(
                self.json, ensure_ascii=False, separators=(",", ":"), indent=None
            ).encode()
            spaces = (4 - (len(json_data) & 3)) & 3
            json_length = len(json_data) + spaces

            bin = bytearray()
            for buffer in self.buffers:
                for _ in range((4 - (len(bin) & 3)) & 3):
                    bin.extend(b"\0")
                bin.extend(buffer)

            alignment = (4 - (len(bin) & 3)) & 3
            bin_length = len(bin) + alignment

            full_length = 12 + (2 * (2 * 4)) + json_length + bin_length

            file.write(SimpleGLB.HEADER_STRUCT.pack(b"glTF", 2, full_length))

            file.write(SimpleGLB.CHUNK_STRUCT.pack(json_length, b"JSON"))
            file.write(json_data)
            file.write(b" " * spaces)

            file.write(SimpleGLB.CHUNK_STRUCT.pack(bin_length, b"BIN\0"))
            file.write(bin)
            file.write(b"\0" * alignment)

    def _refresh_buffer_views(self):
        previous_offset = 0
        for buffer_view in self.json["bufferViews"]:
            length = buffer_view["byteLength"]
            offset = buffer_view["byteOffset"] = previous_offset
            alignment = (4 - (length & 3)) & 3
            previous_offset = length + offset + alignment


class SimpleGLBTextureOptimizer(SimpleGLB):
    def __init__(self, file_path: str | Path) -> None:
        super().__init__(file_path)

    @property
    def textures(self) -> list[dict[str, int | str]]:
        return self.json["images"]

    def optimize_textures(self):
        for tex in self.textures:
            self._optimize_texture(tex)
        self._refresh_buffer_views()
        self._update_file()

    def _optimize_texture(self, tex_obj: dict[str, int | str]):
        def get_jpeg_quality(name: str) -> int:
            name = name.upper()
            if name in ("NM", "NORMAL") or name.endswith("_NM") or "_NM_" in name:
                return 97
            if name == "DF" or name.endswith("_DF") or "_DF_" in name:
                return 95
            if name == "ORM" or name.endswith("_ORM") or "_ORM_" in name:
                return 95
            return 95

        def is_alpha_channel_empty(img: PngImageFile) -> bool:
            assert img.mode == "RGBA"  # cba to handle indexed, etc...

            if img.getextrema()[3] == (255, 255):
                return True
            return False

        if "mimeType" not in tex_obj or "bufferView" not in tex_obj:
            return

        if tex_obj["mimeType"] != "image/png":
            return

        name, buffer_idx = tex_obj.get("name", "Unnamed"), tex_obj["bufferView"]
        orig_size = self.json["bufferViews"][buffer_idx]["byteLength"]
        img = Image.open(BytesIO(self.buffers[buffer_idx]))
        format = "PNG"
        params = {
            "optimize": True,
        }

        if img.mode != "RGBA" or is_alpha_channel_empty(img):
            if img.mode != "RGB":
                img = img.convert("RGB")
            format = "JPEG"
            params |= {
                "quality": get_jpeg_quality(name),
                "subsampling": "4:4:4",
            }

        temp = BytesIO()
        img.save(temp, format=format, **params)
        new_size = temp.getbuffer().nbytes

        diff = orig_size - new_size
        print(
            f"{name}: reduced by {(1 - (new_size / orig_size)) * 100:.3f}%",
            f"({diff} bytes difference)",
        )

        # with open(
        #     rf"C:\Users\floxay\Desktop\ValModels\t_{name}.{format.lower()}",
        #     "wb",
        # ) as out:
        #     out.write(temp.getbuffer())

        tex_obj["mimeType"] = f"image/{format.lower()}"
        self.buffers[buffer_idx] = bytearray(temp.getbuffer())
        self.json["bufferViews"][buffer_idx]["byteLength"] = new_size


if __name__ == "__main__":
    for glb in Path(R"C:\Users\floxay\Desktop\ValModels").glob("*.glb"):
        SimpleGLBTextureOptimizer(glb).optimize_textures()