Add Bcn BC1 encoder

src/PIL/VtfImagePlugin.py

@@ -12,6 +12,8 @@ Full text of the CC0 license:
 
 import struct
 from enum import IntEnum, IntFlag
+from io import BytesIO
+from math import log, ceil
 from typing import NamedTuple
 
 from . import Image, ImageFile
@@ -108,11 +110,9 @@ VTFHeader = NamedTuple(
         ("flags", int),
         ("frames", int),
         ("first_frames", int),
-        ("padding0", int),
         ("reflectivity_r", float),
         ("reflectivity_g", float),
         ("reflectivity_b", float),
-        ("padding1", int),
         ("bumpmap_scale", float),
         ("pixel_format", int),
         ("mipmap_count", int),
@@ -122,10 +122,7 @@ VTFHeader = NamedTuple(
         # V 7.2+
         ("depth", int),
         # V 7.3+
-        ("padding2", int),
-        ("padding2_", int),
         ("resource_count", int),
-        ("padding3", int),
     ],
 )
 RGB_FORMATS = (VtfPF.RGB888,)
@@ -146,6 +143,9 @@ LA_FORMATS = (
 )
 
 SUPPORTED_FORMATS = RGBA_FORMATS + RGB_FORMATS + LA_FORMATS + L_FORMATS
+HEADER_V70_STRUCT = '<I2HI2H4x3f4xfIbI2b'
+HEADER_V72_STRUCT = '<I2HI2H4x3f4xfIbI2bH'
+HEADER_V73_STRUCT = '<I2HI2H4x3f4xfIbI2bH3xI8x'
 
 
 # fmt: off
@@ -163,21 +163,22 @@ def _get_texture_size(pixel_format: VtfPF, width, height):
     raise VTFException(f"Unsupported VTF pixel format: {pixel_format}")
 
 
-def _get_mipmap_count(pixel_format: VtfPF, width: int, height: int):
+def _get_mipmap_count(width: int, height: int):
-    if pixel_format in (VtfPF.DXT1, VtfPF.DXT1_ONEBITALPHA, VtfPF.DXT3, VtfPF.DXT5):
+    mip_count = 1
-        min_size = 4
+    while True:
-    else:
+        mip_width = width >> mip_count
-        min_size = 1
+        mip_height = height >> mip_count
-    mip_count = 0
+        if mip_width < 1 or mip_height < 1:
-    for i in range(8):
-        mip_width = width << i
-        mip_height = height << i
-        if mip_width < min_size or mip_height < min_size:
             break
         mip_count += 1
     return mip_count
 
 
+def closest_power(x):
+    possible_results = round(log(x, 2)), ceil(log(x, 2))
+    return 2 ** min(possible_results, key=lambda z: abs(x - 2 ** z))
+
+
 # fmt: on
 
 
@@ -192,12 +193,16 @@ class VtfImageFile(ImageFile.ImageFile):
         self.fp.seek(4)
         version = struct.unpack("<2I", self.fp.read(8))
         if version <= (7, 2):
-            header = VTFHeader(*struct.unpack("<I2HI2HI3fIfIbI2b", self.fp.read(51)), 0, 0, 0, 0, 0)
+            header = VTFHeader(*struct.unpack(HEADER_V70_STRUCT, self.fp.read(struct.calcsize(HEADER_V70_STRUCT))), 0,
+                               0, 0, 0, 0)
+            self.fp.seek(header.header_size)
         elif (7, 2) <= version < (7, 3):
-            header = VTFHeader(*struct.unpack("<I2HI2HI3fIfIbI2bH", self.fp.read(53)), 0, 0, 0, 0)
+            header = VTFHeader(*struct.unpack(HEADER_V72_STRUCT, self.fp.read(struct.calcsize(HEADER_V72_STRUCT))), 0,
+                               0, 0, 0)
+            self.fp.seek(header.header_size)
         elif (7, 3) <= version < (7, 5):
-            header = VTFHeader(*struct.unpack("<I2HI2HI3fIfIbI2bHHBIQ", self.fp.read(68)))
+            header = VTFHeader(*struct.unpack(HEADER_V73_STRUCT, self.fp.read(struct.calcsize(HEADER_V73_STRUCT))))
-            self.fp.seek(header.resource_count * 8, 1)
+            self.fp.seek(header.header_size)
         else:
             raise VTFException(f"Unsupported VTF version: {version}")
         flags = CompiledVtfFlags(header.flags)
@@ -259,29 +264,66 @@ def _save(im, fp, filename):
         raise OSError(f"cannot write mode {im.mode} as VTF")
     arguments = im.encoderinfo
     pixel_format = VtfPF(arguments.get('pixel_format', VtfPF.RGBA8888))
+    version = arguments.get('version', (7, 4))
     flags = CompiledVtfFlags(0)
     if 'A' in im.mode:
         if pixel_format == VtfPF.DXT1_ONEBITALPHA:
             flags |= CompiledVtfFlags.ONEBITALPHA
+        elif pixel_format == VtfPF.DXT1:
+            im = im.convert('RGB')
         else:
             flags |= CompiledVtfFlags.EIGHTBITALPHA
 
+    im = im.resize((closest_power(im.width), closest_power(im.height)))
     width, height = im.size
 
-    lowres_width = min(16, width)
+    mipmap_count = _get_mipmap_count(width, height)
-    lowres_height = min(16, height)
+
-    mipmap_count = _get_mipmap_count(pixel_format, width, height)
+    thumb_buffer = BytesIO()
-    thumb = im.convert('RGB')
+    thumb = im.convert('RGB')
-    thumb.thumbnail((lowres_width,lowres_height))
+    thumb.thumbnail(((min(16, width)), (min(16, height))))
-    ImageFile._save(thumb, fp, [("bcn", (0, 0, lowres_width, lowres_height), 96, (1, 'DXT1'))])
+    thumb = thumb.resize((closest_power(thumb.width), closest_power(thumb.height)))
+    ImageFile._save(thumb, thumb_buffer, [("bcn", (0, 0) + thumb.size, 0, (1, 'DXT1'))])
+
+    header = VTFHeader(0, width, height, flags, 1, 0, 1.0, 1.0, 1.0, 1.0, pixel_format, mipmap_count,
+                       VtfPF.DXT1, thumb.width, thumb.height, 1, 2)
 
-    header = VTFHeader(96, width, height, flags, 1, 1, 0, 1.0, 1.0, 1.0, 0, 1.0, pixel_format, mipmap_count,
-                       VtfPF.DXT1, lowres_width, lowres_height, 0, 0, 0, 0, 0)
     fp.write(
         b"VTF\x00"
-        + struct.pack('<2I', 7, 2)
+        + struct.pack('<2I', *version)
-        + struct.pack('<I2HI2HI3fIfIbI2b', *header[:17])
     )
+    if version < (7, 2):
+        size = struct.calcsize(HEADER_V70_STRUCT) + 12
+        header = header._replace(header_size=size + (16 - size % 16))
+        fp.write(struct.pack(HEADER_V70_STRUCT, *header[:15]))
+    elif version == (7, 2):
+        size = struct.calcsize(HEADER_V72_STRUCT) + 12
+        header = header._replace(header_size=size + (16 - size % 16))
+        fp.write(struct.pack(HEADER_V72_STRUCT, *header[:16]))
+    elif version > (7, 2):
+        size = struct.calcsize(HEADER_V73_STRUCT) + 12
+        header = header._replace(header_size=size + (16 - size % 16))
+        fp.write(struct.pack(HEADER_V73_STRUCT, *header))
+    else:
+        raise VTFException(f'Unsupported version {version}')
+
+    if version > (7, 2):
+        fp.write(b'\x01\x00\x00\x00')
+        fp.write(struct.pack('<I', header.header_size))
+        fp.write(b'\x30\x00\x00\x00')
+        fp.write(struct.pack('<I', header.header_size + len(thumb_buffer.getbuffer())))
+    else:
+        fp.write(b'\x00' * (16 - fp.tell() % 16))
+
+    fp.write(thumb_buffer.getbuffer())
+    assert fp.tell() == header.header_size + len(thumb_buffer.getbuffer())
+
+    for mip_id in range(mipmap_count - 1, 0, -1):
+        mip_width = max(4, width >> mip_id)
+        mip_height = max(4, height >> mip_id)
+        mip = im.resize((mip_width, mip_height))
+        ImageFile._save(mip, fp, [("bcn", (0, 0) + mip.size, (fp.tell()), (1, 'DXT1'))], mip_width * mip_height // 2)
+    ImageFile._save(im, fp, [("bcn", (0, 0) + im.size, (fp.tell()), (1, 'DXT1'))], im.width * im.height // 2)
 
 
 def _accept(prefix):

src/encode.c

@@ -27,6 +27,7 @@
 
 #include "libImaging/Imaging.h"
 #include "libImaging/Gif.h"
+#include "libImaging/Bcn.h"
 
 #ifdef HAVE_UNISTD_H
 #include <unistd.h> /* write */
@@ -389,16 +390,15 @@ PyImaging_BcnEncoderNew(PyObject *self, PyObject *args) {
     if (!PyArg_ParseTuple(args, "si|s", &mode, &n, &pixel_format)) {
         return NULL;
     }
-    printf("Mode: %s, n: %i, pixel_format: %s", mode, n, pixel_format);
 
-    encoder = PyImaging_EncoderNew(0);
+    encoder = PyImaging_EncoderNew(sizeof(char *));
     if (encoder == NULL) {
         return NULL;
     }
 
     encoder->encode = ImagingBcnEncode;
-    encoder->state.ystep = 4;
+    encoder->state.state = n;
-    encoder->cleanup = NULL;
+    ((BCNSTATE *)encoder->state.context)->pixel_format = pixel_format;
 
     return (PyObject *)encoder;
 }

src/libImaging/BcnEncode.c

@@ -16,8 +16,6 @@
 #include "Bcn.h"
 #include "math.h"
 
-#define PACK_SHORT_565(r, g, b) \
-    ((((r) << 8) & 0xF800) | (((g) << 3) & 0x7E0) | ((b) >> 3))
 #define BIT_MASK(bit_count) ((1 << (bit_count)) - 1)
 #define SET_BITS(target, bit_offset, bit_count, value) \
     target |= (((value)&BIT_MASK(bit_count)) << (bit_offset))
@@ -28,40 +26,44 @@
         B = TMP;         \
     } while (0)
 
+typedef union {
+    struct {
+        UINT16 r : 5;
+        UINT16 g : 6;
+        UINT16 b : 5;
+    } color;
+    UINT16 value;
+} rgb565;
+
+static UINT16
+pack_565(UINT8 r, UINT8 g, UINT8 b) {
+    rgb565 color;
+    color.color.r = r / (255 / 31);
+    color.color.g = g / (255 / 63);
+    color.color.b = b / (255 / 31);
+    return color.value;
+}
+
 static UINT16
 rgb565_diff(UINT16 a, UINT16 b) {
-    UINT8 r0, g0, b0, r1, g1, b1;
+    rgb565 c0;
-    r0 = a & 31;
+    c0.value = a;
-    a >>= 5;
+    rgb565 c1;
-    g0 = a & 63;
+    c1.value = b;
-    a >>= 6;
+    return ((UINT16)abs(c0.color.r - c1.color.r)) + abs(c0.color.g - c1.color.g) +
-    b0 = a & 31;
+           abs(c0.color.b - c1.color.b);
-    r1 = b & 31;
-    b >>= 5;
-    g1 = a & 63;
-    b >>= 6;
-    b1 = b & 31;
-
-    return abs(r0 - r1) + abs(g0 - g1) + abs(b0 - b1);
 }
 
 static inline UINT16
 rgb565_lerp(UINT16 a, UINT16 b, UINT8 a_fac, UINT8 b_fac) {
-    UINT8 r0, g0, b0, r1, g1, b1;
+    rgb565 c0;
-    r0 = a & 31;
+    c0.value = a;
-    a >>= 5;
+    rgb565 c1;
-    g0 = a & 63;
+    c1.value = b;
-    a >>= 6;
+    return pack_565(
-    b0 = a & 31;
+        (c0.color.r * a_fac + c1.color.r * b_fac) / (a_fac + b_fac),
-    r1 = b & 31;
+        (c0.color.g * a_fac + c1.color.g * b_fac) / (a_fac + b_fac),
-    b >>= 5;
+        (c0.color.b * a_fac + c1.color.b * b_fac) / (a_fac + b_fac));
-    g1 = b & 63;
-    b >>= 6;
-    b1 = b & 31;
-    return PACK_SHORT_565(
-        (r0 * a_fac + r1 * b_fac) / (a_fac + b_fac),
-        (g0 * a_fac + g1 * b_fac) / (a_fac + b_fac),
-        (b0 * a_fac + b1 * b_fac) / (a_fac + b_fac));
 }
 
 typedef struct {
@@ -89,24 +91,8 @@ pick_2_major_colors(
     UINT16 color_count,
     UINT16 *color0,
     UINT16 *color1) {
-    Color colors[16] = {
+    Color colors[16];
-        0,
+    memset(colors, 0, sizeof(colors));
-        0,
-        0,
-        0,
-        0,
-        0,
-        0,
-        0,
-        0,
-        0,
-        0,
-        0,
-        0,
-        0,
-        0,
-        0,
-    };
     for (int i = 0; i < color_count; ++i) {
         colors[i].value = unique_colors[i];
         colors[i].frequency = color_freq[i];
@@ -121,8 +107,8 @@ pick_2_major_colors(
 }
 
 static UINT8
-get_closest_color_index(const UINT16 colors[4], UINT16 color) {
+get_closest_color_index(const UINT16 *colors, UINT16 color) {
-    UINT16 color_error = 65535;
+    UINT16 color_error = 0xFFF8;
     UINT16 lowest_id = 0;
 
     for (int color_id = 0; color_id < 4; color_id++) {
@@ -139,26 +125,44 @@ get_closest_color_index(const UINT16 colors[4], UINT16 color) {
 }
 
 int
-ImagingBcnEncode(Imaging im, ImagingCodecState state, UINT8 *buf, int bytes) {
+encode_bc1(Imaging im, ImagingCodecState state, UINT8 *buf, int bytes) {
-    UINT8 *output = buf;
+    bc1_color *blocks = (bc1_color *)buf;
-    BOOL has_alpha = 1;
+    UINT8 no_alpha = 0;
     if (strchr(im->mode, 'A') == NULL)
-        has_alpha = 0;
+        no_alpha = 1;
-    UINT32 row_block_count = im->xsize / 4;
+    UINT32 block_count = (im->xsize * im->ysize) / 16;
+    if (block_count * sizeof(bc1_color) > bytes) {
+        state->errcode = IMAGING_CODEC_MEMORY;
+        return 0;
+    }
 
+    memset(buf, 0, block_count * sizeof(bc1_color));
+    for (int block_index = 0; block_index < block_count; block_index++) {
+        state->x = (block_index % (im->xsize / 4));
+        state->y = (block_index / (im->xsize / 4));
         UINT16 unique_count = 0;
-    UINT16 all_colors[16] = {0};
-    UINT16 unique_colors[16] = {0};
-    UINT8 color_frequency[16] = {0};
 
-    for (int block_x = 0; block_x < row_block_count; ++block_x) {
+        UINT16 all_colors[16];
+        UINT16 unique_colors[16];
+        UINT8 color_frequency[16];
+        UINT8 opaque[16];
+        memset(all_colors, 0, sizeof(all_colors));
+        memset(unique_colors, 0, sizeof(unique_colors));
+        memset(color_frequency, 0, sizeof(color_frequency));
+        memset(opaque, 0, sizeof(opaque));
+
         for (int by = 0; by < 4; ++by) {
-            for (int bx = 0; bx < 16; bx += 4) {
+            for (int bx = 0; bx < 4; ++bx) {
-                UINT8 r = (im->image[state->y + by][bx]);
+                int x = (state->x * 4) + bx;
-                UINT8 g = (im->image[state->y + by][bx + 1]);
+                int y = (state->y * 4) + by;
-                UINT8 b = (im->image[state->y + by][bx] + 2);
+                UINT8 r = im->image[y][x * im->pixelsize + 2];
-                UINT16 color = PACK_SHORT_565(r, g, b);
+                UINT8 g = im->image[y][x * im->pixelsize + 1];
+                UINT8 b = im->image[y][x * im->pixelsize + 0];
+                UINT8 a = im->image[y][x * im->pixelsize + 3];
+                UINT16 color = pack_565(r, g, b);
+                opaque[bx + by * 4] = a >= 127;
                 all_colors[bx + by * 4] = color;
+
                 BOOL new_color = 1;
                 for (UINT16 color_id = 0; color_id < unique_count; color_id++) {
                     if (unique_colors[color_id] == color) {
@@ -175,32 +179,45 @@ ImagingBcnEncode(Imaging im, ImagingCodecState state, UINT8 *buf, int bytes) {
             }
         }
 
-        UINT16 color0 = 0, color1 = 0;
+        UINT16 c0 = 0, c1 = 0;
-        pick_2_major_colors(
+        pick_2_major_colors(unique_colors, color_frequency, unique_count, &c0, &c1);
-            unique_colors, color_frequency, unique_count, &color0, &color1);
+        if (c0 < c1 && no_alpha)
-        if (color0 < color1)
+            SWAP(UINT16, c0, c1);
-            SWAP(UINT16, color0, color1);
 
-        UINT16 output_colors[4] = {color0, color1, 0, 0};
+        UINT16 palette[4] = {c0, c1, 0, 0};
-        if (has_alpha) {
+        if (no_alpha) {
-            output_colors[2] = rgb565_lerp(color0, color1, 1, 1);
+            palette[2] = rgb565_lerp(c0, c1, 2, 1);
-            output_colors[3] = 0;
+            palette[3] = rgb565_lerp(c0, c1, 1, 2);
         } else {
-            output_colors[2] = rgb565_lerp(color0, color1, 2, 1);
+            palette[2] = rgb565_lerp(c0, c1, 1, 1);
-            output_colors[3] = rgb565_lerp(color0, color1, 1, 2);
+            palette[3] = 0;
         }
-        bc1_color *block = &((bc1_color *)output)[block_x];
+        bc1_color *block = &blocks[block_index];
 
-        block->c0 = color0;
+        block->c0 = c0;
-        block->c1 = color1;
+        block->c1 = c1;
         for (UINT32 color_id = 0; color_id < 16; ++color_id) {
-            UINT8 bc_color_id =
+            UINT8 bc_color_id;
-                get_closest_color_index(output_colors, all_colors[color_id]);
+            if (opaque[color_id] || no_alpha)
+                bc_color_id = get_closest_color_index(palette, all_colors[color_id]);
+            else
+                bc_color_id = 3;
             SET_BITS(block->lut, color_id * 2, 2, bc_color_id);
         }
-
+    }
-        output += sizeof(bc1_color);
+    state->errcode = IMAGING_CODEC_END;
+    return block_count * sizeof(bc1_color);
 }
 
-    return output - buf;
+int
+ImagingBcnEncode(Imaging im, ImagingCodecState state, UINT8 *buf, int bytes) {
+    switch (state->state) {
+        case 1: {
+            return encode_bc1(im, state, buf, bytes);
+        }
+        default: {
+            state->errcode = IMAGING_CODEC_CONFIG;
+            return 0;
+        }
+    }
 }