1
0
mirror of synced 2024-11-24 06:20:12 +01:00
bemaniutils/bemani/format/tdxt.py
2024-01-02 02:46:24 +00:00

431 lines
15 KiB
Python

import struct
from PIL import Image
from typing import Optional
from bemani.format.dxt import DXTBuffer
class TDXT:
def __init__(
self,
header_flags1: int,
header_flags2: int,
header_flags3: int,
width: int,
height: int,
fmt: int,
fmtflags: int,
endian: str,
length_fixup: bool,
raw: bytes,
img: Optional[Image.Image],
*,
invert_channels: bool = False,
) -> None:
self.header_flags1 = header_flags1
self.header_flags2 = header_flags2
self.header_flags3 = header_flags3
self.width = width
self.height = height
self.fmt = fmt
self.fmtflags = fmtflags
self.endian = endian
self.length_fixup = length_fixup
self.invert_channels = invert_channels
self.__raw = raw
self.__img = img
@property
def raw(self) -> bytes:
return self.__raw
@raw.setter
def raw(self, newdata: bytes) -> None:
self.__raw = newdata
newimg = self._rawToImg(
self.width,
self.height,
self.fmt,
self.endian,
self.invert_channels,
newdata,
)
width, height = newimg.size
if width != self.width or height != self.height:
raise Exception("Unsupported texture resize operation for TDXT file!")
self.__img = newimg
@property
def img(self) -> Optional[Image.Image]:
return self.__img
@img.setter
def img(self, newimg: Image.Image) -> None:
self.__img = newimg
self.__raw = self._imgToRaw(newimg)
@staticmethod
def fromBytes(raw_data: bytes, *, invert_channels: bool = False) -> "TDXT":
# First, check the endianness.
(magic,) = struct.unpack_from("4s", raw_data)
if magic == b"TDXT":
endian = "<"
elif magic == b"TXDT":
endian = ">"
else:
raise Exception("Unexpected texture format!")
(
magic,
header_flags1,
header_flags2,
raw_length,
width,
height,
fmtflags,
expected_zero1,
expected_zero2,
) = struct.unpack(
f"{endian}4sIIIHHIII",
raw_data[0:32],
)
# Standalone TDXT files don't include the header in the length. Embedded ones
# inside a TXP2 file do. Ugh.
if raw_length == len(raw_data):
length_fixup = False
elif (raw_length + 64) == len(raw_data):
length_fixup = True
else:
raise Exception("Invalid texture length!")
# I have only ever observed the following values across two different games.
# Don't want to keep the chunk around so let's assert our assumptions.
if (expected_zero1 | expected_zero2) != 0:
raise Exception("Found unexpected non-zero value in texture header!")
if raw_data[32:44] != b"\0" * 12:
raise Exception("Found unexpected non-zero value in texture header!")
# This is almost ALWAYS 3, but I've seen it be 1 as well, so I guess we have to
# round-trip it if we want to write files back out. I have no clue what it's for.
# I've seen it be 1 only on files used for fonts so far, but I am not sure there
# is any correlation there.
header_flags3 = struct.unpack(f"{endian}I", raw_data[44:48])[0]
if raw_data[48:64] != b"\0" * 16:
raise Exception("Found unexpected non-zero value in texture header!")
fmt = fmtflags & 0xFF
# Extract flags that the game cares about.
# flags1 = (fmtflags >> 24) & 0xFF
# flags2 = (fmtflags >> 16) & 0xFF
# unk1 = 3 if (flags1 & 0xF == 1) else 1
# unk2 = 3 if ((flags1 >> 4) & 0xF == 1) else 1
# unk3 = 1 if (flags2 & 0xF == 1) else 2
# unk4 = 1 if ((flags2 >> 4) & 0xF == 1) else 2
# Convert texture to image if possible, create structure.
return TDXT(
header_flags1=header_flags1,
header_flags2=header_flags2,
header_flags3=header_flags3,
width=width,
height=height,
fmt=fmt,
fmtflags=fmtflags & 0xFFFFFF00,
endian=endian,
length_fixup=length_fixup,
raw=raw_data[64:],
img=TDXT._rawToImg(width, height, fmt, endian, invert_channels, raw_data[64:]),
invert_channels=invert_channels,
)
@staticmethod
def _rawToImg(
width: int, height: int, fmt: int, endian: str, invert: bool, raw_data: bytes
) -> Optional[Image.Image]:
# Since the AFP file format can be found in both big and little endian, its
# possible that some of these loaders might need byteswapping on some platforms.
# This has been tested on files intended for X86 (little endian) as well as PS3
# (big endian). I've found that the "correct" thing to do is always treat data as
# little-endian instead of the determined endianness of the file. But, this could
# also be broken per-game, so I'm not entirely sure this is fully possible to do
# generically. However, what's here has been tested across a broad range of games
# and does seem to work.
if fmt == 0x01:
# As far as I can tell, this is 8 bit grayscale. Decoding as such results in
# images that are recognizeable and look correct.
img = Image.frombytes(
"L",
(width, height),
raw_data,
"raw",
"L",
)
elif fmt == 0x0B:
# 16-bit 565 color RGB format. Game references D3D9 texture format 23 (R5G6B5).
newdata = []
for i in range(width * height):
pixel = struct.unpack(
"<H",
raw_data[(i * 2) : (2 + (i * 2))],
)[0]
# Extract the raw values
red = ((pixel >> 0) & 0x1F) << 3
green = ((pixel >> 5) & 0x3F) << 2
blue = ((pixel >> 11) & 0x1F) << 3
# Scale the colors so they fill the entire 8 bit range.
red = red | (red >> 5)
green = green | (green >> 6)
blue = blue | (blue >> 5)
newdata.append(struct.pack("<BBB", blue, green, red))
img = Image.frombytes(
"RGB",
(width, height),
b"".join(newdata),
"raw",
"BGR" if invert else "RGB",
)
elif fmt == 0x0E:
# RGB image, no alpha. Game references D3D9 texture format 22 (R8G8B8).
img = Image.frombytes(
"RGB",
(width, height),
raw_data,
"raw",
"BGR" if invert else "RGB",
)
elif fmt == 0x10:
# Seems to be some sort of RGBA with color swapping. Game references D3D9 texture
# format 21 (A8R8B8G8) but does manual byteswapping.
img = Image.frombytes(
"RGBA",
(width, height),
raw_data,
"raw",
"RGBA" if invert else "BGRA",
)
elif fmt == 0x13:
# Some 16-bit texture format. Game references D3D9 texture format 25 (A1R5G5B5).
newdata = []
for i in range(width * height):
pixel = struct.unpack(
"<H",
raw_data[(i * 2) : (2 + (i * 2))],
)[0]
# Extract the raw values
alpha = 255 if ((pixel >> 15) & 0x1) != 0 else 0
red = ((pixel >> 0) & 0x1F) << 3
green = ((pixel >> 5) & 0x1F) << 3
blue = ((pixel >> 10) & 0x1F) << 3
# Scale the colors so they fill the entire 8 bit range.
red = red | (red >> 5)
green = green | (green >> 5)
blue = blue | (blue >> 5)
newdata.append(struct.pack("<BBBB", blue, green, red, alpha))
img = Image.frombytes(
"RGBA",
(width, height),
b"".join(newdata),
"raw",
"BGRA" if invert else "RGBA",
)
elif fmt == 0x15:
# RGBA format. Game references D3D9 texture format 21 (A8R8G8B8).
# Looks like unlike 0x20 below, the game does some endianness swapping.
img = Image.frombytes(
"RGBA",
(width, height),
raw_data,
"raw",
"ABGR" if invert else "ARGB",
)
elif fmt == 0x16:
# DXT1 format. Game references D3D9 DXT1 texture format.
# Konami seems to have screwed up with DDR PS3 where they
# swap every other byte in the format, even though its specified
# as little-endian by all DXT1 documentation.
dxt = DXTBuffer(width, height)
img = Image.frombuffer(
"RGBA",
(width, height),
dxt.DXT1Decompress(raw_data, swap=endian != "<"),
"raw",
"RGBA",
0,
1,
)
elif fmt == 0x1A:
# DXT5 format. Game references D3D9 DXT5 texture format.
# Konami seems to have screwed up with DDR PS3 where they
# swap every other byte in the format, even though its specified
# as little-endian by all DXT5 documentation.
dxt = DXTBuffer(width, height)
img = Image.frombuffer(
"RGBA",
(width, height),
dxt.DXT5Decompress(raw_data, swap=endian != "<"),
"raw",
"RGBA",
0,
1,
)
elif fmt == 0x1E:
# I have no idea what format this is. The game does some byte
# swapping but doesn't actually call any texture create calls.
# This might be leftover from another game.
pass
elif fmt == 0x1F:
# 16-bit 4-4-4-4 RGBA format. Game references D3D9 texture format 26 (A4R4G4B4).
newdata = []
for i in range(width * height):
pixel = struct.unpack(
"<H",
raw_data[(i * 2) : (2 + (i * 2))],
)[0]
# Extract the raw values
blue = ((pixel >> 0) & 0xF) << 4
green = ((pixel >> 4) & 0xF) << 4
red = ((pixel >> 8) & 0xF) << 4
alpha = ((pixel >> 12) & 0xF) << 4
# Scale the colors so they fill the entire 8 bit range.
red = red | (red >> 4)
green = green | (green >> 4)
blue = blue | (blue >> 4)
alpha = alpha | (alpha >> 4)
newdata.append(struct.pack("<BBBB", red, green, blue, alpha))
img = Image.frombytes(
"RGBA",
(width, height),
b"".join(newdata),
"raw",
"BGRA" if invert else "RGBA",
)
elif fmt == 0x20:
# RGBA format. Game references D3D9 surface format 21 (A8R8G8B8).
img = Image.frombytes(
"RGBA",
(width, height),
raw_data,
"raw",
"RGBA" if invert else "BGRA",
)
else:
img = None
return img
def toBytes(self) -> bytes:
# Construct the TDXT texture format from our parsed results.
if self.endian == "<":
magic = b"TDXT"
elif self.endian == ">":
magic = b"TXDT"
else:
raise Exception("Unexpected texture format!")
fmtflags = (self.fmtflags & 0xFFFFFF00) | (self.fmt & 0xFF)
return (
struct.pack(
f"{self.endian}4sIIIHHIII",
magic,
self.header_flags1,
self.header_flags2,
# Some files include the header in the length, some do not. We detect this
# in the header, and round-trip it here.
len(self.raw) if self.length_fixup else (64 + len(self.raw)),
self.width,
self.height,
fmtflags,
0,
0,
)
+ (b"\0" * 12)
+ struct.pack(
f"{self.endian}I",
self.header_flags3,
)
+ (b"\0" * 16)
+ self.raw
)
def _imgToRaw(self, imgdata: Image.Image) -> bytes:
width, height = imgdata.size
if width != self.width or height != self.height:
raise Exception("Unsupported texture resize operation for TDXT file!")
# Ignore alpha, which is basically always in the right place.
if self.invert_channels:
order = (2, 1, 0)
else:
order = (0, 1, 2)
if self.fmt == 0x0B:
# 16-bit 565 color RGB format.
raw = b"".join(
struct.pack(
"<H",
(
(((pixel[order[0]] >> 3) & 0x1F) << 11)
| (((pixel[order[1]] >> 2) & 0x3F) << 5)
| ((pixel[order[2]] >> 3) & 0x1F)
),
)
for pixel in imgdata.getdata()
)
elif self.fmt == 0x13:
# 16-bit A1R5G55 texture format.
raw = b"".join(
struct.pack(
"<H",
(
(0x8000 if pixel[3] >= 128 else 0x0000)
| (((pixel[order[0]] >> 3) & 0x1F) << 10)
| (((pixel[order[1]] >> 3) & 0x1F) << 5)
| ((pixel[order[2]] >> 3) & 0x1F)
),
)
for pixel in imgdata.getdata()
)
elif self.fmt == 0x1F:
# 16-bit 4-4-4-4 RGBA format.
raw = b"".join(
struct.pack(
"<H",
(
((pixel[order[2]] >> 4) & 0xF)
| (((pixel[order[1]] >> 4) & 0xF) << 4)
| (((pixel[order[0]] >> 4) & 0xF) << 8)
| (((pixel[3] >> 4) & 0xF) << 12)
),
)
for pixel in imgdata.getdata()
)
elif self.fmt == 0x20:
# 32-bit RGBA format, stored in BGRA order.
raw = b"".join(
struct.pack(
"<BBBB",
pixel[order[2]],
pixel[order[1]],
pixel[order[0]],
pixel[3],
)
for pixel in imgdata.getdata()
)
else:
raise Exception(f"Unsupported format {hex(self.fmt)} for TDXT file!")
return raw