refactor: Texture2DSwizzler - wrap methods

Author: isHarryh

UnityPy/export/Texture2DConverter.py

@@ -3,7 +3,7 @@
 import struct
 from io import BytesIO
 from threading import Lock
-from typing import TYPE_CHECKING, Callable, Dict, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Callable, Dict, Optional, Sequence, Tuple, Union
 
 import astc_encoder
 import texture2ddecoder
@@ -17,6 +17,9 @@
     from ..classes import Texture2D
 
 
+PlatformBlobType = Union[bytes, Sequence[int]]
+
+
 TEXTURE_FORMAT_BLOCK_SIZE_TABLE: Dict[TF, Optional[Tuple[int, int]]] = {}
 for tf in TF:
     if tf.name.startswith("ASTC"):
@@ -134,7 +137,7 @@ def image_to_texture2d(
     img: Image.Image,
     target_texture_format: Union[TF, int],
     platform: int = 0,
-    platform_blob: Optional[bytes] = None,
+    platform_blob: Optional[PlatformBlobType] = None,
     flip: bool = True,
 ) -> Tuple[bytes, TF]:
     if not isinstance(target_texture_format, TF):
@@ -216,27 +219,18 @@ def image_to_texture2d(
         pil_mode = "RGB"
     # everything else defaulted to RGBA
 
-    switch_swizzle = None
-    if platform == BuildTarget.Switch and platform_blob:
-        gobs_per_block = TextureSwizzler.get_switch_gobs_per_block(platform_blob)
-
+    if TextureSwizzler.is_switch_swizzled(platform, platform_blob):
+        assert platform_blob is not None
         if texture_format == TF.RGB24:
             texture_format = TF.RGBA32
             pil_mode = "RGBA"
         elif texture_format == TF.BGR24:
             texture_format = TF.BGRA32
             pil_mode = "BGRA"
 
-        block_size = TextureSwizzler.TEXTUREFORMAT_BLOCK_SIZE_MAP.get(texture_format)
-        if not block_size:
-            raise NotImplementedError(
-                f"Not implemented swizzle format: {texture_format.name}"
-            )
-
-        width, height = TextureSwizzler.get_padded_texture_size(
-            img.width, img.height, *block_size, gobs_per_block
+        width, height = TextureSwizzler.get_padded_image_size(
+            img.width, img.height, texture_format, platform_blob
         )
-        switch_swizzle = (block_size, gobs_per_block)
     else:
         width, height = get_compressed_image_size(img.width, img.height, texture_format)
 
@@ -248,9 +242,11 @@ def image_to_texture2d(
     else:
         enc_img = img.tobytes("raw", pil_mode)
 
-    if switch_swizzle is not None:
-        block_size, gobs_per_block = switch_swizzle
-        enc_img = TextureSwizzler.swizzle(enc_img, width, height, *block_size, gobs_per_block)
+    if TextureSwizzler.is_switch_swizzled(platform, platform_blob):
+        assert platform_blob is not None
+        enc_img = TextureSwizzler.swizzle(
+            enc_img, width, height, texture_format, platform_blob
+        )
 
     return enc_img, texture_format
 
@@ -285,10 +281,10 @@ def parse_image_data(
     image_data: bytes,
     width: int,
     height: int,
-    texture_format: Union[int, TF],
+    texture_format: Union[TF, int],
     version: Tuple[int, int, int, int],
     platform: int,
-    platform_blob: Optional[bytes] = None,
+    platform_blob: Optional[PlatformBlobType] = None,
     flip: bool = True,
 ) -> Image.Image:
     if not width or not height:
@@ -304,30 +300,19 @@ def parse_image_data(
     if platform == BuildTarget.XBOX360 and texture_format in XBOX_SWAP_FORMATS:
         image_data = swap_bytes_for_xbox(image_data)
 
-    original_width, original_height = width, height
-    switch_swizzle = None
-    if platform == BuildTarget.Switch and platform_blob:
-        gobs_per_block = TextureSwizzler.get_switch_gobs_per_block(platform_blob)
+    ori_width, ori_height = width, height
 
-        pil_mode = "RGBA"
+    if TextureSwizzler.is_switch_swizzled(platform, platform_blob):
+        assert platform_blob is not None
         if texture_format == TF.RGB24:
             texture_format = TF.RGBA32
         elif texture_format == TF.BGR24:
             texture_format = TF.BGRA32
-            pil_mode = "BGRA"
-        elif texture_format == TF.Alpha8:
-            pil_mode = "L"
-
-        block_size = TextureSwizzler.TEXTUREFORMAT_BLOCK_SIZE_MAP.get(texture_format)
-        if not block_size:
-            raise NotImplementedError(
-                f"Not implemented swizzle format: {texture_format.name}"
-            )
 
-        width, height = TextureSwizzler.get_padded_texture_size(
-            width, height, *block_size, gobs_per_block
+        width, height = TextureSwizzler.get_padded_image_size(
+            width, height, texture_format, platform_blob
         )
-        switch_swizzle = (block_size, gobs_per_block, pil_mode)
+        image_data = TextureSwizzler.deswizzle(image_data, width, height, texture_format, platform_blob)
     else:
         width, height = get_compressed_image_size(width, height, texture_format)
 
@@ -342,20 +327,15 @@ def parse_image_data(
         else:
             image_data = texture2ddecoder.unpack_crunch(image_data)
 
-    if switch_swizzle is not None:
-        block_size, gobs_per_block, pil_mode = switch_swizzle
-        image_data = TextureSwizzler.deswizzle(image_data, width, height, *block_size, gobs_per_block)
-
-
     conv_func = CONV_TABLE.get(texture_format)
     if not conv_func:
         raise NotImplementedError(
             f"Not implemented texture format: {texture_format.name}"
         )
     img = conv_func(image_data, width, height)
 
-    if original_width != width or original_height != height:
-        img = img.crop((0, 0, original_width, original_height))
+    if ori_width != width or ori_height != height:
+        img = img.crop((0, 0, ori_width, ori_height))
 
     return img.transpose(Image.FLIP_TOP_BOTTOM) if flip else img
 

UnityPy/helpers/TextureSwizzler.py

@@ -1,8 +1,14 @@
 # based on https://github.com/nesrak1/UABEA/blob/master/TexturePlugin/Texture2DSwitchDeswizzler.cs
-from typing import Dict, Tuple
+from typing import Dict, Optional, Sequence, Tuple, Union
+
+from UnityPy.enums import BuildTarget
 
 from ..enums import TextureFormat as TF
 
+
+PlatformBlobType = Union[bytes, Sequence[int]]
+
+
 GOB_X_TEXEL_COUNT = 4
 GOB_Y_TEXEL_COUNT = 8
 TEXEL_BYTE_SIZE = 16
@@ -13,20 +19,20 @@
 ]
 
 
-def ceil_divide(a: int, b: int) -> int:
+def _ceil_divide(a: int, b: int) -> int:
     return (a + b - 1) // b
 
 
-def deswizzle(
+def _deswizzle(
     data: bytes,
     width: int,
     height: int,
     block_width: int,
     block_height: int,
     texels_per_block: int,
 ) -> bytes:
-    block_count_x = ceil_divide(width, block_width)
-    block_count_y = ceil_divide(height, block_height)
+    block_count_x = _ceil_divide(width, block_width)
+    block_count_y = _ceil_divide(height, block_height)
     gob_count_x = block_count_x // GOB_X_TEXEL_COUNT
     gob_count_y = block_count_y // GOB_Y_TEXEL_COUNT
     new_data = bytearray(len(data))
@@ -50,16 +56,16 @@ def deswizzle(
     return bytes(new_data)
 
 
-def swizzle(
+def _swizzle(
     data: bytes,
     width: int,
     height: int,
     block_width: int,
     block_height: int,
     texels_per_block: int,
 ) -> bytes:
-    block_count_x = ceil_divide(width, block_width)
-    block_count_y = ceil_divide(height, block_height)
+    block_count_x = _ceil_divide(width, block_width)
+    block_count_y = _ceil_divide(height, block_height)
     gob_count_x = block_count_x // GOB_X_TEXEL_COUNT
     gob_count_y = block_count_y // GOB_Y_TEXEL_COUNT
     new_data = bytearray(len(data))
@@ -83,8 +89,31 @@ def swizzle(
     return bytes(new_data)
 
 
+def _get_padded_texture_size(
+    width: int, height: int, block_width: int, block_height: int, texels_per_block: int
+) -> Tuple[int, int]:
+    width = (
+        _ceil_divide(width, block_width * GOB_X_TEXEL_COUNT)
+        * block_width
+        * GOB_X_TEXEL_COUNT
+    )
+    height = (
+        _ceil_divide(height, block_height * GOB_Y_TEXEL_COUNT * texels_per_block)
+        * block_height
+        * GOB_Y_TEXEL_COUNT
+        * texels_per_block
+    )
+    return width, height
+
+
+def _get_texels_per_block(platform_blob: PlatformBlobType) -> int:
+    if not platform_blob:
+        raise ValueError("Given platform_blob is empty")
+    return 1 << int.from_bytes(platform_blob[8:12], "little")
+
+
 # this should be the amount of pixels that can fit 16 bytes
-TEXTUREFORMAT_BLOCK_SIZE_MAP: Dict[TF, Tuple[int, int]] = {
+TEXTURE_FORMAT_BLOCK_SIZE_MAP: Dict[TF, Tuple[int, int]] = {
     TF.Alpha8: (16, 1),  # 1 byte per pixel
     TF.ARGB4444: (8, 1),  # 2 bytes per pixel
     TF.RGBA32: (4, 1),  # 4 bytes per pixel
@@ -117,22 +146,59 @@ def swizzle(
 }
 
 
-def get_padded_texture_size(
-    width: int, height: int, block_width: int, block_height: int, texels_per_block: int
-):
-    width = (
-        ceil_divide(width, block_width * GOB_X_TEXEL_COUNT)
-        * block_width
-        * GOB_X_TEXEL_COUNT
-    )
-    height = (
-        ceil_divide(height, block_height * GOB_Y_TEXEL_COUNT * texels_per_block)
-        * block_height
-        * GOB_Y_TEXEL_COUNT
-        * texels_per_block
-    )
-    return width, height
+def deswizzle(
+    data: bytes,
+    width: int,
+    height: int,
+    texture_format: TF,
+    platform_blob: PlatformBlobType,
+) -> bytes:
+    block_size = TEXTURE_FORMAT_BLOCK_SIZE_MAP.get(texture_format)
+    if not block_size:
+        raise NotImplementedError(
+            f"Not implemented swizzle format: {texture_format.name}"
+        )
+    texels_per_block = _get_texels_per_block(platform_blob)
+    return _deswizzle(data, width, height, *block_size, texels_per_block)
 
 
-def get_switch_gobs_per_block(platform_blob: bytes) -> int:
-    return 1 << int.from_bytes(platform_blob[8:12], "little")
+def swizzle(
+    data: bytes,
+    width: int,
+    height: int,
+    texture_format: TF,
+    platform_blob: PlatformBlobType,
+) -> bytes:
+    block_size = TEXTURE_FORMAT_BLOCK_SIZE_MAP.get(texture_format)
+    if not block_size:
+        raise NotImplementedError(
+            f"Not implemented swizzle format: {texture_format.name}"
+        )
+    texels_per_block = _get_texels_per_block(platform_blob)
+    return _swizzle(data, width, height, *block_size, texels_per_block)
+
+
+def get_padded_image_size(
+    width: int,
+    height: int,
+    texture_format: TF,
+    platform_blob: PlatformBlobType,
+):
+    block_size = TEXTURE_FORMAT_BLOCK_SIZE_MAP.get(texture_format)
+    if not block_size:
+        raise NotImplementedError(
+            f"Not implemented swizzle format: {texture_format.name}"
+        )
+    texels_per_block = _get_texels_per_block(platform_blob)
+    return _get_padded_texture_size(width, height, *block_size, texels_per_block)
+
+
+def is_switch_swizzled(
+    platform: Union[BuildTarget, int], platform_blob: Optional[PlatformBlobType] = None
+) -> bool:
+    if platform != BuildTarget.Switch:
+        return False
+    if not platform_blob or len(platform_blob) < 12:
+        return False
+    gobs_per_block = _get_texels_per_block(platform_blob)
+    return gobs_per_block > 1