evaluate_pointcloud.py

"""Evaluate exported Simam3D points against a reference point cloud."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np


def load_points(path: str | Path) -> np.ndarray:
    """Load Nx3 points from NumPy `.npy` or ASCII PLY files."""
    path = Path(path)
    if path.suffix.lower() == ".npy":
        points = np.load(path)
    elif path.suffix.lower() == ".ply":
        lines = path.read_text(encoding="utf-8").splitlines()
        try:
            end = lines.index("end_header")
            count_line = next(line for line in lines[:end] if line.startswith("element vertex "))
            count = int(count_line.split()[-1])
            if not any(line.startswith("format ascii") for line in lines[:end]):
                raise ValueError("only ASCII PLY is supported")
            points = np.asarray([[float(value) for value in line.split()[:3]] for line in lines[end + 1:end + 1 + count]], dtype=np.float64)
        except (StopIteration, ValueError, IndexError) as exc:
            raise ValueError(f"invalid ASCII PLY point cloud: {path}") from exc
    else:
        raise ValueError("point cloud must be .npy or ASCII .ply")
    points = np.asarray(points, dtype=np.float64)
    if points.ndim != 2 or points.shape[1] < 3:
        raise ValueError("point cloud must have shape Nx3")
    points = points[:, :3]
    points = points[np.isfinite(points).all(axis=1)]
    if not len(points):
        raise ValueError("point cloud has no finite points")
    return points


def _nearest_distances(source: np.ndarray, target: np.ndarray, chunk_size: int = 4096) -> np.ndarray:
    distances = []
    for start in range(0, len(source), chunk_size):
        chunk = source[start:start + chunk_size]
        squared = np.sum((chunk[:, None, :] - target[None, :, :]) ** 2, axis=2)
        distances.append(np.sqrt(np.min(squared, axis=1)))
    return np.concatenate(distances)


def pointcloud_metrics(prediction: np.ndarray, reference: np.ndarray, threshold: float = 0.05) -> dict[str, float | int]:
    """Return symmetric Chamfer, precision, recall, and F-score."""
    if threshold <= 0:
        raise ValueError("threshold must be positive")
    prediction = np.asarray(prediction, dtype=np.float64)
    reference = np.asarray(reference, dtype=np.float64)
    if prediction.ndim != 2 or reference.ndim != 2 or prediction.shape[1] != 3 or reference.shape[1] != 3:
        raise ValueError("prediction and reference must have shape Nx3")
    pred_to_ref = _nearest_distances(prediction, reference)
    ref_to_pred = _nearest_distances(reference, prediction)
    precision = float(np.mean(pred_to_ref <= threshold))
    recall = float(np.mean(ref_to_pred <= threshold))
    fscore = 2.0 * precision * recall / max(precision + recall, 1e-12)
    return {
        "prediction_point_count": int(len(prediction)),
        "reference_point_count": int(len(reference)),
        "threshold": float(threshold),
        "chamfer_l1": float((pred_to_ref.mean() + ref_to_pred.mean()) / 2.0),
        "prediction_to_reference_mean": float(pred_to_ref.mean()),
        "reference_to_prediction_mean": float(ref_to_pred.mean()),
        "precision": precision,
        "recall": recall,
        "fscore": float(fscore),
    }


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("prediction", type=Path)
    parser.add_argument("reference", type=Path)
    parser.add_argument("--threshold", type=float, default=0.05)
    args = parser.parse_args(argv)
    print(json.dumps(pointcloud_metrics(load_points(args.prediction), load_points(args.reference), args.threshold), indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())