sum-parts-test/scripts/coarse_eval.py

#!/usr/bin/env python3
"""Score SUM Parts predictions at the granularity this project actually needs.

The benchmark reports 13 fine classes. The task here is coarser: separate a
drone-photogrammetry mesh into building / vegetation / vehicle / ground. Window
and single-wall detail is explicitly out of scope.

That difference matters for how the numbers read. A fine-grained mIoU averages
in classes like dormer, balcony and roof_installation, which a weak baseline
scores 0 on -- but every one of those collapses into "building" here, so their
individual failure costs nothing. Merging first, then scoring, measures the
thing that is actually wanted.

Mapping (SUM index -> coarse):
    1  terrain            -> ground
    2  high_vegetation    -> vegetation
    3  facade_surface     -> building
    5  car                -> vehicle
    6  boat               -> vehicle
    7  roof_surface       -> building
    8  chimney            -> building
    9  dormer             -> building
    10 balcony            -> building
    11 roof_installation  -> building
    12 wall               -> building
    0  unclassified       -> ignored
    4  water              -> ignored (not a target class here)

Usage:
    python coarse_eval.py pred.ply gt.ply [more pairs ...]
    python coarse_eval.py --pred-dir DIR --gt-dir DIR
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

import numpy as np
from plyfile import PlyData

FINE = ['unclassified', 'terrain', 'high_vegetation', 'facade_surface',
        'water', 'car', 'boat', 'roof_surface', 'chimney', 'dormer',
        'balcony', 'roof_installation', 'wall']

COARSE = ['ignored', 'ground', 'vegetation', 'building', 'vehicle']

# fine index -> coarse index (0 = ignored)
FINE_TO_COARSE = np.array([
    0,  # unclassified
    1,  # terrain           -> ground
    2,  # high_vegetation   -> vegetation
    3,  # facade_surface    -> building
    0,  # water             -> ignored
    4,  # car               -> vehicle
    4,  # boat              -> vehicle
    3,  # roof_surface      -> building
    3,  # chimney           -> building
    3,  # dormer            -> building
    3,  # balcony           -> building
    3,  # roof_installation -> building
    3,  # wall              -> building
], dtype=np.int64)

# main.py's visualization writes colours, not labels; recover the class by
# matching against the palette it used
COLOR_MAP = np.array([
    (0., 0., 0.), (170., 85., 0.), (0., 255., 0.), (255., 255., 0.),
    (0., 255., 255.), (255., 0., 255.), (0., 0., 153.), (85., 85., 127.),
    (255., 50., 50.), (85., 0., 127.), (50., 125., 150.), (50., 0., 50.),
    (215., 160., 140.),
])


def load_labels(path: Path) -> np.ndarray:
    """Fine class per point: from a `label` field, else decoded from colour."""
    v = PlyData.read(str(path))["vertex"]
    props = [p.name for p in v.properties]

    if "label" in props:
        lab = np.asarray(v["label"]).astype(np.int64).ravel()
        # a prediction ply may carry a placeholder label; fall through if so
        if lab.max() >= 0 and not (lab == lab[0]).all():
            return lab
        if lab.max() >= 0 and lab[0] >= 0:
            return lab

    rgb_set = next((s for s in (("red", "green", "blue"), ("r", "g", "b"))
                    if all(c in props for c in s)), None)
    if rgb_set is None:
        raise SystemExit(f"{path}: no label field and no colour to decode")

    rgb = np.stack([np.asarray(v[c], dtype=np.float64) for c in rgb_set], axis=1)
    if rgb.max() <= 1.0:
        rgb *= 255.0
    d = ((rgb[:, None, :] - COLOR_MAP[None, :, :]) ** 2).sum(axis=2)
    return d.argmin(axis=1).astype(np.int64)


def confusion(pred: np.ndarray, true: np.ndarray, n: int) -> np.ndarray:
    k = (true >= 0) & (true < n) & (pred >= 0) & (pred < n)
    return np.bincount(true[k] * n + pred[k], minlength=n * n).reshape(n, n)


def report(cm: np.ndarray, names: list[str], skip: set[int]) -> None:
    tp = np.diag(cm).astype(np.float64)
    actual = cm.sum(axis=1).astype(np.float64)
    predicted = cm.sum(axis=0).astype(np.float64)
    union = actual + predicted - tp

    print(f"  {'class':<12} {'IoU':>7} {'recall':>8} {'points':>12}")
    ious = []
    for i, name in enumerate(names):
        if i in skip:
            continue
        iou = 100.0 * tp[i] / union[i] if union[i] > 0 else 0.0
        rec = 100.0 * tp[i] / actual[i] if actual[i] > 0 else 0.0
        ious.append(iou)
        print(f"  {name:<12} {iou:>6.2f}% {rec:>7.2f}% {int(actual[i]):>12,}")

    scored = [i for i in range(len(names)) if i not in skip]
    oa_tp = tp[scored].sum()
    oa_n = actual[scored].sum()
    print()
    print(f"  mIoU : {np.mean(ious):.2f}%   (over {len(ious)} classes)")
    print(f"  OA   : {100.0 * oa_tp / oa_n if oa_n else 0:.2f}%")


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("pairs", nargs="*", help="pred.ply gt.ply [pred gt ...]")
    ap.add_argument("--pred-dir", type=Path)
    ap.add_argument("--gt-dir", type=Path)
    args = ap.parse_args()

    pairs: list[tuple[Path, Path]] = []
    if args.pred_dir and args.gt_dir:
        for p in sorted(args.pred_dir.glob("*_pred.ply")):
            stem = p.name.replace("_pred.ply", ".ply")
            g = args.gt_dir / stem
            if g.exists():
                pairs.append((p, g))
            else:
                print(f"warn: no ground truth for {p.name}", file=sys.stderr)
    else:
        if len(args.pairs) % 2:
            raise SystemExit("pairs must come as pred gt pred gt ...")
        pairs = [(Path(args.pairs[i]), Path(args.pairs[i + 1]))
                 for i in range(0, len(args.pairs), 2)]

    if not pairs:
        raise SystemExit("nothing to evaluate")

    cm_fine = np.zeros((13, 13), dtype=np.int64)
    cm_coarse = np.zeros((5, 5), dtype=np.int64)

    for pred_p, gt_p in pairs:
        pred = load_labels(pred_p)
        true = load_labels(gt_p)
        if len(pred) != len(true):
            print(f"warn: {pred_p.name} has {len(pred):,} points but "
                  f"{gt_p.name} has {len(true):,} -- skipped", file=sys.stderr)
            continue
        print(f"  + {gt_p.name}  ({len(true):,} pts)")
        cm_fine += confusion(pred, true, 13)
        cm_coarse += confusion(FINE_TO_COARSE[pred], FINE_TO_COARSE[true], 5)

    print()
    print("=== fine (13 SUM classes, benchmark granularity) ===")
    report(cm_fine, FINE, skip={0})

    print()
    print("=== coarse (what this project needs) ===")
    report(cm_coarse, COARSE, skip={0})


if __name__ == "__main__":
    main()