SStrainer3D

SStrainer3D is a tool set for deep-learning based semantic segmentation (SS) for 3D volumes such as X-ray CT image data. You can use the tool set for training SS model and predicting SS. This tool set has been distributed as Python scripts, so you can edit any scripts as you want under the license.

Installation

This tool set requires Python Version 3.10, CUDA and Git. Please install them in advance.

Clone the repository:

$ git clone https://github.com/st707311g/SStrainer3D.git
$ cd SStrainer3D

Install pytorch, and the following dependency:

$ pip install -r requirements.txt

Check whether PyTorch is available.

$ python
>>> import torch
>>> torch.cuda.is_available()
True
>>> exit()

Demo

Training

Make a test volume:

python make_test_volume.py

The test dataset was created in datasets/test_dataset/, conaining 2 directories, namely 00_input, 01_output. The directory 01_output is the ground truth (GT), 400 cylinders segments. The directory 00_input is GT with fractal noise as background.

input	ground truth

Make a dataset for training.

$ python make_dataset.py --root datasets/test_dataset/ --input input --output output

The directory 02_training_dataset is created in datasets/test_dataset/.

Run model training:

$ python train.py --root datasets/test_dataset/ --model_name test --batch_size 16 --epoch 100

The trained model is saved in the directory models/.

Check the loss history:

$ python plot_loss_history.py --model_name test

Then, make additonal test volume:

$ python make_test_volume.py

The test dataset is created in datasets/test_dataset_1/.

Predict the SS using the trained model:

$ python predict.py --root datasets/test_dataset_1/ --model_name test --input input --output predicted

Predicted volume is saved in the directory datasets/test_dataset_1/02_predicted/.

The input volume:

input	ground truth	predicted

Tips

• Multi-GPU training is possible by using --use_multi_gpu option.
• Multi-class training is possible by indicating the class names after --input option.
• Multi-class predicting is possible by indicating the class names after --output option.

License

NARO NON-COMMERCIAL LICENSE Version 1.0

Citations

We are currently submitting a paper on this software. You are free to use this software according to the license, but do not publish the results of using this software before our submitting paper has been accepted and published.