SegMesh: 3D Mesh Segmentation

A fast CUDA-accelerated (GPU) method that uses novel mesh convolutions (spherical harmonics) and neural networks (machine learning/NN) for efficient scene segmentation.

Paper

Based on Mesh Convolution With Continuous Filters for 3-D Surface Parsing (IEEE Transactions on Neural Networks and Learning Systems 2023)

 

Introduction

SegMesh introduces a range of modular operations tailored for 3D triangular mesh segmentation. These include innovative mesh segmentation techniques, GPU-accelerated mesh simplification, and precise mesh pooling/unpooling methods.

Our approach uses spherical harmonics for creating segmentation filters. The GPU-accelerated mesh simplification efficiently processes batched meshes, while our pooling operations cater to varying mesh resolutions. SegMesh also encompasses a neural network designed for 3D surface segmentation, exhibiting exceptional performance in shape analysis and scene segmentation on major 3D benchmarks.

This journal work is a sigificant extension of our original work presented in CVPR 2021.

1. Setup

Tested on Ubuntu 22.04.2 LTS, Cuda 11.7.1, Pytorch 2.0.1. To run within a conda environment, and the following hardware: GPU 3090, 4090 for multiple versions of python /pytorch /cuda /cudnn /ubuntu. Requires GPU support for blocks of 1024 threads.

conda create --prefix .conda python==3.8 -y
conda activate ./.conda
conda install -c "nvidia/label/cuda-11.7.1" cuda -y
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu117
pip install -r requirements.txt
python setup.py install

2. Data Preparation

Dataset(s) supported by Segmesh

Dataset Name	Nature of Dataset	Task	Number of Samples	Categories/Classes	Evaluation (%)	Terms of Usage
ShapeNetCore	Synthetic	Shape Classification	~51,000	55	87.3	Academic
SHREC	Synthetic	Shape Classification	600 (20 samples/class)	30	100	Commercial
CUBE	Synthetic	Shape Classification	4,400 (200 samples/class)	22	100	Commercial
COSEG	Synthetic	Semantic Labelling	Varies per category	3 (aliens, chairs, vases)	98.8, 99.5, 95.6	Commercial
HUMAN	Synthetic	Semantic Labelling	399 (381 training, 18 test)	8 (body parts)	91.5	Commercial
FAUST	Synthetic	3D Manifold Correspondence	100 (80 training, 20 test)	10 subjects with 10 poses each	100	Academic
S3DIS	Real	Shape Segmentation	6 areas (Area 5 for testing)	13 (indoor areas features)	91.3(OA), 77.2(mAcc), 71.0(mIoU)	Academic
ScanNet	Real	Shape Segmentation	1,613 (1,213 training, 300 validation)	40	69.2(mIoU)	Academic

Downloading the Datasets

Datasets can be downloaded following these instructions:

• SHREC, CUBE, COSEG, HUMAN: These datasets are available and can be downloaded from the provided links in the ./data directory.
  • Coseg: Download
  • Human: Download
  • Cubes: Download
  • Shrec16: Download

Alternatively, you can extract these shared datasets of MeshCNN using download_extract_shapes.py script.

  cd data
  python download_extract_shapes.py

• ShapeNetCore, FAUST, S3DIS, ScanNet: These datasets are restricted to academic use. Follow the links to their respective websites for download instructions and adhere to their usage terms.
  • ShapeNetCore: License & Download
  • FAUST: License & Download
  • S3DIS: Terms & Access Form
  • ScanNet: License & Download

Please ensure to follow the respective license terms and agreements when using these datasets.

Unzip the datasets after downloading.

After extracting its compressed file, follow the steps in the Preprocess section to prepare the labels for these datasets.

3. Usage of Segmesh

1. Preprocess

To preprocess datasets for training, run the preprocess script with the specific dataset you want to prepare. Update the main function in preprocess.py if needed. Here are the commands for each supported dataset:

python -m segmesh.datasets.preprocess --dataset cubes --config_file ./config/cubes.yaml
python -m segmesh.datasets.preprocess --dataset shrec --config_file ./config/shrec.yaml
python -m segmesh.datasets.preprocess --dataset coseg_aliens --config_file ./config/coseg_aliens.yaml
python -m segmesh.datasets.preprocess --dataset coseg_chairs --config_file ./config/coseg_chairs.yaml
python -m segmesh.datasets.preprocess --dataset coseg_vases --config_file ./config/coseg_vases.yaml
python -m segmesh.datasets.preprocess --dataset humanseg --config_file ./config/human.yaml
python -m segmesh.datasets.preprocess --dataset faust --config_file ./config/faust_match.yaml

Make sure to update the paths in the config file before you run the command.

2. Training

Shape Classification

python -m train.train_cubes
python -m train.train_shrec
python -m train.train_shapenetcore
python -m train.train_faust_match

Before you start the training for coseg dataset, update the config file path in the train/train_coseg.py to train a specific the coseg dataset (chairs/aliens/vases).

Shape Segmentation

python -m train.train_human
python -m train.train_coseg

Scene Segmentation

python -m train.train_s3dis_render
python -m train.train_scannet_render

3. Support for custom datasets

To add support for a new or custom dataset, follow the steps:-

1. To add custom support for dataloaders and collate functions, create a Dataset class similar to shape_dataset.py and load in the training script
2. To add custom augmentations and normalize functionality, update the augmentations.py and normalize.py with the dataset specific class definitions.
3. Replicate the training code in a new file and update the BaseTrainer and DatasetHandler classes. Load other custom classes.
4. Finally, update the input/outputs and loss functions in train_one_epoch function of fit.py if required.

4. Inference

SegMesh provides a inference framework for various tasks, including scene segmentation, shape classification, and shape segmentation.

Inference Code Structure

The inference code is organized into several Python files, each handling a specific task:

• inference_scene_seg.py: Performs scene segmentation on 3D scenes. It processes 3D scene data and outputs segmented results.
• inference_shape_cls.py: Conducts shape classification on 3D shapes. This script takes a 3D shape and classifies it into predefined categories.
• inference_shape_seg.py: Deals with shape segmentation, where it segments different parts of a given 3D shape.
• __init__.py: Initialization file for the inference module.

Running Inference

To run the inference for a specific task, use the following command structure:

1. Scene Segmentation:

  python -m inference.inference_scene_seg --config path/to/scene_seg_config.yaml --dataset dataset_name --model_path path/to/model --mesh_path path/to/mesh_file.h5 --label_path path/to/label_file.txt

Example command to run inference for a S3DIS dataset test sample:

python -m inference.inference_scene_seg --config ./config/s3dis.yaml --dataset s3dis --model_path ./runs_scenes/s3dis_render_20230913_184719/model_epoch_50 --mesh_file ./data/S3DIS_3cm_hdf5_Rendered/Area_6/conferenceRoom_1.h5 --label_file ./data/S3DIS_3cm_hdf5_Rendered/Area_6/conferenceRoom_1.txt

2. Shape Classification:

python -m inference.inference_shape_cls --config path/to/shape_cls_config.yaml --model_path path/to/model --mesh_file path/to/mesh_file.obj --label_file path/to/label_file.txt

3. Shape Segmentation:

python -m inference.inference_shape_seg --config path/to/shape_seg_config.yaml --model_path path/to/model --mesh_file path/to/mesh_file.obj --label_file path/to/label_file.txt

Example command to run inference for Human dataset sample:

python -m inference.inference_shape_seg --dataset human --model_path './runs_shapes/human_20230922_121456/model_epoch_100' --mesh_file './data/human_seg/test/shrec__10.obj' --label_file './data/human_seg/face_label/shrec__10.txt'

Replace path/to/config_file.yaml with the path to your dataset-specific configuration file, path/to/model with the path to your trained model, and path/to/mesh_file and path/to/label_file with the appropriate paths to your data files.

5. Evaluation

SegMesh provides a evaluation code structure, allowing easy integration and evaluation of various 3D datasets like S3DIS and ScanNet. The evaluation structure is designed to be flexible, enabling the addition of new datasets with minimal changes.

Evaluation Code Structure

The evaluation code is organized into several key subclass files:

• base_evaluator.py: Contains the BaseEvaluator class, which provides common functionalities for all datasets.
• {dataset_name}_evaluator.py: Dataset-specific evaluator classes (e.g., s3dis_evaluator.py, scannet_evaluator.py) that inherit from BaseEvaluator.
• transform_texture.py: Handles dataset-specific data transformations.
• main.py: The main script to select and run the appropriate evaluator based on the dataset.

Running the Evaluation

To run the evaluation for a specific dataset, run the main.py script with the required arguments from the root project directory:

python -m eval.main --config path/to/config_file.yaml --dataset [s3dis | scannet | ...]

Replace path/to/config_file.yaml with the path to your dataset-specific configuration file, and [s3dis | scannet | coseg | ...] with the desired dataset name.

Citation

If you find our work useful in your research, please consider citing:

@article{lei2023mesh,
  title={Mesh Convolution With Continuous Filters for 3-D Surface Parsing},
  author={Lei, Huan and Akhtar, Naveed and Shah, Mubarak and Mian, Ajmal},
  journal={IEEE Transactions on Neural Networks and Learning Systems},
  year={2023},
  publisher={IEEE}
}

@inproceedings{lei2021picasso,
  title={Picasso: A CUDA-based Library for Deep Learning over 3D Meshes},
  author={Lei, Huan and Akhtar, Naveed and Mian, Ajmal},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
  pages={13854--13864},
  year={2021}
} 

Please also cite the original dataset if you use their data or our reprocessed data, and follow their original terms of use.

Licensing

SegMesh is available for non-commercial internal research use by academic institutions or not-for-profit organisations only, free of charge. Please, see the license for further details. To the extent permitted by applicable law, your use is at your own risk and our liability is limited. Interested in a commercial license? For commercial queries, please email [email protected] with subject line "SegMesh Commercial License".

This is an AIML Shop project.