Skip to content

Commit 9defdd7

Browse files
committed
refactor: Texture2DConverter - optimize CONV_TABLE
1 parent b5eab2e commit 9defdd7

File tree

1 file changed

+97
-100
lines changed

1 file changed

+97
-100
lines changed

UnityPy/export/Texture2DConverter.py

Lines changed: 97 additions & 100 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
import struct
44
from io import BytesIO
55
from threading import Lock
6-
from typing import TYPE_CHECKING, Dict, Optional, Tuple, Union
6+
from typing import TYPE_CHECKING, Callable, Dict, Optional, Tuple, Union
77

88
import astc_encoder
99
import texture2ddecoder
@@ -338,11 +338,6 @@ def parse_image_data(
338338
else:
339339
width, height = get_compressed_image_size(width, height, texture_format)
340340

341-
selection = CONV_TABLE[texture_format]
342-
343-
if not selection:
344-
raise NotImplementedError(f"Not implemented texture format: {texture_format}")
345-
346341
if "Crunched" in texture_format.name:
347342
if (
348343
version[0] > 2017
@@ -354,7 +349,11 @@ def parse_image_data(
354349
else:
355350
image_data = texture2ddecoder.unpack_crunch(image_data)
356351

357-
img = selection[0](image_data, width, height, *selection[1:])
352+
conv_func = CONV_TABLE.get(texture_format)
353+
if not conv_func:
354+
raise NotImplementedError(f"Not implemented texture format: {texture_format}")
355+
356+
img = conv_func(image_data, width, height)
358357

359358
if switch_swizzle is not None:
360359
block_size, gobs_per_block, pil_mode = switch_swizzle
@@ -460,18 +459,18 @@ def pvrtc(image_data: bytes, width: int, height: int, fmt: bool) -> Image.Image:
460459
return Image.frombytes("RGBA", (width, height), image_data, "raw", "BGRA")
461460

462461

463-
def etc(image_data: bytes, width: int, height: int, fmt: list) -> Image.Image:
464-
if fmt[0] == 1:
462+
def etc(image_data: bytes, width: int, height: int, fmt: str) -> Image.Image:
463+
if fmt == "ETC1":
465464
image_data = texture2ddecoder.decode_etc1(image_data, width, height)
466-
elif fmt[0] == 2:
467-
if fmt[1] == "RGB":
468-
image_data = texture2ddecoder.decode_etc2(image_data, width, height)
469-
elif fmt[1] == "A1":
470-
image_data = texture2ddecoder.decode_etc2a1(image_data, width, height)
471-
elif fmt[1] == "A8":
472-
image_data = texture2ddecoder.decode_etc2a8(image_data, width, height)
465+
elif fmt == "ETC2_RGB":
466+
image_data = texture2ddecoder.decode_etc2(image_data, width, height)
467+
elif fmt == "ETC2_A1":
468+
image_data = texture2ddecoder.decode_etc2a1(image_data, width, height)
469+
elif fmt == "ETC2_A8":
470+
image_data = texture2ddecoder.decode_etc2a8(image_data, width, height)
473471
else:
474-
raise NotImplementedError("unknown etc mode")
472+
raise NotImplementedError(f"Unknown ETC mode: {fmt}")
473+
475474
return Image.frombytes("RGBA", (width, height), image_data, "raw", "BGRA")
476475

477476

@@ -484,6 +483,9 @@ def eac(image_data: bytes, width: int, height: int, fmt: str) -> Image.Image:
484483
image_data = texture2ddecoder.decode_eacrg(image_data, width, height)
485484
elif fmt == "EAC_RG_SIGNED":
486485
image_data = texture2ddecoder.decode_eacrg_signed(image_data, width, height)
486+
else:
487+
raise NotImplementedError(f"Unknown EAC mode: {fmt}")
488+
487489
return Image.frombytes("RGBA", (width, height), image_data, "raw", "BGRA")
488490

489491

@@ -547,90 +549,85 @@ def rgb9e5float(image_data: bytes, width: int, height: int) -> Image.Image:
547549
return Image.frombytes("RGB", (width, height), rgb, "raw", "RGB")
548550

549551

550-
CONV_TABLE = {
551-
# FORMAT FUNC #ARGS.....
552-
# ----------------------- -------- -------- ------------ ----------------- ------------ ----------
553-
(TF.Alpha8, pillow, "L", "raw", "L"),
554-
(TF.ARGB4444, pillow, "RGBA", "raw", "RGBA;4B", (2, 1, 0, 3)),
555-
(TF.RGB24, pillow, "RGB", "raw", "RGB"),
556-
(TF.RGBA32, pillow, "RGBA", "raw", "RGBA"),
557-
(TF.ARGB32, pillow, "RGBA", "raw", "ARGB"),
558-
(TF.ARGBFloat, pillow, "RGBA", "raw", "RGBAF", (2, 1, 0, 3)),
559-
(TF.RGB565, pillow, "RGB", "raw", "BGR;16"),
560-
(TF.BGR24, pillow, "RGB", "raw", "BGR"),
561-
(TF.R8, pillow, "RGB", "raw", "R"),
562-
(TF.R16, pillow, "RGB", "raw", "R;16"),
563-
(TF.RG16, rg, "RGB", "raw", "RG"),
564-
(TF.DXT1, pillow, "RGBA", "bcn", 1),
565-
(TF.DXT3, pillow, "RGBA", "bcn", 2),
566-
(TF.DXT5, pillow, "RGBA", "bcn", 3),
567-
(TF.RGBA4444, pillow, "RGBA", "raw", "RGBA;4B", (3, 2, 1, 0)),
568-
(TF.BGRA32, pillow, "RGBA", "raw", "BGRA"),
569-
(TF.RHalf, half, "R", "raw", "R"),
570-
(TF.RGHalf, rg, "RGB", "raw", "RGE"),
571-
(TF.RGBAHalf, half, "RGB", "raw", "RGB"),
572-
(TF.RFloat, pillow, "RGB", "raw", "RF"),
573-
(TF.RGFloat, rg, "RGB", "raw", "RGF"),
574-
(TF.RGBAFloat, pillow, "RGBA", "raw", "RGBAF"),
575-
(TF.YUY2,),
576-
(TF.RGB9e5Float, rgb9e5float),
577-
(TF.BC4, pillow, "L", "bcn", 4),
578-
(TF.BC5, pillow, "RGB", "bcn", 5),
579-
(TF.BC6H, pillow, "RGBA", "bcn", 6),
580-
(TF.BC7, pillow, "RGBA", "bcn", 7),
581-
(TF.DXT1Crunched, pillow, "RGBA", "bcn", 1),
582-
(TF.DXT5Crunched, pillow, "RGBA", "bcn", 3),
583-
(TF.PVRTC_RGB2, pvrtc, True),
584-
(TF.PVRTC_RGBA2, pvrtc, True),
585-
(TF.PVRTC_RGB4, pvrtc, False),
586-
(TF.PVRTC_RGBA4, pvrtc, False),
587-
(TF.ETC_RGB4, etc, (1,)),
588-
(TF.ATC_RGB4, atc, False),
589-
(TF.ATC_RGBA8, atc, True),
590-
(TF.EAC_R, eac, "EAC_R"),
591-
(TF.EAC_R_SIGNED, eac, "EAC_R:SIGNED"),
592-
(TF.EAC_RG, eac, "EAC_RG"),
593-
(TF.EAC_RG_SIGNED, eac, "EAC_RG_SIGNED"),
594-
(TF.ETC2_RGB, etc, (2, "RGB")),
595-
(TF.ETC2_RGBA1, etc, (2, "A1")),
596-
(TF.ETC2_RGBA8, etc, (2, "A8")),
597-
(TF.ASTC_RGB_4x4, astc, (4, 4)),
598-
(TF.ASTC_RGB_5x5, astc, (5, 5)),
599-
(TF.ASTC_RGB_6x6, astc, (6, 6)),
600-
(TF.ASTC_RGB_8x8, astc, (8, 8)),
601-
(TF.ASTC_RGB_10x10, astc, (10, 10)),
602-
(TF.ASTC_RGB_12x12, astc, (12, 12)),
603-
(TF.ASTC_RGBA_4x4, astc, (4, 4)),
604-
(TF.ASTC_RGBA_5x5, astc, (5, 5)),
605-
(TF.ASTC_RGBA_6x6, astc, (6, 6)),
606-
(TF.ASTC_RGBA_8x8, astc, (8, 8)),
607-
(TF.ASTC_RGBA_10x10, astc, (10, 10)),
608-
(TF.ASTC_RGBA_12x12, astc, (12, 12)),
609-
(TF.ETC_RGB4_3DS, etc, (1,)),
610-
(TF.ETC_RGBA8_3DS, etc, (1,)),
611-
(TF.ETC_RGB4Crunched, etc, (1,)),
612-
(TF.ETC2_RGBA8Crunched, etc, (2, "A8")),
613-
(TF.ASTC_HDR_4x4, astc, (4, 4)),
614-
(TF.ASTC_HDR_5x5, astc, (5, 5)),
615-
(TF.ASTC_HDR_6x6, astc, (6, 6)),
616-
(TF.ASTC_HDR_8x8, astc, (8, 8)),
617-
(TF.ASTC_HDR_10x10, astc, (10, 10)),
618-
(TF.ASTC_HDR_12x12, astc, (12, 12)),
619-
(TF.RG32, rg, "RGB", "raw", "RG;16"),
620-
(TF.RGB48, pillow, "RGB", "raw", "RGB;16"),
621-
(TF.RGBA64, pillow, "RGBA", "raw", "RGBA;16"),
622-
(TF.R8_SIGNED, pillow, "R", "raw", "R;8s"),
623-
(TF.RG16_SIGNED, rg, "RGB", "raw", "RG;8s"),
624-
(TF.RGB24_SIGNED, pillow, "RGB", "raw", "RGB;8s"),
625-
(TF.RGBA32_SIGNED, pillow, "RGBA", "raw", "RGBA;8s"),
626-
(TF.R16_SIGNED, pillow, "R", "raw", "R;16s"),
627-
(TF.RG32_SIGNED, rg, "RGB", "raw", "RG;16s"),
628-
(TF.RGB48_SIGNED, pillow, "RGB", "raw", "RGB;16s"),
629-
(TF.RGBA64_SIGNED, pillow, "RGBA", "raw", "RGBA;16s"),
552+
CONV_TABLE: Dict[TextureFormat, Callable[[bytes, int, int], Image.Image]] = {
553+
TF.Alpha8: lambda data, w, h: pillow(data, w, h, "L", "raw", "L"),
554+
TF.ARGB4444: lambda data, w, h: pillow(data, w, h, "RGBA", "raw", "RGBA;4B", (2, 1, 0, 3)),
555+
TF.RGB24: lambda data, w, h: pillow(data, w, h, "RGB", "raw", "RGB"),
556+
TF.RGBA32: lambda data, w, h: pillow(data, w, h, "RGBA", "raw", "RGBA"),
557+
TF.ARGB32: lambda data, w, h: pillow(data, w, h, "RGBA", "raw", "ARGB"),
558+
TF.ARGBFloat: lambda data, w, h: pillow(data, w, h, "RGBA", "raw", "RGBAF", (2, 1, 0, 3)),
559+
TF.RGB565: lambda data, w, h: pillow(data, w, h, "RGB", "raw", "BGR;16"),
560+
TF.BGR24: lambda data, w, h: pillow(data, w, h, "RGB", "raw", "BGR"),
561+
TF.R8: lambda data, w, h: pillow(data, w, h, "RGB", "raw", "R"),
562+
TF.R16: lambda data, w, h: pillow(data, w, h, "RGB", "raw", "R;16"),
563+
TF.RG16: lambda data, w, h: rg(data, w, h, "RGB", "raw", "RG"),
564+
TF.DXT1: lambda data, w, h: pillow(data, w, h, "RGBA", "bcn", 1),
565+
TF.DXT3: lambda data, w, h: pillow(data, w, h, "RGBA", "bcn", 2),
566+
TF.DXT5: lambda data, w, h: pillow(data, w, h, "RGBA", "bcn", 3),
567+
TF.RGBA4444: lambda data, w, h: pillow(data, w, h, "RGBA", "raw", "RGBA;4B", (3, 2, 1, 0)),
568+
TF.BGRA32: lambda data, w, h: pillow(data, w, h, "RGBA", "raw", "BGRA"),
569+
TF.RHalf: lambda data, w, h: half(data, w, h, "R", "raw", "R"),
570+
TF.RGHalf: lambda data, w, h: rg(data, w, h, "RGB", "raw", "RGE"),
571+
TF.RGBAHalf: lambda data, w, h: half(data, w, h, "RGB", "raw", "RGB"),
572+
TF.RFloat: lambda data, w, h: pillow(data, w, h, "RGB", "raw", "RF"),
573+
TF.RGFloat: lambda data, w, h: rg(data, w, h, "RGB", "raw", "RGF"),
574+
TF.RGBAFloat: lambda data, w, h: pillow(data, w, h, "RGBA", "raw", "RGBAF"),
575+
# TF.YUY2: lambda data, w, h: NotImplementedError("YUY2 not implemented"),
576+
TF.RGB9e5Float: lambda data, w, h: rgb9e5float(data, w, h),
577+
TF.BC4: lambda data, w, h: pillow(data, w, h, "L", "bcn", 4),
578+
TF.BC5: lambda data, w, h: pillow(data, w, h, "RGB", "bcn", 5),
579+
TF.BC6H: lambda data, w, h: pillow(data, w, h, "RGBA", "bcn", 6),
580+
TF.BC7: lambda data, w, h: pillow(data, w, h, "RGBA", "bcn", 7),
581+
TF.DXT1Crunched: lambda data, w, h: pillow(data, w, h, "RGBA", "bcn", 1),
582+
TF.DXT5Crunched: lambda data, w, h: pillow(data, w, h, "RGBA", "bcn", 3),
583+
TF.PVRTC_RGB2: lambda data, w, h: pvrtc(data, w, h, True),
584+
TF.PVRTC_RGBA2: lambda data, w, h: pvrtc(data, w, h, True),
585+
TF.PVRTC_RGB4: lambda data, w, h: pvrtc(data, w, h, False),
586+
TF.PVRTC_RGBA4: lambda data, w, h: pvrtc(data, w, h, False),
587+
TF.ETC_RGB4: lambda data, w, h: etc(data, w, h, "ETC1"),
588+
TF.ATC_RGB4: lambda data, w, h: atc(data, w, h, False),
589+
TF.ATC_RGBA8: lambda data, w, h: atc(data, w, h, True),
590+
TF.EAC_R: lambda data, w, h: eac(data, w, h, "EAC_R"),
591+
TF.EAC_R_SIGNED: lambda data, w, h: eac(data, w, h, "EAC_R_SIGNED"),
592+
TF.EAC_RG: lambda data, w, h: eac(data, w, h, "EAC_RG"),
593+
TF.EAC_RG_SIGNED: lambda data, w, h: eac(data, w, h, "EAC_RG_SIGNED"),
594+
TF.ETC2_RGB: lambda data, w, h: etc(data, w, h, "ETC2_RGB"),
595+
TF.ETC2_RGBA1: lambda data, w, h: etc(data, w, h, "ETC2_A1"),
596+
TF.ETC2_RGBA8: lambda data, w, h: etc(data, w, h, "ETC2_A8"),
597+
TF.ASTC_RGB_4x4: lambda data, w, h: astc(data, w, h, (4, 4)),
598+
TF.ASTC_RGB_5x5: lambda data, w, h: astc(data, w, h, (5, 5)),
599+
TF.ASTC_RGB_6x6: lambda data, w, h: astc(data, w, h, (6, 6)),
600+
TF.ASTC_RGB_8x8: lambda data, w, h: astc(data, w, h, (8, 8)),
601+
TF.ASTC_RGB_10x10: lambda data, w, h: astc(data, w, h, (10, 10)),
602+
TF.ASTC_RGB_12x12: lambda data, w, h: astc(data, w, h, (12, 12)),
603+
TF.ASTC_RGBA_4x4: lambda data, w, h: astc(data, w, h, (4, 4)),
604+
TF.ASTC_RGBA_5x5: lambda data, w, h: astc(data, w, h, (5, 5)),
605+
TF.ASTC_RGBA_6x6: lambda data, w, h: astc(data, w, h, (6, 6)),
606+
TF.ASTC_RGBA_8x8: lambda data, w, h: astc(data, w, h, (8, 8)),
607+
TF.ASTC_RGBA_10x10: lambda data, w, h: astc(data, w, h, (10, 10)),
608+
TF.ASTC_RGBA_12x12: lambda data, w, h: astc(data, w, h, (12, 12)),
609+
TF.ETC_RGB4_3DS: lambda data, w, h: etc(data, w, h, "ETC1"),
610+
TF.ETC_RGBA8_3DS: lambda data, w, h: etc(data, w, h, "ETC1"),
611+
TF.ETC_RGB4Crunched: lambda data, w, h: etc(data, w, h, "ETC1"),
612+
TF.ETC2_RGBA8Crunched: lambda data, w, h: etc(data, w, h, "ETC2_A8"),
613+
TF.ASTC_HDR_4x4: lambda data, w, h: astc(data, w, h, (4, 4)),
614+
TF.ASTC_HDR_5x5: lambda data, w, h: astc(data, w, h, (5, 5)),
615+
TF.ASTC_HDR_6x6: lambda data, w, h: astc(data, w, h, (6, 6)),
616+
TF.ASTC_HDR_8x8: lambda data, w, h: astc(data, w, h, (8, 8)),
617+
TF.ASTC_HDR_10x10: lambda data, w, h: astc(data, w, h, (10, 10)),
618+
TF.ASTC_HDR_12x12: lambda data, w, h: astc(data, w, h, (12, 12)),
619+
TF.RG32: lambda data, w, h: rg(data, w, h, "RGB", "raw", "RG;16"),
620+
TF.RGB48: lambda data, w, h: pillow(data, w, h, "RGB", "raw", "RGB;16"),
621+
TF.RGBA64: lambda data, w, h: pillow(data, w, h, "RGBA", "raw", "RGBA;16"),
622+
TF.R8_SIGNED: lambda data, w, h: pillow(data, w, h, "R", "raw", "R;8s"),
623+
TF.RG16_SIGNED: lambda data, w, h: rg(data, w, h, "RGB", "raw", "RG;8s"),
624+
TF.RGB24_SIGNED: lambda data, w, h: pillow(data, w, h, "RGB", "raw", "RGB;8s"),
625+
TF.RGBA32_SIGNED: lambda data, w, h: pillow(data, w, h, "RGBA", "raw", "RGBA;8s"),
626+
TF.R16_SIGNED: lambda data, w, h: pillow(data, w, h, "R", "raw", "R;16s"),
627+
TF.RG32_SIGNED: lambda data, w, h: rg(data, w, h, "RGB", "raw", "RG;16s"),
628+
TF.RGB48_SIGNED: lambda data, w, h: pillow(data, w, h, "RGB", "raw", "RGB;16s"),
629+
TF.RGBA64_SIGNED: lambda data, w, h: pillow(data, w, h, "RGBA", "raw", "RGBA;16s"),
630630
}
631631

632-
# format conv_table to a dict
633-
CONV_TABLE = {line[0]: line[1:] for line in CONV_TABLE}
634-
635632
# XBOX Swap Formats
636633
XBOX_SWAP_FORMATS = [TF.RGB565, TF.DXT1, TF.DXT1Crunched, TF.DXT5, TF.DXT5Crunched]

0 commit comments

Comments
 (0)