Image Understands Point Cloud: Weakly Supervised 3D Semantic Segmentation via Association Learning (TIP2024)
Tianfang Sun1, Zhizhong Zhang1, Xin Tan1, Yanyun Qu2, Yuan Xie1
1ECNU, 2XMU
For easy installation, we recommend using conda:
conda create -n iupc python=3.9
conda activate iupc
conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.3 -c pytorch
pip3 install numba nuscenes-devkit
# We need opencv-contrib-python to perform superpixel segment and it is not compatible with opencv-python
pip3 uninstall opencv-python
pip3 install opencv-contrib-python
# We need open3d and hdbscan to perform point cloud pre-segmentation
pip3 install open3d
pip3 install hdbscan
pip3 install tensorboard
Our method is based on torchpack and torchsparse. To install torchpack, we recommend to firstly install openmpi and mpi4py.
conda install -c conda-forge mpi4py openmpi
Install torchpack
pip install git+https://github.com/zhijian-liu/torchpack.git
Before installing torchsparse, install Google Sparse Hash library first.
sudo apt install libsparsehash-dev
Then install torchsparse (v1.4.0) by
pip3 install --upgrade git+https://github.com/mit-han-lab/torchsparse.git@v1.4.0
Please download ImageNet pretrained weight for SwiftNet from Google Drive or BaiduDisk.
Please download the datasets following the official instruction. The official websites of each dataset are listed as following: nuScenes_lidarseg, SemanticKITTI, Waymo open. The color images of SemanticKITTI datasets can be downloaded from KITTI-odometry dataset.
Please run the following command to cluster and generate weak labels for the nuScenes_lidarseg dataset. The whole process will take several hours depended on the CPU.
python3 prepare_sparse_label.py configs/nuscenes/pretrain/preseg.yaml
Statistics | nuScenes | Waymo | SemanticKITTI |
---|---|---|---|
Sparse label | 0.8% | 0.3% | 0.08% |
Propagated label | 48.5% | 21.2% | 22.7% |
Negative label | 44.5% | 76.7% | 70.6% |
Please run the following command to generate superpixels for the nuScenes_lidarseg dataset. The whole process will take several hours depended on the CPU.
python3 prepare_nusc_superpixel.py
- Run the following command to train the model.
CUDA_VISIBLE_DEVICES=0,1,2 torchpack dist-run -np 3 python3 train_lc_sparse_assop2p_nusc.py configs/nuscenes/pretrain/assop2p.yaml --run-dir runs/iupc_nusc/assop2pw0p5_visw0p5_sybn
- For the first round, run the following command to generate pseudo label.
CUDA_VISIBLE_DEVICES=0,1,2 torchpack dist-run -np 3 python3 train_lc_estep_nusc.py configs/nuscenes/estep/it1.yaml --run-dir runs/iupc_nusc/assop2pw0p5_visw0p5_sybn --weight-path runs/iupc_nusc/assop2pw0p5_visw0p5_sybn/checkpoints/max-iou-vox-val.pt
- Run the following command to use the generated pseudo label for training.
CUDA_VISIBLE_DEVICES=0,1,2 torchpack dist-run -np 3 python3 train_lc_sparse_assop2p_mstep_nusc.py configs/nuscenes/mstep/assop2p.yaml --run-dir runs/iupc_nusc/assop2pw0p5_visw0p5_sybn
- For the second round, run the following command to generate pseudo label.
CUDA_VISIBLE_DEVICES=0,1,2 torchpack dist-run -np 3 python3 train_lc_estep_nusc_2.py configs/nuscenes/estep/it2.yaml --run-dir runs/iupc_nusc/assop2pw0p5_visw0p5_sybn --weight-path runs/iupc_nusc/assop2pw0p5_visw0p5_sybn/checkpoints_mstep_0/max-iou-vox-val.pt
- Run the following command to use the generated pseudo label for training.
CUDA_VISIBLE_DEVICES=0,1,2 torchpack dist-run -np 3 python3 train_lc_sparse_assop2p_mstep_nusc.py configs/nuscenes/mstep/assop2p_it2.yaml --run-dir runs/iupc_nusc/assop2pw0p5_visw0p5_sybn
Iter 0 | Iter 1 | Iter 2 | Full | |
---|---|---|---|---|
mIoU(paper) | 77.1 | 77.3 | 77.5 | 78.8 |
mIoU(reprod.) | 76.7 | 76.9 | 77.4 | 78.6 |
- Test and release the codes for nuScenes datasets.
- Test and release the codes for SemanticKITTI datasets.
- Test and release the codes for Waymo datasets.