ivadomed / model_seg_mouse-sc_wm-gm_t1 Goto Github PK

View Code? Open in Web Editor NEW

2.0 1.0 1.0 598 KB

White and grey matter segmentation on T1-weighted exvivo mouse spinal cord

License: MIT License

Python 100.00%

gray-matter spinal-cord white-matter mri segmentation

model_seg_mouse-sc_wm-gm_t1's Introduction

White and grey matter segmentation on T1-weighted exvivo mouse spinal cord

480.mov

Citation

Publication linked to the dataset: Coming soon!

Publication linked to this model: see CITATION.cff

Project description

In this project, we trained a 3D nnU-Net for spinal cord white and grey matter segmentation. The data contains 22 mice with different number of chunks, for a total of 72 MRI 3D images. Each MRI image is T1-weighted, has a size of 200x200x500, with the following resolution: 0.05x0.05x0.05 mm.

Expand this for more information on how we trained the model

In order to train a 3D nnU-Net, the following steps were completed:

First, a total of 161 slices were labelled on various subjects. See Notes for details on the manual labeling.
The slices were then extracted using the extract_slices.py function: it extracted both the slice from the MRI image as well as the mask's slice. These were gathered into a temporary dataset, on which a 2D nnU-Net model was trained to segment spinal cord white and grey matter. The inference was then performed using this model on the full 3D volume from the original dataset.
Then, a 3D nnU-Net was trained on the images, using the results from the previous inference as ground truth as well as using extracted slices (of shape (200x200x1)) and their manual segmentation. The inference, was again performed on the full zurich-mouse dataset. Going from a 2D nnU-Net to a 3D nnU-Net helped improved the continuity of the segmentation on the z-axis.
After that, we selected the best segmentation masks on the dataset totalling 31 images. For each of these images we noted that the top and bottom slices were often poorly annotated. Using the crop_image_and_mask.py script we removed these slices. The objective was to keep only qualitative annotations.
Finally, a 3D nnU-Net was trained on these qualitative image segmentations (31 images) with various dimension as well as annotated slices (161 images). The nnU-Net was trained on 1000 epochs, with "3d_fullres" configuration and on 5 folds. The best Dice score were the following (fold 0 : 0.9135, fold 1: 0.9083, fold 2: 0.9109 , fold 3: 0.9132, fold 4: 0.9173).

For the packaging we decided to keep only fold 4 as it has the best dice score and all performed similarly in terms of final results as well as training evolution (meaning that the dataset is rather homogeneous). The reason for this is to avoid having to upload the full results model which weight around 5 GB and limit ourself to 250 MB. Also, inference is much longer when performed on 5 folds instead of 1 and results are comparable.

For information on how to retrain the same model, refer to this file README.md.

How to use the model

This is the recommended method to use our model.

Install dependencies

Spinal Cord Toolbox (SCT) v6.2 or higher -- follow the installation instructions here
conda
Python

Once the dependencies are installed, download the latest model:

sct_deepseg -install-task seg_mouse_gm_wm_t1w

Getting the WM and GM segmentation

To segment a single image, run the following command:

sct_deepseg -i <INPUT> -o <OUTPUT> -task seg_mouse_gm_wm_t1w

For example:

sct_deepseg -i sub-001_T2w.nii.gz -o sub-001_T2w_wm-gm-seg.nii.gz -task seg_mouse_gm_wm_t1w

Notes

Procedure for ground truth mask creation: https://youtu.be/KVL-JzcSRTo

model_seg_mouse-sc_wm-gm_t1's People

Stargazers

Watchers

Forkers

anriigegliuk

model_seg_mouse-sc_wm-gm_t1's Issues

Intensity scale the image

Try 2D nnUnet

Given the high performance of the nnUnet in general, it would be a good idea to have a benchmark with this architecture.

Todo (update if necessary):

Write a script to convert the current BIDS data to a nnUnet architecture. See https://github.com/ivadomed/data-conversion
Look into MONAI's nnUnetv2 implementation
Write training script with nnUnet that uses wanbd. Ideally, the code would be merged with the existing training and testing codes

Related to #17

Tagging @naga-karthik @louisfb01 @valosekj for help

RuntimeError: Input type (double) and bias type (float) should be the same

When running 93b47a1, I encountered a RuntimeError.

Terminal output

Path to data: /Users/julien/data.neuro/zurich-mouse
wandb: Currently logged in as: jcohenadad. Use `wandb login --relogin` to force relogin
wandb: wandb version 0.14.0 is available!  To upgrade, please run:
wandb:  $ pip install wandb --upgrade
wandb: Tracking run with wandb version 0.13.11
wandb: Run data is saved locally in /Users/julien/code/model_seg_mouse-sc_wm-gm_t1/wandb/run-20230315_164533-084ct9i2
wandb: Run `wandb offline` to turn off syncing.
wandb: Syncing run sweet-pyramid-10
wandb: ⭐️ View project at https://wandb.ai/jcohenadad/mouse-zurich
wandb: 🚀 View run at https://wandb.ai/jcohenadad/mouse-zurich/runs/084ct9i2
----------
epoch 1/100
Backend MacOSX is interactive backend. Turning interactive mode on.
outputs = model(inputs)
PyDev console: starting.
Traceback (most recent call last):
  File "/Applications/PyCharm CE.app/Contents/plugins/python-ce/helpers/pydev/_pydevd_bundle/pydevd_exec2.py", line 3, in Exec
    exec(exp, global_vars, local_vars)
  File "<input>", line 1, in <module>
  File "/Users/julien/code/model_seg_mouse-sc_wm-gm_t1/venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/Users/julien/code/model_seg_mouse-sc_wm-gm_t1/venv/lib/python3.10/site-packages/monai/networks/nets/unet.py", line 303, in forward
    x = self.model(x)
  File "/Users/julien/code/model_seg_mouse-sc_wm-gm_t1/venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/Users/julien/code/model_seg_mouse-sc_wm-gm_t1/venv/lib/python3.10/site-packages/torch/nn/modules/container.py", line 204, in forward
    input = module(input)
  File "/Users/julien/code/model_seg_mouse-sc_wm-gm_t1/venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/Users/julien/code/model_seg_mouse-sc_wm-gm_t1/venv/lib/python3.10/site-packages/monai/networks/blocks/convolutions.py", line 316, in forward
    res: torch.Tensor = self.residual(x)  # create the additive residual from x
  File "/Users/julien/code/model_seg_mouse-sc_wm-gm_t1/venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/Users/julien/code/model_seg_mouse-sc_wm-gm_t1/venv/lib/python3.10/site-packages/torch/nn/modules/conv.py", line 463, in forward
    return self._conv_forward(input, self.weight, self.bias)
  File "/Users/julien/code/model_seg_mouse-sc_wm-gm_t1/venv/lib/python3.10/site-packages/torch/nn/modules/conv.py", line 459, in _conv_forward
    return F.conv2d(input, weight, bias, self.stride,
RuntimeError: Input type (double) and bias type (float) should be the same

Use post-pred transforms with Compose()

See: Project-MONAI/MONAI#415 (comment)

Hyperparameter optimization and wandb sweep

See: https://wandb.ai/gladiator/MONAI_Spleen_3D_Segmentation/reports/3D-Segmentation-with-MONAI-and-PyTorch-Supercharged-by-Weights-Biases---VmlldzoyNDgxNDMz

Use active learning based on uncertainty of ensemble predictions

With ensemble/bagging #27, output uncertainty on the prediction for each slice (eg: compute pixelwise STD and average STD across slice), only keep predictions with high certainty, and re-train model with many more ground truths.

TypeError: Invalid shape (200,) for image data

Commit: 041e2ff

current epoch: 100 current mean dice: 0.0000
best mean dice: 0.0000 at epoch: 10

train completed, best_metric: 0.0000 at epoch: 10
Traceback (most recent call last):
  File "/home/GRAMES.POLYMTL.CA/p101317/code/model_seg_mouse-sc_wm-gm_t1/train.py", line 300, in <module>
    plt.imshow(torch.argmax(
  File "/home/GRAMES.POLYMTL.CA/p101317/code/model_seg_mouse-sc_wm-gm_t1/venv/lib/python3.9/site-packages/matplotlib/pyplot.py", line 2695, in imshow
    __ret = gca().imshow(
  File "/home/GRAMES.POLYMTL.CA/p101317/code/model_seg_mouse-sc_wm-gm_t1/venv/lib/python3.9/site-packages/matplotlib/__init__.py", line 1442, in inner
    return func(ax, *map(sanitize_sequence, args), **kwargs)
  File "/home/GRAMES.POLYMTL.CA/p101317/code/model_seg_mouse-sc_wm-gm_t1/venv/lib/python3.9/site-packages/matplotlib/axes/_axes.py", line 5665, in imshow
    im.set_data(X)
  File "/home/GRAMES.POLYMTL.CA/p101317/code/model_seg_mouse-sc_wm-gm_t1/venv/lib/python3.9/site-packages/matplotlib/image.py", line 710, in set_data
    raise TypeError("Invalid shape {} for image data"
TypeError: Invalid shape (200,) for image data

Compare with/without data augmentation applied to validation set

Improvement of the model - smoothing after 2D nnU-Net inference

Here is one possible improvement:

When :

After training the first 2D nnU-Net model and getting the inference results on the entire dataset
Before training the first 3D nnU-Net model on the previous results

Task:

Perform smoothing of the segmented sipnal cord along the spinal cord centerline
Dilate the segmentation a little bit (because it often forgets to segment a little bit of the spinal cord on the contour) ?

To Do :

Discuss when to do it
Which methods to use
Is it worth it ?

Problem running inference on new data

Problem

According to private discussions between @plbenveniste and Benjamin, running the code version https://github.com/ivadomed/model_seg_mouse-sc_wm-gm_t1/releases/tag/v0.3 on new data produced the following error:

The cause of the issue, identified by @plbenveniste, is that the input data have a corrupted NIfTI header. More specifically, the physical dimensions of the voxels are wrong, which cause memory saturation during resampling as part of the nnUnet preprocessing pipeline.

Solution

A solution is to make sure the data are correct. To do so:

in the first place, the data coming from the scanner should be converted to NIfTI using proper software. Dcm2niix is an example of a good software.
Alternatively (although the previous solution is best), the NIfTI headers could be manually corrected:

Mismatch between image and label qform/sform

image:

qform_name	Aligned Anat
qform_code	2
qto_xyz:1	0.050000 0.000000 0.000000 0.000000 
qto_xyz:2	0.000000 0.050000 0.000000 0.000000 
qto_xyz:3	0.000000 0.000000 0.050000 0.000000 
qto_xyz:4	0.000000 0.000000 0.000000 1.000000 
qform_xorient	Left-to-Right
qform_yorient	Posterior-to-Anterior
qform_zorient	Inferior-to-Superior
sform_name	Scanner Anat
sform_code	1
sto_xyz:1	1.000000 0.000000 0.000000 -99.000000 
sto_xyz:2	0.000000 1.000000 0.000000 -249.000000 
sto_xyz:3	0.000000 0.000000 1.000000 -99.000000 
sto_xyz:4	0.000000 0.000000 0.000000 1.000000

label:

qform_name	Scanner Anat
qform_code	1
qto_xyz:1	0.050000 0.000000 0.000000 -99.000000 
qto_xyz:2	0.000000 0.050000 0.000000 -249.000000 
qto_xyz:3	0.000000 0.000000 0.050000 -99.000000 
qto_xyz:4	0.000000 0.000000 0.000000 1.000000 
qform_xorient	Left-to-Right
qform_yorient	Posterior-to-Anterior
qform_zorient	Inferior-to-Superior
sform_name	Scanner Anat
sform_code	1
sto_xyz:1	0.050000 0.000000 -0.000000 -99.000000 
sto_xyz:2	0.000000 0.050000 -0.000000 -249.000000 
sto_xyz:3	0.000000 0.000000 0.050000 -99.000000 
sto_xyz:4	0.000000 0.000000 0.000000 1.000000

Given that the orientation is already wrong, maybe the best course of action is to:

Fix the orientation in all the datasets
Match qform/sform of label --> image (because the label header is the correct one)

Access to wanb refused

I tried to access the wanb server using the link https://wandb.ai/jcohenadad/mouse-zurich?workspace=user-jcohenadad but I'm not allowed to connect to this project.

Maybe there is a way to be added to the wanb team ?

Thanks !

Metrics measurement based on model output

Using this issue to summarize a discussion with Benjamin.

The objectives are to extract:

Ventro-dorsal diameter (AP diameter) of the spinal cord
Right-left diameter of the spinal cord
Spinal cord area
WM area
GM area

These steps require the use of SCT (these were done on v6.1):

To do so, use the following steps on the predicted file file_pred.nii.gz:

Divide the image in two files, one for GM and one for WM:
For WM : sct_maths -i ~/file_pred.nii.gz -bin 1 -o ~/file_pred_WM.nii.gz
For GM : sct_maths -i ~/file_pred.nii.gz -sub ~/file_pred_WM.nii.gz -o ~/file_pred_GM.nii.gz
sct_maths -i ~/file_pred_GM.nii.gz -sub ~/file_pred_WM.nii.gz -o ~/file_pred_GM.nii.gz (need to do it twice because WM value is 2)
To obtain a Spinal cord segmentation (both WM and GM voxel are equal to 1):
sct_maths -i ~/file_pred.nii.gz -bin 0.5 -o ~/file_pred_sc.nii.gz
To measure AP and RL diameters (for each slice):
sct_process_segmentation -i ~/file_pred_sc.nii.gz -o ~/output_diameters.csv -perslice 1
The information is stored in column MEAN(diameter_AP) and MEAN(diameter_RL)
To measure SC area (for each slice):
sct_process_segmentation -i ~/file_pred_sc.nii.gz -o ~/output_sc_area.csv -perslice 1
The information is stored in MEAN(area)
To measure WM area (for each slice):
sct_process_segmentation -i ~/file_pred_WM.nii.gz -o ~/output_wm_area.csv -perslice 1
The information is stored in MEAN(area)
To measure GM area (for each slice):
sct_process_segmentation -i ~/file_pred_GM.nii.gz -o ~/output_gm_area.csv -perslice 1
The information is stored in MEAN(area)

Furthermore, it is recommanded to do apply sct_proces_segmentation on the original image and not the straightened one, as straightening might bias the results.

Consider using ensemble/bagging for higher performance on prediction

The idea would be to train several models, with randomized train/validation split, and then aggregate the various inferences.

Possibly relevant: https://colab.research.google.com/github/Project-MONAI/tutorials/blob/main/modules/cross_validation_models_ensemble.ipynb#scrollTo=ZWSGnqA12FgX

Compare 2D nnU-Net results with UNet results

Run Julien's model on nnU-Net formatted test folder.
Create gif for visual comparison of nnU-net and U-Net on one full subject (with zero annotation)
Perform visual analysis/comparison of results with @jcohenadad
(if useful)Look into metrics such as Dice to look at similarity/differences between nnU-Net and U-Net results

Related to #30

Do proper 2D slice training from 3D volumes

See this tutorial, notably the 'squeeze dim', which I am currently not doing.

Useful links:

With the current code, I am able to get 2D patches, only those that contain a non-empty label slice:

Output uncertainty with ensemble/bagging

Needed for #31

Add citation for this repos

Consider submitting to zenodo and add a CITATION.cff file.

Find optimal number of epochs

500 (10k steps / 20 images) seems a good number:

Group image/label/prediction on wandb

It would be nice to group the image/label/prediction on wandb so that there is only one slider to move (instead of moving it 3 times):

Create requirements.txt file with proper installation instructions

Try smoothing (esp. along z) to gain performance at test-time

Comparison of image without smoothing, smoothing with 3-slice average and gaussian averaging:

Add artifacts to wandb

Command launched
Logs

Device GPU number not properly displayed

6ab5517 doesn't work as expected

Not enough WM/GM contrast

Some volumes do not show enough contrast between the WM/GM and will be ignored from the training/validation dataset.

Example (sub-mouse8_chunk-4_T1w.nii.gz):

Find out where models are typically stored

Eg: check on various nnUnet-based projects, are they hosted on GH? if so on the release artifacts (what is max size), if not where do ppl usually store these models?

Related to #38

Final nnU-Net model training for zurich-mouse

Here is the followed strategy for the final model training for the zurich-mouse dataset for white and grey matter segmentation.

Extract annotated 2D slices from the dataset
Train a 2D nnU-Net model on the annotated slices
Predict segmentation on full 3D volumes
Train a 3D nnU-Net model on these segmentations and run inference on entire dataset
Select some 3D predictions, improve segmentation (by removing sections which are not annotated) and add to initial training dataset
Train a 3D nnU-Net on the training dataset (slices and selected 3D nnU-Net predictions)
Test on entire zurich-mouse dataset.
Save model and write documentation on how to use it

Related to #32 #37

Create simple test.py inference script

similar to https://github.com/ivadomed/model_seg_mouse-sc_wm-gm_t1/blob/main/test.py, but compatible with nnUnet:

input NIfTI file
create nnUnet compatible folder structure
run nnUnet CLI as system call
reorganize file
delete temp files
done

Related to #38

Error when running on GPU

Commit: 9498ecb

Traceback (most recent call last):
  File "/home/GRAMES.POLYMTL.CA/p101317/code/model_seg_mouse-sc_wm-gm_t1/train.py", line 260, in <module>
    wandb.log({"Validation_Image/Image": wandb.Image(val_inputs, caption=f"Slice: {slice_num}")})
  File "/home/GRAMES.POLYMTL.CA/p101317/code/model_seg_mouse-sc_wm-gm_t1/venv/lib/python3.9/site-packages/wandb/sdk/data_types/image.py", line 161, in __init__
    self._initialize_from_data(data_or_path, mode)
  File "/home/GRAMES.POLYMTL.CA/p101317/code/model_seg_mouse-sc_wm-gm_t1/venv/lib/python3.9/site-packages/wandb/sdk/data_types/image.py", line 296, in _initialize_from_data
    data = data.numpy()
  File "/home/GRAMES.POLYMTL.CA/p101317/code/model_seg_mouse-sc_wm-gm_t1/venv/lib/python3.9/site-packages/monai/data/meta_tensor.py", line 276, in __torch_function__
    ret = super().__torch_function__(func, types, args, kwargs)
  File "/home/GRAMES.POLYMTL.CA/p101317/code/model_seg_mouse-sc_wm-gm_t1/venv/lib/python3.9/site-packages/torch/_tensor.py", line 1279, in __torch_function__
    ret = func(*args, **kwargs)
TypeError: can't convert cuda:0 device type tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first.