From 248d37e99787498fb2f763d164afa68baf4b04f3 Mon Sep 17 00:00:00 2001 From: Harry Huang Date: Sun, 26 Jan 2025 19:26:12 +0800 Subject: [PATCH 1/5] fix: type errors in Texture2DConverter --- UnityPy/export/Texture2DConverter.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/UnityPy/export/Texture2DConverter.py b/UnityPy/export/Texture2DConverter.py index caf9fa0f..9677b856 100644 --- a/UnityPy/export/Texture2DConverter.py +++ b/UnityPy/export/Texture2DConverter.py @@ -1,4 +1,4 @@ -from __future__ import annotations +from __future__ import annotations import struct from copy import copy @@ -177,7 +177,7 @@ def parse_image_data( if len(selection) == 0: raise NotImplementedError( - f"Not implemented texture format: {texture_format.name}" + f"Not implemented texture format: {texture_format}" ) if platform == BuildTarget.XBOX360 and texture_format in XBOX_SWAP_FORMATS: @@ -356,7 +356,7 @@ def rg( padding = bytes(padding_size) rgb_data = b"".join( stream.read(padding_size * 2) + padding - for _ in range(image_data / (2 * padding_size)) + for _ in range(len(image_data) // (2 * padding_size)) ) if codec == "RGE": return half(rgb_data, width, height, mode, "RGB", args) From fed2433f85e82aea4238e1bb09fe5165f4428a2b Mon Sep 17 00:00:00 2001 From: Harry Huang Date: Sun, 26 Jan 2025 19:27:01 +0800 Subject: [PATCH 2/5] fix: incorrect type hint in Texture2DConverter --- UnityPy/export/Texture2DConverter.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/UnityPy/export/Texture2DConverter.py b/UnityPy/export/Texture2DConverter.py index 9677b856..3c7a97d6 100644 --- a/UnityPy/export/Texture2DConverter.py +++ b/UnityPy/export/Texture2DConverter.py @@ -76,7 +76,7 @@ def image_to_texture2d( ) config = astc_encoder.ASTCConfig( - astc_encoder.ASTCProfile.LDR, *block_size, 1, 100 + astc_encoder.ASTCProfile.LDR, *block_size, block_z=1, quality=100 ) context = astc_encoder.ASTCContext(config) raw_img = astc_encoder.ASTCImage( @@ -166,7 +166,7 @@ def parse_image_data( texture_format: Union[int, TextureFormat], version: tuple, platform: int, - platform_blob: bytes = None, + platform_blob: Union[bytes, None] = None, flip: bool = True, ) -> Image.Image: image_data = copy(bytes(image_data)) @@ -230,7 +230,7 @@ def pillow( mode: str, codec: str, args, - swap: tuple = None, + swap: Union[tuple, None] = None, ) -> Image.Image: img = ( Image.frombytes(mode, (width, height), image_data, codec, args) @@ -327,7 +327,7 @@ def half( mode: str, codec: str, args, - swap: tuple = None, + swap: Union[tuple, None] = None, ) -> Image.Image: # convert half-float to int8 stream = BytesIO(image_data) From 830a2bd52a13e1302d12c2f358bb20b735b66c1d Mon Sep 17 00:00:00 2001 From: Harry Huang Date: Sun, 26 Jan 2025 19:43:44 +0800 Subject: [PATCH 3/5] fix: concurrent issue of astc context --- UnityPy/export/Texture2DConverter.py | 77 +++++++++++++++------------- 1 file changed, 42 insertions(+), 35 deletions(-) diff --git a/UnityPy/export/Texture2DConverter.py b/UnityPy/export/Texture2DConverter.py index 3c7a97d6..3239b188 100644 --- a/UnityPy/export/Texture2DConverter.py +++ b/UnityPy/export/Texture2DConverter.py @@ -1,8 +1,9 @@ -from __future__ import annotations +from __future__ import annotations import struct from copy import copy from io import BytesIO +from threading import Lock from typing import TYPE_CHECKING, Dict, Tuple, Union import astc_encoder @@ -74,22 +75,20 @@ def image_to_texture2d( block_size = tuple( map(int, target_texture_format.name.rsplit("_", 1)[1].split("x")) ) - - config = astc_encoder.ASTCConfig( - astc_encoder.ASTCProfile.LDR, *block_size, block_z=1, quality=100 - ) - context = astc_encoder.ASTCContext(config) - raw_img = astc_encoder.ASTCImage( - astc_encoder.ASTCType.U8, img.width, img.height, 1, raw_img - ) - if img.mode == "RGB": - tex_format = getattr(TF, f"ASTC_RGB_{block_size[0]}x{block_size[1]}") - else: - tex_format = getattr(TF, f"ASTC_RGBA_{block_size[0]}x{block_size[1]}") - - swizzle = astc_encoder.ASTCSwizzle.from_str("RGBA") - enc_img = context.compress(raw_img, swizzle) - tex_format = target_texture_format + context, lock = get_astc_context(block_size) + + with lock: + raw_img = astc_encoder.ASTCImage( + astc_encoder.ASTCType.U8, img.width, img.height, 1, raw_img + ) + if img.mode == "RGB": + tex_format = getattr(TF, f"ASTC_RGB_{block_size[0]}x{block_size[1]}") + else: + tex_format = getattr(TF, f"ASTC_RGBA_{block_size[0]}x{block_size[1]}") + + swizzle = astc_encoder.ASTCSwizzle.from_str("RGBA") + enc_img = context.compress(raw_img, swizzle) + tex_format = target_texture_format # A elif target_texture_format == TF.Alpha8: enc_img = img.tobytes("raw", "A") @@ -252,30 +251,38 @@ def atc(image_data: bytes, width: int, height: int, alpha: bool) -> Image.Image: return Image.frombytes("RGBA", (width, height), image_data, "raw", "BGRA") -ASTC_CONTEXTS: Dict[Tuple[int, int], astc_encoder.ASTCContext] = {} +def astc(image_data: bytes, width: int, height: int, block_size: tuple) -> Image.Image: + context, lock = get_astc_context(block_size) + + with lock: + image = astc_encoder.ASTCImage(astc_encoder.ASTCType.U8, width, height, 1) + texture_size = calculate_astc_compressed_size(width, height, block_size) + if len(image_data) < texture_size: + raise ValueError(f"Invalid ASTC data size: {len(image_data)} < {texture_size}") + context.decompress( + image_data[:texture_size], image, astc_encoder.ASTCSwizzle.from_str("RGBA") + ) + + return Image.frombytes("RGBA", (width, height), image.data, "raw", "RGBA") + +ASTC_CONTEXTS: Dict[Tuple[int, int], Tuple[astc_encoder.ASTCContext, Lock]] = {} -def astc(image_data: bytes, width: int, height: int, block_size: tuple) -> Image.Image: - context = ASTC_CONTEXTS.get(block_size) - if context is None: + +def get_astc_context(block_size: tuple): + """Get the ASTC context and its lock using the given `block_size`.""" + if block_size not in ASTC_CONTEXTS: config = astc_encoder.ASTCConfig( astc_encoder.ASTCProfile.LDR, *block_size, - 1, - 100, - astc_encoder.ASTCConfigFlags.USE_DECODE_UNORM8, + block_z=1, + quality=100, + flags=astc_encoder.ASTCConfigFlags.USE_DECODE_UNORM8, ) - context = ASTC_CONTEXTS[block_size] = astc_encoder.ASTCContext(config) - - image = astc_encoder.ASTCImage(astc_encoder.ASTCType.U8, width, height, 1) - texture_size = calculate_astc_compressed_size(width, height, block_size) - if len(image_data) < texture_size: - raise ValueError(f"Invalid ASTC data size: {len(image_data)} < {texture_size}") - context.decompress( - image_data[:texture_size], image, astc_encoder.ASTCSwizzle.from_str("RGBA") - ) - - return Image.frombytes("RGBA", (width, height), image.data, "raw", "RGBA") + context = astc_encoder.ASTCContext(config) + lock = Lock() + ASTC_CONTEXTS[block_size] = (context, lock) + return ASTC_CONTEXTS[block_size] def calculate_astc_compressed_size(width: int, height: int, block_size: tuple) -> int: From b2bd6bbdddc4c969a973c5f9db55c637046a0400 Mon Sep 17 00:00:00 2001 From: Harry Huang Date: Mon, 27 Jan 2025 20:44:31 +0800 Subject: [PATCH 4/5] fix: type hint use Optional --- UnityPy/export/Texture2DConverter.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/UnityPy/export/Texture2DConverter.py b/UnityPy/export/Texture2DConverter.py index 3239b188..e6de3d74 100644 --- a/UnityPy/export/Texture2DConverter.py +++ b/UnityPy/export/Texture2DConverter.py @@ -4,7 +4,7 @@ from copy import copy from io import BytesIO from threading import Lock -from typing import TYPE_CHECKING, Dict, Tuple, Union +from typing import TYPE_CHECKING, Dict, Optional, Tuple, Union import astc_encoder import texture2ddecoder @@ -165,7 +165,7 @@ def parse_image_data( texture_format: Union[int, TextureFormat], version: tuple, platform: int, - platform_blob: Union[bytes, None] = None, + platform_blob: Optional[bytes] = None, flip: bool = True, ) -> Image.Image: image_data = copy(bytes(image_data)) @@ -229,7 +229,7 @@ def pillow( mode: str, codec: str, args, - swap: Union[tuple, None] = None, + swap: Optional[tuple] = None, ) -> Image.Image: img = ( Image.frombytes(mode, (width, height), image_data, codec, args) @@ -334,7 +334,7 @@ def half( mode: str, codec: str, args, - swap: Union[tuple, None] = None, + swap: Optional[tuple] = None, ) -> Image.Image: # convert half-float to int8 stream = BytesIO(image_data) From d39ea26c79dcc37f9282ad2bb7c22ad0ca18edb6 Mon Sep 17 00:00:00 2001 From: Harry Huang Date: Mon, 27 Jan 2025 20:50:09 +0800 Subject: [PATCH 5/5] fix: reduce astc context lock scope --- UnityPy/export/Texture2DConverter.py | 32 ++++++++++++++-------------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/UnityPy/export/Texture2DConverter.py b/UnityPy/export/Texture2DConverter.py index e6de3d74..0bd887ad 100644 --- a/UnityPy/export/Texture2DConverter.py +++ b/UnityPy/export/Texture2DConverter.py @@ -71,24 +71,24 @@ def image_to_texture2d( # ASTC elif target_texture_format.name.startswith("ASTC"): raw_img = img.tobytes("raw", "RGBA") + raw_img = astc_encoder.ASTCImage( + astc_encoder.ASTCType.U8, img.width, img.height, 1, raw_img + ) + if img.mode == "RGB": + tex_format = getattr(TF, f"ASTC_RGB_{block_size[0]}x{block_size[1]}") + else: + tex_format = getattr(TF, f"ASTC_RGBA_{block_size[0]}x{block_size[1]}") + + swizzle = astc_encoder.ASTCSwizzle.from_str("RGBA") block_size = tuple( map(int, target_texture_format.name.rsplit("_", 1)[1].split("x")) ) context, lock = get_astc_context(block_size) - with lock: - raw_img = astc_encoder.ASTCImage( - astc_encoder.ASTCType.U8, img.width, img.height, 1, raw_img - ) - if img.mode == "RGB": - tex_format = getattr(TF, f"ASTC_RGB_{block_size[0]}x{block_size[1]}") - else: - tex_format = getattr(TF, f"ASTC_RGBA_{block_size[0]}x{block_size[1]}") - - swizzle = astc_encoder.ASTCSwizzle.from_str("RGBA") enc_img = context.compress(raw_img, swizzle) - tex_format = target_texture_format + + tex_format = target_texture_format # A elif target_texture_format == TF.Alpha8: enc_img = img.tobytes("raw", "A") @@ -252,13 +252,13 @@ def atc(image_data: bytes, width: int, height: int, alpha: bool) -> Image.Image: def astc(image_data: bytes, width: int, height: int, block_size: tuple) -> Image.Image: - context, lock = get_astc_context(block_size) + image = astc_encoder.ASTCImage(astc_encoder.ASTCType.U8, width, height, 1) + texture_size = calculate_astc_compressed_size(width, height, block_size) + if len(image_data) < texture_size: + raise ValueError(f"Invalid ASTC data size: {len(image_data)} < {texture_size}") + context, lock = get_astc_context(block_size) with lock: - image = astc_encoder.ASTCImage(astc_encoder.ASTCType.U8, width, height, 1) - texture_size = calculate_astc_compressed_size(width, height, block_size) - if len(image_data) < texture_size: - raise ValueError(f"Invalid ASTC data size: {len(image_data)} < {texture_size}") context.decompress( image_data[:texture_size], image, astc_encoder.ASTCSwizzle.from_str("RGBA") )