Upload 2 files

03_infer_halfedge.py

@@ -0,0 +1,198 @@
+# file: 03_infer_halfedge.py
+# -*- coding: utf-8 -*-
+import argparse
+from pathlib import Path
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch_geometric.nn import HeteroConv, SAGEConv, GlobalAttention, JumpingKnowledge, BatchNorm
+from torch_geometric.data import HeteroData
+
+from brep_extractor_utils import load_coedge_arrays, make_heterodata
+
+class HalfEdgeGNN(nn.Module):
+    def __init__(
+        self,
+        coedge_in: int,
+        face_in: int,
+        edge_in: int,
+        global_in: int,
+        hidden=256,
+        layers=6,
+        dropout=0.2,
+        num_classes=3,
+        jk_mode="cat",
+        gating_dim=None,
+    ):
+        super().__init__()
+        self.convs = nn.ModuleList(); self.bns = nn.ModuleList()
+        self.encoders = nn.ModuleDict({
+            "coedge": nn.Sequential(nn.Linear(coedge_in, hidden), nn.ReLU(), nn.Dropout(dropout)),
+            "face": nn.Sequential(nn.Linear(face_in, hidden), nn.ReLU(), nn.Dropout(dropout)),
+            "edge": nn.Sequential(nn.Linear(edge_in, hidden), nn.ReLU(), nn.Dropout(dropout)),
+        })
+        for _ in range(layers):
+            conv = HeteroConv({
+                ('coedge','next','coedge'): SAGEConv((hidden,hidden), hidden),
+                ('coedge','prev','coedge'): SAGEConv((hidden,hidden), hidden),
+                ('coedge','mate','coedge'): SAGEConv((hidden,hidden), hidden),
+                ('coedge','to_face','face'): SAGEConv((hidden, hidden), hidden),
+                ('face','to_coedge','coedge'): SAGEConv((hidden, hidden), hidden),
+                ('coedge','to_edge','edge'): SAGEConv((hidden, hidden), hidden),
+                ('edge','to_coedge','coedge'): SAGEConv((hidden, hidden), hidden),
+                ('face','to_edge','edge'): SAGEConv((hidden, hidden), hidden),
+                ('edge','to_face','face'): SAGEConv((hidden, hidden), hidden),
+            }, aggr='sum')
+            self.convs.append(conv)
+            self.bns.append(nn.ModuleDict({
+                "coedge": BatchNorm(hidden),
+                "face": BatchNorm(hidden),
+                "edge": BatchNorm(hidden),
+            }))
+        self.jk = JumpingKnowledge(mode=jk_mode)
+        self.jk_out = hidden * layers if jk_mode == "cat" else hidden
+        if gating_dim is None:
+            gating_dim = hidden
+        self.gating_dim = gating_dim
+        self.gate = nn.Sequential(
+            nn.Linear(self.jk_out, self.jk_out//2),
+            nn.ReLU(),
+            nn.Linear(self.jk_out//2, 1),
+        )
+        self.pool = GlobalAttention(self.gate)
+        self.proj = nn.Identity() if self.jk_out == gating_dim else nn.Linear(self.jk_out, gating_dim)
+        self.global_mlp = nn.Sequential(
+            nn.Linear(global_in, gating_dim),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.Linear(gating_dim, 2 * gating_dim),
+        )
+        self.head = nn.Sequential(
+            nn.Linear(gating_dim, hidden),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden, num_classes),
+        )
+
+    def forward(self, data: HeteroData):
+        x = {
+            "coedge": self.encoders["coedge"](data["coedge"].x),
+            "face": self.encoders["face"](data["face"].x),
+            "edge": self.encoders["edge"](data["edge"].x),
+        }
+        outs = []
+        for conv, bn in zip(self.convs, self.bns):
+            x_new = conv(x, data.edge_index_dict)
+            x = {k: F.relu(bn[k](x_new[k]) + x[k]) for k in x}
+            outs.append(x["coedge"])
+        xj = self.jk(outs)
+        g = self.pool(xj, data['coedge'].batch)
+        g0 = self.proj(g)
+        global_x = data["global"].x
+        if global_x.dim() == 1:
+            global_x = global_x.view(1, -1)
+        if global_x.size(0) != g0.size(0):
+            raise RuntimeError(
+                f"Global feature batch mismatch: {global_x.size(0)} vs {g0.size(0)}"
+            )
+        gb = self.global_mlp(global_x)
+        gamma, beta = gb.chunk(2, dim=-1)
+        gamma = torch.sigmoid(gamma)
+        g_mod = g0 * gamma + beta
+        return self.head(g_mod)
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model", required=True)
+    ap.add_argument("--npz", required=True, help="Path to a processed BRep extractor npz file")
+    ap.add_argument("--tau", type=float, default=0.0, help="Reject threshold; below this outputs random")
+    ap.add_argument("--min_conf", type=float, default=0.85, help="Hard minimum confidence for known classes")
+    ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    args = ap.parse_args()
+
+    try:
+        ckpt = torch.load(args.model, map_location="cpu", weights_only=False)
+    except TypeError:
+        ckpt = torch.load(args.model, map_location="cpu")
+    if "global_in" not in ckpt or "gating_dim" not in ckpt:
+        raise RuntimeError(
+            "Checkpoint missing gating metadata. Please retrain with global gating enabled."
+        )
+    labels = ckpt["labels"]; inv_labels = {v:k for k,v in labels.items()}
+    random_id = labels.get("random")
+    if (args.tau > 0 or args.min_conf > 0) and random_id is None:
+        raise RuntimeError("Model labels do not include 'random'; retrain a 4-class model.")
+    stats = ckpt["stats"]
+    if not all(k in stats for k in ("coedge", "face", "edge")):
+        raise RuntimeError("Checkpoint missing heterograph stats; retrain required.")
+
+    coedge_in = ckpt.get("coedge_in", ckpt.get("node_in"))
+    face_in = ckpt.get("face_in")
+    edge_in = ckpt.get("edge_in")
+    if coedge_in is None or face_in is None or edge_in is None:
+        raise RuntimeError("Checkpoint missing heterograph input dims; retrain required.")
+
+    graph_data = load_coedge_arrays(Path(args.npz))
+    if int(graph_data["coedge_x"].shape[1]) != int(coedge_in):
+        raise RuntimeError(
+            f"Coedge feature dim mismatch: npz={int(graph_data['coedge_x'].shape[1])} "
+            f"ckpt={int(coedge_in)}"
+        )
+    if int(graph_data["face_x"].shape[1]) != int(face_in):
+        raise RuntimeError(
+            f"Face feature dim mismatch: npz={int(graph_data['face_x'].shape[1])} "
+            f"ckpt={int(face_in)}"
+        )
+    if int(graph_data["edge_x"].shape[1]) != int(edge_in):
+        raise RuntimeError(
+            f"Edge feature dim mismatch: npz={int(graph_data['edge_x'].shape[1])} "
+            f"ckpt={int(edge_in)}"
+        )
+    if int(graph_data["global_x"].shape[0]) != int(ckpt["global_in"]):
+        raise RuntimeError(
+            f"Global feature dim mismatch: npz={int(graph_data['global_x'].shape[0])} "
+            f"ckpt={int(ckpt['global_in'])}"
+        )
+    data = make_heterodata(
+        graph_data["coedge_x"],
+        graph_data["face_x"],
+        graph_data["edge_x"],
+        graph_data["next"],
+        graph_data["mate"],
+        graph_data["coedge_face"],
+        graph_data["coedge_edge"],
+        graph_data["global_x"],
+        label=None,
+        norm_stats=stats,
+    )
+    data['coedge'].batch = torch.zeros(data['coedge'].x.size(0), dtype=torch.long)
+    data["global"].batch = torch.zeros(1, dtype=torch.long)
+    data["face"].batch = torch.zeros(data["face"].x.size(0), dtype=torch.long)
+    data["edge"].batch = torch.zeros(data["edge"].x.size(0), dtype=torch.long)
+
+    global_in = ckpt["global_in"]
+    gating_dim = ckpt["gating_dim"]
+    model = HalfEdgeGNN(coedge_in=coedge_in, face_in=face_in, edge_in=edge_in, global_in=global_in,
+                        hidden=ckpt["hp"]["hidden"],
+                        layers=ckpt["hp"]["layers"], dropout=ckpt["hp"]["dropout"],
+                        num_classes=len(labels), gating_dim=gating_dim).to(args.device)
+    model.load_state_dict(ckpt["state_dict"]); model.eval()
+
+    with torch.no_grad():
+        logits = model(data.to(args.device))
+        probs = F.softmax(logits, dim=-1).cpu().numpy()[0]
+        pred = int(probs.argmax())
+        conf = float(probs[pred])
+        arg_label = inv_labels[pred]
+        effective_tau = max(args.tau, args.min_conf)
+        if conf < effective_tau and random_id is not None:
+            final_label = "random"
+        else:
+            final_label = arg_label
+        print(f"Argmax: {arg_label} (conf={conf:.4f})")
+        print(f"Predicted: {final_label} (tau={effective_tau:.2f})")
+        for i, p in enumerate(probs):
+            print(f"{inv_labels[i]:>6s}: {p:.4f}")
+
+if __name__ == "__main__":
+    main()

brep_extractor_utils.py

@@ -0,0 +1,271 @@
+"""
+Utility helpers for loading BRep extractor-processed STEP data as PyG graphs.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Dict, Iterable, Tuple
+
+import numpy as np
+import torch
+from torch_geometric.data import HeteroData
+
+# Label mapping for the current project
+LABELS: Dict[str, int] = {"pipe": 0, "elbow": 1, "tjoint": 2, "random": 3}
+STEP_EXTS = ("*.step", "*.stp", "*.STEP", "*.STP")
+
+
+def build_label_map(step_root: Path) -> Dict[str, int]:
+    """
+    Scan the STEP directory tree (containing /pipe, /elbow, /tjoint, ...)
+    and build a mapping from file stem to integer label.
+    """
+    mapping: Dict[str, int] = {}
+    for cls, label in LABELS.items():
+        cls_dir = step_root / cls
+        if not cls_dir.exists():
+            continue
+        for ext in STEP_EXTS:
+            for file in cls_dir.glob(ext):
+                mapping[file.stem] = label
+    if not mapping:
+        raise RuntimeError(f"No STEP files found under {step_root} for any of {tuple(LABELS)}")
+    return mapping
+
+
+def _flatten(arr: np.ndarray) -> np.ndarray:
+    return np.asarray(arr, dtype=np.float32).reshape(arr.shape[0], -1)
+
+def _face_grid_stats(face_grids: np.ndarray) -> np.ndarray:
+    """
+    Summarize face point grids into compact stats per face.
+    Returns [F, 10]: xyz_mean (3), xyz_std (3), nrm_mean (3), mask_frac (1).
+    """
+    face_grids = np.asarray(face_grids, dtype=np.float32)
+    f = face_grids.shape[0]
+    xyz = face_grids[:, 0:3, :, :].reshape(f, 3, -1)
+    nrm = face_grids[:, 3:6, :, :].reshape(f, 3, -1)
+    msk = face_grids[:, 6, :, :].reshape(f, -1)
+
+    mask = (msk > 0.5).astype(np.float32)
+    mask_frac = mask.mean(axis=1, keepdims=True)
+    w = mask / (mask.sum(axis=1, keepdims=True) + 1e-6)
+
+    xyz_mean = (xyz * w[:, None, :]).sum(axis=2)
+    xyz_var = (w[:, None, :] * (xyz - xyz_mean[:, :, None]) ** 2).sum(axis=2)
+    xyz_std = np.sqrt(np.maximum(xyz_var, 1e-12))
+    nrm_mean = (nrm * w[:, None, :]).sum(axis=2)
+    return np.concatenate([xyz_mean, xyz_std, nrm_mean, mask_frac], axis=1)
+
+def compute_global_geom_features(data) -> np.ndarray:
+    """
+    Compute compact global geometry descriptors from face/coedge point samples.
+    Returns [5] float32: pca_ev_ratio_1/2/3, line_fit_rmse, plane_fit_rmse.
+    """
+    points = []
+    face_grids = np.asarray(data["face_point_grids"], dtype=np.float32)
+    if face_grids.size:
+        xyz = face_grids[:, 0:3, :, :].transpose(0, 2, 3, 1).reshape(-1, 3)
+        mask = face_grids[:, 6, :, :].reshape(-1) > 0.5
+        if mask.any():
+            points.append(xyz[mask])
+
+    coedge_grids = np.asarray(data["coedge_point_grids"], dtype=np.float32)
+    if coedge_grids.size:
+        co_xyz = coedge_grids[:, 0:3, :].transpose(0, 2, 1).reshape(-1, 3)
+        points.append(co_xyz)
+
+    if not points:
+        return np.zeros(5, dtype=np.float32)
+
+    pts = np.concatenate(points, axis=0)
+    if pts.shape[0] < 3:
+        return np.zeros(5, dtype=np.float32)
+    pts = pts[np.isfinite(pts).all(axis=1)]
+    if pts.shape[0] < 3:
+        return np.zeros(5, dtype=np.float32)
+
+    mean = pts.mean(axis=0, keepdims=True)
+    centered = pts - mean
+    scale = np.sqrt(np.mean(np.sum(centered ** 2, axis=1)))
+    centered = centered / (scale + 1e-6)
+    cov = (centered.T @ centered) / max(1, centered.shape[0])
+    if not np.isfinite(cov).all():
+        return np.zeros(5, dtype=np.float32)
+
+    ev = np.linalg.eigvalsh(cov)
+    ev = np.sort(ev)[::-1]
+    ev = np.maximum(ev, 0.0)
+    total = ev.sum()
+    if not np.isfinite(total) or total <= 0.0:
+        return np.zeros(5, dtype=np.float32)
+
+    ratios = ev / total
+    line_rmse = np.sqrt(max(ev[1] + ev[2], 0.0))
+    plane_rmse = np.sqrt(max(ev[2], 0.0))
+    feats = np.array(
+        [ratios[0], ratios[1], ratios[2], line_rmse, plane_rmse],
+        dtype=np.float32,
+    )
+    if not np.isfinite(feats).all():
+        return np.zeros(5, dtype=np.float32)
+    return feats
+
+def load_coedge_arrays(npz_path: Path) -> Dict[str, np.ndarray]:
+    """
+    Load node features and adjacency indices from a BRep extractor npz.
+    Returns a dict with coedge/face/edge/global features and topology arrays.
+    """
+    with np.load(npz_path) as data:
+        coedge_feats = _flatten(data["coedge_features"])
+        scale = np.asarray(data["coedge_scale_factors"], dtype=np.float32)[:, None]
+        reverse = np.asarray(data["coedge_reverse_flags"], dtype=np.float32)[:, None]
+        point_grids = _flatten(data["coedge_point_grids"])  # [N, 12*U]
+        lcs = _flatten(data["coedge_lcs"])                  # [N, 16]
+
+        face_idx = np.asarray(data["face"], dtype=np.int64)
+        edge_idx = np.asarray(data["edge"], dtype=np.int64)
+        face_feats = np.asarray(data["face_features"], dtype=np.float32)  # [F, 7]
+        edge_feats = np.asarray(data["edge_features"], dtype=np.float32)  # [E, 10]
+
+        face_grid_stats = _face_grid_stats(data["face_point_grids"])
+
+        coedge_x = np.concatenate(
+            [coedge_feats, scale, reverse, point_grids, lcs], axis=1
+        )
+        face_x = np.concatenate([face_feats, face_grid_stats], axis=1)
+        edge_x = edge_feats
+        next_index = np.asarray(data["next"], dtype=np.int64)
+        mate_index = np.asarray(data["mate"], dtype=np.int64)
+        global_features = compute_global_geom_features(data)
+
+    return {
+        "coedge_x": coedge_x,
+        "face_x": face_x,
+        "edge_x": edge_x,
+        "next": next_index,
+        "mate": mate_index,
+        "coedge_face": face_idx,
+        "coedge_edge": edge_idx,
+        "global_x": global_features,
+    }
+
+
+def make_edge_index(source: np.ndarray, target: np.ndarray) -> torch.Tensor:
+    """
+    Build a 2 x E tensor of edge indices (with both directions, deduplicated).
+    """
+    pairs = np.stack([source, target], axis=1)
+    flipped = pairs[:, ::-1]
+    all_pairs = np.concatenate([pairs, flipped], axis=0)
+    all_pairs = np.unique(all_pairs, axis=0)
+    return torch.tensor(all_pairs.T, dtype=torch.long)
+
+def make_directed_edge_index(source: np.ndarray, target: np.ndarray) -> torch.Tensor:
+    """
+    Build a 2 x E tensor of directed edge indices (no deduplication).
+    """
+    return torch.tensor(np.stack([source, target], axis=0), dtype=torch.long)
+
+def make_bipartite_edge_index(source: np.ndarray, target: np.ndarray) -> torch.Tensor:
+    """
+    Build a 2 x E tensor of directed bipartite edge indices (deduplicated).
+    """
+    pairs = np.stack([source, target], axis=1)
+    pairs = np.unique(pairs, axis=0)
+    return torch.tensor(pairs.T, dtype=torch.long)
+
+def make_heterodata(
+    coedge_x: np.ndarray,
+    face_x: np.ndarray,
+    edge_x: np.ndarray,
+    next_index: np.ndarray,
+    mate_index: np.ndarray,
+    coedge_face: np.ndarray,
+    coedge_edge: np.ndarray,
+    global_features: np.ndarray,
+    label: int | None,
+    norm_stats: Dict[str, Dict[str, np.ndarray | torch.Tensor]] | None = None,
+) -> HeteroData:
+    """
+    Create a PyG HeteroData graph for the coedge features/relations.
+    When mean/std are provided the features are normalised element-wise.
+    """
+    def _normalize(x_arr: np.ndarray, stats: Dict[str, np.ndarray | torch.Tensor] | None) -> torch.Tensor:
+        x_t = torch.tensor(x_arr, dtype=torch.float32)
+        if stats is None:
+            return x_t
+        mean = stats.get("mean")
+        std = stats.get("std")
+        if mean is None or std is None:
+            return x_t
+        mean_t = torch.as_tensor(mean, dtype=torch.float32)
+        std_t = torch.as_tensor(std, dtype=torch.float32)
+        return (x_t - mean_t) / std_t
+
+    coedge_stats = norm_stats.get("coedge") if norm_stats else None
+    face_stats = norm_stats.get("face") if norm_stats else None
+    edge_stats = norm_stats.get("edge") if norm_stats else None
+
+    x_coedge = _normalize(coedge_x, coedge_stats)
+    x_face = _normalize(face_x, face_stats)
+    x_edge = _normalize(edge_x, edge_stats)
+
+    idx = np.arange(coedge_x.shape[0], dtype=np.int64)
+    edge_next = make_directed_edge_index(idx, next_index)
+    edge_prev = make_directed_edge_index(next_index, idx)
+    edge_mate = make_edge_index(idx, mate_index)
+    edge_coedge_face = make_directed_edge_index(idx, coedge_face)
+    edge_face_coedge = make_directed_edge_index(coedge_face, idx)
+    edge_coedge_edge = make_directed_edge_index(idx, coedge_edge)
+    edge_edge_coedge = make_directed_edge_index(coedge_edge, idx)
+    edge_face_edge = make_bipartite_edge_index(coedge_face, coedge_edge)
+    edge_edge_face = make_bipartite_edge_index(coedge_edge, coedge_face)
+
+    data = HeteroData()
+    data["coedge"].x = x_coedge
+    data["face"].x = x_face
+    data["edge"].x = x_edge
+    data["global"].x = torch.tensor(global_features, dtype=torch.float32).view(1, -1)
+    data["coedge", "next", "coedge"].edge_index = edge_next
+    data["coedge", "prev", "coedge"].edge_index = edge_prev
+    data["coedge", "mate", "coedge"].edge_index = edge_mate
+    data["coedge", "to_face", "face"].edge_index = edge_coedge_face
+    data["face", "to_coedge", "coedge"].edge_index = edge_face_coedge
+    data["coedge", "to_edge", "edge"].edge_index = edge_coedge_edge
+    data["edge", "to_coedge", "coedge"].edge_index = edge_edge_coedge
+    data["face", "to_edge", "edge"].edge_index = edge_face_edge
+    data["edge", "to_face", "face"].edge_index = edge_edge_face
+    if label is not None:
+        data.y = torch.tensor([int(label)], dtype=torch.long)
+    return data
+
+
+def compute_feature_stats(npz_paths: Iterable[Path]) -> Dict[str, np.ndarray]:
+    """
+    Compute mean and std (per feature dimension) across all node features in the dataset.
+    """
+    totals = {"coedge": 0, "face": 0, "edge": 0}
+    sum_vec: Dict[str, np.ndarray | None] = {"coedge": None, "face": None, "edge": None}
+    sum_sq: Dict[str, np.ndarray | None] = {"coedge": None, "face": None, "edge": None}
+
+    for path in npz_paths:
+        graph = load_coedge_arrays(path)
+        for key, x in (("coedge", graph["coedge_x"]), ("face", graph["face_x"]), ("edge", graph["edge_x"])):
+            if sum_vec[key] is None:
+                sum_vec[key] = np.zeros(x.shape[1], dtype=np.float64)
+                sum_sq[key] = np.zeros(x.shape[1], dtype=np.float64)
+            sum_vec[key] += x.sum(axis=0)
+            sum_sq[key] += (x * x).sum(axis=0)
+            totals[key] += x.shape[0]
+
+    out = {}
+    for key in ("coedge", "face", "edge"):
+        if sum_vec[key] is None or totals[key] == 0:
+            raise RuntimeError(f"Cannot compute feature stats: no {key} features observed.")
+        mean = sum_vec[key] / totals[key]
+        var = sum_sq[key] / totals[key] - mean * mean
+        var = np.maximum(var, 1e-12)
+        std = np.sqrt(var)
+        out[key] = {"mean": mean.astype(np.float32), "std": std.astype(np.float32)}
+    return out