training/edge_patch.py

"""6D cylindrical edge-patch builder and an edge-classifier evaluation check.

Provides:
  - colmap_points_xyz_rgb / build_edge_patch_6d: build the 6D (xyz + RGB)
    cylindrical patch around an edge from COLMAP points (radius 0.5m,
    +0.25m extension at each end). Imported by the dataset generators.

Run as a script to evaluate an edge classifier on a few samples: it scores
each handcrafted edge, splits edges by whether they are ground-truth-positive
(both endpoints near connected GT vertices), and compares the score
distributions.
"""
import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'

import sys
import time
import numpy as np
import torch
from datasets import load_dataset
from scipy.spatial.distance import cdist

CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, CURRENT_DIR)

import hc_helpers as hc
from fast_pointnet_class import (
    load_pointnet_model as load_pnet_class,
    predict_class_from_patch,
)
from hoho2025.example_solutions import read_colmap_rec

NUM_TRIALS = 3
CYL_RADIUS = 0.5
CYL_EXT = 0.25            # extension at each end
GT_VERTEX_THRESH = 0.5    # how close a vertex must be to a GT vertex to count
                          # as "matched" (= HSS vert_thresh)


def colmap_points_xyz_rgb(colmap_rec):
    """Return (xyz, rgb_normalized_0_1) for all COLMAP points."""
    xyz_list, rgb_list = [], []
    for pid, p3D in colmap_rec.points3D.items():
        xyz_list.append(p3D.xyz)
        rgb_list.append(p3D.color / 255.0)
    if not xyz_list:
        return np.empty((0, 3)), np.empty((0, 3))
    return np.array(xyz_list), np.array(rgb_list)


def build_edge_patch_6d(u_xyz, v_xyz, colmap_xyz, colmap_rgb):
    """
    Returns patch dict {'patch_6d': (M, 6)} or None if cylinder too sparse.
    """
    line = v_xyz - u_xyz
    L = float(np.linalg.norm(line))
    if L < 1e-6:
        return None
    direction = line / L

    ext_start = u_xyz - CYL_EXT * direction
    ext_L = L + 2 * CYL_EXT

    rel = colmap_xyz - ext_start[np.newaxis, :]
    proj = rel @ direction
    in_bounds = (proj >= 0) & (proj <= ext_L)

    closest = ext_start[np.newaxis, :] + proj[:, np.newaxis] * direction[np.newaxis, :]
    perp = np.linalg.norm(colmap_xyz - closest, axis=1)
    in_cyl = in_bounds & (perp <= CYL_RADIUS)

    if int(in_cyl.sum()) <= 10:
        return None

    midpoint = (u_xyz + v_xyz) / 2
    pts_centered = colmap_xyz[in_cyl] - midpoint
    rgb_signed = colmap_rgb[in_cyl] * 2.0 - 1.0   # to [-1, 1]
    patch_6d = np.hstack([pts_centered, rgb_signed])
    return {'patch_6d': patch_6d}


def label_user_edges(user_v, user_e, gt_v, gt_e, thresh=GT_VERTEX_THRESH):
    """For each user edge, return True if it matches a GT edge.

    Match := both endpoints have a nearest GT vertex within `thresh`,
    AND those GT vertices are connected in the GT edge set.
    """
    if len(gt_v) == 0 or len(user_v) == 0:
        return [None] * len(user_e)

    d = cdist(user_v, gt_v)
    user_to_gt = {}
    for i in range(len(user_v)):
        j = int(np.argmin(d[i]))
        if d[i, j] < thresh:
            user_to_gt[i] = j

    gt_set = set()
    for a, b in gt_e:
        gt_set.add((int(min(a, b)), int(max(a, b))))

    out = []
    for u, v in user_e:
        gu = user_to_gt.get(int(u))
        gv = user_to_gt.get(int(v))
        if gu is None or gv is None or gu == gv:
            out.append(False)
            continue
        key = (min(gu, gv), max(gu, gv))
        out.append(key in gt_set)
    return out


def smoke_test_one(sample, model, device):
    order_id = sample['order_id']
    print(f"\n=== {order_id} ===")
    t0 = time.time()

    try:
        with hc.suppress_stdout():
            user_v, user_e = hc.hc_predict(sample, {})
    except Exception as e:
        print(f"  user pipeline crashed: {e}")
        return []
    if len(user_v) == 0 or len(user_e) == 0:
        print("  user pipeline empty")
        return []
    print(f"  user: {len(user_v)} vertices, {len(user_e)} edges  ({time.time()-t0:.1f}s)")

    try:
        colmap_rec = read_colmap_rec(sample['colmap'])
    except Exception as e:
        print(f"  colmap parse crashed: {e}")
        return []

    cm_xyz, cm_rgb = colmap_points_xyz_rgb(colmap_rec)
    print(f"  colmap: {len(cm_xyz)} points")
    if len(cm_xyz) == 0:
        return []

    gt_v = np.array(sample['wf_vertices']) if sample.get('wf_vertices') else np.empty((0, 3))
    gt_e = [(int(a), int(b)) for a, b in sample.get('wf_edges', [])]
    labels = label_user_edges(user_v, user_e, gt_v, gt_e)

    results = []
    skipped_sparse = 0
    for (u_idx, v_idx), gt_match in zip(user_e, labels):
        u_xyz = np.asarray(user_v[int(u_idx)])
        v_xyz = np.asarray(user_v[int(v_idx)])
        patch = build_edge_patch_6d(u_xyz, v_xyz, cm_xyz, cm_rgb)
        if patch is None:
            skipped_sparse += 1
            continue
        try:
            cls_label, score = predict_class_from_patch(model, patch, device=device)
        except Exception as e:
            print(f"  inference crashed: {e}")
            continue
        results.append({
            'order_id': order_id,
            'u': int(u_idx),
            'v': int(v_idx),
            'edge_length': float(np.linalg.norm(v_xyz - u_xyz)),
            'patch_n_pts': int(patch['patch_6d'].shape[0]),
            'pred_label': int(cls_label) if cls_label is not None else None,
            'score': float(score),
            'gt_match': bool(gt_match) if gt_match is not None else None,
        })

    if skipped_sparse:
        print(f"  {skipped_sparse} edges skipped (cylinder too sparse)")
    n_gt_pos = sum(1 for r in results if r['gt_match'])
    n_gt_neg = sum(1 for r in results if r['gt_match'] is False)
    print(f"  scored {len(results)} edges: {n_gt_pos} GT-positive, "
          f"{n_gt_neg} GT-negative")
    return results


def main():
    print("Loading pnet_class.pth...")
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = load_pnet_class(
        os.path.join(CURRENT_DIR, '..', 'pnet_class_2026.pth'), device=device)
    print(f"  loaded on {device}")

    print("\nStreaming validation...")
    ds = load_dataset(
        "usm3d/hoho22k_2026_trainval", split="validation",
        streaming=True, trust_remote_code=True,
    )

    all_rows = []
    for i, sample in enumerate(ds):
        if i >= NUM_TRIALS:
            break
        all_rows.extend(smoke_test_one(sample, model, device))

    if not all_rows:
        print("\nNo results -- every sample failed.")
        return

    scores = np.array([r['score'] for r in all_rows])
    pos_scores = np.array([r['score'] for r in all_rows if r['gt_match']])
    neg_scores = np.array([r['score'] for r in all_rows if r['gt_match'] is False])

    print(f"\n=== Aggregate over {len(all_rows)} edges "
          f"({len({r['order_id'] for r in all_rows})} samples) ===")

    print(f"\nScore distribution (all {len(scores)} edges):")
    print(f"  mean={scores.mean():.3f}  median={np.median(scores):.3f}  "
          f"std={scores.std():.3f}  min={scores.min():.3f}  max={scores.max():.3f}")
    print(f"  fraction <0.05: {(scores < 0.05).mean()*100:.0f}%   "
          f">0.65 (paper): {(scores > 0.65).mean()*100:.0f}%   "
          f">0.99: {(scores > 0.99).mean()*100:.0f}%")

    if len(pos_scores) and len(neg_scores):
        diff = float(pos_scores.mean() - neg_scores.mean())
        print(f"\nGT-positive ({len(pos_scores)}): "
              f"mean={pos_scores.mean():.3f}  median={np.median(pos_scores):.3f}  "
              f"std={pos_scores.std():.3f}")
        print(f"GT-negative ({len(neg_scores)}): "
              f"mean={neg_scores.mean():.3f}  median={np.median(neg_scores):.3f}  "
              f"std={neg_scores.std():.3f}")
        print(f"Mean delta (pos-neg): {diff:+.3f}")

        # Quick AUC via ranking
        all_pairs = sorted([(s, 1) for s in pos_scores] + [(s, 0) for s in neg_scores])
        n_pos, n_neg = len(pos_scores), len(neg_scores)
        rank_sum_pos = sum(rank+1 for rank, (_, lab) in enumerate(all_pairs) if lab == 1)
        auc = (rank_sum_pos - n_pos*(n_pos+1)/2) / (n_pos * n_neg) if n_pos and n_neg else 0
        print(f"AUC (pos > neg): {auc:.3f}")
    else:
        print("\nMissing pos or neg group -- can't compute discrimination.")


if __name__ == "__main__":
    main()