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.

Phase	Primitive mIoU	Cluster ARI	Constraint score
Semantic	0.408995	0.494068	0.451532
Geometry	0.422432	0.560268	0.491337
Joint	0.432907	0.525952	0.479429

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.

Downloads last month: 15

FanMingyu
/

cstnet2-stage1-all