linflonet 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- linflonet-0.1.0/LICENSE +21 -0
- linflonet-0.1.0/PKG-INFO +302 -0
- linflonet-0.1.0/README.md +270 -0
- linflonet-0.1.0/linflonet/__init__.py +3 -0
- linflonet-0.1.0/linflonet/__main__.py +3 -0
- linflonet-0.1.0/linflonet/cli.py +210 -0
- linflonet-0.1.0/linflonet/data/__init__.py +0 -0
- linflonet-0.1.0/linflonet/data/template/__init__.py +0 -0
- linflonet-0.1.0/linflonet/data/template/highres_template.vtp +44 -0
- linflonet-0.1.0/linflonet/data/template/highres_template_distance.vtk +0 -0
- linflonet-0.1.0/linflonet/data/template/whole_heart_with_ao.vtp +44 -0
- linflonet-0.1.0/linflonet/paths.py +32 -0
- linflonet-0.1.0/linflonet/predict.py +236 -0
- linflonet-0.1.0/linflonet.egg-info/PKG-INFO +302 -0
- linflonet-0.1.0/linflonet.egg-info/SOURCES.txt +45 -0
- linflonet-0.1.0/linflonet.egg-info/dependency_links.txt +1 -0
- linflonet-0.1.0/linflonet.egg-info/entry_points.txt +2 -0
- linflonet-0.1.0/linflonet.egg-info/requires.txt +11 -0
- linflonet-0.1.0/linflonet.egg-info/top_level.txt +3 -0
- linflonet-0.1.0/pyproject.toml +54 -0
- linflonet-0.1.0/setup.cfg +4 -0
- linflonet-0.1.0/src/__init__.py +0 -0
- linflonet-0.1.0/src/chamfer.py +224 -0
- linflonet-0.1.0/src/dataset.py +68 -0
- linflonet-0.1.0/src/finite_difference.py +39 -0
- linflonet-0.1.0/src/flow.py +116 -0
- linflonet-0.1.0/src/flow_loss.py +62 -0
- linflonet-0.1.0/src/integrator.py +114 -0
- linflonet-0.1.0/src/io_utils.py +204 -0
- linflonet-0.1.0/src/linear_transform.py +208 -0
- linflonet-0.1.0/src/loss.py +109 -0
- linflonet-0.1.0/src/pre_process.py +146 -0
- linflonet-0.1.0/src/random_linear_perturbations.py +29 -0
- linflonet-0.1.0/src/segment_flow.py +516 -0
- linflonet-0.1.0/src/signed_distance_function.py +221 -0
- linflonet-0.1.0/src/template.py +237 -0
- linflonet-0.1.0/src/unet.py +143 -0
- linflonet-0.1.0/src/unet_components.py +193 -0
- linflonet-0.1.0/src/unet_segment.py +57 -0
- linflonet-0.1.0/src/utilities.py +40 -0
- linflonet-0.1.0/test/test_chamfer_normal.py +15 -0
- linflonet-0.1.0/test/test_crop.py +52 -0
- linflonet-0.1.0/test/test_cropped_evaluation.py +32 -0
- linflonet-0.1.0/test/test_linear_transform_distance_function.py +75 -0
- linflonet-0.1.0/test/test_matching_image_label_sizes.py +23 -0
- linflonet-0.1.0/vtk_utils/__init__.py +0 -0
- linflonet-0.1.0/vtk_utils/vtk_utils.py +929 -0
linflonet-0.1.0/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2023 Arjun Narayanan
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
linflonet-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,302 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: linflonet
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Generate simulation-ready 3D heart meshes from CT and MR images
|
|
5
|
+
Author: LinFlo-Net contributors
|
|
6
|
+
License: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/ArjunNarayanan/LinFlo-Net
|
|
8
|
+
Project-URL: Repository, https://github.com/ArjunNarayanan/LinFlo-Net
|
|
9
|
+
Keywords: medical-imaging,mesh-generation,deep-learning,cardiac
|
|
10
|
+
Classifier: Development Status :: 4 - Beta
|
|
11
|
+
Classifier: Intended Audience :: Science/Research
|
|
12
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
13
|
+
Classifier: Programming Language :: Python :: 3
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
17
|
+
Classifier: Topic :: Scientific/Engineering :: Medical Science Apps.
|
|
18
|
+
Requires-Python: >=3.10
|
|
19
|
+
Description-Content-Type: text/markdown
|
|
20
|
+
License-File: LICENSE
|
|
21
|
+
Requires-Dist: numpy<2,>=1.26.0
|
|
22
|
+
Requires-Dist: pandas>=2.1.1
|
|
23
|
+
Requires-Dist: torch>=2.2.0
|
|
24
|
+
Requires-Dist: monai>=1.3.0
|
|
25
|
+
Requires-Dist: SimpleITK>=2.3.1
|
|
26
|
+
Requires-Dist: vtk>=9.3.0
|
|
27
|
+
Requires-Dist: PyYAML>=6.0.1
|
|
28
|
+
Provides-Extra: dev
|
|
29
|
+
Requires-Dist: build; extra == "dev"
|
|
30
|
+
Requires-Dist: twine; extra == "dev"
|
|
31
|
+
Dynamic: license-file
|
|
32
|
+
|
|
33
|
+
# LinFlo-Net
|
|
34
|
+
|
|
35
|
+
**New to LinFlo-Net?** See the [Quick start guide](docs/quick_start.md) for install and prediction in a few minutes.
|
|
36
|
+
|
|
37
|
+
A deep learning package to automatically generate simulation ready 3D meshes of the human heart from biomedical images. [Link to paper](https://asmedigitalcollection.asme.org/biomechanical/article/doi/10.1115/1.4064527/1194613).
|
|
38
|
+
|
|
39
|
+

|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
## Setting up environment on Savio
|
|
43
|
+
|
|
44
|
+
The following instructions can set up a conda environment on the Berkeley Research Computing Savio system. But a similar approach can be used on any SLURM based high-performance computing cluster.
|
|
45
|
+
|
|
46
|
+
Since Savio provides limited space in your home directory, we install all conda packages to our scratch folder.
|
|
47
|
+
|
|
48
|
+
```commandline
|
|
49
|
+
module load cuda/10.2
|
|
50
|
+
module load gcc/5.4.0
|
|
51
|
+
ENVDIR=/global/scratch/users/<your_username>/environments/linflonet
|
|
52
|
+
rm -rf $ENVDIR
|
|
53
|
+
export CONDA_PKGS_DIRS=/global/scratch/users/<your_username>/tmp/.conda
|
|
54
|
+
conda create --prefix $ENVDIR
|
|
55
|
+
```
|
|
56
|
+
|
|
57
|
+
Press `y` when prompted to create your conda environment and then activate your environment,
|
|
58
|
+
|
|
59
|
+
```commandline
|
|
60
|
+
source activate $ENVDIR
|
|
61
|
+
```
|
|
62
|
+
|
|
63
|
+
Next install `pytorch`. Savio does not have the version of `cuda` required for the latest pytorch version, so we will install `pytorch 1.12.1`.
|
|
64
|
+
|
|
65
|
+
```commandline
|
|
66
|
+
conda install pytorch==1.12.1 cudatoolkit=10.2 -c pytorch
|
|
67
|
+
```
|
|
68
|
+
|
|
69
|
+
Press `y` when prompted to start the installation.
|
|
70
|
+
|
|
71
|
+
Next we will install [pytorch3d](https://pytorch3d.org/) which provides several useful routines for dealing with 3D data and mesh data-structures in conjunction with `pytorch`,
|
|
72
|
+
|
|
73
|
+
```commandline
|
|
74
|
+
conda install -c fvcore -c iopath -c conda-forge fvcore iopath
|
|
75
|
+
conda install -c bottler nvidiacub
|
|
76
|
+
conda install pytorch3d -c pytorch3d
|
|
77
|
+
```
|
|
78
|
+
|
|
79
|
+
You can find all the other dependencies in the system generated `requirements.txt` in the repository. You should be able to install these directly with `pip` after installing the above packages.
|
|
80
|
+
|
|
81
|
+
### Test that everything works
|
|
82
|
+
|
|
83
|
+
First request a brief interactive session with a GPU,
|
|
84
|
+
|
|
85
|
+
```commandline
|
|
86
|
+
srun --pty -A <account_name> -p savio3_gpu --nodes=1 --gres=gpu:GTX2080TI:1 --ntasks=1 --cpus-per-task=2 -t 00:30:00 bash -i
|
|
87
|
+
```
|
|
88
|
+
|
|
89
|
+
Once your resources are allocated, load your conda environment and launch python
|
|
90
|
+
|
|
91
|
+
```commandline
|
|
92
|
+
source activate /global/scratch/users/<your_username>/environments/linflonet
|
|
93
|
+
python
|
|
94
|
+
```
|
|
95
|
+
|
|
96
|
+
Now type the following into your python session,
|
|
97
|
+
|
|
98
|
+
```python
|
|
99
|
+
import torch
|
|
100
|
+
from pytorch3d.loss import chamfer_distance
|
|
101
|
+
device = torch.device("cuda")
|
|
102
|
+
|
|
103
|
+
a = torch.rand([5,10000,3]).to(device)
|
|
104
|
+
b = torch.rand([5,10000,3]).to(device)
|
|
105
|
+
loss = chamfer_distance(a, b)
|
|
106
|
+
```
|
|
107
|
+
|
|
108
|
+
If everything runs without error, you are all set!
|
|
109
|
+
|
|
110
|
+
## Install from PyPI
|
|
111
|
+
|
|
112
|
+
For prediction only (Python 3.10+):
|
|
113
|
+
|
|
114
|
+
```commandline
|
|
115
|
+
pip install linflonet
|
|
116
|
+
```
|
|
117
|
+
|
|
118
|
+
`pytorch3d` is required but not listed as a pip dependency because it must be
|
|
119
|
+
built from source on most platforms. Install `torch` first, then:
|
|
120
|
+
|
|
121
|
+
```commandline
|
|
122
|
+
pip install --no-build-isolation "git+https://github.com/facebookresearch/pytorch3d.git@stable"
|
|
123
|
+
```
|
|
124
|
+
|
|
125
|
+
## CLI usage
|
|
126
|
+
|
|
127
|
+
The `linflonet` command generates heart meshes (`.vtp`) and segmentations for
|
|
128
|
+
CT or MR NIfTI images.
|
|
129
|
+
|
|
130
|
+
**Single image:**
|
|
131
|
+
|
|
132
|
+
```commandline
|
|
133
|
+
linflonet predict \
|
|
134
|
+
--image /path/to/scan.nii.gz \
|
|
135
|
+
--model /path/to/best_model.pth \
|
|
136
|
+
--modality ct \
|
|
137
|
+
--output /path/to/output
|
|
138
|
+
```
|
|
139
|
+
|
|
140
|
+
**Folder of images** (flat folder or `image/` subdirectory):
|
|
141
|
+
|
|
142
|
+
```commandline
|
|
143
|
+
linflonet predict \
|
|
144
|
+
--folder /path/to/images \
|
|
145
|
+
--model /path/to/best_model.pth \
|
|
146
|
+
--modality mr \
|
|
147
|
+
--output /path/to/output
|
|
148
|
+
```
|
|
149
|
+
|
|
150
|
+
Template mesh and distance map default to files bundled with the package. Override
|
|
151
|
+
with `--template` and `--template-distance-map` if needed. For linear-transform-only
|
|
152
|
+
models, pass `--linear-transform`.
|
|
153
|
+
|
|
154
|
+
**Using a YAML config** (same format as `config/predict_single_ct.yml`):
|
|
155
|
+
|
|
156
|
+
```commandline
|
|
157
|
+
linflonet predict --config config/predict_single_ct.yml --image /path/to/scan.nii.gz -o /path/to/output
|
|
158
|
+
```
|
|
159
|
+
|
|
160
|
+
Outputs are written to `<output>/meshes/` and `<output>/segmentation/`.
|
|
161
|
+
|
|
162
|
+
You can also run `python -m linflonet predict ...` or install in editable mode
|
|
163
|
+
from a git checkout:
|
|
164
|
+
|
|
165
|
+
```commandline
|
|
166
|
+
pip install -e .
|
|
167
|
+
```
|
|
168
|
+
|
|
169
|
+
## Setting up a local environment with pip (Python 3.12)
|
|
170
|
+
|
|
171
|
+
If you only need to run prediction (not training), you can set up a lightweight
|
|
172
|
+
environment on Python 3.12 using `requirements-py312.txt`.
|
|
173
|
+
|
|
174
|
+
First, initialize the `vtk_utils` submodule and create a virtual environment,
|
|
175
|
+
|
|
176
|
+
```commandline
|
|
177
|
+
git submodule update --init
|
|
178
|
+
python3.12 -m venv .venv
|
|
179
|
+
source .venv/bin/activate
|
|
180
|
+
pip install -r requirements-py312.txt
|
|
181
|
+
```
|
|
182
|
+
|
|
183
|
+
`pytorch3d` does not ship prebuilt wheels for most platforms and must be built
|
|
184
|
+
from source *after* `torch` is installed. Its `setup.py` imports `torch` at build
|
|
185
|
+
time, so you must disable pip's build isolation (otherwise you get
|
|
186
|
+
`ModuleNotFoundError: No module named 'torch'`):
|
|
187
|
+
|
|
188
|
+
```commandline
|
|
189
|
+
pip install --no-build-isolation "git+https://github.com/facebookresearch/pytorch3d.git@stable"
|
|
190
|
+
```
|
|
191
|
+
|
|
192
|
+
On macOS, make sure the Xcode command-line tools are installed first
|
|
193
|
+
(`xcode-select --install`) so the C++ extension can compile.
|
|
194
|
+
|
|
195
|
+
## Dataset Creation
|
|
196
|
+
|
|
197
|
+
We use the multi-modality whole heart segmentation challenge (MMWHS) [dataset](https://zmiclab.github.io/zxh/0/mmwhs/). Download and unzip the data. You should have the following folders,
|
|
198
|
+
|
|
199
|
+
- CT : 2 folders each with 10 images and corresponding segmentations
|
|
200
|
+
- MR : 1 folder with 20 images and corresponding segmentations
|
|
201
|
+
|
|
202
|
+
You can split the data into train and validation as you find appropriate. We chose to use the first 16 samples as training and the remaining 4 samples as validation. Split the data appropriately and place them in separate folders. Make sure to keep the CT and MR data separately as we will need to normalize / scale them differently. We will perform data augmentation on the training data.
|
|
203
|
+
|
|
204
|
+
### Data augmentation
|
|
205
|
+
|
|
206
|
+
We will use the data augmentation procedure available in the [MeshDeformNet](https://github.com/fkong7/MeshDeformNet) package. Clone this package to your system and run `pip install -r requirements.txt` to install package dependencies. (You may want to create a virtual environment first.)
|
|
207
|
+
|
|
208
|
+
To perform augmentation, modify the command below and execute it. The script below launches 16 jobs in parallel (`-np 16`). You can modify that depending on the capacity of the system you are using.
|
|
209
|
+
|
|
210
|
+
```commandline
|
|
211
|
+
mpirun -np 16 python ~/path/to/MeshDeformNet/data/data_augmentation.py \
|
|
212
|
+
--im_dir /path/to/image/directory \
|
|
213
|
+
--seg_dir /path/to/segmentation/directory \
|
|
214
|
+
--out_dir /path/to/output/directory \
|
|
215
|
+
--modality ct or mr \
|
|
216
|
+
--mode train \
|
|
217
|
+
--num number_of_augmentations
|
|
218
|
+
```
|
|
219
|
+
|
|
220
|
+
The output folder will contain two subfolders `modality_train` with the augmented images and `modality_train_seg` with the augmented segmentations where modality is either `ct` or `mr`.
|
|
221
|
+
|
|
222
|
+
### Creating ground-truth meshes
|
|
223
|
+
|
|
224
|
+
We generate ground-truth meshes using marching cubes on the ground-truth segmentations. We can do this using `workflows/prepare_data.py`.
|
|
225
|
+
|
|
226
|
+
```commandline
|
|
227
|
+
python workflows/prepare_data \
|
|
228
|
+
--image /path/to/image/folder \
|
|
229
|
+
--segmentation /path/to/segmentation/folder \
|
|
230
|
+
--output /path/to/output/folder \
|
|
231
|
+
--modality ct # can be either ct or mr
|
|
232
|
+
--ext .nii.gz # input files extension
|
|
233
|
+
```
|
|
234
|
+
|
|
235
|
+
The output folder is going to have 3 subfolders : `seg`, `vtk_image`, `vtk_mesh`. `vtk_image` will be the input to our neural network, and `vtk_mesh` will be the corresponding ground truth meshes. From this point onward, we assume that the folder with the relevant data has the `vtk_image` and `vtk_mesh` subfolders.
|
|
236
|
+
|
|
237
|
+
### Final steps
|
|
238
|
+
|
|
239
|
+
The data set is reasonably large, and we will have to load it from memory. It is useful to store the images as pytorch tensors and the meshes as pytorch3d data structures in pickled files. To do this, we first build a csv index of all the files.
|
|
240
|
+
|
|
241
|
+
```commandline
|
|
242
|
+
python utilities/prepare_train_dataset_csv.py -f /path/to/data/folder
|
|
243
|
+
```
|
|
244
|
+
|
|
245
|
+
Make sure to provide the path to the parent directory containing `vtk_image` and `vtk_mesh` sub-directories. This will create an `index.csv` in the parent folder with the names of all files. Next,
|
|
246
|
+
|
|
247
|
+
```commandline
|
|
248
|
+
python utilities/pickle_image_segmentation_mesh_dataset.py -config /path/to/config/file
|
|
249
|
+
```
|
|
250
|
+
|
|
251
|
+
Look at `config/pickle_dataset.yml` for an example config file. Note that `seg_label` in the config file follows the labelling convention of the MMWHS dataset.
|
|
252
|
+
|
|
253
|
+
The output folder will now contain `.pkl` files which contain the combined image, segmentations, and meshes in a dictionary. This can be used by our dataloader to load the appropriate files during training.
|
|
254
|
+
|
|
255
|
+
|
|
256
|
+
## Training the model
|
|
257
|
+
|
|
258
|
+
Before training, make sure to activate the conda environment that we created earlier. Request a GPU session if you would like to use a GPU for training. Alternatively, submit the below commands as part of a batch job with `sbatch` on a SLURM system. The training workflow will save the best performing model as a checkpoint in the output directory specified in the config file.
|
|
259
|
+
|
|
260
|
+
### Training Linear Transformation module
|
|
261
|
+
|
|
262
|
+
Take a look at the example config file in `config/linear_transform.yml`. Make a copy, and modify it appropriately.
|
|
263
|
+
|
|
264
|
+
Then run the command,
|
|
265
|
+
|
|
266
|
+
```commandline
|
|
267
|
+
python workflows/train_linear_transform.py -config /path/to/config/file
|
|
268
|
+
```
|
|
269
|
+
|
|
270
|
+
### Training the Flow Deformation module
|
|
271
|
+
|
|
272
|
+
Take a look at the example config file in `config/flow_deformation.yml`. Make a copy, and modify it appropriately. In particular, make sure you provide the path to the linear transformation module trained in the previous step.
|
|
273
|
+
|
|
274
|
+
Then run the command,
|
|
275
|
+
|
|
276
|
+
```commandline
|
|
277
|
+
python workflows/train_flow_with_udf.py -config /path/to/config/file
|
|
278
|
+
```
|
|
279
|
+
|
|
280
|
+
## Using trained models on new data
|
|
281
|
+
|
|
282
|
+
The pre-trained model takes as input a CT image in NIFTI format, a template mesh in VTP format and outputs a deformed mesh in VTP format.
|
|
283
|
+
|
|
284
|
+
First, place your image data in a folder named `image`. Let the path to this folder be `/path/to/folder/image`. Make sure that the images have extension `.nii.gz` or `.nii`. Next, run the following command to build an index of the image dataset,
|
|
285
|
+
|
|
286
|
+
```
|
|
287
|
+
python utilities/prepare_test_data_csv.py -f /path/to/folder
|
|
288
|
+
```
|
|
289
|
+
|
|
290
|
+
Note that the argument to `-f` is the path to the **parent** directory of the `image` directory.
|
|
291
|
+
|
|
292
|
+
After generating the index, it's time to execute the model.
|
|
293
|
+
|
|
294
|
+
Take a look at the example config file `config/predict_test_meshes_ct.yml`. Modify the path to the model, path to the image dataset, and the path to your output directory. Next, run the prediction script,
|
|
295
|
+
|
|
296
|
+
```
|
|
297
|
+
python utilities/predict_udf_test_meshes.py -config /path/to/config/file
|
|
298
|
+
```
|
|
299
|
+
|
|
300
|
+
Use the script `utilities/predict_test_meshes.py` if you want to evaluate the Linear Transform as a standalone module.
|
|
301
|
+
|
|
302
|
+
The script will generate output meshes and segmentations for each input image file.
|
|
@@ -0,0 +1,270 @@
|
|
|
1
|
+
# LinFlo-Net
|
|
2
|
+
|
|
3
|
+
**New to LinFlo-Net?** See the [Quick start guide](docs/quick_start.md) for install and prediction in a few minutes.
|
|
4
|
+
|
|
5
|
+
A deep learning package to automatically generate simulation ready 3D meshes of the human heart from biomedical images. [Link to paper](https://asmedigitalcollection.asme.org/biomechanical/article/doi/10.1115/1.4064527/1194613).
|
|
6
|
+
|
|
7
|
+

|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
## Setting up environment on Savio
|
|
11
|
+
|
|
12
|
+
The following instructions can set up a conda environment on the Berkeley Research Computing Savio system. But a similar approach can be used on any SLURM based high-performance computing cluster.
|
|
13
|
+
|
|
14
|
+
Since Savio provides limited space in your home directory, we install all conda packages to our scratch folder.
|
|
15
|
+
|
|
16
|
+
```commandline
|
|
17
|
+
module load cuda/10.2
|
|
18
|
+
module load gcc/5.4.0
|
|
19
|
+
ENVDIR=/global/scratch/users/<your_username>/environments/linflonet
|
|
20
|
+
rm -rf $ENVDIR
|
|
21
|
+
export CONDA_PKGS_DIRS=/global/scratch/users/<your_username>/tmp/.conda
|
|
22
|
+
conda create --prefix $ENVDIR
|
|
23
|
+
```
|
|
24
|
+
|
|
25
|
+
Press `y` when prompted to create your conda environment and then activate your environment,
|
|
26
|
+
|
|
27
|
+
```commandline
|
|
28
|
+
source activate $ENVDIR
|
|
29
|
+
```
|
|
30
|
+
|
|
31
|
+
Next install `pytorch`. Savio does not have the version of `cuda` required for the latest pytorch version, so we will install `pytorch 1.12.1`.
|
|
32
|
+
|
|
33
|
+
```commandline
|
|
34
|
+
conda install pytorch==1.12.1 cudatoolkit=10.2 -c pytorch
|
|
35
|
+
```
|
|
36
|
+
|
|
37
|
+
Press `y` when prompted to start the installation.
|
|
38
|
+
|
|
39
|
+
Next we will install [pytorch3d](https://pytorch3d.org/) which provides several useful routines for dealing with 3D data and mesh data-structures in conjunction with `pytorch`,
|
|
40
|
+
|
|
41
|
+
```commandline
|
|
42
|
+
conda install -c fvcore -c iopath -c conda-forge fvcore iopath
|
|
43
|
+
conda install -c bottler nvidiacub
|
|
44
|
+
conda install pytorch3d -c pytorch3d
|
|
45
|
+
```
|
|
46
|
+
|
|
47
|
+
You can find all the other dependencies in the system generated `requirements.txt` in the repository. You should be able to install these directly with `pip` after installing the above packages.
|
|
48
|
+
|
|
49
|
+
### Test that everything works
|
|
50
|
+
|
|
51
|
+
First request a brief interactive session with a GPU,
|
|
52
|
+
|
|
53
|
+
```commandline
|
|
54
|
+
srun --pty -A <account_name> -p savio3_gpu --nodes=1 --gres=gpu:GTX2080TI:1 --ntasks=1 --cpus-per-task=2 -t 00:30:00 bash -i
|
|
55
|
+
```
|
|
56
|
+
|
|
57
|
+
Once your resources are allocated, load your conda environment and launch python
|
|
58
|
+
|
|
59
|
+
```commandline
|
|
60
|
+
source activate /global/scratch/users/<your_username>/environments/linflonet
|
|
61
|
+
python
|
|
62
|
+
```
|
|
63
|
+
|
|
64
|
+
Now type the following into your python session,
|
|
65
|
+
|
|
66
|
+
```python
|
|
67
|
+
import torch
|
|
68
|
+
from pytorch3d.loss import chamfer_distance
|
|
69
|
+
device = torch.device("cuda")
|
|
70
|
+
|
|
71
|
+
a = torch.rand([5,10000,3]).to(device)
|
|
72
|
+
b = torch.rand([5,10000,3]).to(device)
|
|
73
|
+
loss = chamfer_distance(a, b)
|
|
74
|
+
```
|
|
75
|
+
|
|
76
|
+
If everything runs without error, you are all set!
|
|
77
|
+
|
|
78
|
+
## Install from PyPI
|
|
79
|
+
|
|
80
|
+
For prediction only (Python 3.10+):
|
|
81
|
+
|
|
82
|
+
```commandline
|
|
83
|
+
pip install linflonet
|
|
84
|
+
```
|
|
85
|
+
|
|
86
|
+
`pytorch3d` is required but not listed as a pip dependency because it must be
|
|
87
|
+
built from source on most platforms. Install `torch` first, then:
|
|
88
|
+
|
|
89
|
+
```commandline
|
|
90
|
+
pip install --no-build-isolation "git+https://github.com/facebookresearch/pytorch3d.git@stable"
|
|
91
|
+
```
|
|
92
|
+
|
|
93
|
+
## CLI usage
|
|
94
|
+
|
|
95
|
+
The `linflonet` command generates heart meshes (`.vtp`) and segmentations for
|
|
96
|
+
CT or MR NIfTI images.
|
|
97
|
+
|
|
98
|
+
**Single image:**
|
|
99
|
+
|
|
100
|
+
```commandline
|
|
101
|
+
linflonet predict \
|
|
102
|
+
--image /path/to/scan.nii.gz \
|
|
103
|
+
--model /path/to/best_model.pth \
|
|
104
|
+
--modality ct \
|
|
105
|
+
--output /path/to/output
|
|
106
|
+
```
|
|
107
|
+
|
|
108
|
+
**Folder of images** (flat folder or `image/` subdirectory):
|
|
109
|
+
|
|
110
|
+
```commandline
|
|
111
|
+
linflonet predict \
|
|
112
|
+
--folder /path/to/images \
|
|
113
|
+
--model /path/to/best_model.pth \
|
|
114
|
+
--modality mr \
|
|
115
|
+
--output /path/to/output
|
|
116
|
+
```
|
|
117
|
+
|
|
118
|
+
Template mesh and distance map default to files bundled with the package. Override
|
|
119
|
+
with `--template` and `--template-distance-map` if needed. For linear-transform-only
|
|
120
|
+
models, pass `--linear-transform`.
|
|
121
|
+
|
|
122
|
+
**Using a YAML config** (same format as `config/predict_single_ct.yml`):
|
|
123
|
+
|
|
124
|
+
```commandline
|
|
125
|
+
linflonet predict --config config/predict_single_ct.yml --image /path/to/scan.nii.gz -o /path/to/output
|
|
126
|
+
```
|
|
127
|
+
|
|
128
|
+
Outputs are written to `<output>/meshes/` and `<output>/segmentation/`.
|
|
129
|
+
|
|
130
|
+
You can also run `python -m linflonet predict ...` or install in editable mode
|
|
131
|
+
from a git checkout:
|
|
132
|
+
|
|
133
|
+
```commandline
|
|
134
|
+
pip install -e .
|
|
135
|
+
```
|
|
136
|
+
|
|
137
|
+
## Setting up a local environment with pip (Python 3.12)
|
|
138
|
+
|
|
139
|
+
If you only need to run prediction (not training), you can set up a lightweight
|
|
140
|
+
environment on Python 3.12 using `requirements-py312.txt`.
|
|
141
|
+
|
|
142
|
+
First, initialize the `vtk_utils` submodule and create a virtual environment,
|
|
143
|
+
|
|
144
|
+
```commandline
|
|
145
|
+
git submodule update --init
|
|
146
|
+
python3.12 -m venv .venv
|
|
147
|
+
source .venv/bin/activate
|
|
148
|
+
pip install -r requirements-py312.txt
|
|
149
|
+
```
|
|
150
|
+
|
|
151
|
+
`pytorch3d` does not ship prebuilt wheels for most platforms and must be built
|
|
152
|
+
from source *after* `torch` is installed. Its `setup.py` imports `torch` at build
|
|
153
|
+
time, so you must disable pip's build isolation (otherwise you get
|
|
154
|
+
`ModuleNotFoundError: No module named 'torch'`):
|
|
155
|
+
|
|
156
|
+
```commandline
|
|
157
|
+
pip install --no-build-isolation "git+https://github.com/facebookresearch/pytorch3d.git@stable"
|
|
158
|
+
```
|
|
159
|
+
|
|
160
|
+
On macOS, make sure the Xcode command-line tools are installed first
|
|
161
|
+
(`xcode-select --install`) so the C++ extension can compile.
|
|
162
|
+
|
|
163
|
+
## Dataset Creation
|
|
164
|
+
|
|
165
|
+
We use the multi-modality whole heart segmentation challenge (MMWHS) [dataset](https://zmiclab.github.io/zxh/0/mmwhs/). Download and unzip the data. You should have the following folders,
|
|
166
|
+
|
|
167
|
+
- CT : 2 folders each with 10 images and corresponding segmentations
|
|
168
|
+
- MR : 1 folder with 20 images and corresponding segmentations
|
|
169
|
+
|
|
170
|
+
You can split the data into train and validation as you find appropriate. We chose to use the first 16 samples as training and the remaining 4 samples as validation. Split the data appropriately and place them in separate folders. Make sure to keep the CT and MR data separately as we will need to normalize / scale them differently. We will perform data augmentation on the training data.
|
|
171
|
+
|
|
172
|
+
### Data augmentation
|
|
173
|
+
|
|
174
|
+
We will use the data augmentation procedure available in the [MeshDeformNet](https://github.com/fkong7/MeshDeformNet) package. Clone this package to your system and run `pip install -r requirements.txt` to install package dependencies. (You may want to create a virtual environment first.)
|
|
175
|
+
|
|
176
|
+
To perform augmentation, modify the command below and execute it. The script below launches 16 jobs in parallel (`-np 16`). You can modify that depending on the capacity of the system you are using.
|
|
177
|
+
|
|
178
|
+
```commandline
|
|
179
|
+
mpirun -np 16 python ~/path/to/MeshDeformNet/data/data_augmentation.py \
|
|
180
|
+
--im_dir /path/to/image/directory \
|
|
181
|
+
--seg_dir /path/to/segmentation/directory \
|
|
182
|
+
--out_dir /path/to/output/directory \
|
|
183
|
+
--modality ct or mr \
|
|
184
|
+
--mode train \
|
|
185
|
+
--num number_of_augmentations
|
|
186
|
+
```
|
|
187
|
+
|
|
188
|
+
The output folder will contain two subfolders `modality_train` with the augmented images and `modality_train_seg` with the augmented segmentations where modality is either `ct` or `mr`.
|
|
189
|
+
|
|
190
|
+
### Creating ground-truth meshes
|
|
191
|
+
|
|
192
|
+
We generate ground-truth meshes using marching cubes on the ground-truth segmentations. We can do this using `workflows/prepare_data.py`.
|
|
193
|
+
|
|
194
|
+
```commandline
|
|
195
|
+
python workflows/prepare_data \
|
|
196
|
+
--image /path/to/image/folder \
|
|
197
|
+
--segmentation /path/to/segmentation/folder \
|
|
198
|
+
--output /path/to/output/folder \
|
|
199
|
+
--modality ct # can be either ct or mr
|
|
200
|
+
--ext .nii.gz # input files extension
|
|
201
|
+
```
|
|
202
|
+
|
|
203
|
+
The output folder is going to have 3 subfolders : `seg`, `vtk_image`, `vtk_mesh`. `vtk_image` will be the input to our neural network, and `vtk_mesh` will be the corresponding ground truth meshes. From this point onward, we assume that the folder with the relevant data has the `vtk_image` and `vtk_mesh` subfolders.
|
|
204
|
+
|
|
205
|
+
### Final steps
|
|
206
|
+
|
|
207
|
+
The data set is reasonably large, and we will have to load it from memory. It is useful to store the images as pytorch tensors and the meshes as pytorch3d data structures in pickled files. To do this, we first build a csv index of all the files.
|
|
208
|
+
|
|
209
|
+
```commandline
|
|
210
|
+
python utilities/prepare_train_dataset_csv.py -f /path/to/data/folder
|
|
211
|
+
```
|
|
212
|
+
|
|
213
|
+
Make sure to provide the path to the parent directory containing `vtk_image` and `vtk_mesh` sub-directories. This will create an `index.csv` in the parent folder with the names of all files. Next,
|
|
214
|
+
|
|
215
|
+
```commandline
|
|
216
|
+
python utilities/pickle_image_segmentation_mesh_dataset.py -config /path/to/config/file
|
|
217
|
+
```
|
|
218
|
+
|
|
219
|
+
Look at `config/pickle_dataset.yml` for an example config file. Note that `seg_label` in the config file follows the labelling convention of the MMWHS dataset.
|
|
220
|
+
|
|
221
|
+
The output folder will now contain `.pkl` files which contain the combined image, segmentations, and meshes in a dictionary. This can be used by our dataloader to load the appropriate files during training.
|
|
222
|
+
|
|
223
|
+
|
|
224
|
+
## Training the model
|
|
225
|
+
|
|
226
|
+
Before training, make sure to activate the conda environment that we created earlier. Request a GPU session if you would like to use a GPU for training. Alternatively, submit the below commands as part of a batch job with `sbatch` on a SLURM system. The training workflow will save the best performing model as a checkpoint in the output directory specified in the config file.
|
|
227
|
+
|
|
228
|
+
### Training Linear Transformation module
|
|
229
|
+
|
|
230
|
+
Take a look at the example config file in `config/linear_transform.yml`. Make a copy, and modify it appropriately.
|
|
231
|
+
|
|
232
|
+
Then run the command,
|
|
233
|
+
|
|
234
|
+
```commandline
|
|
235
|
+
python workflows/train_linear_transform.py -config /path/to/config/file
|
|
236
|
+
```
|
|
237
|
+
|
|
238
|
+
### Training the Flow Deformation module
|
|
239
|
+
|
|
240
|
+
Take a look at the example config file in `config/flow_deformation.yml`. Make a copy, and modify it appropriately. In particular, make sure you provide the path to the linear transformation module trained in the previous step.
|
|
241
|
+
|
|
242
|
+
Then run the command,
|
|
243
|
+
|
|
244
|
+
```commandline
|
|
245
|
+
python workflows/train_flow_with_udf.py -config /path/to/config/file
|
|
246
|
+
```
|
|
247
|
+
|
|
248
|
+
## Using trained models on new data
|
|
249
|
+
|
|
250
|
+
The pre-trained model takes as input a CT image in NIFTI format, a template mesh in VTP format and outputs a deformed mesh in VTP format.
|
|
251
|
+
|
|
252
|
+
First, place your image data in a folder named `image`. Let the path to this folder be `/path/to/folder/image`. Make sure that the images have extension `.nii.gz` or `.nii`. Next, run the following command to build an index of the image dataset,
|
|
253
|
+
|
|
254
|
+
```
|
|
255
|
+
python utilities/prepare_test_data_csv.py -f /path/to/folder
|
|
256
|
+
```
|
|
257
|
+
|
|
258
|
+
Note that the argument to `-f` is the path to the **parent** directory of the `image` directory.
|
|
259
|
+
|
|
260
|
+
After generating the index, it's time to execute the model.
|
|
261
|
+
|
|
262
|
+
Take a look at the example config file `config/predict_test_meshes_ct.yml`. Modify the path to the model, path to the image dataset, and the path to your output directory. Next, run the prediction script,
|
|
263
|
+
|
|
264
|
+
```
|
|
265
|
+
python utilities/predict_udf_test_meshes.py -config /path/to/config/file
|
|
266
|
+
```
|
|
267
|
+
|
|
268
|
+
Use the script `utilities/predict_test_meshes.py` if you want to evaluate the Linear Transform as a standalone module.
|
|
269
|
+
|
|
270
|
+
The script will generate output meshes and segmentations for each input image file.
|