Skip to content

Commit d8f7159

Browse files
committed
refactor: Texture2DSwizzler - wrap methods
1 parent 67e1961 commit d8f7159

File tree

2 files changed

+100
-67
lines changed

2 files changed

+100
-67
lines changed

UnityPy/export/Texture2DConverter.py

Lines changed: 20 additions & 42 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, Callable, Dict, Optional, Tuple, Union
6+
from typing import TYPE_CHECKING, Callable, Dict, Optional, Sequence, Tuple, Union
77

88
import astc_encoder
99
import texture2ddecoder
@@ -17,6 +17,9 @@
1717
from ..classes import Texture2D
1818

1919

20+
PlatformBlobType = Union[bytes, Sequence[int]]
21+
22+
2023
TEXTURE_FORMAT_BLOCK_SIZE_TABLE: Dict[TF, Optional[Tuple[int, int]]] = {}
2124
for tf in TF:
2225
if tf.name.startswith("ASTC"):
@@ -134,7 +137,7 @@ def image_to_texture2d(
134137
img: Image.Image,
135138
target_texture_format: Union[TF, int],
136139
platform: int = 0,
137-
platform_blob: Optional[bytes] = None,
140+
platform_blob: Optional[PlatformBlobType] = None,
138141
flip: bool = True,
139142
) -> Tuple[bytes, TF]:
140143
if not isinstance(target_texture_format, TF):
@@ -216,27 +219,17 @@ def image_to_texture2d(
216219
pil_mode = "RGB"
217220
# everything else defaulted to RGBA
218221

219-
switch_swizzle = None
220222
if platform == BuildTarget.Switch and platform_blob:
221-
gobs_per_block = TextureSwizzler.get_switch_gobs_per_block(platform_blob)
222-
223223
if texture_format == TF.RGB24:
224224
texture_format = TF.RGBA32
225225
pil_mode = "RGBA"
226226
elif texture_format == TF.BGR24:
227227
texture_format = TF.BGRA32
228228
pil_mode = "BGRA"
229229

230-
block_size = TextureSwizzler.TEXTUREFORMAT_BLOCK_SIZE_MAP.get(texture_format)
231-
if not block_size:
232-
raise NotImplementedError(
233-
f"Not implemented swizzle format: {texture_format.name}"
234-
)
235-
236-
width, height = TextureSwizzler.get_padded_texture_size(
237-
img.width, img.height, *block_size, gobs_per_block
230+
width, height = TextureSwizzler.get_padded_image_size(
231+
img.width, img.height, texture_format, platform_blob
238232
)
239-
switch_swizzle = (block_size, gobs_per_block)
240233
else:
241234
width, height = get_compressed_image_size(img.width, img.height, texture_format)
242235

@@ -248,9 +241,10 @@ def image_to_texture2d(
248241
else:
249242
enc_img = img.tobytes("raw", pil_mode)
250243

251-
if switch_swizzle is not None:
252-
block_size, gobs_per_block = switch_swizzle
253-
enc_img = TextureSwizzler.swizzle(enc_img, width, height, *block_size, gobs_per_block)
244+
if platform == BuildTarget.Switch and platform_blob:
245+
enc_img = TextureSwizzler.swizzle(
246+
enc_img, width, height, texture_format, platform_blob
247+
)
254248

255249
return enc_img, texture_format
256250

@@ -285,10 +279,10 @@ def parse_image_data(
285279
image_data: bytes,
286280
width: int,
287281
height: int,
288-
texture_format: Union[int, TF],
282+
texture_format: Union[TF, int],
289283
version: Tuple[int, int, int, int],
290284
platform: int,
291-
platform_blob: Optional[bytes] = None,
285+
platform_blob: Optional[PlatformBlobType] = None,
292286
flip: bool = True,
293287
) -> Image.Image:
294288
if not width or not height:
@@ -304,31 +298,20 @@ def parse_image_data(
304298
if platform == BuildTarget.XBOX360 and texture_format in XBOX_SWAP_FORMATS:
305299
image_data = swap_bytes_for_xbox(image_data)
306300

307-
original_width, original_height = width, height
308-
switch_swizzle = None
309-
if platform == BuildTarget.Switch and platform_blob:
310-
gobs_per_block = TextureSwizzler.get_switch_gobs_per_block(platform_blob)
301+
ori_width, ori_height = width, height
311302

312-
pil_mode = "RGBA"
303+
if platform == BuildTarget.Switch and platform_blob:
313304
if texture_format == TF.RGB24:
314305
texture_format = TF.RGBA32
315306
elif texture_format == TF.BGR24:
316307
texture_format = TF.BGRA32
317-
pil_mode = "BGRA"
318308
elif texture_format == TF.Alpha8:
319309
texture_format = texture_format
320-
pil_mode = "L"
321310

322-
block_size = TextureSwizzler.TEXTUREFORMAT_BLOCK_SIZE_MAP.get(texture_format)
323-
if not block_size:
324-
raise NotImplementedError(
325-
f"Not implemented swizzle format: {texture_format.name}"
326-
)
327-
328-
width, height = TextureSwizzler.get_padded_texture_size(
329-
width, height, *block_size, gobs_per_block
311+
width, height = TextureSwizzler.get_padded_image_size(
312+
width, height, texture_format, platform_blob
330313
)
331-
switch_swizzle = (block_size, gobs_per_block, pil_mode)
314+
image_data = TextureSwizzler.deswizzle(image_data, width, height, texture_format, platform_blob)
332315
else:
333316
width, height = get_compressed_image_size(width, height, texture_format)
334317

@@ -343,20 +326,15 @@ def parse_image_data(
343326
else:
344327
image_data = texture2ddecoder.unpack_crunch(image_data)
345328

346-
if switch_swizzle is not None:
347-
block_size, gobs_per_block, pil_mode = switch_swizzle
348-
image_data = TextureSwizzler.deswizzle(image_data, width, height, *block_size, gobs_per_block)
349-
350-
351329
conv_func = CONV_TABLE.get(texture_format)
352330
if not conv_func:
353331
raise NotImplementedError(
354332
f"Not implemented texture format: {texture_format.name}"
355333
)
356334
img = conv_func(image_data, width, height)
357335

358-
if original_width != width or original_height != height:
359-
img = img.crop((0, 0, original_width, original_height))
336+
if ori_width != width or ori_height != height:
337+
img = img.crop((0, 0, ori_width, ori_height))
360338

361339
return img.transpose(Image.FLIP_TOP_BOTTOM) if flip else img
362340

UnityPy/helpers/TextureSwizzler.py

Lines changed: 80 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,12 @@
11
# based on https://github.com/nesrak1/UABEA/blob/master/TexturePlugin/Texture2DSwitchDeswizzler.cs
2-
from typing import Dict, Tuple
2+
from typing import Dict, Sequence, Tuple, Union
33

44
from ..enums import TextureFormat as TF
55

6+
7+
PlatformBlobType = Union[bytes, Sequence[int]]
8+
9+
610
GOB_X_TEXEL_COUNT = 4
711
GOB_Y_TEXEL_COUNT = 8
812
TEXEL_BYTE_SIZE = 16
@@ -13,20 +17,20 @@
1317
]
1418

1519

16-
def ceil_divide(a: int, b: int) -> int:
20+
def _ceil_divide(a: int, b: int) -> int:
1721
return (a + b - 1) // b
1822

1923

20-
def deswizzle(
24+
def _deswizzle(
2125
data: bytes,
2226
width: int,
2327
height: int,
2428
block_width: int,
2529
block_height: int,
2630
texels_per_block: int,
2731
) -> bytes:
28-
block_count_x = ceil_divide(width, block_width)
29-
block_count_y = ceil_divide(height, block_height)
32+
block_count_x = _ceil_divide(width, block_width)
33+
block_count_y = _ceil_divide(height, block_height)
3034
gob_count_x = block_count_x // GOB_X_TEXEL_COUNT
3135
gob_count_y = block_count_y // GOB_Y_TEXEL_COUNT
3236
new_data = bytearray(len(data))
@@ -50,16 +54,16 @@ def deswizzle(
5054
return bytes(new_data)
5155

5256

53-
def swizzle(
57+
def _swizzle(
5458
data: bytes,
5559
width: int,
5660
height: int,
5761
block_width: int,
5862
block_height: int,
5963
texels_per_block: int,
6064
) -> bytes:
61-
block_count_x = ceil_divide(width, block_width)
62-
block_count_y = ceil_divide(height, block_height)
65+
block_count_x = _ceil_divide(width, block_width)
66+
block_count_y = _ceil_divide(height, block_height)
6367
gob_count_x = block_count_x // GOB_X_TEXEL_COUNT
6468
gob_count_y = block_count_y // GOB_Y_TEXEL_COUNT
6569
new_data = bytearray(len(data))
@@ -83,8 +87,31 @@ def swizzle(
8387
return bytes(new_data)
8488

8589

90+
def _get_padded_texture_size(
91+
width: int, height: int, block_width: int, block_height: int, texels_per_block: int
92+
) -> Tuple[int, int]:
93+
width = (
94+
_ceil_divide(width, block_width * GOB_X_TEXEL_COUNT)
95+
* block_width
96+
* GOB_X_TEXEL_COUNT
97+
)
98+
height = (
99+
_ceil_divide(height, block_height * GOB_Y_TEXEL_COUNT * texels_per_block)
100+
* block_height
101+
* GOB_Y_TEXEL_COUNT
102+
* texels_per_block
103+
)
104+
return width, height
105+
106+
107+
def _get_texels_per_block(platform_blob: PlatformBlobType) -> int:
108+
if not platform_blob:
109+
raise ValueError("Given platform_blob is empty")
110+
return 1 << int.from_bytes(platform_blob[8:12], "little")
111+
112+
86113
# this should be the amount of pixels that can fit 16 bytes
87-
TEXTUREFORMAT_BLOCK_SIZE_MAP: Dict[TF, Tuple[int, int]] = {
114+
TEXTURE_FORMAT_BLOCK_SIZE_MAP: Dict[TF, Tuple[int, int]] = {
88115
TF.Alpha8: (16, 1), # 1 byte per pixel
89116
TF.ARGB4444: (8, 1), # 2 bytes per pixel
90117
TF.RGBA32: (4, 1), # 4 bytes per pixel
@@ -117,22 +144,50 @@ def swizzle(
117144
}
118145

119146

120-
def get_padded_texture_size(
121-
width: int, height: int, block_width: int, block_height: int, texels_per_block: int
122-
):
123-
width = (
124-
ceil_divide(width, block_width * GOB_X_TEXEL_COUNT)
125-
* block_width
126-
* GOB_X_TEXEL_COUNT
127-
)
128-
height = (
129-
ceil_divide(height, block_height * GOB_Y_TEXEL_COUNT * texels_per_block)
130-
* block_height
131-
* GOB_Y_TEXEL_COUNT
132-
* texels_per_block
147+
def deswizzle(
148+
data: bytes,
149+
width: int,
150+
height: int,
151+
texture_format: TF,
152+
platform_blob: PlatformBlobType,
153+
) -> bytes:
154+
block_size = TEXTURE_FORMAT_BLOCK_SIZE_MAP.get(texture_format)
155+
if not block_size:
156+
raise NotImplementedError(
157+
f"Not implemented swizzle format: {texture_format.name}"
158+
)
159+
texels_per_block = _get_texels_per_block(platform_blob)
160+
return _deswizzle(
161+
data, width, height, *block_size, texels_per_block
133162
)
134-
return width, height
135163

136164

137-
def get_switch_gobs_per_block(platform_blob: bytes) -> int:
138-
return 1 << int.from_bytes(platform_blob[8:12], "little")
165+
def swizzle(
166+
data: bytes,
167+
width: int,
168+
height: int,
169+
texture_format: TF,
170+
platform_blob: PlatformBlobType,
171+
) -> bytes:
172+
block_size = TEXTURE_FORMAT_BLOCK_SIZE_MAP.get(texture_format)
173+
if not block_size:
174+
raise NotImplementedError(
175+
f"Not implemented swizzle format: {texture_format.name}"
176+
)
177+
texels_per_block = _get_texels_per_block(platform_blob)
178+
return _swizzle(data, width, height, *block_size, texels_per_block)
179+
180+
181+
def get_padded_image_size(
182+
width: int,
183+
height: int,
184+
texture_format: TF,
185+
platform_blob: PlatformBlobType,
186+
):
187+
block_size = TEXTURE_FORMAT_BLOCK_SIZE_MAP.get(texture_format)
188+
if not block_size:
189+
raise NotImplementedError(
190+
f"Not implemented swizzle format: {texture_format.name}"
191+
)
192+
texels_per_block = _get_texels_per_block(platform_blob)
193+
return _get_padded_texture_size(width, height, *block_size, texels_per_block)

0 commit comments

Comments
 (0)