dreamcast/gen/k_means/python/decode_pvrt.py

142 lines
4.0 KiB
Python

from dataclasses import dataclass
import sys
import struct
from PIL import Image
codebook_size = 256 * 2 * 4
@dataclass
class PVRT:
texture_data_size: int
texture_type: int
width: int
height: int
data: list[int]
def parse_pvrt_header(buf):
header = buf[0:16]
#codebook = buf[16:codebook_size + 16]
#indices = buf[codebook_size + 16:]
data = buf[16:]
assert len(header) == 16
assert header[0:4] == b"PVRT"
unpacked = struct.unpack('<LLHH', header[4:])
texture_data_size, texture_type, width, height = unpacked
#print(texture_data_size)
#print("texture type", hex(texture_type))
#print(width)
#print(height)
#assert len(indices) + len(codebook) == texture_data_size - 8, (len(indices) + len(codebook), texture_data_size - 8)
#assert len(indices) == width * height / 4, (len(indices), width * height / 4)
return PVRT(
texture_data_size,
texture_type,
width,
height,
data,
)
def rgb24(color):
r = (color >> 11) & 31
g = (color >> 5) & 63
b = (color >> 0) & 31
return r << 3, g << 2, b << 3
def get_colors(buf, codebook_ix):
codeword = buf[codebook_ix * 2 * 4:][:2 * 4]
assert len(codeword) == 2 * 4
colors = struct.unpack('<HHHH', codeword)
return list(map(rgb24, colors))
def log2(n):
if n == 8:
return 3
if n == 16:
return 4
if n == 32:
return 5
if n == 64:
return 6
if n == 128:
return 7
if n == 256:
return 8
if n == 512:
return 9
if n == 1024:
return 10
assert False, n
def from_xy(x, y, width, height):
# maximum texture size : 1024x1024
# maximum 1-dimensional index: 0xfffff
# bits : 19-0
# y bits: 0, 2, 4, 6, 8, 10, 12, 14, 16, 18
# x bits: 1, 3, 5, 7, 9, 11, 13, 15, 17, 19
width_max = log2(width);
height_max = log2(height);
twiddle_ix = 0
i = 0
while i < (20 / 2):
if i < width_max and i < height_max:
twiddle_ix |= ((y >> i) & 1) << (i * 2 + 0)
twiddle_ix |= ((x >> i) & 1) << (i * 2 + 1)
elif i < width_max:
twiddle_ix |= ((x >> i) & 1) << (i + height_max)
elif i < height_max:
twiddle_ix |= ((y >> i) & 1) << (i + width_max)
else:
break
i += 1
return twiddle_ix
def decode_vq_indices(data, width, height):
codebook = data[:codebook_size]
indices = data[codebook_size:]
canvas = [0] * width * height
for ty in range(height // 2):
for tx in range(width // 2):
codebook_ix = indices[from_xy(tx, ty, width, height)]
codeword = get_colors(codebook, codebook_ix)
ai = ((ty * 2) + 0) * width + ((tx * 2) + 0)
bi = ((ty * 2) + 1) * width + ((tx * 2) + 0)
ci = ((ty * 2) + 0) * width + ((tx * 2) + 1)
di = ((ty * 2) + 1) * width + ((tx * 2) + 1)
#print(width, height, ai, ty, tx)
canvas[ai] = codeword[0]
canvas[bi] = codeword[1]
canvas[ci] = codeword[2]
canvas[di] = codeword[3]
return canvas
def decode_twiddled(data, width, height):
canvas = [0] * width * height
for y in range(height):
for x in range(width):
ix = from_xy(x, y, width, height) * 2
color, = struct.unpack("<H", data[ix:ix+2])
canvas[y * width + x] = rgb24(color)
return canvas
in_filename = sys.argv[1]
out_filename = sys.argv[2]
with open(in_filename, 'rb') as f:
buf = f.read()
pvrt = parse_pvrt_header(buf)
print(pvrt.texture_data_size, hex(pvrt.texture_type), pvrt.width, pvrt.height)
if (pvrt.texture_type & 0xff00) == 0x300: # vq
canvas = decode_vq_indices(pvrt.data, pvrt.width, pvrt.height)
elif (pvrt.texture_type & 0xff00) == 0x100: # twiddled
canvas = decode_twiddled(pvrt.data, pvrt.width, pvrt.height)
else:
assert False, ("unsupported texture type:", hex(pvrt.texture_type))
palimage = Image.new('RGB', (pvrt.width, pvrt.height))
palimage.putdata(canvas)
palimage.save(out_filename)