"""Image file I/O for textures: loading and saving PNG/JPG."""
import logging
import struct
import zlib
from pathlib import Path
from typing import Any
import numpy as np
import vulkan as vk
from PIL import Image
from simvx.graphics.gpu.memory import upload_image_data
log = logging.getLogger(__name__)
[docs]
def premultiply_alpha_rgba(pixels: np.ndarray) -> np.ndarray:
"""Return a new RGBA uint8 array with RGB channels premultiplied by alpha.
Fixes halo artefacts on alpha-blended PNGs where transparent pixels still
carry stale RGB values (e.g. ``(255, 255, 255, 0)`` border pixels that
bleed white into anti-aliased sprite edges). The transform is
``rgb' = rgb * a / 255``; for fully transparent pixels the output RGB is
zero, regardless of the source RGB. Engine code paths use
straight-alpha blending today, so this is opt-in: pass
``premultiply_alpha=True`` to ``TextureManager.resolve`` / ``.load`` /
``.load_from_bytes`` to enable it per-texture.
Args:
pixels: ``(H, W, 4)`` uint8 RGBA array.
Returns:
A *new* ``(H, W, 4)`` uint8 array; the input is not mutated.
"""
if pixels.ndim != 3 or pixels.shape[2] != 4 or pixels.dtype != np.uint8:
raise ValueError(
f"premultiply_alpha_rgba: expected RGBA uint8 (H, W, 4); got {pixels.shape}/{pixels.dtype}"
)
out = pixels.astype(np.uint16, copy=True)
alpha = out[..., 3:4] # broadcast across RGB
out[..., :3] = (out[..., :3] * alpha + 127) // 255
return out.astype(np.uint8)
[docs]
def load_texture_from_file(
device: Any,
physical_device: Any,
queue: Any,
cmd_pool: Any,
file_path: str,
*,
premultiply_alpha: bool = False,
) -> tuple[Any, Any, int, int]:
"""Load PNG/JPG texture from disk → device-local VkImage.
Returns: (image, memory, width, height)
Args:
premultiply_alpha: When True, multiply RGB by alpha before upload so
anti-aliased edges don't bleed garbage RGB through transparent
pixels. Default False keeps existing snapshot tests stable.
"""
img = Image.open(file_path).convert("RGBA")
width, height = img.size
pixels = np.ascontiguousarray(np.array(img, dtype=np.uint8))
if premultiply_alpha:
pixels = premultiply_alpha_rgba(pixels)
image, memory = upload_image_data(
device,
physical_device,
queue,
cmd_pool,
pixels,
width,
height,
vk.VK_FORMAT_R8G8B8A8_UNORM,
)
return image, memory, width, height
# ---------------------------------------------------------------------------
# PNG I/O (pure Python, no Pillow)
# ---------------------------------------------------------------------------
def _png_chunk(chunk_type: bytes, data: bytes) -> bytes:
"""Build a single PNG chunk: length + type + data + CRC."""
crc = zlib.crc32(chunk_type + data) & 0xFFFFFFFF
return struct.pack(">I", len(data)) + chunk_type + data + struct.pack(">I", crc)
[docs]
def save_png(path: str | Path, pixels: np.ndarray) -> None:
"""Save RGBA uint8 pixels (H, W, 4) as a PNG file. Pure Python, no Pillow."""
h, w = pixels.shape[:2]
channels = pixels.shape[2] if pixels.ndim == 3 else 1
if channels not in (3, 4):
raise ValueError(f"Expected 3 or 4 channels, got {channels}")
colour_type = 6 if channels == 4 else 2 # RGBA or RGB
raw = bytearray()
row_bytes = pixels[:, :, :channels].reshape(h, -1)
for y in range(h):
raw.append(0) # filter type: None
raw.extend(row_bytes[y].tobytes())
ihdr = struct.pack(">IIBBBBB", w, h, 8, colour_type, 0, 0, 0)
compressed = zlib.compress(bytes(raw), 9)
p = Path(path)
with open(p, "wb") as f:
f.write(b"\x89PNG\r\n\x1a\n")
f.write(_png_chunk(b"IHDR", ihdr))
f.write(_png_chunk(b"IDAT", compressed))
f.write(_png_chunk(b"IEND", b""))
def _unfilter_scanlines(raw: bytes, height: int, width: int, channels: int) -> np.ndarray:
"""Reverse the PNG per-scanline filters (RFC 2083 §6) into an (H, W, C) array.
Handles all five filter types (0 None, 1 Sub, 2 Up, 3 Average, 4 Paeth); only
None was supported before, so any PNG not written by our own ``save_png`` (e.g.
a Playwright/Chrome screenshot, which uses Up/Paeth) failed to load. ``bpp`` is
``channels`` because bit depth is fixed at 8. None/Sub/Up are vectorised;
Average/Paeth carry an intra-row left dependency so they iterate per pixel
(channel-vectorised), which is fine for the occasional image load.
"""
stride = width * channels
prev = np.zeros((width, channels), dtype=np.int32)
out = np.empty((height, width, channels), dtype=np.uint8)
pos = 0
for y in range(height):
ft = raw[pos]
pos += 1
filt = np.frombuffer(raw, np.uint8, stride, pos).astype(np.int32).reshape(width, channels)
pos += stride
if ft == 0: # None
recon = filt % 256
elif ft == 1: # Sub: predictor is the pixel to the left -> per-channel running sum
recon = np.cumsum(filt, axis=0) % 256
elif ft == 2: # Up: predictor is the pixel above
recon = (filt + prev) % 256
elif ft == 3: # Average: floor((left + up) / 2)
recon = np.empty((width, channels), dtype=np.int32)
for x in range(width):
a = recon[x - 1] if x > 0 else 0
recon[x] = (filt[x] + ((a + prev[x]) >> 1)) % 256
elif ft == 4: # Paeth
recon = np.empty((width, channels), dtype=np.int32)
zero = np.zeros(channels, dtype=np.int32)
for x in range(width):
a = recon[x - 1] if x > 0 else zero
b = prev[x]
c = prev[x - 1] if x > 0 else zero
p = a + b - c
pa, pb, pc = np.abs(p - a), np.abs(p - b), np.abs(p - c)
pred = np.where((pa <= pb) & (pa <= pc), a, np.where(pb <= pc, b, c))
recon[x] = (filt[x] + pred) % 256
else:
raise ValueError(f"Unsupported PNG filter type {ft} at row {y}")
prev = recon
out[y] = recon.astype(np.uint8)
return out
def _load_png(path: str | Path) -> np.ndarray:
"""Load a PNG written by save_png() back to an RGBA (H, W, 4) uint8 ndarray."""
data = Path(path).read_bytes()
if data[:8] != b"\x89PNG\r\n\x1a\n":
raise ValueError(f"Not a PNG file: {path}")
pos = 8
width = height = 0
channels = 4
idat_parts: list[bytes] = []
while pos < len(data):
length = struct.unpack(">I", data[pos : pos + 4])[0]
chunk_type = data[pos + 4 : pos + 8]
chunk_data = data[pos + 8 : pos + 8 + length]
pos += 12 + length
if chunk_type == b"IHDR":
width, height, bit_depth, colour_type = struct.unpack(">IIBB", chunk_data[:10])
if bit_depth != 8:
raise ValueError(f"Unsupported bit depth: {bit_depth}")
channels = 4 if colour_type == 6 else 3
elif chunk_type == b"IDAT":
idat_parts.append(chunk_data)
elif chunk_type == b"IEND":
break
raw = zlib.decompress(b"".join(idat_parts))
pixels = _unfilter_scanlines(raw, height, width, channels)
if channels == 3:
alpha = np.full((height, width, 1), 255, dtype=np.uint8)
pixels = np.concatenate([pixels, alpha], axis=2)
return pixels