| """ |
| 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 |
|
|
| |
| 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"]) |
| lcs = _flatten(data["coedge_lcs"]) |
|
|
| 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) |
| edge_feats = np.asarray(data["edge_features"], dtype=np.float32) |
|
|
| 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 |
|
|
