185 lines
7.6 KiB
Python
185 lines
7.6 KiB
Python
"""Reading, writing and resampling heightmaps with nothing but numpy, so the authoring scripts run on the
|
|
engine's own Python (which has no PIL). Greyscale PNG in 8 or 16 bit, raw 16-bit little-endian (.r16 / .raw,
|
|
what World Machine, Gaea and the engine's own exporter write), bilinear resampling and a centred square crop:
|
|
enough to take a real heightmap from any of the usual sources and put it on the landscape.
|
|
"""
|
|
import math
|
|
import struct
|
|
import zlib
|
|
|
|
import numpy as np
|
|
|
|
PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
|
|
|
|
|
|
def write_png(path, data):
|
|
"""Greyscale PNG, 8 or 16 bit from the array dtype. Row filter 0, one zlib stream."""
|
|
if data.dtype == np.uint16:
|
|
depth, payload = 16, data.astype(">u2")
|
|
else:
|
|
depth, payload = 8, data.astype(np.uint8)
|
|
height, width = data.shape
|
|
raw = b"".join(b"\x00" + payload[y].tobytes() for y in range(height))
|
|
|
|
def chunk(kind, body):
|
|
return struct.pack(">I", len(body)) + kind + body + struct.pack(">I", zlib.crc32(kind + body) & 0xFFFFFFFF)
|
|
|
|
ihdr = struct.pack(">IIBBBBB", width, height, depth, 0, 0, 0, 0)
|
|
with open(path, "wb") as f:
|
|
f.write(PNG_SIGNATURE + chunk(b"IHDR", ihdr) + chunk(b"IDAT", zlib.compress(raw, 6)) + chunk(b"IEND", b""))
|
|
|
|
|
|
def read_png_header(path):
|
|
"""(width, height, bit_depth, colour_type) without decoding the image."""
|
|
with open(path, "rb") as f:
|
|
head = f.read(8 + 8 + 13)
|
|
if head[:8] != PNG_SIGNATURE or head[12:16] != b"IHDR":
|
|
raise ValueError(f"{path}: not a PNG")
|
|
width, height, depth, colour_type = struct.unpack(">IIBB", head[16:26])
|
|
return width, height, depth, colour_type
|
|
|
|
|
|
def _unfilter_sequential(filter_type, row, prev, bpp):
|
|
"""Average and Paeth depend on the byte just decoded, so they go pixel by pixel. Rare in practice; a
|
|
4081x4081 16-bit file with every row Paeth-filtered takes some tens of seconds, once, on import."""
|
|
out = bytearray(row)
|
|
n = len(out)
|
|
if filter_type == 3:
|
|
for i in range(n):
|
|
left = out[i - bpp] if i >= bpp else 0
|
|
out[i] = (out[i] + ((left + prev[i]) >> 1)) & 0xFF
|
|
else:
|
|
for i in range(n):
|
|
if i >= bpp:
|
|
a, c = out[i - bpp], prev[i - bpp]
|
|
else:
|
|
a, c = 0, 0
|
|
b = prev[i]
|
|
p = a + b - c
|
|
pa, pb, pc = abs(p - a), abs(p - b), abs(p - c)
|
|
if pa <= pb and pa <= pc:
|
|
predictor = a
|
|
elif pb <= pc:
|
|
predictor = b
|
|
else:
|
|
predictor = c
|
|
out[i] = (out[i] + predictor) & 0xFF
|
|
return out
|
|
|
|
|
|
def read_png(path):
|
|
"""The first channel of a non-interlaced PNG as a 2D uint8 or uint16 array (greyscale, grey+alpha, RGB
|
|
and RGBA are accepted; palette and interlaced files are not)."""
|
|
with open(path, "rb") as f:
|
|
blob = f.read()
|
|
if blob[:8] != PNG_SIGNATURE:
|
|
raise ValueError(f"{path}: not a PNG")
|
|
pos, idat, ihdr = 8, [], None
|
|
while pos + 8 <= len(blob):
|
|
length, kind = struct.unpack(">I4s", blob[pos:pos + 8])
|
|
body = blob[pos + 8:pos + 8 + length]
|
|
pos += 12 + length
|
|
if kind == b"IHDR":
|
|
ihdr = struct.unpack(">IIBBBBB", body)
|
|
elif kind == b"IDAT":
|
|
idat.append(body)
|
|
elif kind == b"IEND":
|
|
break
|
|
if ihdr is None:
|
|
raise ValueError(f"{path}: no IHDR")
|
|
width, height, depth, colour_type, _, _, interlace = ihdr
|
|
channels = {0: 1, 2: 3, 4: 2, 6: 4}.get(colour_type)
|
|
if channels is None or depth not in (8, 16) or interlace != 0:
|
|
raise ValueError(f"{path}: unsupported PNG (colour type {colour_type}, {depth} bit, interlace {interlace}); "
|
|
"use a non-interlaced 8 or 16 bit greyscale or RGB file")
|
|
bytes_per_sample = depth // 8
|
|
bpp = channels * bytes_per_sample
|
|
stride = width * bpp
|
|
data = zlib.decompress(b"".join(idat))
|
|
if len(data) != height * (stride + 1):
|
|
raise ValueError(f"{path}: PNG data is {len(data)} bytes, expected {height * (stride + 1)}")
|
|
|
|
rows = np.empty((height, stride), dtype=np.uint8)
|
|
prev = np.zeros(stride, dtype=np.uint8)
|
|
for y in range(height):
|
|
start = y * (stride + 1)
|
|
filter_type = data[start]
|
|
row = np.frombuffer(data, dtype=np.uint8, count=stride, offset=start + 1)
|
|
if filter_type == 0:
|
|
out = row.copy()
|
|
elif filter_type == 1:
|
|
out = (np.cumsum(row.reshape(width, bpp), axis=0, dtype=np.uint64) & 0xFF).astype(np.uint8).reshape(stride)
|
|
elif filter_type == 2:
|
|
out = ((row.astype(np.uint16) + prev) & 0xFF).astype(np.uint8)
|
|
elif filter_type in (3, 4):
|
|
out = np.frombuffer(bytes(_unfilter_sequential(filter_type, bytes(row), bytes(prev), bpp)), dtype=np.uint8)
|
|
else:
|
|
raise ValueError(f"{path}: bad PNG filter {filter_type} on row {y}")
|
|
rows[y] = out
|
|
prev = rows[y]
|
|
|
|
dtype = ">u2" if depth == 16 else np.uint8
|
|
samples = rows.reshape(height, width * channels * bytes_per_sample).view(dtype).reshape(height, width, channels)
|
|
first = samples[:, :, 0]
|
|
return first.astype(np.uint16) if depth == 16 else first.astype(np.uint8)
|
|
|
|
|
|
def read_r16(path, width=None):
|
|
"""Raw 16-bit little-endian samples, square unless a width is given."""
|
|
values = np.fromfile(path, dtype="<u2")
|
|
if width is None:
|
|
width = math.isqrt(len(values))
|
|
if width * width != len(values):
|
|
raise ValueError(f"{path}: {len(values)} samples is not a square; give the width in the manifest")
|
|
if len(values) % width != 0:
|
|
raise ValueError(f"{path}: {len(values)} samples do not divide by width {width}")
|
|
return values.reshape(len(values) // width, width).astype(np.uint16)
|
|
|
|
|
|
def read_heightmap(path, width=None):
|
|
"""Any supported file as a 2D uint16 array with the full 0..65535 range (8-bit files are widened)."""
|
|
lower = path.lower()
|
|
if lower.endswith(".png"):
|
|
values = read_png(path)
|
|
return values.astype(np.uint16) * 257 if values.dtype == np.uint8 else values
|
|
if lower.endswith((".r16", ".raw")):
|
|
return read_r16(path, width)
|
|
raise ValueError(f"{path}: unknown heightmap format; use 16-bit PNG or raw .r16")
|
|
|
|
|
|
def center_crop_square(values):
|
|
height, width = values.shape
|
|
side = min(height, width)
|
|
y0, x0 = (height - side) // 2, (width - side) // 2
|
|
return values[y0:y0 + side, x0:x0 + side]
|
|
|
|
|
|
def block_mean(values, factor):
|
|
"""Downsample by an integer factor with a box filter, trimming the edge that does not divide."""
|
|
height, width = values.shape
|
|
height, width = height // factor * factor, width // factor * factor
|
|
trimmed = values[:height, :width].astype(np.float32)
|
|
return trimmed.reshape(height // factor, factor, width // factor, factor).mean(axis=(1, 3))
|
|
|
|
|
|
def resample(values, size):
|
|
"""Bilinear resample of a 2D array to size x size, box-filtered first when shrinking by 2x or more."""
|
|
source = values.astype(np.float32)
|
|
factor = min(source.shape) // size
|
|
if factor >= 2:
|
|
source = block_mean(source, factor)
|
|
src_h, src_w = source.shape
|
|
if (src_h, src_w) == (size, size):
|
|
return source
|
|
ys = np.linspace(0.0, src_h - 1, size, dtype=np.float32)
|
|
xs = np.linspace(0.0, src_w - 1, size, dtype=np.float32)
|
|
y0 = np.floor(ys).astype(np.int64)
|
|
x0 = np.floor(xs).astype(np.int64)
|
|
y1 = np.minimum(y0 + 1, src_h - 1)
|
|
x1 = np.minimum(x0 + 1, src_w - 1)
|
|
ty = (ys - y0)[:, None]
|
|
tx = (xs - x0)[None, :]
|
|
top = source[np.ix_(y0, x0)] * (1 - tx) + source[np.ix_(y0, x1)] * tx
|
|
bottom = source[np.ix_(y1, x0)] * (1 - tx) + source[np.ix_(y1, x1)] * tx
|
|
return (top * (1 - ty) + bottom * ty).astype(np.float32)
|