|
| 1 | +# Sparse Keypoint Segmentation of Lung Fissures: Efficient Geometric Deep Learning for Abstracting Volumetric Images |
| 2 | +This repository contains the code for segmenting pulmonary fissures in CT images using keypoint extraction and geometric deep learning. |
| 3 | + |
| 4 | + |
| 5 | + |
| 6 | +# Environment |
| 7 | +Install the required packages using Anaconda and the provided environment file: |
| 8 | +```bash |
| 9 | +conda env create -f environment.yml |
| 10 | +conda activate fissure-segmentation |
| 11 | +``` |
| 12 | + |
| 13 | +# Data |
| 14 | +Datasets used in the evaluation are openly available for research purposes. Please refer to the original sources |
| 15 | +detailed below. |
| 16 | + |
| 17 | +## Training / cross-validation data |
| 18 | +The TotalSegmentator data set is the main data set used for this work. |
| 19 | +Download it from here: https://doi.org/10.5281/zenodo.6802614 (v1.0). |
| 20 | +The data has to be preprocessed with `preprocess_totalsegmentator_dataset.py`, make sure to point to the correct data |
| 21 | +path `TS_RAW_DATA_PATH` in `constants.py`. The script selects applicable images (showing the entire thorax), extracts fissure |
| 22 | +segmentations from the lobe masks and crops the images to the thoracic region. Results will be written to |
| 23 | +`IMG_DIR_TS_PREPROC`. |
| 24 | + |
| 25 | +## COPD validation data set |
| 26 | +The COPD data set can be accessed here: https://med.emory.edu/departments/radiation-oncology/research-laboratories/deformable-image-registration/downloads-and-reference-data/copdgene.html |
| 27 | +(citation: Castillo et al., Phys Med Biol 2013 https://doi.org/10.1088/0031-9155/58/9/2861) |
| 28 | +The manual fissure annotations are available from here: http://www.mpheinrich.de/research.html#COPD (citation: |
| 29 | +Rühaak et al., IEEE TMI 2017 https://doi.org/10.1109/TMI.2017.2691259). |
| 30 | + |
| 31 | +The images should be named `<patid>_img_<fixed/moving>.nii.gz` and be in the same folder (`data/images/COPD`) as the |
| 32 | +fissure annotations `<patid>_fissures_<fixed/moving>.nii.gz` and lung masks `<patid>_mask_<fixed/moving>.nii.gz`. |
| 33 | +Masks can be generated automatically by, e.g., using the `lungmask` tool https://github.com/JoHof/lungmask. |
| 34 | + |
| 35 | +Then, preprocessing can be applied by running `preprocess_copd_dataset.py`. |
| 36 | + |
| 37 | +## Pre-computing keypoints and features |
| 38 | +Keypoints and features are pre-computed from the image data. Run the following scripts: |
| 39 | +1. `python -m data_processing.keypint_extraction` |
| 40 | +2. `python -m data_processing.point_features` |
| 41 | +Results will be written into `data/points` by default. |
| 42 | + |
| 43 | +# Train models |
| 44 | +The different models can be trained using the following scripts: |
| 45 | +- Point Cloud segmentation networks (DGCNN, PointNet, PointTransformer): `train_point_segmentation.py` |
| 46 | +- Point Cloud Autoencoder (PC-AE): `train_pc_ae.py` |
| 47 | +- Keypoint CNN (MobileNetV3 + LR-ASPP): `train_keypoint_cnn.py` (pre-trained weights are available in `results/lraspp_recall_loss`) |
| 48 | +Run the scripts with the `-h` flag to see available options. |
| 49 | + |
| 50 | +# Evaluate models |
| 51 | +The given training scripts already contain test functions for all models. |
| 52 | + |
| 53 | +Evaluation of the PC-AE reconstruction of fissures can be done using the `evaluate_pc_ae.py` script. This will take a |
| 54 | +trained point cloud segmentation network and apply a trained PC-AE for mesh reconstruction (instead of Poisson surface |
| 55 | +reconstruction). |
| 56 | + |
| 57 | +The nnU-Net baseline was evaluated using `evaluate_nnunet.py`. |
| 58 | + |
| 59 | +# Citation |
| 60 | +If you used code in this repository or found it useful in any way for your research, please cite the following papers: |
| 61 | + |
| 62 | +1. Journal paper for Point Cloud Autoencoder (PC-AE) and comparison of geometric segmentation networks (under review): |
| 63 | +``` |
| 64 | +@article{Kaftan2024_FissureSegmentation_IJCARS, |
| 65 | + title={Sparse Keypoint Segmentation of Lung Fissures: Efficient Geometric Deep Learning for Abstracting Volumetric Images}, |
| 66 | + author={Kaftan, Paul and Heinrich, Mattias P and Hansen, Lasse and Rasche, Volker and Kestler, Hans A and Bigalke, Alexander}, |
| 67 | + journal={International Journal of Computer Assisted Radiology and Surgery}, |
| 68 | + year={tbd}, |
| 69 | + volume={tbd}, |
| 70 | + number={tbd}, |
| 71 | + pages={tbd}, |
| 72 | + doi={tbd} |
| 73 | +} |
| 74 | +``` |
| 75 | + |
| 76 | +2. Conference paper for comparison of more keypoints and features: |
| 77 | +``` |
| 78 | +@InProceedings{Kaftan2024_FissureSegmentation_BVM, |
| 79 | + author = {Kaftan, Paul and Heinrich, Mattias P. and Hansen, Lasse and Rasche, Volker and Kestler, Hans A. and Bigalke, Alexander}, |
| 80 | + editor = {Maier, Andreas and Deserno, Thomas M. and Handels, Heinz and Maier-Hein, Klaus and Palm, Christoph and Tolxdorff, Thomas}, |
| 81 | + title = {Abstracting Volumetric Medical Images with Sparse Keypoints for Efficient Geometric Segmentation of Lung Fissures with a Graph CNN}, |
| 82 | + booktitle = {Bildverarbeitung f{\"u}r die Medizin 2024}, |
| 83 | + year = {2024}, |
| 84 | + publisher = {Springer Fachmedien Wiesbaden}, |
| 85 | + address = {Wiesbaden}, |
| 86 | + pages = {60--65}, |
| 87 | + isbn = {978-3-658-44037-4}, |
| 88 | + doi = {10.1007/978-3-658-44037-4_19} |
| 89 | +} |
| 90 | +``` |
0 commit comments