jpmschweitzerandClaude Opus 4.6
6286cf5d2c
Remove unused imports (os, sys, math, ImageFilter, gaussian_filter, MAX_BIOME_ID) flagged by ruff F401. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
485 lines
18 KiB
Python
485 lines
18 KiB
Python
"""
|
||
render_heightmap.py
|
||
-------------------
|
||
Renders a 4096×2048 annotated equirectangular heightmap PNG from a terrain dict.
|
||
|
||
This is the PRIMARY output of the planet generator pipeline.
|
||
The globe render is a separate downstream step that reads the same terrain dict.
|
||
|
||
Equirectangular projection:
|
||
X axis: longitude 0°→360° (left to right)
|
||
Y axis: latitude +90°→-90° (top to bottom, north pole at row 0)
|
||
|
||
Each terrain grid cell maps to a block of output pixels via bicubic upscale.
|
||
All rendering is in float32; final conversion to uint8 at save time.
|
||
|
||
Output layers (composited in order):
|
||
1. Biome colour — smooth-blended from Whittaker grid, not hard-snapped
|
||
2. Elevation shading — subtle darkening in valleys, lightening on peaks
|
||
3. Hillshade — surface normal lighting pass (makes terrain 3D-readable)
|
||
4. Coastline — 1px dark border at sea level threshold
|
||
5. Rivers — anti-aliased polylines from river list
|
||
6. Lat/lon grid — every 30°, semi-transparent
|
||
7. Title panel — body metadata strip at top
|
||
8. Legend — biome colour swatches at bottom
|
||
|
||
Geographic only. No settlements, roads, or cultural data.
|
||
Those live in a separate JSON sidecar and are overlaid by the atlas app.
|
||
|
||
Usage:
|
||
from render_heightmap import render_heightmap
|
||
from planet_simulation import simulate
|
||
from body_definition_parser import parse_system
|
||
|
||
defs = parse_system("index.md")
|
||
terrain = simulate(defs[0])
|
||
img = render_heightmap(defs[0], terrain)
|
||
img.save("GJ144d_heightmap.png")
|
||
"""
|
||
|
||
import numpy as np
|
||
from PIL import Image, ImageDraw, ImageFont
|
||
from scipy.ndimage import binary_dilation
|
||
|
||
from biome_config import (
|
||
BIOME_PALETTE as _BIOME_PALETTE_CFG,
|
||
RIVER_RGB as _RIVER_RGB_CFG,
|
||
COAST_RGB as _COAST_RGB_CFG,
|
||
build_biome_rgb,
|
||
)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Output resolution
|
||
# ---------------------------------------------------------------------------
|
||
|
||
OUT_W = 4096
|
||
OUT_H = 2048
|
||
UI_SCALE = OUT_W / 1024 # 4.0 — all pixel sizes scale with this
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Biome colour palette
|
||
# Indices match planet_simulation.WHITTAKER_TABLE class IDs.
|
||
# Extended exotic classes appended at end.
|
||
# ---------------------------------------------------------------------------
|
||
|
||
# Biome palette loaded from biomes.toml via biome_config.
|
||
# Per-planet overrides can patch _BIOME_PALETTE_CFG before rendering.
|
||
BIOME_PALETTE = _BIOME_PALETTE_CFG
|
||
|
||
def _build_biome_rgb(mode: str = "cartographic") -> dict:
|
||
return build_biome_rgb(mode)
|
||
|
||
RENDER_MODE = "cartographic"
|
||
BIOME_RGB = _build_biome_rgb(RENDER_MODE)
|
||
|
||
def _ocean_arrays(mode: str = "cartographic"):
|
||
return (
|
||
np.array(BIOME_PALETTE[0][mode], dtype=np.float32),
|
||
np.array(BIOME_PALETTE[1][mode], dtype=np.float32),
|
||
np.array(BIOME_PALETTE[2][mode], dtype=np.float32),
|
||
)
|
||
|
||
OCEAN_DEEP, OCEAN_MID, OCEAN_SHALLOW = _ocean_arrays(RENDER_MODE)
|
||
|
||
RIVER_RGB = _RIVER_RGB_CFG
|
||
COAST_RGB = _COAST_RGB_CFG
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Helpers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _upscale(grid: np.ndarray, order: int = 1) -> np.ndarray:
|
||
"""
|
||
Upscale a (GRID_H, GRID_W) float32 grid to (OUT_H, OUT_W).
|
||
order=1 → bilinear (smooth, good for continuous fields)
|
||
order=0 → nearest (sharp, good for integer class grids)
|
||
"""
|
||
from scipy.ndimage import zoom
|
||
zy = OUT_H / grid.shape[0]
|
||
zx = OUT_W / grid.shape[1]
|
||
return zoom(grid.astype(np.float32), (zy, zx), order=order).astype(np.float32)
|
||
|
||
|
||
def _upscale_int(grid: np.ndarray) -> np.ndarray:
|
||
"""Nearest-neighbour upscale for integer class grids (biome, etc)."""
|
||
from scipy.ndimage import zoom
|
||
zy = OUT_H / grid.shape[0]
|
||
zx = OUT_W / grid.shape[1]
|
||
return zoom(grid.astype(np.int32), (zy, zx), order=0).astype(np.int8)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Layer 1 + 2 + 3: Biome colour + elevation shading + hillshade
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _render_surface(terrain: dict) -> np.ndarray:
|
||
"""
|
||
Returns (OUT_H, OUT_W, 3) float32 RGB in [0, 1].
|
||
|
||
Compositing order:
|
||
biome_colour × elevation_shade × hillshade_factor
|
||
"""
|
||
elevation = _upscale(terrain["elevation"], order=1)
|
||
hillshade = _upscale(terrain["hillshade"], order=1)
|
||
biome_up = _upscale_int(terrain["biome"])
|
||
surf_water = _upscale(terrain["surface_water"].astype(np.float32),
|
||
order=0) > 0.5
|
||
sea_level = terrain["sea_level"]
|
||
|
||
H, W = elevation.shape
|
||
|
||
# ── Biome base colour ─────────────────────────────────────────────────
|
||
# Clamp biome index, look up palette
|
||
# Build lookup array from active BIOME_RGB dict for vectorised indexing
|
||
max_id = max(BIOME_RGB.keys())
|
||
pal_arr = np.zeros((max_id + 1, 3), dtype=np.float32)
|
||
for k, v in BIOME_RGB.items():
|
||
pal_arr[k] = v
|
||
biome_clamped = np.clip(biome_up, 0, max_id)
|
||
rgb = pal_arr[biome_clamped].astype(np.float32) / 255.0
|
||
|
||
# ── Ocean depth blending ───────────────────────────────────────────────
|
||
# Override flat ocean biome with smooth depth gradient
|
||
if surf_water.any():
|
||
depth = np.clip((sea_level - elevation) / (sea_level + 1e-9), 0, 1)
|
||
deep_col = OCEAN_DEEP / 255.0
|
||
mid_col = OCEAN_MID / 255.0
|
||
shallow_col = OCEAN_SHALLOW / 255.0
|
||
|
||
# Three-stop blend: 0=shallow, 0.5=mid, 1=deep
|
||
t1 = np.clip(depth * 2.0, 0, 1) # 0→0.5 depth: shallow→mid
|
||
t2 = np.clip((depth - 0.5) * 2.0, 0, 1) # 0.5→1 depth: mid→deep
|
||
ocean_rgb = (shallow_col * (1 - t1)[..., None]
|
||
+ mid_col * (t1 * (1 - t2))[..., None]
|
||
+ deep_col * t2[..., None])
|
||
rgb = np.where(surf_water[..., None], ocean_rgb, rgb)
|
||
|
||
# ── Elevation shading on land ──────────────────────────────────────────
|
||
# Slight darkening in lowlands, brightening on ridges
|
||
elev_norm = np.where(
|
||
~surf_water,
|
||
np.clip((elevation - sea_level) / (1.0 - sea_level + 1e-9), 0, 1),
|
||
0.0)
|
||
elev_shade = 0.88 + 0.18 * elev_norm # [0.88, 1.06] — clamp below
|
||
rgb = np.where(~surf_water[..., None],
|
||
np.clip(rgb * elev_shade[..., None], 0, 1),
|
||
rgb)
|
||
|
||
# ── Hillshade ──────────────────────────────────────────────────────────
|
||
# Apply only on land — ocean gets its own depth shading
|
||
# Blend factor: 0.55 hillshade + 0.45 flat (keeps colours readable)
|
||
hs_blend = 0.55 * hillshade + 0.45
|
||
rgb = np.where(~surf_water[..., None],
|
||
np.clip(rgb * hs_blend[..., None], 0, 1),
|
||
rgb)
|
||
|
||
return rgb.astype(np.float32)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Layer 4: Coastline
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _render_coastline(terrain: dict,
|
||
rgb: np.ndarray) -> np.ndarray:
|
||
"""Draw a 1–2px dark border at the sea level threshold."""
|
||
surf_water = _upscale(terrain["surface_water"].astype(np.float32),
|
||
order=0) > 0.5
|
||
|
||
# Dilate water mask by 1px, XOR with original → coastline ring
|
||
dilated = binary_dilation(surf_water, iterations=2)
|
||
coastline = dilated & ~surf_water
|
||
|
||
coast_col = np.array(COAST_RGB, dtype=np.float32) / 255.0
|
||
out = rgb.copy()
|
||
out[coastline] = coast_col
|
||
return out
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Layer 5: Rivers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _render_rivers(terrain: dict,
|
||
rgb: np.ndarray) -> np.ndarray:
|
||
"""
|
||
Draw rivers as anti-aliased polylines.
|
||
River list is in simulation grid coords (row, col) at GRID_H×GRID_W.
|
||
Scale to output pixels, draw with PIL.
|
||
"""
|
||
rivers = terrain.get("rivers", [])
|
||
if not rivers:
|
||
return rgb
|
||
|
||
GRID_H, GRID_W = terrain["_grid_h"], terrain["_grid_w"]
|
||
scale_y = OUT_H / GRID_H
|
||
scale_x = OUT_W / GRID_W
|
||
|
||
# Work on a PIL image for anti-aliased line drawing
|
||
img = Image.fromarray((rgb * 255).clip(0, 255).astype(np.uint8), mode="RGB")
|
||
draw = ImageDraw.Draw(img)
|
||
|
||
river_col = RIVER_RGB
|
||
|
||
for path in rivers:
|
||
if len(path) < 2:
|
||
continue
|
||
# Scale grid coords to output pixels
|
||
pts = [(int(c * scale_x), int(r * scale_y)) for r, c in path]
|
||
# Line width scales with path length — longer rivers are wider.
|
||
# Base width doubled for readability at high output resolutions.
|
||
width = max(2, min(6, len(path) // 40))
|
||
draw.line(pts, fill=river_col, width=width, joint="curve")
|
||
|
||
return np.array(img).astype(np.float32) / 255.0
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Layer 6: Lat/lon grid
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _render_grid(rgb: np.ndarray) -> np.ndarray:
|
||
"""Draw lat/lon lines every 30° as semi-transparent overlays."""
|
||
out = rgb.copy()
|
||
col = np.array([255, 255, 255], dtype=np.float32) / 255.0
|
||
alpha = 0.12 # very subtle
|
||
|
||
# Latitude lines (horizontal) every 30°: at 1/6, 2/6, 3/6, 4/6, 5/6 of height
|
||
for frac in [1/6, 2/6, 3/6, 4/6, 5/6]:
|
||
y = int(frac * OUT_H)
|
||
y0 = max(0, y - 1); y1 = min(OUT_H - 1, y + 1)
|
||
out[y0:y1, :] = out[y0:y1, :] * (1 - alpha) + col * alpha
|
||
|
||
# Longitude lines (vertical) every 30°
|
||
for frac in [1/6, 2/6, 3/6, 4/6, 5/6]:
|
||
x = int(frac * OUT_W)
|
||
x0 = max(0, x - 1); x1 = min(OUT_W - 1, x + 1)
|
||
out[:, x0:x1] = out[:, x0:x1] * (1 - alpha) + col * alpha
|
||
|
||
return out
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Layer 7: Title panel
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _load_font(size: int):
|
||
try:
|
||
return ImageFont.load_default(size=size)
|
||
except TypeError:
|
||
return ImageFont.load_default()
|
||
|
||
|
||
def _render_title(img: Image.Image, body_def: dict) -> Image.Image:
|
||
"""Draw metadata strip at top of image."""
|
||
panel_h = int(52 * UI_SCALE)
|
||
panel = Image.new("RGBA", (OUT_W, panel_h), (12, 15, 22, 210))
|
||
|
||
img_rgba = img.convert("RGBA")
|
||
img_rgba.paste(panel, (0, 0), panel)
|
||
img_out = img_rgba.convert("RGB")
|
||
draw = ImageDraw.Draw(img_out)
|
||
|
||
name = body_def.get("name") or body_def.get("id", "Unknown")
|
||
bid = body_def.get("id", "")
|
||
pclass = body_def.get("planet_class", "").replace("_ringed", "")
|
||
star = body_def.get("star", {})
|
||
orbit = body_def.get("orbit", {})
|
||
phys = body_def.get("physical", {})
|
||
env = body_def.get("environment", {})
|
||
|
||
star_str = f"{star.get('type','?')}-type"
|
||
dist_str = f"{orbit.get('distance_au', 0):.2f} AU"
|
||
grav_str = f"{phys.get('gravity_g', '?')}g"
|
||
atmo_str = phys.get("atmosphere", "?")
|
||
hydro_str = env.get("hydrosphere", "?")
|
||
|
||
px = int(14 * UI_SCALE)
|
||
py = int(7 * UI_SCALE)
|
||
lh = int(17 * UI_SCALE)
|
||
|
||
title_col = (200, 210, 228)
|
||
sub_col = (130, 145, 168)
|
||
dim_col = (75, 88, 110)
|
||
|
||
line1 = f"{name.upper()} · {bid} · {pclass}"
|
||
line2 = f"{star_str} · {dist_str} · {grav_str} · atmo: {atmo_str} · hydro: {hydro_str}"
|
||
line3 = "HEIGHTMAP · Settled Reach"
|
||
|
||
draw.text((px, py), line1, fill=title_col, font=_load_font(int(14 * UI_SCALE)))
|
||
draw.text((px, py + lh), line2, fill=sub_col, font=_load_font(int(12 * UI_SCALE)))
|
||
draw.text((px, py + lh*2), line3, fill=dim_col, font=_load_font(int(11 * UI_SCALE)))
|
||
|
||
return img_out
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Layer 8: Legend
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _biome_legend_items(terrain: dict) -> list:
|
||
"""
|
||
Return list of (label, RGB) for biome classes actually present
|
||
in this terrain — no phantom legend entries.
|
||
"""
|
||
biome = terrain["biome"]
|
||
present = set(np.unique(biome).tolist())
|
||
|
||
LABELS = {
|
||
0: "ocean deep", 1: "ocean", 2: "coastal water",
|
||
3: "coast", 5: "rainforest", 6: "trop. forest",
|
||
7: "savanna", 8: "grassland", 9: "forest",
|
||
10: "rainforest", 11: "boreal", 12: "shrubland",
|
||
13: "temperate desert", 14: "desert", 15: "hot desert",
|
||
16: "tundra", 17: "ice / snow", 18: "mountain rock",
|
||
19: "lava field", 20: "chemosyn. mat", 21: "thermophilic",
|
||
22: "sulfuric scrub", 23: "crypto. crust", 25: "ash field",
|
||
}
|
||
|
||
items = []
|
||
# Fixed display order — most common first, exotic last
|
||
order = [0, 1, 2, 3, 7, 8, 5, 6, 9, 10, 11, 12, 13, 14, 15, 16, 17,
|
||
18, 19, 20, 21, 22, 23, 25]
|
||
for cls_id in order:
|
||
if cls_id in present and cls_id in LABELS:
|
||
rgb = BIOME_RGB.get(cls_id, (128, 128, 128))
|
||
items.append((LABELS[cls_id], rgb))
|
||
|
||
# Always include river swatch if rivers exist
|
||
if terrain.get("rivers"):
|
||
items.append(("river", RIVER_RGB))
|
||
|
||
return items
|
||
|
||
|
||
def _render_legend(img: Image.Image, terrain: dict) -> Image.Image:
|
||
"""Draw biome legend strip at bottom of image."""
|
||
items = _biome_legend_items(terrain)
|
||
if not items:
|
||
return img
|
||
|
||
draw = ImageDraw.Draw(img)
|
||
sw = int(14 * UI_SCALE) # swatch width
|
||
sh = int(12 * UI_SCALE) # swatch height
|
||
pad_x = int(14 * UI_SCALE)
|
||
leg_y = OUT_H - int(34 * UI_SCALE)
|
||
font = _load_font(int(10 * UI_SCALE))
|
||
gap = int(6 * UI_SCALE)
|
||
step = int(108 * UI_SCALE)
|
||
|
||
lx = pad_x
|
||
for label, rgb in items:
|
||
if lx + step > OUT_W - pad_x:
|
||
break
|
||
draw.rectangle([(lx, leg_y), (lx + sw, leg_y + sh)], fill=rgb)
|
||
draw.text((lx + sw + gap, leg_y), label,
|
||
fill=(185, 192, 205), font=font)
|
||
lx += step
|
||
|
||
return img
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Main entry point
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def render_heightmap(body_def: dict,
|
||
terrain: dict,
|
||
out_w: int = OUT_W,
|
||
out_h: int = OUT_H,
|
||
render_mode: str = "cartographic",
|
||
chrome: bool = True) -> Image.Image:
|
||
"""
|
||
Render a 4096×2048 annotated equirectangular heightmap PNG.
|
||
|
||
Parameters
|
||
----------
|
||
body_def : dict — body definition from body_definition_parser
|
||
terrain : dict — terrain dict from planet_simulation.simulate()
|
||
out_w, out_h — output resolution (default 4096×2048)
|
||
|
||
Returns
|
||
-------
|
||
PIL.Image.Image RGB
|
||
"""
|
||
global OUT_W, OUT_H, UI_SCALE
|
||
OUT_W = out_w
|
||
OUT_H = out_h
|
||
UI_SCALE = out_w / 1024
|
||
|
||
# Set active colour mode for this render
|
||
global BIOME_RGB, OCEAN_DEEP, OCEAN_MID, OCEAN_SHALLOW, RENDER_MODE
|
||
RENDER_MODE = render_mode
|
||
BIOME_RGB = _build_biome_rgb(render_mode)
|
||
OCEAN_DEEP, OCEAN_MID, OCEAN_SHALLOW = _ocean_arrays(render_mode)
|
||
|
||
# Guard: require simulation data
|
||
required = ("elevation", "biome", "surface_water", "hillshade", "sea_level")
|
||
missing = [k for k in required if k not in terrain]
|
||
if missing:
|
||
raise ValueError(f"terrain dict missing keys: {missing}")
|
||
|
||
# 1+2+3: surface colour with elevation shading and hillshade
|
||
rgb = _render_surface(terrain)
|
||
|
||
# 4: coastline
|
||
rgb = _render_coastline(terrain, rgb)
|
||
|
||
# 5: rivers
|
||
rgb = _render_rivers(terrain, rgb)
|
||
|
||
# 6: lat/lon grid
|
||
rgb = _render_grid(rgb)
|
||
|
||
# Convert to PIL for text rendering
|
||
img = Image.fromarray(
|
||
(rgb * 255).clip(0, 255).astype(np.uint8), mode="RGB")
|
||
|
||
if chrome:
|
||
# 7: title panel
|
||
img = _render_title(img, body_def)
|
||
# 8: legend
|
||
img = _render_legend(img, terrain)
|
||
|
||
return img
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# CLI
|
||
# ---------------------------------------------------------------------------
|
||
|
||
if __name__ == "__main__":
|
||
import sys, json, time
|
||
|
||
if len(sys.argv) < 2:
|
||
print("Usage: python3 render_heightmap.py body_def.json [--small]")
|
||
sys.exit(1)
|
||
|
||
with open(sys.argv[1]) as f:
|
||
bd = json.load(f)
|
||
|
||
# --small flag renders at 1024×512 for fast iteration
|
||
small = "--small" in sys.argv
|
||
w, h = (1024, 512) if small else (OUT_W, OUT_H)
|
||
|
||
from planet_simulation import simulate
|
||
|
||
print(f"Simulating: {bd['id']} ({bd['planet_class']})")
|
||
t0 = time.time()
|
||
terrain = simulate(bd)
|
||
sim_t = time.time() - t0
|
||
|
||
if not terrain:
|
||
print("Gas giant — no heightmap.")
|
||
sys.exit(0)
|
||
|
||
print(f"Rendering heightmap {w}×{h}…")
|
||
t1 = time.time()
|
||
img = render_heightmap(bd, terrain, out_w=w, out_h=h)
|
||
ren_t = time.time() - t1
|
||
|
||
out = f"/mnt/user-data/outputs/{bd['id']}_heightmap.png"
|
||
img.save(out, format="PNG")
|
||
print(f"Saved: {out}")
|
||
print(f" simulate={sim_t:.1f}s render={ren_t:.1f}s total={sim_t+ren_t:.1f}s")
|