Milestone 4+5: UV unwrap (xatlas), symmetry, and normal map baking (nvdiffrast)

Milestone 4 — Stage 2D-2E:
- stage2_symmetry.py: bilateral-X/Y symmetry via half-mesh mirror + merge
- stage2_uv.py: xatlas atlas generation with configurable texels_per_unit
and padding; exports unwrapped.glb with TEXCOORD_0

Milestone 5 — Stage 2F:
- stage2_bake_normal.py: full tangent-space normal map baker
- Rasterizes low-poly in UV space via nvdiffrast
- Interpolates per-pixel world position + TBN frame
- Queries high-poly nearest-surface normals via trimesh.proximity
- Transforms to tangent space, packs RGB, dilates past UV islands
- Saves both GL (Y-up) and DX/UE5 (Y-down) variants
- @spaces.GPU(duration=120)

app.py: wire M4/M5 steps in run_post_process (removed from stubs)
requirements.txt: add xatlas, scipy, nvdiffrast (PyPI version)

app.py

@@ -125,11 +125,39 @@ def run_post_process(
         except Exception as e:
             log.append(f"⚠️ Decimation error: {e}")
 
+    if do_symmetry:
+        try:
+            from src.stages.stage2_symmetry import apply_symmetry
+            sym_axis = "bilateral-X"  # default; exposed via Advanced in M4 polish
+            current_glb, msg = apply_symmetry(current_glb, sym_axis)
+            log.append(f"✅ {msg}")
+        except Exception as e:
+            log.append(f"⚠️ Symmetry error: {e}")
+
+    if do_unwrap:
+        try:
+            from src.stages.stage2_uv import unwrap_uvs
+            current_glb, msg = unwrap_uvs(current_glb)
+            log.append(f"✅ {msg}")
+        except Exception as e:
+            log.append(f"⚠️ UV unwrap error: {e}")
+
+    if do_normal_bake:
+        try:
+            from src.stages.stage2_bake_normal import bake_normal_map
+            st = workspace.get_state()
+            hp = st.high_poly_glb
+            lo = st.unwrapped_glb or current_glb
+            if not hp or not hp.exists():
+                log.append("⚠️ Normal bake: no high-poly GLB. Generate first.")
+            else:
+                _gl, _dx, msg = bake_normal_map(hp, lo, map_size=2048, dx_format=True)
+                log.append(f"✅ {msg}")
+        except Exception as e:
+            log.append(f"⚠️ Normal bake error: {e}")
+
     # Later milestones — stubs with informative messages
     pending = []
-    if do_symmetry:    pending.append("Symmetry (M4)")
-    if do_unwrap:      pending.append("UV unwrap (M4)")
-    if do_normal_bake: pending.append("Normal bake (M5)")
     if do_albedo:      pending.append("Albedo bake (M6)")
     if do_material:    pending.append("Material bake (M6)")
     if do_ao:          pending.append("AO bake (M6)")

requirements.txt

@@ -32,7 +32,14 @@ trimesh[easy]
 pymeshfix
 pymeshlab
 fast-simplification
-# xatlas        — Milestone 4
+
+# ===== Milestone 4: Stage 2D-2E UV + Symmetry =====
+xatlas
+scipy               # normal map dilation
+
+# ===== Milestone 5: Stage 2F Normal Baking =====
+nvdiffrast          # PyPI version — compiles against installed CUDA at runtime
+
 # coacd==1.0.4  — Milestone 8
 
 # ===== Milestone 4: Stage 2 GPU baking =====

src/stages/stage2_bake_normal.py

@@ -0,0 +1,189 @@
+"""
+Stage 2F — Normal map baking with nvdiffrast.
+
+Pipeline:
+  1. Convert low-poly UV coords to clip space.
+  2. Rasterize the low-poly in UV-space → per-pixel triangle + bary coords.
+  3. Interpolate world positions and per-vertex TBN frame at each pixel.
+  4. Query the high-poly mesh for the nearest-surface normal at each pixel's
+     world position (trimesh proximity).
+  5. Transform that world-space normal into the low-poly's tangent space.
+  6. Pack as RGB [0, 255], dilate past UV island edges.
+  7. Save both GL (Y-up) and DX (Y-down / UE5) variants.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import spaces
+
+from src import workspace
+from src.workspace import CURRENT
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _vertex_tangents(verts: np.ndarray, faces: np.ndarray, uvs: np.ndarray) -> np.ndarray:
+    """Compute per-vertex tangent vectors from UV parametrisation."""
+    v0, v1, v2 = verts[faces[:, 0]], verts[faces[:, 1]], verts[faces[:, 2]]
+    u0, u1, u2 = uvs[faces[:, 0]], uvs[faces[:, 1]], uvs[faces[:, 2]]
+
+    e1, e2 = v1 - v0, v2 - v0
+    du1, dv1 = u1[:, 0] - u0[:, 0], u1[:, 1] - u0[:, 1]
+    du2, dv2 = u2[:, 0] - u0[:, 0], u2[:, 1] - u0[:, 1]
+
+    denom = du1 * dv2 - du2 * dv1
+    r = np.where(np.abs(denom) > 1e-10, 1.0 / denom, 0.0)
+
+    T = (dv2[:, None] * e1 - dv1[:, None] * e2) * r[:, None]
+
+    tangents = np.zeros_like(verts, dtype=np.float64)
+    np.add.at(tangents, faces[:, 0], T)
+    np.add.at(tangents, faces[:, 1], T)
+    np.add.at(tangents, faces[:, 2], T)
+
+    norms = np.linalg.norm(tangents, axis=1, keepdims=True)
+    return (tangents / np.where(norms < 1e-8, 1.0, norms)).astype(np.float32)
+
+
+def _dilate(img: np.ndarray, mask: np.ndarray, iterations: int = 8) -> np.ndarray:
+    """Push pixel values outward from UV islands using nearest-neighbour dilation."""
+    from scipy.ndimage import binary_dilation, convolve
+
+    result = img.astype(np.float32)
+    valid = mask.astype(bool)
+    kernel = np.array([[0, 1, 0], [1, 0, 1], [0, 1, 0]], dtype=np.float32)
+
+    for _ in range(iterations):
+        expanded = binary_dilation(valid)
+        border = expanded & ~valid
+        if not border.any():
+            break
+        for c in range(3):
+            weighted = convolve(result[:, :, c] * valid, kernel)
+            weight_n = convolve(valid.astype(np.float32), kernel)
+            blended = np.where(weight_n > 0, weighted / weight_n, 0.0)
+            result[:, :, c] = np.where(border, blended, result[:, :, c])
+        valid = expanded
+
+    return result.clip(0, 255).astype(np.uint8)
+
+
+# ---------------------------------------------------------------------------
+# GPU bake
+# ---------------------------------------------------------------------------
+
+@spaces.GPU(duration=120)
+def bake_normal_map(
+    high_poly_path: Path,
+    unwrapped_glb_path: Path,
+    map_size: int = 2048,
+    dx_format: bool = True,
+) -> tuple[Path, Path, str]:
+    """Bake tangent-space normal maps. Returns (gl_path, dx_path, message)."""
+    import torch
+    import trimesh
+    import nvdiffrast.torch as dr
+    from PIL import Image
+
+    device = torch.device("cuda")
+    ctx = dr.RasterizeCudaContext()
+
+    # ---- Load meshes -------------------------------------------------------
+    hi = trimesh.load(str(high_poly_path), force="mesh")
+
+    lo = trimesh.load(str(unwrapped_glb_path), force="mesh")
+    if not isinstance(lo.visual, trimesh.visual.TextureVisuals) or lo.visual.uv is None:
+        raise ValueError("Low-poly has no UV coordinates — run UV unwrap first.")
+
+    uvs = np.array(lo.visual.uv, dtype=np.float32)
+    lo_v = np.array(lo.vertices, dtype=np.float32)
+    lo_f = np.array(lo.faces, dtype=np.int32)
+    lo_n = np.array(lo.vertex_normals, dtype=np.float32)
+    lo_t = _vertex_tangents(lo_v.astype(np.float64), lo_f, uvs)
+    lo_b = np.cross(lo_n, lo_t).astype(np.float32)
+    bn = np.linalg.norm(lo_b, axis=1, keepdims=True)
+    lo_b /= np.where(bn < 1e-8, 1.0, bn)
+
+    # ---- UV → clip space ---------------------------------------------------
+    # UV (u, v) in [0,1] → NDC (2u-1, 1-2v, 0, 1)
+    pos_clip = np.zeros((len(uvs), 4), dtype=np.float32)
+    pos_clip[:, 0] = uvs[:, 0] * 2.0 - 1.0
+    pos_clip[:, 1] = (1.0 - uvs[:, 1]) * 2.0 - 1.0
+    pos_clip[:, 3] = 1.0
+
+    v_clip = torch.tensor(pos_clip, device=device).unsqueeze(0)   # [1, N, 4]
+    faces_t = torch.tensor(lo_f, device=device, dtype=torch.int32) # [M, 3]
+
+    # ---- Rasterize ---------------------------------------------------------
+    rast, _ = dr.rasterize(ctx, v_clip, faces_t, resolution=[map_size, map_size])
+    # rast: [1, H, W, 4]; last channel = tri_id + 1 (0 = background)
+
+    # ---- Interpolate world pos + TBN frame ---------------------------------
+    lo_v_t = torch.tensor(lo_v, device=device).unsqueeze(0)
+    T_t = torch.tensor(lo_t, device=device).unsqueeze(0)
+    B_t = torch.tensor(lo_b, device=device).unsqueeze(0)
+    N_t = torch.tensor(lo_n, device=device).unsqueeze(0)
+
+    world_pos, _ = dr.interpolate(lo_v_t, rast, faces_t)  # [1, H, W, 3]
+    T_px, _ = dr.interpolate(T_t, rast, faces_t)
+    B_px, _ = dr.interpolate(B_t, rast, faces_t)
+    N_px, _ = dr.interpolate(N_t, rast, faces_t)
+
+    mask = (rast[..., 3] > 0).squeeze(0).cpu().numpy()   # [H, W]
+    wp = world_pos.squeeze(0).cpu().numpy()               # [H, W, 3]
+    T_np = T_px.squeeze(0).cpu().numpy()
+    B_np = B_px.squeeze(0).cpu().numpy()
+    N_np = N_px.squeeze(0).cpu().numpy()
+
+    # ---- Query high-poly normals at pixel world positions ------------------
+    ys, xs = np.where(mask)
+    if len(ys) == 0:
+        raise ValueError("No UV-covered pixels — check that the mesh has valid UVs.")
+
+    query_pts = wp[ys, xs]
+    prox = trimesh.proximity.ProximityQuery(hi)
+    _, _, tri_ids = prox.on_surface(query_pts)
+    world_normals = hi.face_normals[tri_ids].astype(np.float32)
+
+    # ---- Transform to tangent space ----------------------------------------
+    T_q = T_np[ys, xs]
+    B_q = B_np[ys, xs]
+    N_q = N_np[ys, xs]
+
+    for v in (T_q, B_q, N_q):
+        norms = np.linalg.norm(v, axis=1, keepdims=True)
+        v /= np.where(norms < 1e-8, 1.0, norms)
+
+    ts_x = (world_normals * T_q).sum(axis=1)
+    ts_y = (world_normals * B_q).sum(axis=1)
+    ts_z = (world_normals * N_q).sum(axis=1)
+    ts = np.stack([ts_x, ts_y, ts_z], axis=1)
+
+    # ---- Pack + dilate -----------------------------------------------------
+    normal_map = np.full((map_size, map_size, 3), 127, dtype=np.uint8)
+    packed = np.clip(ts * 127.5 + 127.5, 0, 255).astype(np.uint8)
+    normal_map[ys, xs] = packed
+    normal_map = _dilate(normal_map, mask)
+
+    # GL variant (Y-up, default for most renderers)
+    gl_map = normal_map
+    # DX variant (Y-down, UE5)
+    dx_map = normal_map.copy()
+    dx_map[:, :, 1] = 255 - dx_map[:, :, 1]
+
+    tex_dir = CURRENT / "textures"
+    tex_dir.mkdir(parents=True, exist_ok=True)
+    gl_path = tex_dir / "normal_gl.png"
+    dx_path = tex_dir / "normal_dx.png"
+    Image.fromarray(gl_map).save(str(gl_path))
+    Image.fromarray(dx_map).save(str(dx_path))
+
+    state = workspace.get_state()
+    state.normal_gl_png = gl_path
+    state.normal_dx_png = dx_path
+
+    return gl_path, dx_path, f"Normal maps baked: {map_size}×{map_size}"

src/stages/stage2_symmetry.py

@@ -0,0 +1,65 @@
+"""Stage 2D — Bilateral symmetry enforcement."""
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+
+from src import workspace
+from src.workspace import CURRENT
+
+
+def apply_symmetry(input_glb: Path, axis: str) -> tuple[Path, str]:
+    """
+    Enforce bilateral symmetry. Takes faces from the negative-axis half
+    (e.g. x<=0 for bilateral-X) and mirrors them to fill the other half.
+
+    Works best when the mesh is already roughly centred at the symmetry plane.
+    """
+    if axis == "off":
+        return input_glb, "Symmetry: skipped"
+
+    import trimesh
+
+    mesh = trimesh.load(str(input_glb), force="mesh")
+    v = np.array(mesh.vertices, dtype=np.float64)
+    f = np.array(mesh.faces, dtype=np.int32)
+
+    ax_map = {"bilateral-X": 0, "bilateral": 0, "bilateral-Y": 1}
+    if axis not in ax_map:
+        return input_glb, f"Symmetry: '{axis}' not supported"
+    ax = ax_map[axis]
+
+    # Pick faces whose centroid is on the negative side
+    centroids = v[f].mean(axis=1)
+    keep = centroids[:, ax] <= 0
+    if not keep.any():
+        keep = centroids[:, ax] >= 0  # fallback: use positive side
+
+    kept_f = f[keep]
+    used = np.unique(kept_f)
+    remap = np.full(len(v), -1, dtype=np.int32)
+    remap[used] = np.arange(len(used))
+
+    half_v = v[used]
+    half_f = remap[kept_f]
+
+    mirror_v = half_v.copy()
+    mirror_v[:, ax] *= -1
+    mirror_f = half_f[:, [0, 2, 1]]  # flip winding for mirrored half
+
+    combined_v = np.vstack([half_v, mirror_v])
+    combined_f = np.vstack([half_f, mirror_f + len(half_v)])
+
+    sym = trimesh.Trimesh(vertices=combined_v, faces=combined_f, process=False)
+    sym.merge_vertices(merge_tex=False)
+    sym.remove_duplicate_faces()
+    sym.remove_unreferenced_vertices()
+
+    out_path = CURRENT / "symmetric.glb"
+    sym.export(str(out_path))
+
+    state = workspace.get_state()
+    state.cleaned_glb = out_path  # replaces the pre-UV working mesh
+
+    return out_path, f"Symmetry ({axis}): {len(sym.faces):,} faces"

src/stages/stage2_uv.py

@@ -0,0 +1,60 @@
+"""Stage 2E — UV unwrap with xatlas."""
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+
+from src import workspace
+from src.workspace import CURRENT
+
+
+def unwrap_uvs(
+    input_glb: Path,
+    texels_per_unit: float = 256.0,
+    padding: int = 2,
+) -> tuple[Path, str]:
+    """
+    Generate a UV atlas with xatlas and attach the coordinates to the mesh.
+    Output GLB has TEXCOORD_0 set — required before any texture baking step.
+    """
+    import trimesh
+    import xatlas
+
+    mesh = trimesh.load(str(input_glb), force="mesh")
+
+    positions = np.array(mesh.vertices, dtype=np.float32)
+    indices = np.array(mesh.faces, dtype=np.uint32)
+    normals = np.array(mesh.vertex_normals, dtype=np.float32)
+
+    atlas = xatlas.Atlas()
+    atlas.add_mesh(positions, indices, normals)
+
+    pack_opts = xatlas.PackOptions()
+    pack_opts.padding = padding
+    pack_opts.texels_per_unit = texels_per_unit
+
+    atlas.generate(chart_options=xatlas.ChartOptions(), pack_options=pack_opts)
+
+    vmapping, new_indices, uvs = atlas[0]
+    # vmapping: [new_N] — original vertex index for each new vertex
+    # new_indices: [new_F * 3] flat
+    # uvs: [new_N, 2] in [0, 1]
+
+    new_verts = positions[vmapping]
+    new_faces = new_indices.reshape(-1, 3).astype(np.int32)
+
+    unwrapped = trimesh.Trimesh(vertices=new_verts, faces=new_faces, process=False)
+    unwrapped.visual = trimesh.visual.TextureVisuals(uv=uvs.astype(np.float32))
+
+    out_path = CURRENT / "unwrapped.glb"
+    unwrapped.export(str(out_path))
+
+    state = workspace.get_state()
+    state.unwrapped_glb = out_path
+
+    atlas_w, atlas_h = atlas.width, atlas.height
+    return out_path, (
+        f"UV unwrap: {len(new_faces):,} faces · "
+        f"atlas {atlas_w}×{atlas_h} · {texels_per_unit:.0f} texels/unit"
+    )