refactor(texture): texture loading accelerated

Author: danila-schelkov

xcoder/console.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/images.py

@@ -20,49 +20,45 @@
 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.frombytes(
+            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),
+                "raw",
+            )
 
-                    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 +126,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 +135,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(bytearray(pixel))
 
-            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 +151,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/objects/texture.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