README.md

LRGNet: Learnable Region Growing for Point Clouds

A PyTorch implementation of LRGNet — Learnable Region Growing for Class-Agnostic Point Cloud Segmentation — originally described in:

LRGNet: Learnable Region Growing for Class-Agnostic Point Cloud Segmentation
Jingdao Chen, Zsolt Kira, Yong K. Cho
IEEE Robotics and Automation Letters (RAL), 2021
arXiv:2103.09160
Original code: jingdao/learn_region_grow

---

What this repository does

This repo provides a clean, modern PyTorch port of the LRGNet algorithm that can be applied to any PLY or PCD point cloud without being tied to specific benchmark datasets (S3DIS, ScanNet, etc.).

• Load .ply or .pcd point clouds.
• Preprocess (voxel equalization, normal/curvature estimation).
• Train an LrgNet model on labeled point clouds.
• Segment unlabeled point clouds into class-agnostic object instances.

---

Quick start

1. Install

pip install -r requirements.txt

2. Train on a labeled point cloud

You need a PLY/PCD file and a .npy file with shape (N,) containing integer instance labels per point.

# Stage training data
python scripts/stage_data.py \
    --input my_scene.ply \
    --labels my_scene_labels.npy \
    --output my_scene_staged.h5 \
    --resolution 0.1

# Train
python scripts/train.py \
    --data_dir . \
    --epochs 50 --batch_size 16 --lr 1e-3 \
    --save_dir checkpoints

3. Segment a new point cloud

python scripts/segment.py \
    --input new_scene.ply \
    --ckpt checkpoints/best_model.pt \
    --output segmented.ply \
    --device cuda

Open segmented.ply in CloudCompare or MeshLab — each object instance is colored differently.

---

How the algorithm works internally

Classical region growing vs. LRGNet

	Classical region growing (PCL)	LRGNet
Similarity rule	Hand-crafted threshold on normals / color / curvature	Learned by a neural network
Direction	Can only add points to the region	Can add AND remove points ("morphing")
Class knowledge	Often class-specific or needs tuning	Class-agnostic — works on any object
Recovery	None — one bad seed ruins the segment	Network learns to recover from mistakes

Step-by-step pipeline

Raw point cloud (.ply / .pcd)
        │
        ▼
┌─────────────────────┐
│  Voxel equalization │  ← keep 1 point per voxel (default 0.1 m)
│  (removes density   │     prevents oversampled regions from dominating
│   bias)               │
└─────────────────────┘
        │
        ▼
┌─────────────────────┐
│ Normal + curvature  │  ← PCA on 3×3×3 voxel neighborhood
│ estimation          │     flat areas → low curvature (good seeds)
└─────────────────────┘
        │
        ▼
┌─────────────────────┐
│ Build 13-D feature  │  ← [x,y,z, room_x, room_y, room_z,
│ vector per point    │      r,g,b, nx,ny,nz, curvature]
└─────────────────────┘
        │
        ▼
┌─────────────────────┐
│ Region growing loop │  ← for each seed (sorted by curvature):
│                     │     1. Find boundary neighbors (6-connected voxels)
│   ┌───────────────┐ │     2. Sample 512 inliers + 512 neighbors
│   │   LrgNet      │ │     3. Center coordinates (translation invariance)
│   │  dual-branch  │ │     4. Forward pass → add logits + remove logits
│   │  1D PointNet  │ │     5. Stochastic decision: accept / reject points
│   └───────────────┘ │     6. Repeat until "stuck" or max steps
│                     │
└─────────────────────┘
        │
        ▼
   Colored PLY output

Network architecture (lrg_net.py)

LrgNet is a dual-branch 1D PointNet:

• Inlier branch — processes the current region points.
• Neighbor branch — processes candidate boundary points.

Both branches run a stack of 1D convolutions:

Conv1D(13 -> 64) → Conv1D(64 -> 64) → Conv1D(64 -> 64)
      → Conv1D(64 -> 128) → Conv1D(128 -> 512)

After the deepest layer, a global max-pool extracts one feature vector per branch. These global vectors are tiled back to the point counts and concatenated with per-point features. Two classification heads then predict:

1. add_head — for each neighbor point: probability it should join the region.
2. remove_head — for each inlier point: probability it should leave the region.

Both heads use binary cross-entropy loss and are trained jointly.

Why two heads?
Classical region growing is monotonic — it can only expand. LRGNet can shrink the region too. This is critical for recovering from an early bad seed or an over-segmentation.

Training data generation (stage_data.py)

Training is fully supervised from ground-truth instance labels:

1. For each object instance, pick k random seeds inside it.
2. Simulate region growing using the true labels.
3. At every step, record:
   • the current (possibly noisy) inlier set
   • the boundary neighbor set
   • ground-truth binary labels: add / don't add, remove / don't remove
4. Inject controlled noise (add_mistake_prob ≈ 0.2, remove_mistake_prob ≈ 0.2)
   → the network sees "messy" regions and learns to clean them up.

Centering (translation invariance)

At both training and inference time, the XYZ coordinates of inliers and neighbors are centered by subtracting the median XY of the inlier set, and feature channels (RGB, normals, curvature) are centered by subtracting their median over the inlier set. This makes the network translation invariant — it only cares about local geometry and appearance, not absolute room coordinates.

---

Repository layout

learn_region_grow/
├── __init__.py          # Package metadata
├── io.py                # PLY / PCD loaders and savers + color maps
├── preprocess.py         # Voxel equalization, normal/curvature, feature vector
├── lrg_net.py           # PyTorch LrgNet model
├── growing.py           # Region growing inference engine
├── stage_data.py        # Training data generator from labeled clouds
├── train.py             # PyTorch training loop
└── utils.py             # Sampling, centering, clustering metrics

scripts/
├── segment.py           # CLI: run inference on a new point cloud
├── train.py             # CLI: train from staged H5 files
└── stage_data.py        # CLI: generate staged H5 from labeled cloud

---

Parameters & tuning tips

Parameter	Default	Effect
resolution	0.1 m	Voxel grid size. Smaller = more detail, slower.
lite	0	Model size: 0=full, 1=half, 2=quarter channels.
add_threshold / remove_threshold	0.5	Confidence cutoffs for deterministic mode.
stochastic	True	Random sampling weighted by confidence (paper default).
cluster_threshold	10	Minimum points to keep a segment.
seeds_per_instance	5	Random seeds per object during training staging.
add_mistake_prob	0.2	Noise injected into "add" labels during staging.
remove_mistake_prob	0.2	Noise injected into "remove" labels during staging.

When to change resolution

• Indoor scenes (S3DIS, ScanNet): 0.1 m works well.
• Outdoor LiDAR (Semantic KITTI): use 0.3 m because objects are larger and sparser.
• Small objects / high detail: try 0.05 m.

When to use lite

• lite=0: best accuracy, use on GPU.
• lite=1 or 2: faster inference, lower memory, slight accuracy drop.

---

Differences from the original TensorFlow 1.x code

Aspect	Original	This port
Framework	TensorFlow 1.x (tf.compat.v1)	PyTorch 2.x
Input format	H5 dumps from S3DIS / ScanNet	Any PLY or PCD
Model branches	Shared conv weights	Independent branches (clearer, no leakage)
Normals	3×3×3 voxel window	Same, with k-NN fallback for sparse clouds
Seed ordering	Curvature ascending	Same
Loss	TF cross-entropy	PyTorch BCEWithLogitsLoss
Output	Colored PLY only	PLY + PCD, with label IDs

---

Citation

If you use this code, please cite the original paper:

@article{chen2021lrgnet,
  title={LRGNet: Learnable Region Growing for Class-Agnostic Point Cloud Segmentation},
  author={Chen, Jingdao and Kira, Zsolt and Cho, Yong K},
  journal={IEEE Robotics and Automation Letters},
  volume={6},
  number={3},
  pages={5205--5212},
  year={2021},
  publisher={IEEE}
}

---

License

This port is released under the MIT license. The original repository is available at jingdao/learn_region_grow.