update readme

README.md

@@ -1,14 +1,11 @@
-# Implementation of Paper: A simulation-friendly auto segmentation method for individual tooth and bones
-##  Backbone model: Multi-Planar U-Net
+#  Backbone model: Multi-Planar U-Net (modified)
 Most of the repository is based on the official implemenation of original Multi-Planar U-Net.
 Please follow https://github.com/perslev/MultiPlanarUNet for the general installation and usage of Multi-Planar U-Net
 
 But make sure you clone from the current repository.
 
 ---
 
-#  Teeth-bone Specific Pipeline with Transfer learning and Instance segmentation post-processing
-
 ##  Installation
 ```
 git clone https://github.com/diku-dk/AutoJawSegment
@@ -25,9 +22,9 @@ pip install -e AutoJawSegment
 ## Initializing two projects, one for pretraining on inaccurate data, and one for fine-tuning as follows:
 
 ```
-# Initialize a project for pretraining
+# Initialize a project for pretraining/or just training without a pretrained model
 mp init_project --name my_pretraining_project --data_dir ./pretraining_data_folder
-# Initialize a project for pretraining
+# Initialize a project for finetuing
 mp init_project --name my_finetune_project --data_dir ./fineune_data_folder
 ```
 
@@ -50,92 +47,19 @@ The standard fine-tuning process without disatance weight map can be applied by
 ```
 # make sure you are at root dir
 cd my_finetune_project
-mp train --num_GPUs=1 --overwrite --transfer_last_layer --force_GPU=0 \
---initialize_from '../my_pretraining_project/model/@epoch_xxx_val_dice_xxxx.h5'
-```
-Make sure you change ```xxx``` according to the .h5 file of the model generated from pre-training step 
-
-### Fine-tuning with weight_map loss
-
-![](figs/training_pipeline.jpg)
-
-
-
-To fine-tune with the weight_map loss for emphasizing the gap between teeth and bones in the loss function, 
-make sure you first compute the weight_map based on the label maps of 
-your training data. You can compute and store them by invoking the following command, but make sure you 
-have modified the path to your fine-tuning dataset. You can also change the hyperparameters 
-(mean and std in the ```__main__``` function) 
-```
-python mpunet/preprocessing/weight_map.py
-```
-
-After the weight maps are stored on disk, you can fine-tune with weight_map loss by invoking
-```
-# make sure you are at root dir
-cd my_finetune_project
-mp train --num_GPUs=1 --overwrite \
---distance_loss --transfer_last_layer \
---distance_map_folder my_weight_map_folder \
---initialize_from '../my_pretraining_project/model/@epoch_xxx_val_dice_xxxx.h5'
-```
-Make sure you change ```xxx``` according to the .h5 file of the model generated from pre-training step 
-and change ```my_weight_map_folder``` to the folder that the previous command generated.
-
-### Fine-tuning with double weight_map loss
-In order to enforce not only on gaps between teeth and bones, but also between neighboring teeth
-to the loss function, 
-we need to first label each tooth with a different integer from the ground truth. 
-To do this, run 
-
-```
-python mpunet/preprocessing/instance_labelling.py
-```
-Note that this process can be very memory-consuming, so make sure you have enough memory and disk for memory-swap. 
-Then, run the same command to compute the weigh map.
-```
-python mpunet/preprocessing/weight_map.py
-```
-
-![](figs/double_weight.jpg)
-
-### Interative Fine-tuning
-To do Interative fine-tuning with the weight_map loss, make sure you first put the new training dataset to the data_dir,
-then compute the weight_map based on the labels of 
-your new training data. You can compute and store them by invoking the following command, but make sure you 
-have modified the path to your fine-tuning dataset. You can also change the hyperparameters 
-(std and erosion radius in the ```__main__``` function) 
-
-```
-python mpunet/preprocessing/weight_map.py
-```
-
-After the weight maps are stored on disk, you can fine-tune with weight_map loss by invoking the following, basically
-by removing ```--initialize_from``` and ```--transfer_last_layer```  tag and changing ```--overwrite``` to ```--continue``` 
-
-```
-# make sure you are at root dir
-cd my_finetune_project
-mp train --num_GPUs=1 --continue --distance_loss \
---distance_map_folder my_weight_map_folder
+mp train --num_GPUs=1 --overwrite --transfer_last_layer  \
+--force_GPU=0 --initialize_from '../my_pretraining_project/model/@epoch_xxx_val_dice_xxxx.h5'
 ```
-Make sure you change ```xxx``` according to the .h5 file of the model generated from pre-training step 
-and change ```my_weight_map_folder``` to the folder that the previous command generated.
-
+Make sure you change ```xxx``` according to the .h5 file of the model generated from pre-training step
 
+Also make sure to comment out ```transfer_last_layer``` if not working on the same class
 
 ## Predict
 The trained model can now be evaluated on the testing data in 
 ```data_folder/test``` by invoking:
 
 ```
-mp predict --num_GPUs=1 \
---force_GPU=0 \
---sum_fusion \
---overwrite \
---no_eval \
---by_radius \
---out_dir predictions 
+mp predict --num_GPUs=1 --force_GPU=0 --sum_fusion --overwrite --no_eval --by_radius --out_dir predictions 
 ```
 note that the ```--sum_fusion``` tag is essential since we are not training 
 another fusion model. ```--by_radius``` tage is also essential in ensuring different sampling strategy during training
@@ -144,51 +68,16 @@ and testing to avoid incomplete boundaries, as mentioned in the paper.
 This will create a folder ```my_project/predictions``` storing the predicted 
 images along with dice coefficient performance metrics.
 
-
-## Individual Tooth Segmention with Post-Processing
-
-![](figs/watershed.jpg)
-
-The proposed watershed instance segmentation pipeline works on the UNet output probability map (for the teeth class). 
-In order to get the fused (averaged) probability map (the softmax score) before argmax of MPUNet, 
-run the following prediction command.
-Note that ```save_single_class``` is the integer label for the teeth class, which is 2 in our case.
-
-```
-mp predict --num_GPUs=1 --save_single_class=2 \
---sum_fusion \
---overwrite --by_radius \
---no_eval --out_dir prob_maps \
---no_argmax
-```
-
-
-After you get the probability maps, run the following to get different labels for each tooth. 
-
-```
-python instance_seg/watershed
-```
-
-Default hyperprameters are given in the main function 
-but will probably need to be adapted for specific use case, e.g., the ```open``` radius.
-Note that they denote number of voxels, so will probably also need to adapt to new voxel sizes.
-
-
-The ```watershed``` will only work on the teeth class probability map until the very end, where 
-it will be combined with the bone class that was generated by the previous ```mp predict``` command 
-without the argument ```--no_argmax```. Therefore, please also make sure that 
-the integer results are stored
-
-The bone class (labeled with 1 is our case) consists of mandible and maxilla. Unlike teeth, 
-the mandible and maxilla (upper and lower bone jaws) are always disconnected by a large gap. 
-Therefore, a straightforward Connected Component Decomposition while keeping the largest two
-components to remove small disconnected false positives suffice to further separate maxilla and mandibles.
-You can do this by running
-
+## Post-Processing
+Usually, each class should be a single connected component. Run the following to keep the largest component
+for each class
 ```
 python postprocessing/ccd
 ```
 
+Other post-processing such as opening to remove small floating parts and closing to remove holes
+are often also helpful, but would need to specify the size.
+
 
 ## Evaluate/Numerical Validation
 The results can be evaluated on the testing data to get DICE, GapDICE, and Hausdoff Distance, and ASSD by running
@@ -199,25 +88,3 @@ python mpunet/evaluate/compute_metrics.py
 
 This will print the mean and standard deviation for the 4 metrics, as well as generate .csv files for the score for each image. 
 You can use simple a script or Excel functions to get the mean and std of the overall performance
-
-### Sample Result
-
-```results``` folder includes some prediction results of our pipeline, with the following structure:
-
-* ```UNet```: The direct output from MPUNet (with argmax). Bone class will be labeled ```1``` and Teeth will be labeled ```2```
-* ```Watershed```: The instance segmentation of each individual tooth from proposed method, where each tooth will be labeled from ```1``` to ```num_teeth``` of each scan
-* ```Combined```: The Watershed result combined with the UNet Bone class result, where bone class will be labeled 1 and each tooth will be labeled from ```2``` to ```num_teeth + 1``` of each san
-* ```Combined_binary```: The result from combined where each individual tooth is mapped back to label ```2```.
-
-## Final Finite Element Model
-
-Much of the motivation of this project is related to a larger project named ["Open-Full-Jaw: An open-access dataset and pipeline for finite element models of human jaw"](https://www.sciencedirect.com/science/article/pii/S0169260722003911)
-
-available at https://github.com/diku-dk/Open-Full-Jaw
-
-
-![Fig6](https://user-images.githubusercontent.com/30265621/215486601-7702284b-639e-47ff-a787-15cdc5b95e91.jpg)
-
-
-
-The auto segmentation of teeth and bones is the essential first step to generate patient-specific finite element models of the human jaw. As mentioned in [Open-Full-Jaw](https://www.sciencedirect.com/science/article/pii/S0169260722003911), this step involves time-consuming and labor-intensive segmentation task of fine geometries, i.e., teeth, PDL, and bone structures. A potential use case for this study is to use the output of our model as the initial step. Then, improve the quality of the reconstructed geometries as well as the generated gap regions. Finally, generate the required PDL geometries for the simulation purposes. 

mpunet/bashes.sh

@@ -0,0 +1,25 @@
+
+## pretraining on OWIA dataset
+mp init_project --name pretraining_IWOAI --data_dir ./datasets/pretrain_IWOAI
+mp train --overwrite --num_GPUs=1 --force_GPU=0
+
+## finetune
+mp init_project --name fine_tune --data_dir ./datasets/rasmus
+
+# first training own dataset using the model pretrained on OWIA dataset
+
+mp train --overwrite --force_GPU=0 --epochs 100 --initialize_from "C:\\Users\\Administrator\\PycharmProjects\\AutoJawSegment\\pretraining_IWOAI\\model\\@epoch_170_val_dice_0.84136.h5"
+
+# predictions
+mp predict --num_GPUs=1 --force_GPU=0 --sum_fusion --overwrite --no_eval --by_radius --out_dir predictions
+
+### continue training when having more data
+
+mp train --continue --force_GPU=0 --epochs 80
+
+### miscellaneous
+mp predict --sum_fusion --overwrite --by_radius --no_eval --force_GPU=0 --test_dir "C:\Users\Administrator\PycharmProjects\AutoJawSegment\datasets\rasmus/train"
+--out_dir predictions --use_diag_dim
+
+mp predict --sum_fusion --overwrite --by_radius --no_eval  --out_dir predictions_no_arg --use_diag_dim --num_GPUs=0
+mp predict --num_GPUs=1 --force_GPU=0 --sum_fusion --overwrite --no_eval --by_radius --out_dir predictions_2

mpunet/postprocessing/symmetric_separator_show.py

@@ -0,0 +1,93 @@
+from scipy.ndimage import label
+import numpy as np
+
+import nibabel as nib
+import os
+
+from connected_component_3D import connected_component_3D
+from morphologies import binary_closing_label
+
+def get_max_label(labels):
+    values, counts = np.unique(labels, return_counts=True)
+    inds = np.argsort(counts)
+
+    for i in reversed(inds):
+        if values[i] == 0:
+            continue
+        else:
+            res = values[i]
+            return res
+
+
+def symmetric_separator(file_path, save_path=None, connectivity=26, portion_foreground=0.01,
+                                    bg_val=0, pairs=None, no_save=False, morph=False):
+
+    img_func = nib.load(file_path)
+    affine = img_func.affine
+    img = img_func.get_fdata().astype(np.uint8)
+    img = np.squeeze(img)
+    # res = np.empty(img.shape)
+
+    for (i, j) in pairs:
+        img_i = np.copy(img == i)
+        img_j = np.copy(img == j)
+
+        img[img_i] = 0
+        img[img_j] = 0
+
+        labels_out_i, _ = label(img_i)
+        max_label = get_max_label(labels_out_i)
+
+        print(f'max_label of class {i} is {max_label}')
+
+        img[np.logical_and(labels_out_i > 0, labels_out_i != max_label)] = j
+        img[labels_out_i == max_label] = i
+
+        labels_out_j, _ = label(img_j)
+        max_label = get_max_label(labels_out_j)
+
+        print(f'max_label of class {j} is {max_label}')
+
+        img[np.logical_and(labels_out_j > 0, labels_out_j != max_label)] = i
+        img[labels_out_j == max_label] = j
+
+    img = connected_component_3D(img, connectivity=26, portion_foreground=portion_foreground)
+
+    if morph:
+        img = binary_closing_label(img, radius=3)
+
+    if no_save:
+        return img
+
+    ni_img = nib.Nifti1Image(img.astype(np.uint8)
+                             , affine)
+
+    nib.save(ni_img, save_path)
+
+
+if __name__ == '__main__':
+
+    root_path = '/Users/px/GoogleDrive/MultiPlanarUNet/my_hip_project_weight_map/predictions_0902/nii_files'
+
+    ccd_path = os.path.join(root_path, 'ccd_symmetric')
+
+    if not os.path.exists(ccd_path):
+        os.mkdir(ccd_path)
+
+    for i in reversed(sorted(os.listdir(root_path))):
+        print(f'working on {i}')
+        if i.startswith('.') or not (i.endswith('nii') or i.endswith('gz')):
+            continue
+
+        img_path = os.path.join(root_path, i)
+
+        save_path = os.path.join(ccd_path, i)
+
+        portion_foreground = 0.3
+
+        symmetric_separator(img_path, save_path=save_path,
+                            connectivity=6, portion_foreground=portion_foreground,
+                            pairs=[(1, 2), (3, 4)],
+                            morph=True
+                            )
+

requirements.txt

@@ -1,8 +1,12 @@
+## need to do this conda install for gpu to work
+## conda install -c conda-forge cudatoolkit=11.2 cudnn=8.1.0
+
+
 matplotlib>=3.3.0
 scipy>=1.4.0
 # numpy>=1.18.5, <1.19.0
 numpy>1.19.0
-nibabel>=3.1.0
+nibabel==3.1.0
 pandas>=1.0.0
 ruamel.yaml>=0.16.0
 scikit_learn>=0.23.0
@@ -21,4 +25,6 @@ PyMCubes
 PyQt5
 # pyqtgraph
 # change made to h5py, no < 2.11.0 previously
-# , <1.19
+# , <1.19
+### check tg gpu tf.test.is_gpu_available()
+#### python -c "import tensorflow as tf; print(tf.config.list_physical_devices('GPU'))"