Add support for other common pixel formats

This commit is contained in:
REDxEYE 2022-08-13 22:37:42 +03:00
parent c6c0749a16
commit 7285b74a61

View File

@ -12,7 +12,7 @@ Full text of the CC0 license:
import struct import struct
from enum import IntEnum, IntFlag from enum import IntEnum, IntFlag
from io import BytesIO from io import BytesIO, BufferedIOBase
from math import ceil, log from math import ceil, log
from typing import NamedTuple from typing import NamedTuple
@ -125,7 +125,7 @@ VTFHeader = NamedTuple(
("resource_count", int), ("resource_count", int),
], ],
) )
RGB_FORMATS = (VtfPF.RGB888,) RGB_FORMATS = (VtfPF.RGB888, VtfPF.BGR888, VtfPF.UV88,)
RGBA_FORMATS = ( RGBA_FORMATS = (
VtfPF.DXT1, VtfPF.DXT1,
VtfPF.DXT1_ONEBITALPHA, VtfPF.DXT1_ONEBITALPHA,
@ -139,7 +139,6 @@ L_FORMATS = (
) )
LA_FORMATS = ( LA_FORMATS = (
VtfPF.IA88, VtfPF.IA88,
VtfPF.UV88,
) )
BLOCK_COMPRESSED = (VtfPF.DXT1, VtfPF.DXT1_ONEBITALPHA, VtfPF.DXT3, VtfPF.DXT5) BLOCK_COMPRESSED = (VtfPF.DXT1, VtfPF.DXT1_ONEBITALPHA, VtfPF.DXT3, VtfPF.DXT5)
@ -153,18 +152,22 @@ def _get_texture_size(pixel_format: VtfPF, width, height):
if pixel_format in (VtfPF.DXT1, VtfPF.DXT1_ONEBITALPHA): if pixel_format in (VtfPF.DXT1, VtfPF.DXT1_ONEBITALPHA):
return width * height // 2 return width * height // 2
elif ( elif (
pixel_format pixel_format
in ( in (
VtfPF.DXT3, VtfPF.DXT3,
VtfPF.DXT5, VtfPF.DXT5,
) )
+ L_FORMATS + L_FORMATS
): ):
return width * height return width * height
elif pixel_format == VtfPF.UV88:
return width * height * 2
elif pixel_format in LA_FORMATS: elif pixel_format in LA_FORMATS:
return width * height * 2 return width * height * 2
elif pixel_format == VtfPF.RGB888: elif pixel_format == VtfPF.RGB888:
return width * height * 3 return width * height * 3
elif pixel_format == VtfPF.BGR888:
return width * height * 3
elif pixel_format == VtfPF.RGBA8888: elif pixel_format == VtfPF.RGBA8888:
return width * height * 4 return width * height * 4
raise VTFException(f"Unsupported VTF pixel format: {pixel_format}") raise VTFException(f"Unsupported VTF pixel format: {pixel_format}")
@ -181,9 +184,47 @@ def _get_mipmap_count(width: int, height: int):
return mip_count return mip_count
def closest_power(x): def _write_image(fp: BufferedIOBase, im: Image.Image, pixel_format: VtfPF):
extents = (0, 0) + im.size
if pixel_format == VtfPF.DXT1:
encoder = 'bcn'
encoder_args = (1, "DXT1")
elif pixel_format == VtfPF.DXT3:
encoder = 'bcn'
encoder_args = (3, "DXT3")
elif pixel_format == VtfPF.DXT5:
encoder = 'bcn'
encoder_args = (5, "DXT5")
elif pixel_format == VtfPF.RGB888:
encoder = 'raw'
encoder_args = ("RGB", 0, 0)
elif pixel_format == VtfPF.BGR888:
encoder = 'raw'
encoder_args = ("BGR", 0, 0)
elif pixel_format == VtfPF.RGBA8888:
encoder = 'raw'
encoder_args = ("RGBA", 0, 0)
elif pixel_format == VtfPF.I8:
encoder = 'raw'
encoder_args = ("L", 0, 0)
elif pixel_format == VtfPF.IA88:
encoder = 'raw'
encoder_args = ("LA", 0, 0)
elif pixel_format == VtfPF.UV88:
encoder = 'raw'
r, g, *_ = im.split()
im = Image.merge('LA', (r, g))
encoder_args = ("LA", 0, 0)
else:
raise VTFException(f"Unsupported pixel format: {pixel_format!r}")
tile = [(encoder, extents, fp.tell(), encoder_args)]
ImageFile._save(im, fp, tile, _get_texture_size(pixel_format, *im.size))
def _closest_power(x):
possible_results = round(log(x, 2)), ceil(log(x, 2)) possible_results = round(log(x, 2)), ceil(log(x, 2))
return 2 ** min(possible_results, key=lambda z: abs(x - 2**z)) return 2 ** min(possible_results, key=lambda z: abs(x - 2 ** z))
class VtfImageFile(ImageFile.ImageFile): class VtfImageFile(ImageFile.ImageFile):
@ -268,12 +309,23 @@ class VtfImageFile(ImageFile.ImageFile):
def _save(im, fp, filename): def _save(im, fp, filename):
im: Image.Image im: Image.Image
if im.mode not in ("RGB", "RGBA"): if im.mode not in ("RGB", "RGBA", 'L', 'LA'):
raise OSError(f"cannot write mode {im.mode} as VTF") raise OSError(f"cannot write mode {im.mode} as VTF")
encoderinfo = im.encoderinfo encoderinfo = im.encoderinfo
pixel_format = VtfPF(encoderinfo.get("pixel_format", VtfPF.RGBA8888)) pixel_format = VtfPF(encoderinfo.get("pixel_format", VtfPF.RGBA8888))
version = encoderinfo.get("version", (7, 4)) version = encoderinfo.get("version", (7, 4))
generate_mips = encoderinfo.get('generate_mips', True)
flags = CompiledVtfFlags(0) flags = CompiledVtfFlags(0)
if pixel_format in RGBA_FORMATS:
im = im.convert('RGBA')
if pixel_format in RGB_FORMATS:
im = im.convert('RGB')
if pixel_format in L_FORMATS:
im = im.convert('L')
if pixel_format in LA_FORMATS:
im = im.convert('LA')
if "A" in im.mode: if "A" in im.mode:
if pixel_format == VtfPF.DXT1_ONEBITALPHA: if pixel_format == VtfPF.DXT1_ONEBITALPHA:
flags |= CompiledVtfFlags.ONEBITALPHA flags |= CompiledVtfFlags.ONEBITALPHA
@ -282,16 +334,18 @@ def _save(im, fp, filename):
else: else:
flags |= CompiledVtfFlags.EIGHTBITALPHA flags |= CompiledVtfFlags.EIGHTBITALPHA
im = im.resize((closest_power(im.width), closest_power(im.height))) im = im.resize((_closest_power(im.width), _closest_power(im.height)))
width, height = im.size width, height = im.size
mipmap_count = _get_mipmap_count(width, height) mipmap_count = 0
if generate_mips:
mipmap_count = _get_mipmap_count(width, height)
thumb_buffer = BytesIO() thumb_buffer = BytesIO()
thumb = im.convert("RGB") thumb = im.convert("RGB")
thumb.thumbnail(((min(16, width)), (min(16, height)))) thumb.thumbnail(((min(16, width)), (min(16, height))))
thumb = thumb.resize((closest_power(thumb.width), closest_power(thumb.height))) thumb = thumb.resize((_closest_power(thumb.width), _closest_power(thumb.height)))
ImageFile._save(thumb, thumb_buffer, [("bcn", (0, 0) + thumb.size, 0, (1, "DXT1"))]) _write_image(thumb_buffer, thumb, VtfPF.DXT1)
header = VTFHeader( header = VTFHeader(
0, 0,
@ -338,19 +392,17 @@ def _save(im, fp, filename):
fp.write(b"\x00" * (16 - fp.tell() % 16)) fp.write(b"\x00" * (16 - fp.tell() % 16))
fp.write(thumb_buffer.getbuffer()) fp.write(thumb_buffer.getbuffer())
if pixel_format in BLOCK_COMPRESSED:
min_size = 4
else:
min_size = 1
for mip_id in range(mipmap_count - 1, 0, -1): for mip_id in range(mipmap_count - 1, 0, -1):
mip_width = max(4, width >> mip_id) mip_width = max(min_size, width >> mip_id)
mip_height = max(4, height >> mip_id) mip_height = max(min_size, height >> mip_id)
mip = im.resize((mip_width, mip_height)) mip = im.resize((mip_width, mip_height))
buffer_size = mip_width * mip_height // 2 _write_image(fp, mip, pixel_format)
extents = (0, 0) + mip.size _write_image(fp, im, pixel_format)
ImageFile._save(
mip, fp, [("bcn", extents, fp.tell(), (1, "DXT1"))], buffer_size
)
buffer_size = im.width * im.height // 2
ImageFile._save(
im, fp, [("bcn", (0, 0) + im.size, fp.tell(), (1, "DXT1"))], buffer_size
)
def _accept(prefix): def _accept(prefix):