fix(Texture2DConverter): decode ASTC_HDR formats with the HDR profile

UnityPy/export/Texture2DConverter.py

@@ -353,38 +353,49 @@ def atc(image_data: bytes, width: int, height: int, alpha: bool) -> Image.Image:
     return Image.frombytes("RGBA", (width, height), image_data, "raw", "BGRA")
 
 
-def astc(image_data: bytes, width: int, height: int, block_size: tuple) -> Image.Image:
+def astc(image_data: bytes, width: int, height: int, block_size: tuple, hdr: bool = False) -> Image.Image:
+    """Decode ASTC data to an RGBA image.
+
+    `hdr` selects the HDR decode profile, needed for the ASTC_HDR formats: the LDR profile does not decode HDR
+    blocks and returns magenta for blocks with HDR endpoint modes and white for HDR constant-colour blocks.
+    HDR values are clamped to [0, 1] in the 8-bit image.
+    """
     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 = get_astc_context(block_size)
+    context = get_astc_context(block_size, hdr)
     context.decompress(image_data[:texture_size], image, astc_encoder.ASTCSwizzle.from_str("RGBA"))
     assert image.data is not None, "Decompression failed, image data is None"
 
     return Image.frombytes("RGBA", (width, height), image.data, "raw", "RGBA")
 
 
 @lru_cache(maxsize=128)
-def _get_astc_context(ident: int, block_size: tuple):
-    config = astc_encoder.ASTCConfig(
-        astc_encoder.ASTCProfile.LDR,
-        *block_size,
-        block_z=1,
-        quality=100,
-        flags=astc_encoder.ASTCConfigFlags.USE_DECODE_UNORM8,
-    )
+def _get_astc_context(ident: int, block_size: tuple, hdr: bool = False):
+    if hdr:
+        # the unorm8 decode mode is only valid for the LDR profiles
+        config = astc_encoder.ASTCConfig(astc_encoder.ASTCProfile.HDR, *block_size, block_z=1, quality=100)
+    else:
+        config = astc_encoder.ASTCConfig(
+            astc_encoder.ASTCProfile.LDR,
+            *block_size,
+            block_z=1,
+            quality=100,
+            flags=astc_encoder.ASTCConfigFlags.USE_DECODE_UNORM8,
+        )
     context = astc_encoder.ASTCContext(config)
     return context
 
 
-def get_astc_context(block_size: tuple):
-    """Get the ASTC context for the current thread using the given `block_size`.
+def get_astc_context(block_size: tuple, hdr: bool = False):
+    """Get the ASTC context for the current thread using the given `block_size`,
+    with the HDR profile if `hdr` is set, else the LDR profile.
     Created contexts belong to and only to the calling thread, and may be cached.
     This function is thread safe.
     """
-    return _get_astc_context(get_ident(), block_size)
+    return _get_astc_context(get_ident(), block_size, hdr)
 
 
 def calculate_astc_compressed_size(width: int, height: int, block_size: tuple) -> int:
@@ -546,12 +557,12 @@ def rgb9e5float(image_data: bytes, width: int, height: int) -> Image.Image:
     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.ASTC_HDR_4x4: (astc, ((4, 4), True)),
+    TF.ASTC_HDR_5x5: (astc, ((5, 5), True)),
+    TF.ASTC_HDR_6x6: (astc, ((6, 6), True)),
+    TF.ASTC_HDR_8x8: (astc, ((8, 8), True)),
+    TF.ASTC_HDR_10x10: (astc, ((10, 10), True)),
+    TF.ASTC_HDR_12x12: (astc, ((12, 12), True)),
     TF.RG32: (rg, ("RGB", "raw", "RG;16")),
     TF.RGB48: (pillow, ("RGB", "raw", "RGB;16")),
     TF.RGBA64: (pillow, ("RGBA", "raw", "RGBA;16")),

tests/test_astc.py

@@ -0,0 +1,65 @@
+import struct
+
+import astc_encoder
+import pytest
+
+from UnityPy.enums import TextureFormat as TF
+from UnityPy.export.Texture2DConverter import CONV_TABLE
+
+# the colour astcenc returns for a block it cannot decode in the selected profile
+ERROR_COLOUR = (255, 0, 255, 255)
+WIDTH = HEIGHT = 24
+
+
+def compress(profile, block_size, pixels):
+    image = astc_encoder.ASTCImage(
+        astc_encoder.ASTCType.F32, WIDTH, HEIGHT, 1, struct.pack(f"<{len(pixels)}f", *pixels)
+    )
+    config = astc_encoder.ASTCConfig(profile, *block_size, block_z=1, quality=100)
+    return astc_encoder.ASTCContext(config).compress(image, astc_encoder.ASTCSwizzle.from_str("RGBA"))
+
+
+def gradient(peak):
+    """RGBA floats rising from left to right up to `peak` in red, half of that in green."""
+    pixels = []
+    for _ in range(HEIGHT):
+        for x in range(WIDTH):
+            v = peak * (x + 1) / WIDTH
+            pixels += [v, v / 2, 0.1, 1.0]
+    return pixels
+
+
+def decode(texture_format, data):
+    func, args = CONV_TABLE[texture_format]
+    image = func(data, WIDTH, HEIGHT, *args)
+    raw = image.tobytes()
+    return [tuple(raw[i : i + 4]) for i in range(0, len(raw), 4)]
+
+
+def expected(pixels):
+    """The 8-bit values of the source floats, clamped to [0, 1]."""
+    return [tuple(round(min(max(c, 0.0), 1.0) * 255) for c in pixels[i : i + 4]) for i in range(0, len(pixels), 4)]
+
+
+def assert_close(decoded, source, tolerance):
+    for got, want in zip(decoded, expected(source)):
+        assert all(abs(g - w) <= tolerance for g, w in zip(got, want)), (got, want)
+
+
+@pytest.mark.parametrize("block_size", [(4, 4), (5, 5), (6, 6), (8, 8), (10, 10), (12, 12)])
+def test_astc_hdr(block_size):
+    source = gradient(peak=2.0)
+    data = compress(astc_encoder.ASTCProfile.HDR, block_size, source)
+    decoded = decode(getattr(TF, f"ASTC_HDR_{block_size[0]}x{block_size[1]}"), data)
+    assert ERROR_COLOUR not in decoded
+    assert_close(decoded, source, tolerance=16)
+    # values above 1.0 are clamped in the 8-bit image
+    assert decoded[WIDTH - 1][0] == 255
+
+
+@pytest.mark.parametrize("block_size", [(4, 4), (6, 6), (8, 8)])
+def test_astc_ldr(block_size):
+    source = gradient(peak=1.0)
+    data = compress(astc_encoder.ASTCProfile.LDR, block_size, source)
+    decoded = decode(getattr(TF, f"ASTC_RGBA_{block_size[0]}x{block_size[1]}"), data)
+    assert_close(decoded, source, tolerance=4)