Default to pixel format specific to mode

This commit is contained in:
Andrew Murray 2025-05-06 22:17:15 +10:00
parent 487c7b2be1
commit b40aa6888c
2 changed files with 107 additions and 123 deletions

View File

@ -1,5 +1,7 @@
from __future__ import annotations from __future__ import annotations
from pathlib import Path
import pytest import pytest
from PIL import Image from PIL import Image
@ -71,54 +73,59 @@ def test_get_texture_size(
@pytest.mark.parametrize( @pytest.mark.parametrize(
"etalon_path, file_path, expected_mode, epsilon", "etalon_path, file_path, expected_mode, epsilon",
[ [
("Tests/images/vtf_i8.png", "Tests/images/vtf_i8.vtf", "L", 0.0), ("Tests/images/vtf_i8.png", "Tests/images/vtf_i8.vtf", "L", 0),
("Tests/images/vtf_a8.png", "Tests/images/vtf_a8.vtf", "RGBA", 0.0), ("Tests/images/vtf_a8.png", "Tests/images/vtf_a8.vtf", "RGBA", 0),
("Tests/images/vtf_ia88.png", "Tests/images/vtf_ia88.vtf", "LA", 0.0), ("Tests/images/vtf_ia88.png", "Tests/images/vtf_ia88.vtf", "LA", 0),
("Tests/images/vtf_uv88.png", "Tests/images/vtf_uv88.vtf", "RGB", 0.0), ("Tests/images/vtf_uv88.png", "Tests/images/vtf_uv88.vtf", "RGB", 0),
("Tests/images/vtf_rgb888.png", "Tests/images/vtf_rgb888.vtf", "RGB", 0.0), ("Tests/images/vtf_rgb888.png", "Tests/images/vtf_rgb888.vtf", "RGB", 0),
("Tests/images/vtf_bgr888.png", "Tests/images/vtf_bgr888.vtf", "RGB", 0.0), ("Tests/images/vtf_bgr888.png", "Tests/images/vtf_bgr888.vtf", "RGB", 0),
("Tests/images/vtf_dxt1.png", "Tests/images/vtf_dxt1.vtf", "RGBA", 3.0), ("Tests/images/vtf_dxt1.png", "Tests/images/vtf_dxt1.vtf", "RGBA", 3),
("Tests/images/vtf_dxt1A.png", "Tests/images/vtf_dxt1A.vtf", "RGBA", 8.0), ("Tests/images/vtf_dxt1A.png", "Tests/images/vtf_dxt1A.vtf", "RGBA", 8),
("Tests/images/vtf_rgba8888.png", "Tests/images/vtf_rgba8888.vtf", "RGBA", 0), ("Tests/images/vtf_rgba8888.png", "Tests/images/vtf_rgba8888.vtf", "RGBA", 0),
], ],
) )
def test_vtf_read( def test_vtf_read(
etalon_path: str, file_path: str, expected_mode: str, epsilon: float etalon_path: str, file_path: str, expected_mode: str, epsilon: int
) -> None: ) -> None:
with Image.open(file_path) as f: with Image.open(file_path) as f:
assert f.mode == expected_mode assert f.mode == expected_mode
with Image.open(etalon_path) as e: with Image.open(etalon_path) as e:
converted_e = e.convert(expected_mode) converted_e = e.convert(expected_mode)
if epsilon == 0: if epsilon:
assert_image_equal(converted_e, f)
else:
assert_image_similar(converted_e, f, epsilon) assert_image_similar(converted_e, f, epsilon)
else:
assert_image_equal(converted_e, f)
@pytest.mark.parametrize( @pytest.mark.parametrize(
"pixel_format, file_path, expected_mode, epsilon", "pixel_format, file_path, expected_mode, epsilon",
[ [
(VtfPF.I8, "Tests/images/vtf_i8.png", "L", 0.0), (VtfPF.I8, "Tests/images/vtf_i8.png", "L", 0),
(VtfPF.A8, "Tests/images/vtf_a8.png", "RGBA", 0.0), (VtfPF.A8, "Tests/images/vtf_a8.png", "RGBA", 0),
(VtfPF.IA88, "Tests/images/vtf_ia88.png", "LA", 0.0), (VtfPF.IA88, "Tests/images/vtf_ia88.png", "LA", 0),
(VtfPF.UV88, "Tests/images/vtf_uv88.png", "RGB", 0.0), (VtfPF.UV88, "Tests/images/vtf_uv88.png", "RGB", 0),
(VtfPF.RGB888, "Tests/images/vtf_rgb888.png", "RGB", 0.0), (VtfPF.RGB888, "Tests/images/vtf_rgb888.png", "RGB", 0),
(VtfPF.BGR888, "Tests/images/vtf_bgr888.png", "RGB", 0.0), (VtfPF.BGR888, "Tests/images/vtf_bgr888.png", "RGB", 0),
(VtfPF.DXT1, "Tests/images/vtf_dxt1.png", "RGBA", 3.0), (VtfPF.DXT1, "Tests/images/vtf_dxt1.png", "RGBA", 3),
(VtfPF.RGBA8888, "Tests/images/vtf_rgba8888.png", "RGBA", 0), (VtfPF.RGBA8888, "Tests/images/vtf_rgba8888.png", "RGBA", 0),
], ],
) )
def test_vtf_save( def test_vtf_save(
pixel_format: VtfPF, file_path: str, expected_mode: str, epsilon: float, tmp_path pixel_format: VtfPF,
file_path: str,
expected_mode: str,
epsilon: int,
tmp_path: Path,
) -> None: ) -> None:
f: Image.Image = Image.open(file_path) im: Image.Image
out = (tmp_path / "tmp.vtf").as_posix() with Image.open(file_path) as im:
f.save(out, pixel_format=pixel_format) out = tmp_path / "tmp.vtf"
if pixel_format == VtfPF.DXT1: im.save(out, pixel_format=pixel_format)
f = f.convert("RGBA") if pixel_format == VtfPF.DXT1:
e = Image.open(out) im = im.convert("RGBA")
assert e.mode == expected_mode with Image.open(out) as expected:
if epsilon == 0: assert expected.mode == expected_mode
assert_image_equal(e, f) if epsilon:
else: assert_image_similar(im, expected, epsilon)
assert_image_similar(e, f, epsilon) else:
assert_image_equal(im, expected)

View File

@ -119,8 +119,8 @@ class VTFHeader(NamedTuple):
low_pixel_format: int low_pixel_format: int
low_width: int low_width: int
low_height: int low_height: int
depth: int depth: int = 0
resource_count: int resource_count: int = 0
BLOCK_COMPRESSED = (VtfPF.DXT1, VtfPF.DXT1_ONEBITALPHA, VtfPF.DXT3, VtfPF.DXT5) BLOCK_COMPRESSED = (VtfPF.DXT1, VtfPF.DXT1_ONEBITALPHA, VtfPF.DXT3, VtfPF.DXT5)
@ -129,7 +129,7 @@ HEADER_V72 = "<I2HI2H4x3f4xfIbI2bH"
HEADER_V73 = "<I2HI2H4x3f4xfIbI2bH3xI8x" HEADER_V73 = "<I2HI2H4x3f4xfIbI2bH3xI8x"
def _get_texture_size(pixel_format: VtfPF, width: int, height: int) -> int: def _get_texture_size(pixel_format: int, width: int, height: int) -> int:
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 pixel_format in ( elif pixel_format in (
@ -160,28 +160,21 @@ def _get_mipmap_count(width: int, height: int) -> int:
def _write_image(fp: IO[bytes], im: Image.Image, pixel_format: VtfPF) -> None: def _write_image(fp: IO[bytes], im: Image.Image, pixel_format: VtfPF) -> None:
encoder_args: tuple[int,] | tuple[str, int, int] args: dict[VtfPF, tuple[int,] | str] = {
if pixel_format == VtfPF.DXT1: VtfPF.DXT1: (1,),
encoder_args = (1,) VtfPF.DXT3: (3,),
elif pixel_format == VtfPF.DXT3: VtfPF.DXT5: (5,),
encoder_args = (3,) VtfPF.RGB888: "RGB",
elif pixel_format == VtfPF.DXT5: VtfPF.BGR888: "BGR",
encoder_args = (5,) VtfPF.RGBA8888: "RGBA",
elif pixel_format == VtfPF.RGB888: VtfPF.A8: "A",
encoder_args = ("RGB", 0, 0) VtfPF.I8: "L",
elif pixel_format == VtfPF.BGR888: VtfPF.IA88: "LA",
encoder_args = ("BGR", 0, 0) VtfPF.UV88: "RG",
elif pixel_format == VtfPF.RGBA8888: }
encoder_args = ("RGBA", 0, 0) try:
elif pixel_format == VtfPF.A8: encoder_args = args[pixel_format]
encoder_args = ("A", 0, 0) except KeyError:
elif pixel_format == VtfPF.I8:
encoder_args = ("L", 0, 0)
elif pixel_format == VtfPF.IA88:
encoder_args = ("LA", 0, 0)
elif pixel_format == VtfPF.UV88:
encoder_args = ("RG", 0, 0)
else:
msg = f"Unsupported pixel format: {pixel_format!r}" msg = f"Unsupported pixel format: {pixel_format!r}"
raise VTFException(msg) raise VTFException(msg)
@ -206,35 +199,16 @@ class VtfImageFile(ImageFile.ImageFile):
self.fp.seek(4) self.fp.seek(4)
version = struct.unpack("<2I", self.fp.read(8)) version = struct.unpack("<2I", self.fp.read(8))
if version <= (7, 2): if version > (7, 4):
header = VTFHeader(
*struct.unpack(HEADER_V70, self.fp.read(struct.calcsize(HEADER_V70))),
0,
0,
0,
0,
0,
)
self.fp.seek(header.header_size)
elif version < (7, 3):
header = VTFHeader(
*struct.unpack(HEADER_V72, self.fp.read(struct.calcsize(HEADER_V72))),
0,
0,
0,
0,
)
self.fp.seek(header.header_size)
elif version < (7, 5):
header = VTFHeader(
*struct.unpack(HEADER_V73, self.fp.read(struct.calcsize(HEADER_V73)))
)
self.fp.seek(header.header_size)
else:
msg = f"Unsupported VTF version: {version}" msg = f"Unsupported VTF version: {version}"
raise VTFException(msg) raise VTFException(msg)
pixel_format = VtfPF(header.pixel_format) header = VTFHeader(
*struct.unpack(HEADER_V70, self.fp.read(struct.calcsize(HEADER_V70)))
)
self.fp.seek(header.header_size)
pixel_format = header.pixel_format
if pixel_format in ( if pixel_format in (
VtfPF.DXT1_ONEBITALPHA, VtfPF.DXT1_ONEBITALPHA,
VtfPF.DXT1, VtfPF.DXT1,
@ -255,38 +229,37 @@ class VtfImageFile(ImageFile.ImageFile):
msg = f"Unsupported VTF pixel format: {pixel_format}" msg = f"Unsupported VTF pixel format: {pixel_format}"
raise VTFException(msg) raise VTFException(msg)
if pixel_format in (VtfPF.DXT1, VtfPF.DXT1_ONEBITALPHA): if pixel_format in BLOCK_COMPRESSED:
args = (1, "DXT1") codec_name = "bcn"
elif pixel_format == VtfPF.DXT3: if pixel_format in (VtfPF.DXT1, VtfPF.DXT1_ONEBITALPHA):
args = (2, "DXT3") args = (1, "DXT1")
elif pixel_format == VtfPF.DXT5: elif pixel_format == VtfPF.DXT3:
args = (3, "DXT5") args = (2, "DXT3")
elif pixel_format == VtfPF.RGBA8888: else:
args = ("RGBA", 0, 1) args = (3, "DXT5")
elif pixel_format == VtfPF.RGB888:
args = ("RGB", 0, 1)
elif pixel_format == VtfPF.BGR888:
args = ("BGR", 0, 1)
elif pixel_format == VtfPF.BGRA8888:
args = ("BGRA", 0, 1)
elif pixel_format == VtfPF.UV88:
args = ("RG", 0, 1)
elif pixel_format == VtfPF.I8:
args = ("L", 0, 1)
elif pixel_format == VtfPF.A8:
args = ("A", 0, 1)
elif pixel_format == VtfPF.IA88:
args = ("LA", 0, 1)
else: else:
msg = f"Unsupported VTF pixel format: {pixel_format}" codec_name = "raw"
raise VTFException(msg) try:
args = {
VtfPF.RGBA8888: "RGBA",
VtfPF.RGB888: "RGB",
VtfPF.BGR888: "BGR",
VtfPF.BGRA8888: "BGRA",
VtfPF.UV88: "RG",
VtfPF.I8: "L",
VtfPF.A8: "A",
VtfPF.IA88: "LA",
}[pixel_format]
except KeyError:
msg = f"Unsupported VTF pixel format: {pixel_format}"
raise VTFException(msg)
self._size = (header.width, header.height) self._size = (header.width, header.height)
low_format = VtfPF(header.low_pixel_format)
data_start = self.fp.tell() data_start = self.fp.tell()
data_start += _get_texture_size(low_format, header.low_width, header.low_height) data_start += _get_texture_size(
header.low_pixel_format, header.low_width, header.low_height
)
min_res = 4 if pixel_format in BLOCK_COMPRESSED else 1 min_res = 4 if pixel_format in BLOCK_COMPRESSED else 1
for mip_id in range(header.mipmap_count - 1, 0, -1): for mip_id in range(header.mipmap_count - 1, 0, -1):
mip_width = max(header.width >> mip_id, min_res) mip_width = max(header.width >> mip_id, min_res)
@ -294,7 +267,6 @@ class VtfImageFile(ImageFile.ImageFile):
data_start += _get_texture_size(pixel_format, mip_width, mip_height) data_start += _get_texture_size(pixel_format, mip_width, mip_height)
codec_name = "bcn" if pixel_format in BLOCK_COMPRESSED else "raw"
self.tile = [ImageFile._Tile(codec_name, (0, 0) + self.size, data_start, args)] self.tile = [ImageFile._Tile(codec_name, (0, 0) + self.size, data_start, args)]
@ -302,9 +274,22 @@ def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None:
if im.mode not in ("RGB", "RGBA", "L", "LA"): if im.mode not in ("RGB", "RGBA", "L", "LA"):
msg = f"cannot write mode {im.mode} as VTF" msg = f"cannot write mode {im.mode} as VTF"
raise OSError(msg) raise OSError(msg)
encoderinfo = im.encoderinfo encoderinfo = im.encoderinfo
pixel_format = VtfPF(encoderinfo.get("pixel_format", VtfPF.RGBA8888))
version = encoderinfo.get("version", (7, 4)) version = encoderinfo.get("version", (7, 4))
if version > (7, 4):
msg = f"Unsupported VTF version: {version}"
raise VTFException(msg)
pixel_format = (
encoderinfo.get("pixel_format")
or {
"RGB": VtfPF.RGB888,
"RGBA": VtfPF.RGBA8888,
"L": VtfPF.I8,
"LA": VtfPF.IA88,
}[im.mode]
)
generate_mips = encoderinfo.get("generate_mips", True) generate_mips = encoderinfo.get("generate_mips", True)
flags = CompiledVtfFlags(0) flags = CompiledVtfFlags(0)
@ -321,13 +306,11 @@ def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None:
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 = 0 mipmap_count = _get_mipmap_count(width, height) if generate_mips else 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)))
_write_image(thumb_buffer, thumb, VtfPF.DXT1) _write_image(thumb_buffer, thumb, VtfPF.DXT1)
@ -360,13 +343,10 @@ def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None:
size = struct.calcsize(HEADER_V72) + 12 size = struct.calcsize(HEADER_V72) + 12
header = header._replace(header_size=size + (16 - size % 16)) header = header._replace(header_size=size + (16 - size % 16))
fp.write(struct.pack(HEADER_V72, *header[:16])) fp.write(struct.pack(HEADER_V72, *header[:16]))
elif version > (7, 2): else:
size = struct.calcsize(HEADER_V73) + 12 size = struct.calcsize(HEADER_V73) + 12
header = header._replace(header_size=size + (16 - size % 16)) header = header._replace(header_size=size + (16 - size % 16))
fp.write(struct.pack(HEADER_V73, *header)) fp.write(struct.pack(HEADER_V73, *header))
else:
msg = f"Unsupported version {version}"
raise VTFException(msg)
if version > (7, 2): if version > (7, 2):
fp.write(b"\x01\x00\x00\x00") fp.write(b"\x01\x00\x00\x00")
@ -377,10 +357,7 @@ def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None:
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 if pixel_format in BLOCK_COMPRESSED else 1
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(min_size, width >> mip_id) mip_width = max(min_size, width >> mip_id)