Merge pull request #31 from xcoder-tool/texture-loading

refactor(texture): texture loading accelerated

Author: danila-schelkov

poetry.lock

@@ -11,7 +11,7 @@ sc-compression = "0.6.5"
 colorama = "0.4.6"
 pylzham = "^0.1.3"
 zstandard = "^0.23.0"
-pillow = "11.2.0"
+pillow = "~11.2.1"
 loguru = "0.7.3"
 
 

pyproject.toml

@@ -1,36 +1,45 @@
 import io
+import struct
 from typing import Literal
 
 
-class Reader(io.BytesIO):
+class Reader:
     def __init__(
         self,
         initial_buffer: bytes = b"",
         endian: Literal["little", "big"] = "little",
     ):
-        super().__init__(initial_buffer)
+        self._internal_reader = io.BytesIO(initial_buffer)
+
         self.endian: Literal["little", "big"] = endian
+        self.endian_sign: Literal["<", ">"] = "<" if endian == "little" else ">"
+
+    def seek(self, position: int) -> None:
+        self._internal_reader.seek(position)
 
-    def read_integer(self, length: int, signed=False) -> int:
-        return int.from_bytes(self.read(length), self.endian, signed=signed)
+    def tell(self) -> int:
+        return self._internal_reader.tell()
+
+    def read(self, size: int) -> bytes:
+        return self._internal_reader.read(size)
 
     def read_uchar(self) -> int:
-        return self.read_integer(1)
+        return struct.unpack("B", self.read(1))[0]
 
     def read_char(self) -> int:
-        return self.read_integer(1, True)
+        return struct.unpack("b", self.read(1))[0]
 
     def read_ushort(self) -> int:
-        return self.read_integer(2)
+        return struct.unpack(f"{self.endian_sign}H", self.read(2))[0]
 
     def read_short(self) -> int:
-        return self.read_integer(2, True)
+        return struct.unpack(f"{self.endian_sign}h", self.read(2))[0]
 
     def read_uint(self) -> int:
-        return self.read_integer(4)
+        return struct.unpack(f"{self.endian_sign}I", self.read(4))[0]
 
     def read_int(self) -> int:
-        return self.read_integer(4, True)
+        return struct.unpack(f"{self.endian_sign}i", self.read(4))[0]
 
     def read_twip(self) -> float:
         return self.read_int() / 20
@@ -72,6 +81,7 @@ def write_string(self, string: str | None = None):
         if string is None:
             self.write_byte(0xFF)
             return
+
         encoded = string.encode()
         self.write_byte(len(encoded))
         self.write(encoded)

xcoder/bytestream.py

@@ -1,15 +1,19 @@
 class Console:
-    @staticmethod
-    def progress_bar(message, current, total):
+    previous_percentage: int = -1
+
+    @classmethod
+    def progress_bar(cls, message: str, current: int, total: int) -> None:
         percentage = (current + 1) * 100 // total
-        print("\r", end="")
-        print(f"[{percentage}%] {message}", end="")
-        if current + 1 == total:
-            print()
+        if percentage == cls.previous_percentage:
+            return
 
-    @staticmethod
-    def percent(current: int, total: int) -> int:
-        return (current + 1) * 100 // total
+        print(f"\r[{percentage}%] {message}", end="")
+
+        if percentage == 100:
+            print()
+            cls.previous_percentage = -1
+        else:
+            cls.previous_percentage = percentage
 
     @staticmethod
     def ask_integer(message: str):
@@ -20,13 +24,16 @@ def ask_integer(message: str):
                 pass
 
     @staticmethod
-    def question(message):
+    def question(message: str) -> bool:
         while True:
             answer = input(f"[????] {message} [Y/n] ").lower()
+            if not answer:
+                return True
+
             if answer in "ny":
                 break
 
-        return "ny".index(answer)
+        return bool("ny".index(answer))
 
 
 if __name__ == "__main__":

xcoder/console.py

@@ -20,49 +20,44 @@
 CHUNK_SIZE = 32
 
 
-def load_image_from_buffer(img: Image.Image) -> None:
-    width, height = img.size
-    loaded_image = img.load()
-    if loaded_image is None:
-        raise Exception("loaded_image is None")
-
+def load_image_from_buffer(pixel_type: int, width: int, height: int) -> Image.Image:
     with open("pixel_buffer", "rb") as pixel_buffer:
-        channel_count = int.from_bytes(pixel_buffer.read(1), "little")
+        pixel_buffer.read(1)
+
+        return Image.frombuffer(
+            get_format_by_pixel_type(pixel_type), (width, height), pixel_buffer.read()
+        )
 
-        for y in range(height):
-            for x in range(width):
-                loaded_image[x, y] = tuple(pixel_buffer.read(channel_count))
 
+def join_image(pixel_type: int, width: int, height: int) -> Image.Image:
+    mode = get_format_by_pixel_type(pixel_type)
+    image = Image.new(mode, (width, height))
 
-def join_image(img: Image.Image) -> None:
     with open("pixel_buffer", "rb") as pixel_buffer:
         channel_count = int.from_bytes(pixel_buffer.read(1), "little")
 
-        width, height = img.size
-        loaded_image = img.load()
-        if loaded_image is None:
-            raise Exception("loaded_image is None")
+        chunk_count_x = math.ceil(width / CHUNK_SIZE)
+        chunk_count_y = math.ceil(height / CHUNK_SIZE)
+        chunk_count = chunk_count_x * chunk_count_y
 
-        x_chunks_count = width // CHUNK_SIZE
-        y_chunks_count = height // CHUNK_SIZE
+        for chunk_index in range(chunk_count):
+            chunk_x = chunk_index % chunk_count_x
+            chunk_y = chunk_index // chunk_count_x
 
-        for y_chunk in range(y_chunks_count + 1):
-            for x_chunk in range(x_chunks_count + 1):
-                for y in range(CHUNK_SIZE):
-                    pixel_y = y_chunk * CHUNK_SIZE + y
-                    if pixel_y >= height:
-                        break
+            chunk_width = min(width - chunk_x * CHUNK_SIZE, CHUNK_SIZE)
+            chunk_height = min(height - chunk_y * CHUNK_SIZE, CHUNK_SIZE)
+
+            sub_image = Image.frombuffer(
+                mode,
+                (chunk_width, chunk_height),
+                pixel_buffer.read(channel_count * chunk_width * chunk_height),
+            )
 
-                    for x in range(CHUNK_SIZE):
-                        pixel_x = x_chunk * CHUNK_SIZE + x
-                        if pixel_x >= width:
-                            break
+            image.paste(sub_image, (chunk_x * CHUNK_SIZE, chunk_y * CHUNK_SIZE))
 
-                        loaded_image[pixel_x, pixel_y] = tuple(
-                            pixel_buffer.read(channel_count)
-                        )
+            Console.progress_bar(locale.join_pic, chunk_index, chunk_count)
 
-            Console.progress_bar(locale.join_pic, y_chunk, y_chunks_count + 1)
+    return image
 
 
 def _add_pixel(
@@ -130,7 +125,7 @@ def get_format_by_pixel_type(pixel_type: int) -> str:
     raise Exception(locale.unknown_pixel_type % pixel_type)
 
 
-def load_texture(reader: Reader, pixel_type: int, img: Image.Image) -> None:
+def load_texture(reader: Reader, pixel_type: int, width: int, height: int) -> None:
     channel_count = get_channel_count_by_pixel_type(pixel_type)
     read_pixel = get_read_function(pixel_type)
     if read_pixel is None:
@@ -139,18 +134,12 @@ def load_texture(reader: Reader, pixel_type: int, img: Image.Image) -> None:
     with open("pixel_buffer", "wb") as pixel_buffer:
         pixel_buffer.write(channel_count.to_bytes(1, "little"))
 
-        width, height = img.size
-        point = -1
         for y in range(height):
-            for x in range(width):
-                pixel = read_pixel(reader)
-                for channel in pixel:
-                    pixel_buffer.write(channel.to_bytes(1, "little"))
+            pixel_buffer.write(
+                b"".join([bytearray(read_pixel(reader)) for _ in range(width)])
+            )
 
-            curr = Console.percent(y, height)
-            if curr > point:
-                Console.progress_bar(locale.crt_pic, y, height)
-                point = curr
+            Console.progress_bar(locale.crt_pic, y, height)
 
 
 def save_texture(writer: Writer, image: Image.Image, pixel_type: int) -> None:
@@ -161,16 +150,12 @@ def save_texture(writer: Writer, image: Image.Image, pixel_type: int) -> None:
     width, height = image.size
 
     pixels = image.getdata()
-    point = -1
     for y in range(height):
         for x in range(width):
             # noinspection PyTypeChecker
             writer.write(encode_pixel(pixels[y * width + x]))
 
-        curr = Console.percent(y, height)
-        if curr > point:
-            Console.progress_bar(locale.writing_pic, y, height)
-            point = curr
+        Console.progress_bar(locale.writing_pic, y, height)
 
 
 def transform_image(

xcoder/images.py

@@ -6,12 +6,7 @@
 import zstandard
 from PIL import Image
 
-from xcoder.images import (
-    get_format_by_pixel_type,
-    join_image,
-    load_image_from_buffer,
-    load_texture,
-)
+from xcoder.images import join_image, load_image_from_buffer, load_texture
 from xcoder.pvr_tex_tool import get_image_from_ktx_data
 
 if TYPE_CHECKING:
@@ -59,17 +54,14 @@ def load(self, swf: SupercellSWF, tag: int, has_texture: bool):
             )
             return
 
-        image = Image.new(
-            get_format_by_pixel_type(self.pixel_type), (self.width, self.height)
-        )
+        try:
+            load_texture(swf.reader, self.pixel_type, self.width, self.height)
 
-        load_texture(swf.reader, self.pixel_type, image)
-
-        if tag in (27, 28, 29):
-            join_image(image)
-        else:
-            load_image_from_buffer(image)
-
-        os.remove("pixel_buffer")
+            if tag in (27, 28, 29):
+                image = join_image(self.pixel_type, self.width, self.height)
+            else:
+                image = load_image_from_buffer(self.pixel_type, self.width, self.height)
+        finally:
+            os.remove("pixel_buffer")
 
         self.image = image

xcoder/objects/texture.py

@@ -27,12 +27,7 @@ def get_channel_count_by_pixel_type(pixel_type: int) -> int:
 
 
 def _read_rgba8(reader: Reader) -> PixelChannels:
-    return (
-        reader.read_uchar(),
-        reader.read_uchar(),
-        reader.read_uchar(),
-        reader.read_uchar(),
-    )
+    return tuple(reader.read(4))
 
 
 def _read_rgba4(reader: Reader) -> PixelChannels:
@@ -61,11 +56,11 @@ def _read_rgb565(reader: Reader) -> PixelChannels:
 
 
 def _read_luminance8_alpha8(reader: Reader) -> PixelChannels:
-    return (reader.read_uchar(), reader.read_uchar())[::-1]
+    return tuple(reader.read(2))[::-1]
 
 
 def _read_luminance8(reader: Reader) -> PixelChannels:
-    return (reader.read_uchar(),)
+    return tuple(reader.read(1))
 
 
 def _write_rgba8(pixel: PixelChannels) -> bytes: