cstnet2-stage1-all

CSTNet2 Stage 1 — Three-Phase Constraint Extractor

This repository contains the trained CSTNet2 Stage 1 Attn3DGCN checkpoints. The model predicts per-point primitive type and instance embedding, then uses clustering and primitive fitting to construct the final geometric constraint representation.

Source code: xcheng-tsinghua/cstnet2

Training dataset: ZXCCHENGXI/cstnet2_stage1_mini

Checkpoints

semantic/
  last.pth
  best_pmt_miou.pth
  best_cluster_ari.pth
  best_constraint_score.pth
geometry/
  last.pth
  best_pmt_miou.pth
  best_cluster_ari.pth
  best_constraint_score.pth
joint/
  last.pth
  best_pmt_miou.pth
  best_cluster_ari.pth
  best_constraint_score.pth

For downstream constraint extraction, start with joint/best_constraint_score.pth. The other phases are included for reproducibility and phase-specific analysis.

Training configuration

• Backbone: Attn3DGCN
• Points per sample: 2,048
• Input: XYZ plus curvature and local-density features
• Hardware: 4 GPUs with DDP
• Precision: BF16 AMP
• Per-GPU batch size: 16 (global batch size 64)
• Semantic: 30 epochs
• Geometry: 20 epochs
• Joint: 30 epochs

Training-set metrics

These are best training-set metrics stored in each phase checkpoint; no held-out validation split was used for checkpoint selection.

Loading

import torch

from functional.point_features import stage1_feature_dim
from networks.cst_pred_wrapper import CstPredWrapper

checkpoint = torch.load(
    "joint/best_constraint_score.pth",
    map_location="cpu",
    weights_only=False,
)
args = checkpoint["args"]
model = CstPredWrapper(
    args["model"],
    channel_fea=stage1_feature_dim(args.get("use_extra_features", True)),
)
model.load_state_dict(checkpoint["model"], strict=True)
model.eval()

The repository source code is required for feature construction, clustering, and geometric primitive fitting.

Limitations

• Metrics above are training metrics and should not be interpreted as held-out generalization results.
• The model is trained on mechanical CAD-derived point clouds and may not generalize to arbitrary scans or noisy real-world data.
• Checkpoint files use Python/PyTorch serialization. Load only files from trusted sources.