Merge 207843f6cf into 3e5df07b34

Tests/test_file_vtf.py

@@ -0,0 +1,167 @@
+from __future__ import annotations
+
+import struct
+from io import BytesIO
+from pathlib import Path
+
+import pytest
+
+from PIL import Image
+from PIL.VtfImagePlugin import (
+    VTFException,
+    VtfImageFile,
+    VtfPF,
+    _closest_power,
+    _get_mipmap_count,
+    _get_texture_size,
+)
+
+from .helper import assert_image_equal, assert_image_similar
+
+
+@pytest.mark.parametrize(
+    "size, expected_size",
+    [
+        (8, 8),
+        (7, 8),
+        (9, 8),
+        (192, 256),
+        (1, 1),
+        (2000, 2048),
+    ],
+)
+def test_closest_power(size: int, expected_size: int) -> None:
+    assert _closest_power(size) == expected_size
+
+
+@pytest.mark.parametrize(
+    "size, expected_count",
+    [
+        ((1, 1), 1),
+        ((2, 2), 2),
+        ((4, 4), 3),
+        ((8, 8), 4),
+        ((128, 128), 8),
+        ((256, 256), 9),
+        ((512, 512), 10),
+        ((1024, 1024), 11),
+        ((1024, 1), 11),
+    ],
+)
+def test_get_mipmap_count(size: tuple[int, int], expected_count: int) -> None:
+    assert _get_mipmap_count(size) == expected_count
+
+
+@pytest.mark.parametrize(
+    "pixel_format, size, expected_size",
+    [
+        (VtfPF.DXT1, (16, 16), (16 * 16) // 2),
+        (VtfPF.DXT1_ONEBITALPHA, (16, 16), (16 * 16) // 2),
+        (VtfPF.DXT3, (16, 16), 16 * 16),
+        (VtfPF.DXT5, (16, 16), 16 * 16),
+        (VtfPF.BGR888, (16, 16), 16 * 16 * 3),
+        (VtfPF.RGB888, (16, 16), 16 * 16 * 3),
+        (VtfPF.RGBA8888, (16, 16), 16 * 16 * 4),
+        (VtfPF.UV88, (16, 16), 16 * 16 * 2),
+        (VtfPF.A8, (16, 16), 16 * 16),
+        (VtfPF.I8, (16, 16), 16 * 16),
+        (VtfPF.IA88, (16, 16), 16 * 16 * 2),
+    ],
+)
+def test_get_texture_size(
+    pixel_format: VtfPF, size: tuple[int, int], expected_size: int
+) -> None:
+    assert _get_texture_size(pixel_format, size) == expected_size
+
+
+@pytest.mark.parametrize(
+    "file_path, expected_mode, epsilon",
+    [
+        ("Tests/images/vtf_i8.vtf", "L", 0),
+        ("Tests/images/vtf_a8.vtf", "RGBA", 0),
+        ("Tests/images/vtf_ia88.vtf", "LA", 0),
+        ("Tests/images/vtf_uv88.vtf", "RGB", 0),
+        ("Tests/images/vtf_rgb888.vtf", "RGB", 0),
+        ("Tests/images/vtf_bgr888.vtf", "RGB", 0),
+        ("Tests/images/vtf_dxt1.vtf", "RGBA", 3),
+        ("Tests/images/vtf_dxt1A.vtf", "RGBA", 8),
+        ("Tests/images/vtf_rgba8888.vtf", "RGBA", 0),
+    ],
+)
+def test_vtf_read(file_path: str, expected_mode: str, epsilon: int) -> None:
+    with Image.open(file_path) as im:
+        assert im.mode == expected_mode
+        with Image.open(file_path.replace(".vtf", ".png")) as expected:
+            if epsilon:
+                assert_image_similar(im, expected, epsilon)
+            else:
+                assert_image_equal(im, expected)
+
+
+def test_invalid_file() -> None:
+    invalid_file = "Tests/images/flower.jpg"
+
+    with pytest.raises(SyntaxError, match="not a VTF file"):
+        VtfImageFile(invalid_file)
+
+
+def test_vtf_read_unsupported_version() -> None:
+    b = BytesIO(b"VTF\x00" + struct.pack("<2I", 7, 5))
+    with pytest.raises(VTFException, match=r"Unsupported VTF version: \(7, 5\)"):
+        with Image.open(b):
+            pass
+
+
+def test_vtf_read_unsupported_pixel_format() -> None:
+    b = BytesIO(b"VTF\x00" + struct.pack("<2I40xI7x", 7, 2, 1))
+    with pytest.raises(VTFException, match="Unsupported VTF pixel format: 1"):
+        with Image.open(b):
+            pass
+
+
+@pytest.mark.parametrize(
+    "pixel_format, file_path, expected_mode, epsilon",
+    [
+        (VtfPF.I8, "Tests/images/vtf_i8.png", "L", 0),
+        (VtfPF.A8, "Tests/images/vtf_a8.png", "RGBA", 0),
+        (VtfPF.IA88, "Tests/images/vtf_ia88.png", "LA", 0),
+        (VtfPF.UV88, "Tests/images/vtf_uv88.png", "RGB", 0),
+        (VtfPF.RGB888, "Tests/images/vtf_rgb888.png", "RGB", 0),
+        (VtfPF.BGR888, "Tests/images/vtf_bgr888.png", "RGB", 0),
+        (VtfPF.DXT1, "Tests/images/vtf_dxt1.png", "RGBA", 3),
+        (VtfPF.RGBA8888, "Tests/images/vtf_rgba8888.png", "RGBA", 0),
+    ],
+)
+@pytest.mark.parametrize("version", ((7, 1), (7, 2), (7, 3)))
+def test_vtf_save(
+    pixel_format: VtfPF,
+    file_path: str,
+    expected_mode: str,
+    epsilon: int,
+    version: tuple[int, int],
+    tmp_path: Path,
+) -> None:
+    im: Image.Image
+    with Image.open(file_path) as im:
+        out = tmp_path / "tmp.vtf"
+        im.save(out, pixel_format=pixel_format, version=version)
+        with Image.open(out) as reloaded:
+            assert reloaded.mode == expected_mode
+            if epsilon:
+                assert_image_similar(im, reloaded, epsilon)
+            else:
+                assert_image_equal(im, reloaded)
+
+
+def test_vtf_save_unsupported_mode(tmp_path: Path) -> None:
+    out = tmp_path / "temp.vtf"
+    im = Image.new("HSV", (1, 1))
+    with pytest.raises(OSError, match="cannot write mode HSV as VTF"):
+        im.save(out)
+
+
+def test_vtf_save_unsupported_version(tmp_path: Path) -> None:
+    out = tmp_path / "temp.vtf"
+    im = Image.new("L", (1, 1))
+    with pytest.raises(VTFException, match=r"Unsupported VTF version: \(7, 5\)"):
+        im.save(out, version=(7, 5))

src/PIL/VtfImagePlugin.py

@@ -0,0 +1,367 @@
+"""
+A Pillow loader for .vtf files (aka Valve Texture Format)
+REDxEYE <med45c@gmail.com>
+
+Documentation:
+  https://developer.valvesoftware.com/wiki/Valve_Texture_Format
+
+The contents of this file are hereby released in the public domain (CC0)
+Full text of the CC0 license:
+  https://creativecommons.org/publicdomain/zero/1.0/
+"""
+
+from __future__ import annotations
+
+import struct
+from enum import IntEnum, IntFlag
+from math import ceil, log
+from typing import IO, NamedTuple, cast
+
+from . import Image, ImageFile
+
+
+class VTFException(Exception):
+    pass
+
+
+class CompiledVtfFlags(IntFlag):
+    # Flags from the *.txt config file
+    POINTSAMPLE = 0x00000001
+    TRILINEAR = 0x00000002
+    CLAMPS = 0x00000004
+    CLAMPT = 0x00000008
+    ANISOTROPIC = 0x00000010
+    HINT_DXT5 = 0x00000020
+    PWL_CORRECTED = 0x00000040
+    NORMAL = 0x00000080
+    NOMIP = 0x00000100
+    NOLOD = 0x00000200
+    ALL_MIPS = 0x00000400
+    PROCEDURAL = 0x00000800
+
+    # These are automatically generated by vtex from the texture data.
+    ONEBITALPHA = 0x00001000
+    EIGHTBITALPHA = 0x00002000
+
+    # Newer flags from the *.txt config file
+    ENVMAP = 0x00004000
+    RENDERTARGET = 0x00008000
+    DEPTHRENDERTARGET = 0x00010000
+    NODEBUGOVERRIDE = 0x00020000
+    SINGLECOPY = 0x00040000
+    PRE_SRGB = 0x00080000
+
+    UNUSED_00100000 = 0x00100000
+    UNUSED_00200000 = 0x00200000
+    UNUSED_00400000 = 0x00400000
+
+    NODEPTHBUFFER = 0x00800000
+
+    UNUSED_01000000 = 0x01000000
+
+    CLAMPU = 0x02000000
+    VERTEXTEXTURE = 0x04000000
+    SSBUMP = 0x08000000
+
+    UNUSED_10000000 = 0x10000000
+
+    BORDER = 0x20000000
+
+    UNUSED_40000000 = 0x40000000
+    UNUSED_80000000 = 0x80000000
+
+
+class VtfPF(IntEnum):
+    # NONE = -1
+    RGBA8888 = 0
+    # ABGR8888 = 1
+    RGB888 = 2
+    BGR888 = 3
+    # RGB565 = 4
+    I8 = 5
+    IA88 = 6
+    # P8 = 7
+    A8 = 8
+    # RGB888_BLUESCREEN = 9
+    # BGR888_BLUESCREEN = 10
+    # ARGB8888 = 11
+    BGRA8888 = 12
+    DXT1 = 13
+    DXT3 = 14
+    DXT5 = 15
+    # BGRX8888 = 16
+    # BGR565 = 17
+    # BGRX5551 = 18
+    # BGRA4444 = 19
+    DXT1_ONEBITALPHA = 20
+    # BGRA5551 = 21
+    UV88 = 22
+    # UVWQ8888 = 23
+    # RGBA16161616F = 24
+    # RGBA16161616 = 25
+    # UVLX8888 = 26
+
+
+class VTFHeader(NamedTuple):
+    header_size: int
+    width: int
+    height: int
+    flags: int
+    frames: int
+    first_frames: int
+    reflectivity_r: float
+    reflectivity_g: float
+    reflectivity_b: float
+    bumpmap_scale: float
+    pixel_format: int
+    mipmap_count: int
+    low_pixel_format: int
+    low_width: int
+    low_height: int
+    depth: int = 0
+    resource_count: int = 0
+
+
+BLOCK_COMPRESSED = (VtfPF.DXT1, VtfPF.DXT1_ONEBITALPHA, VtfPF.DXT3, VtfPF.DXT5)
+HEADER_V70 = "2HI2H4x3f4xfIbI2b"
+
+
+def _get_texture_size(pixel_format: int, size: tuple[int, int]) -> int:
+    for factor, pixel_formats in (
+        (0.5, (VtfPF.DXT1, VtfPF.DXT1_ONEBITALPHA)),
+        (1, (VtfPF.DXT3, VtfPF.DXT5, VtfPF.A8, VtfPF.I8)),
+        (2, (VtfPF.UV88, VtfPF.IA88)),
+        (3, (VtfPF.RGB888, VtfPF.BGR888)),
+        (4, (VtfPF.RGBA8888,)),
+    ):
+        if pixel_format in pixel_formats:
+            return int(size[0] * size[1] * factor)
+    msg = f"Unsupported VTF pixel format: {pixel_format}"
+    raise VTFException(msg)
+
+
+def _get_mipmap_count(size: tuple[int, int]) -> int:
+    mip_count = 1
+    while True:
+        mip_width = size[0] >> mip_count
+        mip_height = size[1] >> mip_count
+        if mip_width == 0 and mip_height == 0:
+            return mip_count
+        mip_count += 1
+
+
+def _write_image(fp: IO[bytes], im: Image.Image, pixel_format: VtfPF) -> None:
+    args: dict[VtfPF, tuple[int,] | str] = {
+        VtfPF.DXT1: (1,),
+        VtfPF.DXT3: (3,),
+        VtfPF.DXT5: (5,),
+        VtfPF.RGB888: "RGB",
+        VtfPF.BGR888: "BGR",
+        VtfPF.RGBA8888: "RGBA",
+        VtfPF.A8: "A",
+        VtfPF.I8: "L",
+        VtfPF.IA88: "LA",
+        VtfPF.UV88: "RG",
+    }
+    try:
+        encoder_args = args[pixel_format]
+    except KeyError:
+        msg = f"Unsupported pixel format: {pixel_format!r}"
+        raise VTFException(msg)
+
+    codec_name = "bcn" if pixel_format in BLOCK_COMPRESSED else "raw"
+    tile = [ImageFile._Tile(codec_name, (0, 0) + im.size, fp.tell(), encoder_args)]
+    ImageFile._save(im, fp, tile, _get_texture_size(pixel_format, im.size))
+
+
+def _closest_power(x: int) -> int:
+    possible_results = round(log(x, 2)), ceil(log(x, 2))
+    return 2 ** min(possible_results, key=lambda z: abs(x - 2**z))
+
+
+class VtfImageFile(ImageFile.ImageFile):
+    format = "VTF"
+    format_description = "Valve Texture Format"
+
+    def _open(self):
+        if not _accept(self.fp.read(12)):
+            msg = "not a VTF file"
+            raise SyntaxError(msg)
+
+        self.fp.seek(4)
+        version = struct.unpack("<2I", self.fp.read(8))
+        if version > (7, 4):
+            msg = f"Unsupported VTF version: {version}"
+            raise VTFException(msg)
+
+        header = VTFHeader(
+            *struct.unpack(
+                "<I" + HEADER_V70, self.fp.read(struct.calcsize("<I" + HEADER_V70))
+            )
+        )
+
+        pixel_format = header.pixel_format
+        if pixel_format in (
+            VtfPF.DXT1_ONEBITALPHA,
+            VtfPF.DXT1,
+            VtfPF.DXT3,
+            VtfPF.DXT5,
+            VtfPF.RGBA8888,
+            VtfPF.BGRA8888,
+            VtfPF.A8,
+        ):
+            self._mode = "RGBA"
+        elif pixel_format in (VtfPF.RGB888, VtfPF.BGR888, VtfPF.UV88):
+            self._mode = "RGB"
+        elif pixel_format == VtfPF.I8:
+            self._mode = "L"
+        elif pixel_format == VtfPF.IA88:
+            self._mode = "LA"
+        else:
+            msg = f"Unsupported VTF pixel format: {pixel_format}"
+            raise VTFException(msg)
+
+        if pixel_format in BLOCK_COMPRESSED:
+            codec_name = "bcn"
+            if pixel_format in (VtfPF.DXT1, VtfPF.DXT1_ONEBITALPHA):
+                args = (1, "DXT1")
+            elif pixel_format == VtfPF.DXT3:
+                args = (2, "DXT3")
+            else:
+                args = (3, "DXT5")
+        else:
+            codec_name = "raw"
+            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)
+
+        data_start = header.header_size
+        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
+        for mip_id in range(header.mipmap_count - 1, 0, -1):
+            mip_width = max(header.width >> mip_id, min_res)
+            mip_height = max(header.height >> mip_id, min_res)
+
+            data_start += _get_texture_size(pixel_format, (mip_width, mip_height))
+
+        self.tile = [ImageFile._Tile(codec_name, (0, 0) + self.size, data_start, args)]
+
+
+def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None:
+    if im.mode not in ("RGB", "RGBA", "L", "LA"):
+        msg = f"cannot write mode {im.mode} as VTF"
+        raise OSError(msg)
+
+    encoderinfo = im.encoderinfo
+    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)
+
+    flags = CompiledVtfFlags(0)
+
+    if (
+        pixel_format in (VtfPF.DXT1, VtfPF.DXT3, VtfPF.DXT5)
+        and im.mode in ("RGBA", "LA")
+    ) or pixel_format in (VtfPF.RGBA8888, VtfPF.BGRA8888, VtfPF.A8, VtfPF.IA88):
+        flags |= CompiledVtfFlags.EIGHTBITALPHA
+    im = im.resize((_closest_power(im.width), _closest_power(im.height)))
+
+    mipmap_count = _get_mipmap_count(im.size) if generate_mips else 0
+
+    thumb = im.convert("RGBA")
+    thumb.thumbnail((16, 16))
+    thumb = thumb.resize((_closest_power(thumb.width), _closest_power(thumb.height)))
+
+    header = VTFHeader(
+        0,
+        im.width,
+        im.height,
+        flags,
+        1,
+        0,
+        1,
+        1,
+        1,
+        1,
+        pixel_format,
+        mipmap_count,
+        VtfPF.DXT1,
+        thumb.width,
+        thumb.height,
+        1,
+        2,
+    )
+
+    header_bytes = struct.pack("<" + HEADER_V70, *header[1:15])
+    if version >= (7, 2):
+        header_bytes += struct.pack("<H", header.depth)
+    if version >= (7, 3):
+        header_bytes += struct.pack("<3xI8x", header.resource_count)
+    # Align header size to 16 bytes
+    if len(header_bytes) % 16:
+        header_bytes += b"\x00" * (16 - len(header_bytes) % 16)
+    header_length = 16 + len(header_bytes)
+    if version >= (7, 3):
+        header_length += 16  # Resource entries
+    fp.write(b"VTF\x00" + struct.pack("<2II", *version, header_length) + header_bytes)
+
+    if version > (7, 2):
+        # Resource entries
+        for tag, offset in {
+            b"\x01\x00\x00": header_length,  # Low-res
+            b"\x30\x00\x00": header_length
+            + _get_texture_size(VtfPF.DXT1, thumb.size),  # High-res
+        }.items():
+            fp.write(tag + b"\x00")  # Tag, flags
+            fp.write(struct.pack("<I", offset))
+
+    _write_image(fp, thumb, VtfPF.DXT1)
+
+    min_size = 4 if pixel_format in BLOCK_COMPRESSED else 1
+
+    for mip_id in range(mipmap_count - 1, 0, -1):
+        size = cast(
+            tuple[int, int],
+            tuple(max(min_size, dimension >> mip_id) for dimension in im.size),
+        )
+        mip = im.resize(size)
+        _write_image(fp, mip, pixel_format)
+    _write_image(fp, im, pixel_format)
+
+
+def _accept(prefix: bytes) -> bool:
+    if not prefix.startswith(b"VTF\x00"):
+        return False
+    return struct.unpack_from("<2I", prefix, 4) >= (7, 0)
+
+
+Image.register_open(VtfImageFile.format, VtfImageFile, _accept)
+Image.register_save(VtfImageFile.format, _save)
+Image.register_extension(VtfImageFile.format, ".vtf")

src/PIL/__init__.py

@@ -72,6 +72,7 @@ _plugins = [
     "XbmImagePlugin",
     "XpmImagePlugin",
     "XVThumbImagePlugin",
+    "VtfImagePlugin",
 ]
 
 

src/libImaging/Pack.c

@@ -242,6 +242,17 @@ packLA(UINT8 *out, const UINT8 *in, int pixels) {
         in += 4;
     }
 }
+static void
+packRG(UINT8 *out, const UINT8 *in, int pixels) {
+    int i;
+    /* RG, pixel interleaved */
+    for (i = 0; i < pixels; i++) {
+        out[0] = in[R];
+        out[1] = in[G];
+        out += 2;
+        in += 4;
+    }
+}
 
 static void
 packLAL(UINT8 *out, const UINT8 *in, int pixels) {
@@ -474,7 +485,7 @@ copy2(UINT8 *out, const UINT8 *in, int pixels) {
 static void
 copy4(UINT8 *out, const UINT8 *in, int pixels) {
     /* RGBA, CMYK quadruples */
-    memcpy(out, in, 4 * pixels);
+    memcpy(out, in, pixels * 4);
 }
 
 static void
@@ -563,6 +574,7 @@ static struct {
     {"RGB", "BGRX", 32, ImagingPackBGRX},
     {"RGB", "XBGR", 32, ImagingPackXBGR},
     {"RGB", "RGB;L", 24, packRGBL},
+    {"RGB", "RG", 16, packRG},
     {"RGB", "R", 8, band0},
     {"RGB", "G", 8, band1},
     {"RGB", "B", 8, band2},
@@ -575,6 +587,7 @@ static struct {
     {"RGBA", "BGRA", 32, ImagingPackBGRA},
     {"RGBA", "ABGR", 32, ImagingPackABGR},
     {"RGBA", "BGRa", 32, ImagingPackBGRa},
+    {"RGBA", "RG", 16, packRG},
     {"RGBA", "R", 8, band0},
     {"RGBA", "G", 8, band1},
     {"RGBA", "B", 8, band2},

src/libImaging/Unpack.c

@@ -426,6 +426,18 @@ unpackLA(UINT8 *_out, const UINT8 *in, int pixels) {
     }
 }
 
+static void
+unpackRG(UINT8 *_out, const UINT8 *in, int pixels) {
+    int i;
+    /* RG, pixel interleaved */
+    for (i = 0; i < pixels; i++) {
+        _out[R] = in[0];
+        _out[G] = in[1];
+        in += 2;
+        _out += 4;
+    }
+}
+
 static void
 unpackLAL(UINT8 *_out, const UINT8 *in, int pixels) {
     int i;
@@ -1606,6 +1618,9 @@ static struct {
     {"PA", "PA;L", 16, unpackLAL},
     {"PA", "LA", 16, unpackLA},
 
+    /* 2 channel to RGB/RGBA */
+    {"RGB", "RG", 16, unpackRG},
+    {"RGBA", "RG", 16, unpackRG},
     /* true colour */
     {"RGB", "RGB", 24, ImagingUnpackRGB},
     {"RGB", "RGB;L", 24, unpackRGBL},