astc_encoder.pil_codec

  1"""PIL codec for ASTC images.
  2
  3Importing this module provides an ASTC codec for PIL for encoding and decoding ASTC images.
  4"""
  5
  6import struct
  7from typing import Any, List, Union
  8
  9from PIL import Image, ImageFile
 10
 11from astc_encoder import (
 12    ASTCConfig,
 13    ASTCConfigFlags,
 14    ASTCContext,
 15    ASTCImage,
 16    ASTCProfile,
 17    ASTCSwizzle,
 18    ASTCType,
 19)
 20
 21
 22class ASTCEncoder(ImageFile.PyEncoder):  # noqa: D101
 23    _pushes_fd: bool = True
 24    context: ASTCContext
 25
 26    def init(  # noqa: D102
 27        self,
 28        args: List[Any],
 29    ):
 30        assert len(args) in (4, 5), "Invalid number of arguments"
 31        profile: int = args[0]
 32        quality: float = args[1]
 33        block_width: int = args[2]
 34        block_height: int = args[3]
 35        block_depth: int = args[4] if len(args) > 4 else 1
 36
 37        assert block_depth == 1, "Cannot handle 3D textures"
 38
 39        config = ASTCConfig(
 40            ASTCProfile(profile),
 41            block_width,
 42            block_height,
 43            block_depth,
 44            quality=quality,
 45        )
 46        self.context = ASTCContext(config)
 47
 48    def encode(self, bufsize: int) -> tuple[int, int, bytes]:  # noqa: D102
 49        assert self.im is not None, "No image set"  # type: ignore
 50
 51        mode: str = self.mode
 52        pack_struct: struct.Struct
 53        if mode == "RGBA":
 54            swizzle = ASTCSwizzle.from_str("rgba")
 55            pack_struct = struct.Struct("4B")
 56        elif mode == "RGB":
 57            swizzle = ASTCSwizzle.from_str("rgb1")
 58            pack_struct = struct.Struct("3Bx")
 59        else:
 60            raise ValueError(f"Unsupported mode: {mode}")
 61
 62        pa = self.im.pixel_access()
 63        data = b"".join(
 64            pack_struct.pack(*pa[x, y])  # type: ignore
 65            for y in range(self.state.ysize)
 66            for x in range(self.state.xsize)
 67        )
 68
 69        astc_img = ASTCImage(
 70            ASTCType.U8,
 71            self.state.xsize,
 72            self.state.ysize,
 73            data=data,
 74        )
 75        comp = self.context.compress(astc_img, swizzle)
 76        return len(comp), 1, comp
 77
 78
 79Image.register_encoder("astc", ASTCEncoder)
 80
 81
 82class ASTCDecoder(ImageFile.PyDecoder):  # noqa: D101
 83    _pull_fd: bool = True
 84    context: ASTCContext
 85
 86    def init(self, args: List[Any]):  # noqa: D102
 87        assert len(args) in (3, 4), "Invalid number of arguments"
 88        profile: int = ASTCProfile(args[0])
 89        block_width: int = args[1]
 90        block_height: int = args[2]
 91        block_depth: int = args[3] if len(args) > 3 else 1
 92        assert block_depth == 1, "Cannot handle 3D textures"
 93
 94        config = ASTCConfig(
 95            profile,
 96            block_width,
 97            block_height,
 98            block_depth,
 99            flags=ASTCConfigFlags.DECOMPRESS_ONLY,
100        )
101        self.context = ASTCContext(config)
102
103    def decode(
104        self, buffer: Union[bytes, Image.SupportsArrayInterface]
105    ) -> tuple[int, int]:  # noqa: D102
106        assert self.state.xoff == 0 and self.state.yoff == 0, "Cannot handle offsets"
107
108        config = self.context.config
109        assert self.state.xsize % config.block_x == 0, (
110            "Invalid width, must be multiple of block width"
111        )
112        assert self.state.ysize % config.block_y == 0, (
113            "Invalid height, must be multiple of block height"
114        )
115
116        block_count_x = (self.state.xsize + config.block_x - 1) // config.block_x
117        block_count_y = (self.state.ysize + config.block_y - 1) // config.block_y
118        expected_size = block_count_x * block_count_y * 16
119
120        if len(buffer) != expected_size and self.fd is not None:
121            buffer = self.fd.read(expected_size)
122
123        if len(buffer) != expected_size:
124            raise ValueError("Not enough data")
125
126        astc_img = ASTCImage(
127            ASTCType.U8,
128            self.state.xsize,
129            self.state.ysize,
130        )
131
132        mode: str = self.mode
133        if mode == "RGBA":
134            swizzle = ASTCSwizzle.from_str("rgba")
135        elif mode == "RGB":
136            swizzle = ASTCSwizzle.from_str("rgb1")
137        else:
138            # TODO: LA, L
139            raise ValueError(f"Unsupported mode: {mode}")
140
141        self.context.decompress(buffer, astc_img, swizzle)
142
143        self.set_as_raw(astc_img.data, "RGBA")
144        return -1, 0
145
146
147Image.register_decoder("astc", ASTCDecoder)