sum-parts-test/scripts/mesh_to_ply.py

#!/usr/bin/env python3
"""Convert a textured mesh into the PLY point cloud SUM Parts expects.

Target schema, from openpoints/dataset/sumv2_triangle/sumv2_triangle.py
(read_ply_with_plyfilelib) and confirmed against the shipped demo tile:

    element vertex N
    property float x / y / z
    property uchar red / green / blue     (the demo uses r/g/b; the loader
                                           accepts either spelling)
    property uchar label                  0 == unclassified, which cfg ignores
    property int   object_index           optional

Written for ContextCapture-style aerial OBJ exports, which are the shape the
Seosan Myeongcheon blocks come in:

  * multi-material -- one texture atlas per material (17-21 per block), so the
    mesh has to be sampled per material, not as one merged blob. Merging with
    force='mesh' silently drops every texture and the output comes out grey.
  * no vertex colours -- the texture is the ONLY colour source. There is no
    fallback to degrade to.
  * coordinates already local and metric, with the true origin recorded in
    metadata.xml as SRSOrigin (EPSG:5186+9999). Same convention SUM Parts
    itself ships, so no reprojection is needed -- but never feed raw easting/
    northing or degrees, since every cfg radius and voxel size is in metres.

--bbox crops before sampling, which is how you get one ~252 m SUM-sized tile
out of a 465 x 487 m block.

Usage:
    python mesh_to_ply.py Block.obj out.ply --points 500000
    python mesh_to_ply.py Block.obj tile.ply --bbox 300 700 552 952
"""
from __future__ import annotations

import argparse
import sys
import time
from pathlib import Path

import numpy as np


def die(msg: str) -> None:
    print(f"error: {msg}", file=sys.stderr)
    raise SystemExit(1)


def log(msg: str) -> None:
    print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)


def load_scene(path: Path):
    """Load as a Scene so each material keeps its own texture."""
    try:
        import trimesh
    except ImportError:
        die("trimesh not installed. run: pip install trimesh pillow")

    log(f"loading {path.name} ({path.stat().st_size / 1e6:,.0f} MB) ...")
    scene = trimesh.load(path, process=False)

    if isinstance(scene, trimesh.Trimesh):
        geoms = [scene]
    else:
        geoms = list(scene.geometry.values())

    geoms = [g for g in geoms if getattr(g, "faces", None) is not None and len(g.faces)]
    if not geoms:
        die(f"{path} has no faces")

    log(f"loaded {len(geoms)} geometry group(s), "
        f"{sum(len(g.faces) for g in geoms):,} faces total")
    return geoms


def crop(geom, bbox):
    """Keep faces whose centroid falls inside the XY bbox. Returns a mask."""
    if bbox is None:
        return np.ones(len(geom.faces), dtype=bool)
    x0, y0, x1, y1 = bbox
    c = geom.vertices[geom.faces].mean(axis=1)
    return (c[:, 0] >= x0) & (c[:, 0] < x1) & (c[:, 1] >= y0) & (c[:, 1] < y1)


def sample_geom(geom, face_mask, n_points, seed):
    """Area-weighted sampling restricted to the masked faces, with colour."""
    import trimesh

    idx = np.flatnonzero(face_mask)
    if idx.size == 0 or n_points <= 0:
        return np.empty((0, 3)), np.empty((0, 3), dtype=np.uint8)

    tri = geom.vertices[geom.faces[idx]]                       # (m, 3, 3)
    area = trimesh.triangles.area(tri)
    total = area.sum()
    if total <= 0:
        return np.empty((0, 3)), np.empty((0, 3), dtype=np.uint8)

    rng = np.random.default_rng(seed)
    pick = rng.choice(idx.size, size=n_points, p=area / total)

    # uniform barycentric coordinates over each chosen triangle
    r1, r2 = rng.random(n_points), rng.random(n_points)
    s = np.sqrt(r1)
    bary = np.stack([1 - s, s * (1 - r2), s * r2], axis=1)     # (n, 3)

    xyz = (tri[pick] * bary[:, :, None]).sum(axis=1)
    rgb = colour(geom, idx[pick], bary)
    return xyz, rgb


def colour(geom, face_idx, bary) -> np.ndarray:
    """Per-point RGB from the geometry's texture; grey if it has none."""
    n = len(face_idx)
    visual = getattr(geom, "visual", None)

    uv = getattr(visual, "uv", None)
    material = getattr(visual, "material", None)
    image = getattr(material, "image", None) if material is not None else None

    if uv is not None and image is not None:
        try:
            tri_uv = np.asarray(uv)[geom.faces[face_idx]]       # (n, 3, 2)
            point_uv = (tri_uv * bary[:, :, None]).sum(axis=1)
            img = np.asarray(image.convert("RGB"))
            h, w = img.shape[:2]
            # OBJ UV origin is bottom-left; image rows run top-down
            px = np.clip((point_uv[:, 0] % 1.0) * (w - 1), 0, w - 1).astype(np.int32)
            py = np.clip((1.0 - point_uv[:, 1] % 1.0) * (h - 1), 0, h - 1).astype(np.int32)
            return img[py, px].astype(np.uint8)
        except Exception as e:  # noqa: BLE001 - exporters vary wildly
            print(f"warn: texture sampling failed ({e})", file=sys.stderr)

    vc = getattr(visual, "vertex_colors", None)
    if vc is not None and len(vc) == len(geom.vertices):
        vc = np.asarray(vc)[:, :3].astype(np.float64)
        return (vc[geom.faces[face_idx]] * bary[:, :, None]).sum(axis=1).astype(np.uint8)

    return np.full((n, 3), 128, dtype=np.uint8)


def write_ply(path: Path, xyz, rgb, label: int, object_index: int | None,
              colour_style: str = "sum") -> None:
    """Write the point cloud.

    colour_style='sum' reproduces exactly what SUM Parts ships: fields named
    r/g/b holding float32 in [0, 1]. That match is not cosmetic. The loader
    normalises nothing --

        rgb = np.stack([...'red','green','blue'...]).astype(np.uint8)
        except ValueError:
            rgb = np.stack([...'r','g','b'...]).astype(np.float32)
        if np.max(rgb) > 1:
            rgb = rgb            # <- no-op, despite how it reads

    -- and the active datatransforms in cfgs/sumv2_triangle/default.yaml are
    [PointsToTensor, PointCloudScaling, PointCloudRotation, PointCloudJitter];
    NumpyChromaticNormalize is commented out. So whatever is in the file reaches
    the network unscaled. Writing 0-255 where the training data was 0-1 hands
    the model colour features 255x too large.

    colour_style='uint8' writes red/green/blue as uint8 instead, for viewers
    that expect the conventional PLY spelling.
    """
    try:
        from plyfile import PlyData, PlyElement
    except ImportError:
        die("plyfile not installed. run: pip install plyfile")

    if colour_style == "sum":
        cols = [("r", "f4"), ("g", "f4"), ("b", "f4")]
        vals = (rgb.astype(np.float32) / 255.0)
    else:
        cols = [("red", "u1"), ("green", "u1"), ("blue", "u1")]
        vals = rgb

    # demo tiles store label as int32
    dtype = [("x", "f4"), ("y", "f4"), ("z", "f4"), *cols, ("label", "i4")]
    if object_index is not None:
        dtype.append(("object_index", "i4"))

    arr = np.empty(len(xyz), dtype=dtype)
    arr["x"], arr["y"], arr["z"] = xyz[:, 0], xyz[:, 1], xyz[:, 2]
    for i, (name, _) in enumerate(cols):
        arr[name] = vals[:, i]
    arr["label"] = label
    if object_index is not None:
        arr["object_index"] = object_index

    path.parent.mkdir(parents=True, exist_ok=True)
    PlyData([PlyElement.describe(arr, "vertex")], text=False).write(str(path))


def main() -> None:
    ap = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("input", type=Path)
    ap.add_argument("output", type=Path)
    ap.add_argument("--points", type=int, default=500_000,
                    help="total points to sample (default: 500000)")
    ap.add_argument("--bbox", type=float, nargs=4, metavar=("X0", "Y0", "X1", "Y1"),
                    help="crop to this XY box in the mesh's own local coords")
    ap.add_argument("--label", type=int, default=0,
                    help="constant label to stamp; 0 = unclassified (default: 0)")
    ap.add_argument("--object-index", type=int, default=None)
    ap.add_argument("--colour-style", choices=("sum", "uint8"), default="sum",
                    help="'sum' (default): r/g/b float32 0-1, byte-for-byte what "
                         "SUM Parts ships. 'uint8': red/green/blue 0-255 for viewers")
    ap.add_argument("--keep-origin", action="store_true",
                    help="do NOT translate the result to a local origin")
    ap.add_argument("--seed", type=int, default=0)
    args = ap.parse_args()

    if not args.input.exists():
        die(f"{args.input} not found")
    if not 0 <= args.label <= 255:
        die("--label must fit in a uint8 (0-255)")

    geoms = load_scene(args.input)

    # Budget points per geometry by cropped surface area, so a material that
    # covers half the tile gets half the points.
    import trimesh
    masks, areas = [], []
    for g in geoms:
        m = crop(g, args.bbox)
        masks.append(m)
        a = trimesh.triangles.area(g.vertices[g.faces[m]]).sum() if m.any() else 0.0
        areas.append(a)

    total_area = float(sum(areas))
    if total_area <= 0:
        die("nothing left after cropping -- check --bbox against the mesh bbox")
    log(f"cropped surface area: {total_area:,.0f} m^2 across "
        f"{sum(int(m.sum()) for m in masks):,} faces")

    parts_xyz, parts_rgb = [], []
    for i, (g, m, a) in enumerate(zip(geoms, masks, areas)):
        if a <= 0:
            continue
        n = int(round(args.points * a / total_area))
        if n <= 0:
            continue
        xyz, rgb = sample_geom(g, m, n, args.seed + i)
        if len(xyz):
            parts_xyz.append(xyz)
            parts_rgb.append(rgb)
        log(f"  geom {i:>3}: {int(m.sum()):>8,} faces  {a:>12,.0f} m^2  -> {len(xyz):>8,} pts")

    if not parts_xyz:
        die("no points sampled")

    xyz = np.vstack(parts_xyz)
    rgb = np.vstack(parts_rgb)

    lo, hi = xyz.min(axis=0), xyz.max(axis=0)
    log(f"extent: {np.round(hi - lo, 2)} m   origin: {np.round(lo, 2)}")

    if not args.keep_origin:
        xyz = xyz - lo
        log("translated to local origin (min corner -> 0,0,0)")

    if (hi - lo).max() < 0.5:
        print("warn: extent under 0.5 units -- looks like degrees, not metres. "
              "Reproject to a metric CRS (e.g. EPSG:5186) first.", file=sys.stderr)

    grey = int((rgb == 128).all(axis=1).sum())
    if grey:
        print(f"warn: {grey:,}/{len(rgb):,} points fell back to grey "
              "(missing texture)", file=sys.stderr)

    write_ply(args.output, xyz, rgb, args.label, args.object_index, args.colour_style)
    log(f"wrote {args.output}  ({len(xyz):,} points, label={args.label}, "
        f"colour={args.colour_style})")


if __name__ == "__main__":
    main()