benchmark.py

#!/usr/bin/env python3
"""
Benchmark all wood classification models against manual labels in data/test/.

Ground truth: scalar_alpha field (1 = wood, 0 = leaf)
Predictions: model output remapped to 1 (leaf) / 255 (wood)
"""

import os
import sys
import time
import numpy as np
import torch
from pathlib import Path

sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'modules', 'filter'))

from run_models_inference import TreeAIBoxInference

TEST_DIR = Path(__file__).parent / "data" / "test"
MODEL_DIR = Path(__file__).parent / "models"

# Only wood classification models are relevant for this benchmark
WOODCLS_MODELS = sorted(MODEL_DIR.glob("woodcls_*.pth"))


def compute_metrics(gt, pred):
    """Compute classification metrics.

    Args:
        gt: ground truth array (1 = wood, 0 = leaf)
        pred: prediction array (255 = wood, 1 = leaf)

    Returns:
        dict of metrics
    """
    # Convert prediction encoding: 255 -> 1 (wood), 1 -> 0 (leaf)
    pred_binary = (pred == 255).astype(int)

    tp = np.sum((pred_binary == 1) & (gt == 1))
    tn = np.sum((pred_binary == 0) & (gt == 0))
    fp = np.sum((pred_binary == 1) & (gt == 0))
    fn = np.sum((pred_binary == 0) & (gt == 1))

    total = tp + tn + fp + fn
    accuracy = (tp + tn) / total if total > 0 else 0
    precision = tp / (tp + fp) if (tp + fp) > 0 else 0
    recall = tp / (tp + fn) if (tp + fn) > 0 else 0
    f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0
    iou_wood = tp / (tp + fp + fn) if (tp + fp + fn) > 0 else 0
    iou_leaf = tn / (tn + fp + fn) if (tn + fp + fn) > 0 else 0
    miou = (iou_wood + iou_leaf) / 2

    return {
        "accuracy": accuracy,
        "precision": precision,
        "recall": recall,
        "f1": f1,
        "iou_wood": iou_wood,
        "iou_leaf": iou_leaf,
        "miou": miou,
        "tp": int(tp),
        "tn": int(tn),
        "fp": int(fp),
        "fn": int(fn),
    }


def main():
    test_files = sorted(TEST_DIR.glob("*.ply"))
    if not test_files:
        print(f"No PLY files found in {TEST_DIR}")
        return 1

    if not WOODCLS_MODELS:
        print(f"No woodcls models found in {MODEL_DIR}")
        return 1

    print(f"Test files: {len(test_files)}")
    print(f"Models: {len(WOODCLS_MODELS)}")
    for m in WOODCLS_MODELS:
        print(f"  - {m.name}")
    print()

    engine = TreeAIBoxInference(use_gpu=True)

    # results[model_name][file_name] = metrics dict
    results = {}

    for model_path in WOODCLS_MODELS:
        model_name = model_path.stem
        wood_type = "branch" if "branch" in model_name else "stem"
        results[model_name] = {}

        print(f"\n{'='*70}")
        print(f"MODEL: {model_path.name}")
        print(f"{'='*70}")

        for test_file in test_files:
            file_label = test_file.stem
            print(f"\n  File: {file_label}")

            # Load point cloud
            try:
                pc = engine.load_ply_file(str(test_file))
            except Exception as e:
                print(f"    ERROR loading: {e}")
                results[model_name][file_label] = None
                continue

            # Get ground truth from scalar_alpha
            if "scalar_alpha" not in pc._scalar_fields:
                print(f"    ERROR: no scalar_alpha field")
                results[model_name][file_label] = None
                continue

            gt = pc._scalar_fields["scalar_alpha"].asArray().astype(int)

            # Run inference
            t0 = time.time()
            result_pc = engine.run_wood_classification(
                pc, str(model_path), wood_type=wood_type
            )
            elapsed = time.time() - t0

            if result_pc is None:
                print(f"    ERROR: inference failed")
                results[model_name][file_label] = None
                continue

            # Get predictions from the classification field
            field_name = f"{wood_type}_classification"
            if field_name not in result_pc._scalar_fields:
                print(f"    ERROR: no {field_name} field in output")
                results[model_name][file_label] = None
                continue

            pred = result_pc._scalar_fields[field_name].asArray()
            metrics = compute_metrics(gt, pred)
            metrics["time_s"] = elapsed
            metrics["n_points"] = len(gt)
            results[model_name][file_label] = metrics

            print(f"    Accuracy: {metrics['accuracy']:.4f}  "
                  f"F1: {metrics['f1']:.4f}  "
                  f"mIoU: {metrics['miou']:.4f}  "
                  f"Time: {elapsed:.1f}s")

    # Print summary table
    print(f"\n\n{'='*70}")
    print("SUMMARY")
    print(f"{'='*70}")

    header = f"{'Model':<55} {'Acc':>6} {'F1':>6} {'mIoU':>6} {'Prec':>6} {'Rec':>6} {'Time':>6}"
    print(header)
    print("-" * len(header))

    for model_name, file_results in results.items():
        valid = [m for m in file_results.values() if m is not None]
        if not valid:
            print(f"{model_name:<55} {'FAIL':>6}")
            continue

        avg_acc = np.mean([m["accuracy"] for m in valid])
        avg_f1 = np.mean([m["f1"] for m in valid])
        avg_miou = np.mean([m["miou"] for m in valid])
        avg_prec = np.mean([m["precision"] for m in valid])
        avg_rec = np.mean([m["recall"] for m in valid])
        avg_time = np.mean([m["time_s"] for m in valid])

        print(f"{model_name:<55} {avg_acc:>6.3f} {avg_f1:>6.3f} {avg_miou:>6.3f} "
              f"{avg_prec:>6.3f} {avg_rec:>6.3f} {avg_time:>5.1f}s")

    # Per-file breakdown
    print(f"\n\n{'='*70}")
    print("PER-FILE BREAKDOWN")
    print(f"{'='*70}")

    for model_name, file_results in results.items():
        print(f"\n{model_name}:")
        print(f"  {'File':<35} {'Acc':>6} {'F1':>6} {'mIoU':>6} {'Points':>10} {'Time':>6}")
        print(f"  {'-'*75}")
        for file_label, metrics in file_results.items():
            if metrics is None:
                print(f"  {file_label:<35} {'FAIL':>6}")
            else:
                print(f"  {file_label:<35} {metrics['accuracy']:>6.3f} {metrics['f1']:>6.3f} "
                      f"{metrics['miou']:>6.3f} {metrics['n_points']:>10,} {metrics['time_s']:>5.1f}s")

    return 0


if __name__ == "__main__":
    sys.exit(main())