refactor(Texture2DConverter): converter table

Author: isHarryh

UnityPy/export/Texture2DConverter.py

@@ -1,4 +1,4 @@
-from __future__ import annotations
+from __future__ import annotations
 
 import struct
 from io import BytesIO
@@ -8,7 +8,6 @@
 import astc_encoder
 import texture2ddecoder
 from PIL import Image
-from typing_extensions import Unpack
 
 from ..enums import BuildTarget, TextureFormat
 from ..helpers import TextureSwizzler
@@ -235,11 +234,11 @@ def image_to_texture2d(
             assert platform_blob is not None
         gobsPerBlock = TextureSwizzler.get_switch_gobs_per_block(platform_blob)
         s_tex_format = tex_format
-        if tex_format == TextureFormat.RGB24:
-            s_tex_format = TextureFormat.RGBA32
+        if tex_format == TF.RGB24:
+            s_tex_format = TF.RGBA32
             pil_mode = "RGBA"
-        elif tex_format == TextureFormat.BGR24:
-            s_tex_format = TextureFormat.BGRA32
+        elif tex_format == TF.BGR24:
+            s_tex_format = TF.BGRA32
             pil_mode = "BGRA"
         block_size = TextureSwizzler.TEXTUREFORMAT_BLOCK_SIZE_MAP[s_tex_format]
         width, height = TextureSwizzler.get_padded_texture_size(
@@ -311,10 +310,10 @@ def parse_image_data(
             # due to earlier check platform_blob can't be None
             assert platform_blob is not None
         gobsPerBlock = TextureSwizzler.get_switch_gobs_per_block(platform_blob)
-        if texture_format == TextureFormat.RGB24:
-            texture_format = TextureFormat.RGBA32
-        elif texture_format == TextureFormat.BGR24:
-            texture_format = TextureFormat.BGRA32
+        if texture_format == TF.RGB24:
+            texture_format = TF.RGBA32
+        elif texture_format == TF.BGR24:
+            texture_format = TF.BGRA32
         block_size = TextureSwizzler.TEXTUREFORMAT_BLOCK_SIZE_MAP[texture_format]
         width, height = TextureSwizzler.get_padded_texture_size(
             width, height, *block_size, gobsPerBlock
@@ -325,11 +324,6 @@ def parse_image_data(
     else:
         width, height = get_compressed_image_size(width, height, texture_format)
 
-    selection = CONV_TABLE[texture_format]
-
-    if not selection:
-        raise NotImplementedError(f"Not implemented texture format: {texture_format}")
-
     image_data = bytes(image_data)
 
     if "Crunched" in texture_format.name:
@@ -343,7 +337,13 @@ def parse_image_data(
         else:
             image_data = texture2ddecoder.unpack_crunch(image_data)
 
-    img = selection[0](image_data, width, height, *selection[1:])
+    if texture_format not in CONV_TABLE:
+        raise NotImplementedError(
+            f"Not implemented texture format: {texture_format.name}"
+        )
+    conv_func, conv_args = CONV_TABLE[texture_format]
+
+    img = conv_func(image_data, width, height, *conv_args)
 
     if original_width != width or original_height != height:
         img = img.crop((0, 0, original_width, original_height))
@@ -443,22 +443,22 @@ def pvrtc(image_data: bytes, width: int, height: int, fmt: bool) -> Image.Image:
     return Image.frombytes("RGBA", (width, height), image_data, "raw", "BGRA")
 
 
-def etc(image_data: bytes, width: int, height: int, fmt: list) -> Image.Image:
-    if fmt[0] == 1:
+def etc(image_data: bytes, width: int, height: int, fmt: str) -> Image.Image:
+    if fmt == "ETC1":
         image_data = texture2ddecoder.decode_etc1(image_data, width, height)
-    elif fmt[0] == 2:
-        if fmt[1] == "RGB":
-            image_data = texture2ddecoder.decode_etc2(image_data, width, height)
-        elif fmt[1] == "A1":
-            image_data = texture2ddecoder.decode_etc2a1(image_data, width, height)
-        elif fmt[1] == "A8":
-            image_data = texture2ddecoder.decode_etc2a8(image_data, width, height)
+    elif fmt == "ETC2_RGB":
+        image_data = texture2ddecoder.decode_etc2(image_data, width, height)
+    elif fmt == "ETC2_A1":
+        image_data = texture2ddecoder.decode_etc2a1(image_data, width, height)
+    elif fmt == "ETC2_A8":
+        image_data = texture2ddecoder.decode_etc2a8(image_data, width, height)
     else:
-        raise NotImplementedError("unknown etc mode")
+        raise NotImplementedError(f"Unknown ETC mode: {fmt}")
+
     return Image.frombytes("RGBA", (width, height), image_data, "raw", "BGRA")
 
 
-def eac(image_data: bytes, width: int, height: int, fmt: list) -> Image.Image:
+def eac(image_data: bytes, width: int, height: int, fmt: str) -> Image.Image:
     if fmt == "EAC_R":
         image_data = texture2ddecoder.decode_eacr(image_data, width, height)
     elif fmt == "EAC_R_SIGNED":
@@ -467,6 +467,9 @@ def eac(image_data: bytes, width: int, height: int, fmt: list) -> Image.Image:
         image_data = texture2ddecoder.decode_eacrg(image_data, width, height)
     elif fmt == "EAC_RG_SIGNED":
         image_data = texture2ddecoder.decode_eacrg_signed(image_data, width, height)
+    else:
+        raise NotImplementedError(f"Unknown EAC mode: {fmt}")
+
     return Image.frombytes("RGBA", (width, height), image_data, "raw", "BGRA")
 
 
@@ -530,87 +533,85 @@ def rgb9e5float(image_data: bytes, width: int, height: int) -> Image.Image:
     return Image.frombytes("RGB", (width, height), rgb, "raw", "RGB")
 
 
-# The type for the conversion table
-CONV_TABLE: Dict[TextureFormat, Optional[Tuple[Callable, Unpack[Tuple[Any, ...]]]]] = {
-    #  FORMAT                  FUNC     #ARGS.....
-    # ----------------------- -------- -------- ------------ ----------------- ------------ ----------
-    TF.Alpha8: (pillow, "RGBA", "raw", "A"),
-    TF.ARGB4444: (pillow, "RGBA", "raw", "RGBA;4B", (2, 1, 0, 3)),
-    TF.RGB24: (pillow, "RGB", "raw", "RGB"),
-    TF.RGBA32: (pillow, "RGBA", "raw", "RGBA"),
-    TF.ARGB32: (pillow, "RGBA", "raw", "ARGB"),
-    TF.ARGBFloat: (pillow, "RGBA", "raw", "RGBAF", (2, 1, 0, 3)),
-    TF.RGB565: (pillow, "RGB", "raw", "BGR;16"),
-    TF.BGR24: (pillow, "RGB", "raw", "BGR"),
-    TF.R8: (pillow, "RGB", "raw", "R"),
-    TF.R16: (pillow, "RGB", "raw", "R;16"),
-    TF.RG16: (rg, "RGB", "raw", "RG"),
-    TF.DXT1: (pillow, "RGBA", "bcn", 1),
-    TF.DXT3: (pillow, "RGBA", "bcn", 2),
-    TF.DXT5: (pillow, "RGBA", "bcn", 3),
-    TF.RGBA4444: (pillow, "RGBA", "raw", "RGBA;4B", (3, 2, 1, 0)),
-    TF.BGRA32: (pillow, "RGBA", "raw", "BGRA"),
-    TF.RHalf: (half, "R", "raw", "R"),
-    TF.RGHalf: (rg, "RGB", "raw", "RGE"),
-    TF.RGBAHalf: (half, "RGB", "raw", "RGB"),
-    TF.RFloat: (pillow, "RGB", "raw", "RF"),
-    TF.RGFloat: (rg, "RGB", "raw", "RGF"),
-    TF.RGBAFloat: (pillow, "RGBA", "raw", "RGBAF"),
-    TF.YUY2: None,
-    TF.RGB9e5Float: (rgb9e5float,),
-    TF.BC4: (pillow, "L", "bcn", 4),
-    TF.BC5: (pillow, "RGB", "bcn", 5),
-    TF.BC6H: (pillow, "RGBA", "bcn", 6),
-    TF.BC7: (pillow, "RGBA", "bcn", 7),
-    TF.DXT1Crunched: (pillow, "RGBA", "bcn", 1),
-    TF.DXT5Crunched: (pillow, "RGBA", "bcn", 3),
-    TF.PVRTC_RGB2: (pvrtc, True),
-    TF.PVRTC_RGBA2: (pvrtc, True),
-    TF.PVRTC_RGB4: (pvrtc, False),
-    TF.PVRTC_RGBA4: (pvrtc, False),
-    TF.ETC_RGB4: (etc, (1,)),
-    TF.ATC_RGB4: (atc, False),
-    TF.ATC_RGBA8: (atc, True),
-    TF.EAC_R: (eac, "EAC_R"),
-    TF.EAC_R_SIGNED: (eac, "EAC_R:SIGNED"),
-    TF.EAC_RG: (eac, "EAC_RG"),
-    TF.EAC_RG_SIGNED: (eac, "EAC_RG_SIGNED"),
-    TF.ETC2_RGB: (etc, (2, "RGB")),
-    TF.ETC2_RGBA1: (etc, (2, "A1")),
-    TF.ETC2_RGBA8: (etc, (2, "A8")),
-    TF.ASTC_RGB_4x4: (astc, (4, 4)),
-    TF.ASTC_RGB_5x5: (astc, (5, 5)),
-    TF.ASTC_RGB_6x6: (astc, (6, 6)),
-    TF.ASTC_RGB_8x8: (astc, (8, 8)),
-    TF.ASTC_RGB_10x10: (astc, (10, 10)),
-    TF.ASTC_RGB_12x12: (astc, (12, 12)),
-    TF.ASTC_RGBA_4x4: (astc, (4, 4)),
-    TF.ASTC_RGBA_5x5: (astc, (5, 5)),
-    TF.ASTC_RGBA_6x6: (astc, (6, 6)),
-    TF.ASTC_RGBA_8x8: (astc, (8, 8)),
-    TF.ASTC_RGBA_10x10: (astc, (10, 10)),
-    TF.ASTC_RGBA_12x12: (astc, (12, 12)),
-    TF.ETC_RGB4_3DS: (etc, (1,)),
-    TF.ETC_RGBA8_3DS: (etc, (1,)),
-    TF.ETC_RGB4Crunched: (etc, (1,)),
-    TF.ETC2_RGBA8Crunched: (etc, (2, "A8")),
-    TF.ASTC_HDR_4x4: (astc, (4, 4)),
-    TF.ASTC_HDR_5x5: (astc, (5, 5)),
-    TF.ASTC_HDR_6x6: (astc, (6, 6)),
-    TF.ASTC_HDR_8x8: (astc, (8, 8)),
-    TF.ASTC_HDR_10x10: (astc, (10, 10)),
-    TF.ASTC_HDR_12x12: (astc, (12, 12)),
-    TF.RG32: (rg, "RGB", "raw", "RG;16"),
-    TF.RGB48: (pillow, "RGB", "raw", "RGB;16"),
-    TF.RGBA64: (pillow, "RGBA", "raw", "RGBA;16"),
-    TF.R8_SIGNED: (pillow, "R", "raw", "R;8s"),
-    TF.RG16_SIGNED: (rg, "RGB", "raw", "RG;8s"),
-    TF.RGB24_SIGNED: (pillow, "RGB", "raw", "RGB;8s"),
-    TF.RGBA32_SIGNED: (pillow, "RGBA", "raw", "RGBA;8s"),
-    TF.R16_SIGNED: (pillow, "R", "raw", "R;16s"),
-    TF.RG32_SIGNED: (rg, "RGB", "raw", "RG;16s"),
-    TF.RGB48_SIGNED: (pillow, "RGB", "raw", "RGB;16s"),
-    TF.RGBA64_SIGNED: (pillow, "RGBA", "raw", "RGBA;16s"),
+# Mapping TextureFormat -> (converter function, (additional args, ...))
+CONV_TABLE: Dict[TF, Tuple[Callable[..., Image.Image], Tuple[Any, ...]]] = {
+    TF.Alpha8: (pillow, ("RGBA", "raw", "A")),
+    TF.ARGB4444: (pillow, ("RGBA", "raw", "RGBA;4B", (2, 1, 0, 3))),
+    TF.RGB24: (pillow, ("RGB", "raw", "RGB")),
+    TF.RGBA32: (pillow, ("RGBA", "raw", "RGBA")),
+    TF.ARGB32: (pillow, ("RGBA", "raw", "ARGB")),
+    TF.ARGBFloat: (pillow, ("RGBA", "raw", "RGBAF", (2, 1, 0, 3))),
+    TF.RGB565: (pillow, ("RGB", "raw", "BGR;16")),
+    TF.BGR24: (pillow, ("RGB", "raw", "BGR")),
+    TF.R8: (pillow, ("RGB", "raw", "R")),
+    TF.R16: (pillow, ("RGB", "raw", "R;16")),
+    TF.RG16: (rg, ("RGB", "raw", "RG")),
+    TF.DXT1: (pillow, ("RGBA", "bcn", 1)),
+    TF.DXT3: (pillow, ("RGBA", "bcn", 2)),
+    TF.DXT5: (pillow, ("RGBA", "bcn", 3)),
+    TF.RGBA4444: (pillow, ("RGBA", "raw", "RGBA;4B", (3, 2, 1, 0))),
+    TF.BGRA32: (pillow, ("RGBA", "raw", "BGRA")),
+    TF.RHalf: (half, ("R", "raw", "R")),
+    TF.RGHalf: (rg, ("RGB", "raw", "RGE")),
+    TF.RGBAHalf: (half, ("RGB", "raw", "RGB")),
+    TF.RFloat: (pillow, ("RGB", "raw", "RF")),
+    TF.RGFloat: (rg, ("RGB", "raw", "RGF")),
+    TF.RGBAFloat: (pillow, ("RGBA", "raw", "RGBAF")),
+    # TF.YUY2: NotImplementedError("YUY2 not implemented"),
+    TF.RGB9e5Float: (rgb9e5float, ()),
+    TF.BC4: (pillow, ("L", "bcn", 4)),
+    TF.BC5: (pillow, ("RGB", "bcn", 5)),
+    TF.BC6H: (pillow, ("RGBA", "bcn", 6)),
+    TF.BC7: (pillow, ("RGBA", "bcn", 7)),
+    TF.DXT1Crunched: (pillow, ("RGBA", "bcn", 1)),
+    TF.DXT5Crunched: (pillow, ("RGBA", "bcn", 3)),
+    TF.PVRTC_RGB2: (pvrtc, (True,)),
+    TF.PVRTC_RGBA2: (pvrtc, (True,)),
+    TF.PVRTC_RGB4: (pvrtc, (False,)),
+    TF.PVRTC_RGBA4: (pvrtc, (False,)),
+    TF.ETC_RGB4: (etc, ("ETC1",)),
+    TF.ATC_RGB4: (atc, (False,)),
+    TF.ATC_RGBA8: (atc, (True,)),
+    TF.EAC_R: (eac, ("EAC_R",)),
+    TF.EAC_R_SIGNED: (eac, ("EAC_R_SIGNED",)),
+    TF.EAC_RG: (eac, ("EAC_RG",)),
+    TF.EAC_RG_SIGNED: (eac, ("EAC_RG_SIGNED",)),
+    TF.ETC2_RGB: (etc, ("ETC2_RGB",)),
+    TF.ETC2_RGBA1: (etc, ("ETC2_A1",)),
+    TF.ETC2_RGBA8: (etc, ("ETC2_A8",)),
+    TF.ASTC_RGB_4x4: (astc, ((4, 4),)),
+    TF.ASTC_RGB_5x5: (astc, ((5, 5),)),
+    TF.ASTC_RGB_6x6: (astc, ((6, 6),)),
+    TF.ASTC_RGB_8x8: (astc, ((8, 8),)),
+    TF.ASTC_RGB_10x10: (astc, ((10, 10),)),
+    TF.ASTC_RGB_12x12: (astc, ((12, 12),)),
+    TF.ASTC_RGBA_4x4: (astc, ((4, 4),)),
+    TF.ASTC_RGBA_5x5: (astc, ((5, 5),)),
+    TF.ASTC_RGBA_6x6: (astc, ((6, 6),)),
+    TF.ASTC_RGBA_8x8: (astc, ((8, 8),)),
+    TF.ASTC_RGBA_10x10: (astc, ((10, 10),)),
+    TF.ASTC_RGBA_12x12: (astc, ((12, 12),)),
+    TF.ETC_RGB4_3DS: (etc, ("ETC1",)),
+    TF.ETC_RGBA8_3DS: (etc, ("ETC1",)),
+    TF.ETC_RGB4Crunched: (etc, ("ETC1",)),
+    TF.ETC2_RGBA8Crunched: (etc, ("ETC2_A8",)),
+    TF.ASTC_HDR_4x4: (astc, ((4, 4),)),
+    TF.ASTC_HDR_5x5: (astc, ((5, 5),)),
+    TF.ASTC_HDR_6x6: (astc, ((6, 6),)),
+    TF.ASTC_HDR_8x8: (astc, ((8, 8),)),
+    TF.ASTC_HDR_10x10: (astc, ((10, 10),)),
+    TF.ASTC_HDR_12x12: (astc, ((12, 12),)),
+    TF.RG32: (rg, ("RGB", "raw", "RG;16")),
+    TF.RGB48: (pillow, ("RGB", "raw", "RGB;16")),
+    TF.RGBA64: (pillow, ("RGBA", "raw", "RGBA;16")),
+    TF.R8_SIGNED: (pillow, ("R", "raw", "R;8s")),
+    TF.RG16_SIGNED: (rg, ("RGB", "raw", "RG;8s")),
+    TF.RGB24_SIGNED: (pillow, ("RGB", "raw", "RGB;8s")),
+    TF.RGBA32_SIGNED: (pillow, ("RGBA", "raw", "RGBA;8s")),
+    TF.R16_SIGNED: (pillow, ("R", "raw", "R;16s")),
+    TF.RG32_SIGNED: (rg, ("RGB", "raw", "RG;16s")),
+    TF.RGB48_SIGNED: (pillow, ("RGB", "raw", "RGB;16s")),
+    TF.RGBA64_SIGNED: (pillow, ("RGBA", "raw", "RGBA;16s")),
 }
 
 # XBOX Swap Formats