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.
Files changed (47) hide show
  1. linflonet-0.1.0/LICENSE +21 -0
  2. linflonet-0.1.0/PKG-INFO +302 -0
  3. linflonet-0.1.0/README.md +270 -0
  4. linflonet-0.1.0/linflonet/__init__.py +3 -0
  5. linflonet-0.1.0/linflonet/__main__.py +3 -0
  6. linflonet-0.1.0/linflonet/cli.py +210 -0
  7. linflonet-0.1.0/linflonet/data/__init__.py +0 -0
  8. linflonet-0.1.0/linflonet/data/template/__init__.py +0 -0
  9. linflonet-0.1.0/linflonet/data/template/highres_template.vtp +44 -0
  10. linflonet-0.1.0/linflonet/data/template/highres_template_distance.vtk +0 -0
  11. linflonet-0.1.0/linflonet/data/template/whole_heart_with_ao.vtp +44 -0
  12. linflonet-0.1.0/linflonet/paths.py +32 -0
  13. linflonet-0.1.0/linflonet/predict.py +236 -0
  14. linflonet-0.1.0/linflonet.egg-info/PKG-INFO +302 -0
  15. linflonet-0.1.0/linflonet.egg-info/SOURCES.txt +45 -0
  16. linflonet-0.1.0/linflonet.egg-info/dependency_links.txt +1 -0
  17. linflonet-0.1.0/linflonet.egg-info/entry_points.txt +2 -0
  18. linflonet-0.1.0/linflonet.egg-info/requires.txt +11 -0
  19. linflonet-0.1.0/linflonet.egg-info/top_level.txt +3 -0
  20. linflonet-0.1.0/pyproject.toml +54 -0
  21. linflonet-0.1.0/setup.cfg +4 -0
  22. linflonet-0.1.0/src/__init__.py +0 -0
  23. linflonet-0.1.0/src/chamfer.py +224 -0
  24. linflonet-0.1.0/src/dataset.py +68 -0
  25. linflonet-0.1.0/src/finite_difference.py +39 -0
  26. linflonet-0.1.0/src/flow.py +116 -0
  27. linflonet-0.1.0/src/flow_loss.py +62 -0
  28. linflonet-0.1.0/src/integrator.py +114 -0
  29. linflonet-0.1.0/src/io_utils.py +204 -0
  30. linflonet-0.1.0/src/linear_transform.py +208 -0
  31. linflonet-0.1.0/src/loss.py +109 -0
  32. linflonet-0.1.0/src/pre_process.py +146 -0
  33. linflonet-0.1.0/src/random_linear_perturbations.py +29 -0
  34. linflonet-0.1.0/src/segment_flow.py +516 -0
  35. linflonet-0.1.0/src/signed_distance_function.py +221 -0
  36. linflonet-0.1.0/src/template.py +237 -0
  37. linflonet-0.1.0/src/unet.py +143 -0
  38. linflonet-0.1.0/src/unet_components.py +193 -0
  39. linflonet-0.1.0/src/unet_segment.py +57 -0
  40. linflonet-0.1.0/src/utilities.py +40 -0
  41. linflonet-0.1.0/test/test_chamfer_normal.py +15 -0
  42. linflonet-0.1.0/test/test_crop.py +52 -0
  43. linflonet-0.1.0/test/test_cropped_evaluation.py +32 -0
  44. linflonet-0.1.0/test/test_linear_transform_distance_function.py +75 -0
  45. linflonet-0.1.0/test/test_matching_image_label_sizes.py +23 -0
  46. linflonet-0.1.0/vtk_utils/__init__.py +0 -0
  47. linflonet-0.1.0/vtk_utils/vtk_utils.py +929 -0
@@ -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.
@@ -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
+ ![image](figures/flow-deformation-no-encoder.png)
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
+ ![image](figures/flow-deformation-no-encoder.png)
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.
@@ -0,0 +1,3 @@
1
+ """LinFlo-Net: heart mesh generation from biomedical images."""
2
+
3
+ __version__ = "0.1.0"
@@ -0,0 +1,3 @@
1
+ from linflonet.cli import main
2
+
3
+ main()