Source code: Multi-colour 3D-printed terrain with Python: one STL per colour for AMS and MMU printers

"""Multi-colour 3D-printed terrain: split a model of Yr Wyddfa (Snowdon) into one STL per colour.

Copernicus GLO-30 terrain and ESA WorldCover land cover are put on one 20 m British National Grid,
classified into four printable colours (water, woodland, open ground below 700 m, high ground
above 700 m), cleaned so that no colour region is too small to print, and cut into four watertight
parts that fit together exactly. The parts are written as separate STLs and as one 3MF for
multi-material printers such as a Bambu Lab AMS or a Prusa MMU.

TechnicalGIS tutorial: https://www.technicalgis.com/tutorials/multi-colour-3d-print-terrain/

Data: Copernicus DEM GLO-30 (© DLR e.V. 2010-2014 and © Airbus Defence and Space GmbH 2014-2018,
provided under COPERNICUS by the European Union and ESA); ESA WorldCover 10 m 2021 v200
(© ESA WorldCover project 2021, contains modified Copernicus Sentinel data (2021) processed by the
ESA WorldCover consortium, CC BY 4.0).
"""
import sys
import time
import zipfile
import xml.etree.ElementTree as ET
from importlib.metadata import version
from pathlib import Path

import numpy as np
import rasterio
from rasterio import features
from rasterio.transform import from_origin
from rasterio.warp import reproject, Resampling
from scipy import ndimage
from shapely.geometry import shape
from shapely.geometry.polygon import orient
import manifold3d as mf
import pyvista as pv
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap, LightSource, BoundaryNorm
from matplotlib.patches import Patch

plt.rcParams.update({"figure.dpi": 150, "savefig.bbox": "tight", "font.size": 9,
                     "figure.facecolor": "white", "axes.facecolor": "white"})
pv.OFF_SCREEN = True
OUT = Path("output"); OUT.mkdir(exist_ok=True)

DEM_URL = ("https://copernicus-dem-30m.s3.amazonaws.com/Copernicus_DSM_COG_10_N53_00_W005_00_DEM/"
           "Copernicus_DSM_COG_10_N53_00_W005_00_DEM.tif")
WC_URL = ("https://esa-worldcover.s3.eu-central-1.amazonaws.com/v200/2021/map/"
          "ESA_WorldCover_10m_2021_v200_N51W006_Map.tif")
AOI = (256_500, 350_000, 264_500, 357_000)      # 8 x 7 km around Yr Wyddfa, British National Grid
RES = 20.0
X0, Y0, X1, Y1 = AOI
DST = from_origin(X0, Y1, RES, RES)
ROWS, COLS = int((Y1 - Y0) / RES), int((X1 - X0) / RES)

MODEL_LONG_MM = 150.0
Z_EXAGGERATION = 1.5
BASE_MM = 3.0
HIGH_M = 700.0               # open ground above this is "high ground"
WATER_DROP_MM = 0.4          # lakes sit a little below their shores
MIN_REGION_MM2 = 1.5         # smallest colour island worth printing (about 3 x 3 nozzle widths)
CAP_MM = 1.2                 # thickness of the coloured skin in design B
LAYER_MM = 0.2

# 0 water, 1 woodland, 2 open ground, 3 high ground; colours from the Spectral palette
NAMES = ["water", "woodland", "open ground", "high ground"]
COLOURS = ["#3288bd", "#1a9850", "#abdda4", "#f2efe6"]
EDGE = "#555555"


def section(title):
    print(f"\n=== {title} ===")


# ------------------------------------------------------------------ mesh helpers
def grid_triangles(r, c):
    idx = np.arange(r * c).reshape(r, c)
    a, b = idx[:-1, :-1].ravel(), idx[:-1, 1:].ravel()
    d, e = idx[1:, :-1].ravel(), idx[1:, 1:].ravel()
    return np.concatenate([np.column_stack([a, d, b]), np.column_stack([b, d, e])])


def heightfield_solid(Z, cell_mm):
    rows, cols = Z.shape
    X, Y = np.meshgrid(np.arange(cols) * cell_mm, (rows - 1 - np.arange(rows)) * cell_mm)
    top = np.column_stack([X.ravel(), Y.ravel(), Z.ravel()])
    bot = top.copy()
    bot[:, 2] = 0.0
    n = rows * cols
    tri_top = grid_triangles(rows, cols)
    idx = np.arange(n).reshape(rows, cols)
    ring = np.concatenate([idx[0, :], idx[1:, -1], idx[-1, -2::-1], idx[-2:0:-1, 0]])
    p, q = ring, np.roll(ring, -1)
    walls = np.concatenate([np.column_stack([p, q, p + n]), np.column_stack([q, q + n, p + n])])
    return np.vstack([top, bot]), np.vstack([tri_top, tri_top[:, ::-1] + n, walls])


def to_manifold(verts, tris):
    m = mf.Manifold(mf.Mesh(vert_properties=np.ascontiguousarray(verts, dtype=np.float32),
                            tri_verts=np.ascontiguousarray(tris, dtype=np.uint32)))
    assert m.status() == mf.Error.NoError, m.status()
    return m


def from_manifold(m):
    mesh = m.to_mesh()
    return np.asarray(mesh.vert_properties)[:, :3].astype("float64"), np.asarray(mesh.tri_verts).astype("int64")


def check_mesh(name, verts, tris):
    e = np.concatenate([tris[:, [0, 1]], tris[:, [1, 2]], tris[:, [2, 0]]])
    n = len(verts)
    und = np.unique(np.sort(e, axis=1) @ np.array([n, 1]), return_counts=True)[1]
    drc = np.unique(e @ np.array([n, 1]), return_counts=True)[1]
    t = verts[tris]
    vol = np.einsum("ij,ij->i", t[:, 0], np.cross(t[:, 1], t[:, 2])).sum() / 6
    ok = bool((und == 2).all() and (drc == 1).all() and vol > 0)
    print(f"{name:<12s} {len(tris):>8,} triangles | watertight {bool((und == 2).all())} | "
          f"oriented {bool((drc == 1).all())} | volume {vol / 1000:6.2f} cm3")
    return ok, vol


def write_stl(path, verts, tris, header):
    t = verts[tris]
    nrm = np.cross(t[:, 1] - t[:, 0], t[:, 2] - t[:, 0])
    ln = np.linalg.norm(nrm, axis=1, keepdims=True)
    rec = np.zeros(len(tris), dtype=[("normal", "<f4", 3), ("v", "<f4", (3, 3)), ("attr", "<u2")])
    rec["normal"] = np.divide(nrm, ln, out=np.zeros_like(nrm), where=ln > 0)
    rec["v"] = t
    with open(path, "wb") as f:
        f.write(header.encode("ascii")[:80].ljust(80, b" "))
        f.write(np.uint32(len(rec)).tobytes())
        f.write(rec.tobytes())


def write_3mf(path, parts):
    """Minimal 3MF: one mesh object per part, each with its own display colour."""
    ns = "http://schemas.microsoft.com/3dmanufacturing/core/2015/02"
    xml = [f'<?xml version="1.0" encoding="UTF-8"?>\n<model unit="millimeter" xml:lang="en-GB" xmlns="{ns}">',
           "<resources>", '<basematerials id="1">']
    xml += [f'<base name="{name}" displaycolor="{colour}FF"/>' for name, _, _, colour in parts]
    xml.append("</basematerials>")
    for i, (name, v, t, _) in enumerate(parts):
        xml.append(f'<object id="{i + 2}" name="{name}" type="model" pid="1" pindex="{i}"><mesh><vertices>')
        xml.append("".join(f'<vertex x="{x:.4f}" y="{y:.4f}" z="{z:.4f}"/>' for x, y, z in v))
        xml.append("</vertices><triangles>")
        xml.append("".join(f'<triangle v1="{a}" v2="{b}" v3="{c}"/>' for a, b, c in t))
        xml.append("</triangles></mesh></object>")
    xml += ["</resources>", "<build>"] + [f'<item objectid="{i + 2}"/>' for i in range(len(parts))]
    xml += ["</build>", "</model>"]
    with zipfile.ZipFile(path, "w", zipfile.ZIP_DEFLATED) as z:
        z.writestr("[Content_Types].xml",
                   '<?xml version="1.0" encoding="UTF-8"?><Types xmlns="http://schemas.openxmlformats.org/'
                   'package/2006/content-types"><Default Extension="rels" ContentType="application/'
                   'vnd.openxmlformats-package.relationships+xml"/><Default Extension="model" '
                   'ContentType="application/vnd.ms-package.3dmanufacturing-3dmodel+xml"/></Types>')
        z.writestr("_rels/.rels",
                   '<?xml version="1.0" encoding="UTF-8"?><Relationships xmlns="http://schemas.openxmlformats.org/'
                   'package/2006/relationships"><Relationship Target="/3D/3dmodel.model" Id="rel0" '
                   'Type="http://schemas.microsoft.com/3dmanufacturing/2013/01/3dmodel"/></Relationships>')
        z.writestr("3D/3dmodel.model", "\n".join(xml))


def count_regions(lab):
    return {NAMES[k]: int(ndimage.label(lab == k)[1]) for k in range(len(NAMES))}


section("versions")
print(f"python {sys.version.split()[0]} | numpy {np.__version__} | rasterio {rasterio.__version__} | "
      f"shapely {version('shapely')} | manifold3d {version('manifold3d')} | pyvista {pv.__version__}")

# ------------------------------------------------------------------ Step 1: two rasters, one grid
section("Step 1: DEM and land cover on one 20 m grid")
t0 = time.perf_counter()
dem = np.full((ROWS, COLS), np.nan, dtype="float32")
wc = np.zeros((ROWS, COLS), dtype="uint8")
with rasterio.Env(GDAL_HTTP_MAX_RETRY="4", GDAL_HTTP_RETRY_DELAY="5"):
    with rasterio.open(DEM_URL) as src:
        reproject(rasterio.band(src, 1), dem, dst_transform=DST, dst_crs="EPSG:27700",
                  resampling=Resampling.bilinear, dst_nodata=np.nan)
    with rasterio.open(WC_URL) as src:
        reproject(rasterio.band(src, 1), wc, dst_transform=DST, dst_crs="EPSG:27700",
                  resampling=Resampling.mode)
print(f"grid {COLS} x {ROWS} cells of {RES:.0f} m | read and reprojected in {time.perf_counter() - t0:.1f} s")
print(f"DEM: {np.nanmin(dem):.1f} - {np.nanmax(dem):.1f} m, NaNs {int(np.isnan(dem).sum())}")
wc_names = {10: "tree cover", 20: "shrubland", 30: "grassland", 40: "cropland", 50: "built-up",
            60: "bare / sparse", 80: "permanent water", 90: "herbaceous wetland", 95: "mangroves", 100: "moss and lichen"}
vals, cnt = np.unique(wc, return_counts=True)
print("WorldCover classes: " + ", ".join(f"{wc_names.get(v, v)} {c / wc.size * 100:.1f}%" for v, c in zip(vals, cnt)))

# ------------------------------------------------------------------ Step 2: four printable colours
section("Step 2: classify into four colours and remove unprintable specks")
lab = np.where(dem >= HIGH_M, 3, 2).astype("uint8")
lab[wc == 10] = 1
lab[wc == 80] = 0
raw = lab.copy()
scale = MODEL_LONG_MM / ((max(ROWS, COLS) - 1) * RES)           # mm per metre
cell_mm = RES * scale
min_px = int(np.ceil(MIN_REGION_MM2 / cell_mm ** 2))
lab = features.sieve(lab, size=min_px, connectivity=4)
print(f"cell on the print {cell_mm:.3f} mm -> minimum region {MIN_REGION_MM2} mm2 = {min_px} cells")
before, after = count_regions(raw), count_regions(lab)
for k, name in enumerate(NAMES):
    print(f"{name:<12s} regions {before[name]:>4d} -> {after[name]:>3d} | "
          f"area {np.mean(raw == k) * 100:5.1f}% -> {np.mean(lab == k) * 100:5.1f}%")
print(f"cells changed by the sieve: {int((raw != lab).sum()):,} ({np.mean(raw != lab) * 100:.2f}%)")

# ------------------------------------------------------------------ Step 3: the terrain solid
section("Step 3: terrain solid")
Z = BASE_MM + (dem - dem.min()) * scale * Z_EXAGGERATION
water = lab == 0
Z[water] -= WATER_DROP_MM
W_mm, D_mm = (COLS - 1) * cell_mm, (ROWS - 1) * cell_mm
print(f"horizontal scale 1:{1000 / scale:,.0f} | model {W_mm:.1f} x {D_mm:.1f} x {Z.max():.1f} mm")
print(f"lakes lowered by {WATER_DROP_MM} mm at {int(water.sum()):,} vertices")
z_high = BASE_MM + (HIGH_M - dem.min()) * scale * Z_EXAGGERATION
print(f"single-extruder alternative: change filament at z = {z_high:.2f} mm for ground above {HIGH_M:.0f} m")
terrain = to_manifold(*heightfield_solid(Z, cell_mm))

# ------------------------------------------------------------------ Step 4: cut one part per colour
section("Step 4: one watertight part per colour")
# pixel (r, c) is centred on the mesh vertex at x = c * cell, y = (ROWS - 1 - r) * cell
px_to_mm = rasterio.Affine(cell_mm, 0, -cell_mm / 2, 0, -cell_mm, (ROWS - 0.5) * cell_mm)


def class_prism(k):
    contours = []
    for geom, _ in features.shapes(lab, mask=lab == k, transform=px_to_mm, connectivity=4):
        poly = orient(shape(geom), 1.0)              # outer rings anticlockwise, holes clockwise
        contours += [np.array(poly.exterior.coords)[:-1]] + [np.array(h.coords)[:-1] for h in poly.interiors]
    return mf.Manifold.extrude(mf.CrossSection(contours), Z.max() + 10).translate((0, 0, -5))


BODY = NAMES.index("open ground")                    # the most common colour carries the whole model
prisms = {k: class_prism(k) for k in range(len(NAMES)) if k != BODY}
skin = terrain - to_manifold(*heightfield_solid(Z - CAP_MM, cell_mm))   # the top CAP_MM of the terrain

# design A: every colour runs from the bed to the surface
full = {NAMES[k]: terrain ^ pr for k, pr in prisms.items()}
# design B: colours are a CAP_MM skin over a single-colour body
caps = {NAMES[k]: skin ^ pr for k, pr in prisms.items()}
for design in (full, caps):    # the body is whatever the other colours leave behind
    design[NAMES[BODY]] = terrain - mf.Manifold.batch_boolean(
        [m for name, m in design.items() if name != NAMES[BODY]], mf.OpType.Add)


def layers_with_swaps(parts, layer=LAYER_MM):
    zs = np.arange(layer / 2, Z.max(), layer)
    n_colours = np.array([sum(m.slice(z).area() > 0.01 for m in parts.values()) for z in zs])
    return len(zs), int((n_colours > 1).sum())


for label, design in [("A: full-depth colours", full), (f"B: {CAP_MM} mm colour caps", caps)]:
    n, multi = layers_with_swaps(design)
    accent = sum(m.volume() for name, m in design.items() if name != NAMES[BODY]) / 1000
    print(f"{label:<24s} accent colours {accent:5.2f} cm3 | {multi:>3d} of {n} layers ({LAYER_MM} mm) "
          f"need more than one filament")


def drop_slivers(m, min_mm3=0.05):
    """Booleans can leave zero-volume slivers where two regions touch only at a corner."""
    pieces = m.decompose()
    keep = [p for p in pieces if p.volume() >= min_mm3]
    return mf.Manifold.batch_boolean(keep, mf.OpType.Add), len(pieces) - len(keep), sum(p.volume() for p in pieces) - sum(p.volume() for p in keep)


parts = {}
for name in NAMES:
    parts[name], n_drop, v_drop = drop_slivers(caps[name])
    print(f"{name:<12s} dropped {n_drop:>3d} slivers totalling {v_drop:.4f} mm3")

meshes = {}
for name, m in parts.items():
    v, t = from_manifold(m)
    ok, vol = check_mesh(name, v, t)
    assert ok, name
    meshes[name] = (v, t, vol)
total = sum(m[2] for m in meshes.values())
names = list(parts)
overlap = sum((parts[a] ^ parts[b]).volume() for i, a in enumerate(names) for b in names[i + 1:])
print(f"sum of parts {total / 1000:.3f} cm3 | whole terrain {terrain.volume() / 1000:.3f} cm3 | "
      f"pairwise overlap {overlap:.4f} mm3")
for name, m in parts.items():
    print(f"{name:<12s} {len(m.decompose()):>3d} separate piece(s)")

# ------------------------------------------------------------------ Step 5: files
section("Step 5: write STLs and a 3MF")
for k, name in enumerate(NAMES):
    v, t, _ = meshes[name]
    fn = f"yr_wyddfa_{k + 1}_{name.replace(' ', '_')}.stl"
    write_stl(OUT / fn, v, t, f"TechnicalGIS Yr Wyddfa {name} - Copernicus GLO-30, ESA WorldCover")
    print(f"{fn:<34s} {(OUT / fn).stat().st_size / 1e6:5.1f} MB")
write_3mf(OUT / "yr_wyddfa_multicolour.3mf",
          [(name, *meshes[name][:2], COLOURS[k]) for k, name in enumerate(NAMES)])
with zipfile.ZipFile(OUT / "yr_wyddfa_multicolour.3mf") as z:
    model = ET.fromstring(z.read("3D/3dmodel.model"))
objs = model.findall(".//{http://schemas.microsoft.com/3dmanufacturing/core/2015/02}object")
print(f"yr_wyddfa_multicolour.3mf          {(OUT / 'yr_wyddfa_multicolour.3mf').stat().st_size / 1e6:5.1f} MB | "
      f"parts: {', '.join(o.get('name') for o in objs)}")

# ------------------------------------------------------------------ Figures
section("figures")
ext = (X0 / 1000, X1 / 1000, Y0 / 1000, Y1 / 1000)
hs = LightSource(azdeg=315, altdeg=45).hillshade(dem, vert_exag=1.5, dx=RES, dy=RES)
wc_cols = {10: "#1a9850", 20: "#a6d96a", 30: "#d9ef8b", 60: "#bababa", 80: "#3288bd", 90: "#66c2a5",
           50: "#d53e4f", 40: "#fee08b", 100: "#e6f598"}
present = [v for v in vals if v in wc_cols]
wc_idx = np.searchsorted(present, wc)
fig, axes = plt.subplots(1, 2, figsize=(10, 4.4))
im = axes[0].imshow(dem, cmap="Spectral_r", extent=ext)
axes[0].imshow(hs, cmap="gray", extent=ext, alpha=0.35)
plt.colorbar(im, ax=axes[0], shrink=0.8, label="Elevation (m)")
axes[0].set_title("Copernicus GLO-30, reprojected to 20 m")
axes[1].imshow(wc_idx, cmap=ListedColormap([wc_cols[v] for v in present]), extent=ext, interpolation="nearest")
axes[1].imshow(hs, cmap="gray", extent=ext, alpha=0.25)
axes[1].legend(handles=[Patch(color=wc_cols[v], label=wc_names[v]) for v in present], fontsize=7,
               loc="upper left", bbox_to_anchor=(1.01, 1))
axes[1].set_title("ESA WorldCover 2021, mode-resampled to 20 m")
for ax in axes:
    ax.set_xlabel("Easting (km, BNG)")
axes[0].set_ylabel("Northing (km)")
fig.savefig(OUT / "01-dem-and-land-cover.png"); plt.close(fig)

cmap4 = ListedColormap(COLOURS)
norm4 = BoundaryNorm(np.arange(-0.5, 4.5), 4)
fig, axes = plt.subplots(1, 2, figsize=(10, 4.4))
for ax, arr, title, counts in [(axes[0], raw, "Before: raw classes", before),
                               (axes[1], lab, f"After: regions under {MIN_REGION_MM2} mm² on the print removed", after)]:
    ax.imshow(arr, cmap=cmap4, norm=norm4, extent=ext, interpolation="nearest")
    ax.imshow(hs, cmap="gray", extent=ext, alpha=0.2)
    ax.set_title(f"{title}\n{sum(counts.values()):,} colour regions")
    ax.set_xlabel("Easting (km, BNG)")
axes[0].set_ylabel("Northing (km)")
axes[1].legend(handles=[Patch(color=c, label=n) for n, c in zip(NAMES, COLOURS)], fontsize=8,
               loc="upper left", bbox_to_anchor=(1.01, 1))
fig.savefig(OUT / "02-colour-regions-sieve.png"); plt.close(fig)


def stl_path(k):
    return OUT / f"yr_wyddfa_{k + 1}_{NAMES[k].replace(' ', '_')}.stl"


def place_camera(p, elev_deg=35, azim_deg=-60, dist_factor=2.0):
    c = np.array((W_mm / 2, D_mm / 2, Z.max() / 3))
    el, az = np.radians(elev_deg), np.radians(azim_deg)
    eye = c + dist_factor * max(W_mm, D_mm) * np.array([np.cos(el) * np.cos(az), np.cos(el) * np.sin(az), np.sin(el)])
    p.camera_position = [tuple(eye), tuple(c), (0, 0, 1)]
    p.camera.view_angle = 30


def save_cropped(p, png, pad=20):
    img = p.screenshot(return_img=True)
    ink = np.where((img < 250).any(axis=2))
    r0, r1 = max(ink[0].min() - pad, 0), min(ink[0].max() + pad, img.shape[0])
    c0, c1 = max(ink[1].min() - pad, 0), min(ink[1].max() + pad, img.shape[1])
    plt.imsave(OUT / png, img[r0:r1, c0:c1])
    p.close()


# hero: the four STLs read back from disk and shown together, as the slicer will see them
p = pv.Plotter(off_screen=True, window_size=(1800, 1300))
p.set_background("white")
for k in range(len(NAMES)):
    p.add_mesh(pv.read(stl_path(k)), color=COLOURS[k], ambient=0.25, diffuse=0.75, specular=0.1)
p.enable_lightkit()
place_camera(p)
save_cropped(p, "03-stl-preview-multicolour.png")

# one panel per part, each read from its own STL, inside the outline of the full model
box = pv.Box(bounds=(0, W_mm, 0, D_mm, 0, float(Z.max()))).extract_feature_edges()
p = pv.Plotter(off_screen=True, window_size=(1600, 1200), shape=(2, 2), border=False)
for k, name in enumerate(NAMES):
    p.subplot(k // 2, k % 2)
    p.set_background("white")
    p.add_mesh(pv.read(stl_path(k)), color=COLOURS[k], ambient=0.3, diffuse=0.7)
    p.add_mesh(box, color="#aaaaaa", line_width=1)
    p.add_text(f"{k + 1}. {name}", font_size=14, color="#222222", position="upper_left")
    p.enable_lightkit()
    place_camera(p, dist_factor=2.1)
save_cropped(p, "04-stl-parts.png")
print("wrote 01-dem-and-land-cover, 02-colour-regions-sieve, 03-stl-preview-multicolour, 04-stl-parts")