Train a DL model for spinal cord centerline detection/prediction
Opened this issue · 18 comments
This idea was floated across other projects as well, especially, in model_seg_sci and canproco projects and thus is gaining more traction. This issue tries to unify those ideas into a common place so that all the brain storming in the future will be done here (thank you @valosekj for the suggestion!).
This is what we have in mind for testing this idea (taken from the issue linked above in the model_seg_sci project):
We could formulate this as a segmentation problem and make the model output the spinal cord (SC) centerlines as segmentation niftis.
Rationale? - We already get decent outputs from sct_get_centerline as binary nifti files containing discretized SC centerlines. Moreover, these are of the same shape as the input image, thereby setting up conveniently as a segmentation problem. Hence, the input to the model will be a T2w sagittal image (initially), and it is trained to output a binarized prediction of the centerline.
Implementational details that need to be ironed out:
- Usually, the input image is cropped to a manageable size using the (dilated) centerline or SC mask as the reference. How will the input be cropped if the goal itself is to predict the centerline?
- Should orientations other than sagittal be considered? If so, how? What additional information can they provide?
Alternative - If the above idea does not work for some reason, here's an alternative. Instead of posing it as a segmentation problem, pose it as a regression problem. This is because sct_get_centerline also outputs a .csv file containing the voxel coordinates of the centerline along with binarized nifti. Then, a model is trained to regress these coordinates (3 values) for each slice in the sagittal image (hoping that it would be continuous). A centerline could then be constructed based the predicted voxel coordinates.
EDIT:
Also tagging @NathanMolinier for getting up to speed on this latest project idea and possibly getting your thoughts on this!
EDIT:
Repo: https://github.com/ivadomed/model_sc_centerline
Answering @jcohenadad's question in the issue(comment) here:
Yes! I did mean "other than sagittal", missed a word there - corrected that in my description above. I do not want to enforce this but I said sagittal to begin with because in this orientation we could see the centerline as a whole as opposed to other orientations where it is a single voxel. That said, this shouldn't matter if we are training 3D models but if we are looking into 2D models then maybe using the sagittal slices for centerline prediction should be better.
if we are looking into 2D models then maybe using the sagittal slices for centerline prediction should be better.
yes, good point. In fact, I am thinking of a possibly good approach: average ~10-15 sagittal slices in the middle of the FOV (medial plane), the spinal cord should appear in most cases (because FOVs are usually centered in the medial plane), and train a 2D model to detect the centerline in this plane.
Then, do the exact same thing for the coronal plane, and then combine the sagittal and coronal centerline coordinates.
yes, good point. In fact, I am thinking of a possibly good approach: average ~10-15 sagittal slices in the middle of the FOV (medial plane), the spinal cord should appear in most cases (because FOVs are usually centered in the medial plane), and train a 2D model to detect the centerline in this plane.
Thank you for this suggestion! We tried this with @naga-karthik and @NathanMolinier.
This is what the cropped (sct_crop_image) and averaged (sct_maths -mean x) sagittal image looks like.
Note that this is only the average of 15 slices in the sagittal plane, which results in a 2D image.
Then, do the exact same thing for the coronal plane, and then combine the sagittal and coronal centerline coordinates.
And this is what the averaged (sct_maths -mean y) coronal image looks like.
Note that this is the average in the coronal plane across the entire image (i.e., no cropping applied), which results in a 2D image. We did not do the cropping because it will be highly heterogeneous across subjects (due to the curvate of the spinal cord).
With these two images, we should be able to extract all three coordinates along the spinal cord.
Maybe, for the sagittal plane, we could also avoid the cropping, and we can just average the entire image? To avoid issues with subjects with "abnormal" SC anatomy such as scoliosis. The tradeoff would be the loss of the contrast, though (as in the case of averaging in coronal plane).
Maybe, for the sagittal plane, we could also avoid the cropping, and we can just average the entire image? To avoid issues with subjects with "abnormal" SC anatomy such as scoliosis.
if it works, fine, but if not let's crop in the middle as suggested
Here are a few preliminary updates from the 3D model we trained.
Background: We went ahead with our initial hypothesis mentioned in the comment:
We could formulate this as a segmentation problem and make the model output the spinal cord (SC) centerlines as segmentation niftis.
As far as these questions are concerned:
Usually, the input image is cropped to a manageable size using the (dilated) centerline or SC mask as the reference. How will the input be cropped if the goal itself is to predict the centerline?
We did not crop the image using SCT before training the model, rather, we directly used the CenterCrop in ivadomed to roughly crop around the spinal cord. More details below.
Should orientations other than sagittal be considered? If so, how? What additional information can they provide?
Since, this is a 3D model, all the planes are considered.
Config file used:
{
"command": "train",
"gpu_ids": [2],
"path_output": "outputs/3d-model_ccrop=320-224-32_len=320-224-32_str=160-112-32_N=250",
"model_name": "modified_unet",
"debugging": true,
"object_detection_params": {
"object_detection_path": null,
"safety_factor": [1.0, 1.0, 1.0]
},
"wandb": {
"wandb_api_key": "bf043c64b6c6b4abcc1ee7d8501c300b19945028",
"project_name": "centerline-detection",
"group_name": "ivadomed-3d-unet",
"run_name": "ccrop=320-224-32_len=320-224-32_str=160-112-32_N=250",
"log_grads_every": 1000
},
"loader_parameters": {
"path_data": ["/home/GRAMES.POLYMTL.CA/u114716/duke/temp/janvalosek/canproco_T2w_centerline_2023-01-23/data_processed"],
"target_suffix": ["_seg_centerline"],
"extensions": [".nii.gz"],
"roi_params": {
"suffix": null,
"slice_filter_roi": null
},
"contrast_params": {
"training_validation": ["T2w"],
"testing": ["T2w"],
"balance": {}
},
"slice_filter_params": {
"filter_empty_mask": false,
"filter_empty_input": false
},
"subject_selection": {
"n": [50, 200],
"metadata": ["pathology", "pathology"],
"value": ["HC", "MS"]
},
"slice_axis": "sagittal",
"multichannel": false,
"soft_gt": false,
"bids_validate": false
},
"split_dataset": {
"fname_split": null,
"random_seed": 100,
"split_method" : "participant_id",
"data_testing": {"data_type": null, "data_value":[]},
"balance": null,
"train_fraction": 0.6,
"test_fraction": 0.2
},
"training_parameters": {
"batch_size": 4,
"loss": {
"name": "DiceLoss"
},
"training_time": {
"num_epochs": 200,
"early_stopping_patience": 50,
"early_stopping_epsilon": 0.001
},
"scheduler": {
"initial_lr": 1e-4,
"lr_scheduler": {
"name": "CosineAnnealingLR",
"base_lr": 1e-5,
"max_lr": 1e-3
}
},
"balance_samples": {"applied": false, "type": "gt"},
"transfer_learning": {
"retrain_model": null,
"retrain_fraction": 1.0,
"reset": true
}
},
"default_model": {
"name": "Unet",
"dropout_rate": 0.25,
"bn_momentum": 0.1,
"is_2d": false,
"final_activation": "relu"
},
"uncertainty": {
"epistemic": false,
"aleatoric": false,
"n_it": 0
},
"postprocessing": {
"remove_noise": {"thr": -1},
"keep_largest": {},
"binarize_prediction": {"thr": 0.5},
"uncertainty": {"thr": -1, "suffix": "_unc-vox.nii.gz"},
"fill_holes": {},
"remove_small": {"unit": "vox", "thr": 3}
},
"evaluation_parameters": {},
"Modified3DUNet": {
"applied": true,
"length_3D": [320, 224, 32],
"stride_3D": [160, 112, 32],
"attention": false,
"n_filters": 8
},
"transformation": {
"Resample": {
"hspace": 0.8,
"wspace": 0.8,
"dspace": 0.8
},
"CenterCrop": {
"size": [320, 224, 32]
},
"RandomAffine": {
"degrees": 5,
"scale": [0.15, 0.15, 0.15],
"translate": [0.1, 0.1, 0.1],
"applied_to": ["im", "gt"],
"dataset_type": ["training"]
},
"RandomBiasField": {
"coefficients": 0.5,
"order": 3,
"p": 0.25,
"applied_to": ["im"],
"dataset_type": ["training"]
},
"HistogramClipping": {
"min_percentile": 3,
"max_percentile": 97,
"applied_to": ["im"]
},
"NormalizeInstance": {"applied_to": ["im"]}
}
}
Notable parameters here are: CenterCrop - [320, 224, 32], length_3D - [320, 224, 32], stride_3D - [160, 112, 32]. These were chosen by looking at how cropped inputs look on WandB (see loss curves and image patches here). The actual size of the inputs are [320, 320, 56]. Note that for this experiment, we used the original (un-averaged) input image.
Results:
In Green is the GT centerline (generated by sct_get_centerline) and in Red are the model's prediction.
The good part
We observe that the model is able to predict the centerline to a good extent for the cervical spine (mainly because they were the most exposed during training). It also gets some parts of the centerline that were missed by sct_get_centerline.
The bad part
We observe here that for most of the thoracic part, there is no prediction at all. Why I think this is the case is because we use 32 in the depth dimension during CenterCrop which is not considering the slices containing the centerline in the thoracic part. I'm currently another experiment which increases the CenterCrop dimensions so that more input is included.
My Thoughts: Our preliminary hypothesis of testing out the centerline prediction as a segmentation problem definitely looks promising. The average Dice based on 50 test subjects is about 0.49. As mentioned above, the cropping might still need to optimized along with other standard hyperparameters by running more experiments.
Great work! You were so fast in getting results, this is exciting! Few suggestions:
- Open a repository under the ivadomed organization, anticipating that the model will be hosted there alongside all the issues/discussions). The earlier we do this, the earlier we can centralize all the related information in that single repository.
- The CenterCrop is a major issue. Some images will be larger than
[320, 224, 32]at test time, and if cropping is applied the centerline will also be cropped. However, I'm thinking maybe you could try to remove the CenterCrop at test time. - From #15 (comment), looking at the 'The bad part', the GIF anim does not show any red, but there seem to be two green overlays. So, is it possible that the prediction was mistakenly colored in green?
Yes, it is exciting indeed!
Open a repository under the ivadomed organization, anticipating that the model will be hosted there alongside all the issues/discussions).
Sounds good, how about we call it model_sc_centerline?
The CenterCrop is a major issue. Some images will be larger than
[320, 224, 32]at test time, and if cropping is applied the centerline will also be cropped. However, I'm thinking maybe you could try to remove the CenterCrop at test time.
Yes, I'll try removing this parameter during test time in the next run
Looking at the 'The bad part', the GIF anim does not show any red, but there seem to be two green overlays. So, is it possible that the prediction was mistakenly colored in green?
Aha! I purposely did it this way to show that there were indeed no predictions for that part of the spinal cord. The model only predicted the cervical SC parts
Sounds good, how about we call it model_sc_centerline?
👍
Aha! I purposely did it this way to show that there were indeed no predictions for that part of the spinal cord. The model only predicted the cervical SC parts
I still don't understand. Your image shows two green centerlines around thoracic 4 (T4). So, if one green centerline is the ground truth, what is the other green centerline?
Sorry for not being clear enough. Notice that I am overlaying the prediction over the GT in the 1st gif. So, the 3rd image in the 1st gif (in red) is the prediction overlayed on the GT. Now, for the 2nd gif, since there is no prediction, it appears that there are 2 green centerlines (because there is no 3rd, or, it is empty).
updated 2nd gif with filenames for clarity
Note that when we come to the 3rd image, we still see the green centerline (which is the actual GT). There is no red centerline (because no prediction)
ok, let me clarify: in your GIF animation, when going from "_seg_centerline" to "_pred", some green pixels are displaced (look closely at the screen, you will see it). Now: if there was no prediction, i would not expect any green pixel to be displaced (because it is all zeros).
That's a good catch! But, it is very strange though, I am now showing the prediction and the GT in separate gifs, and we can see that there is indeed no prediction! Idk what's the problem here.
Input image and GT
Input image and Prediction
Even by looking at this gif very closely, it seems that the input image has moved (in the middle of the image, near the cervical region). How can this be the case if they essentially the same screenshots? It has got to do something with the gif creation process
It has got to do something with the gif creation process
Yup, highly possible. Do you use JPEG? I always use PNG instead of JPEG to avoid compression errors.
Anyway, to confirm this issue, you can simply zoom heavily on this region, and re-do your GIF.
BTW: in the future it is better to show a single GIF switching between GT and prediction, instead of showing two separate GIF (as here: #15 (comment)). The problem of showing two separate GIFs is that you might be selecting another sagittal plane, and the reason the centerline is not appearing could be because it is located in another plane.
So, instead of doing:
FRAME 1:
- pred: 0
- GT: 1
- image: 1
FRAME 2: - pred: 1
- GT: 1
- image: 1
Do this:
FRAME 1:
- pred: 0
- GT: 1
- image: 1
FRAME 2: - pred: 1
- GT: 0
- image: 1
Also, to overcome the possibility that the predicted centerline does appear in another sagittal slice, I suggest to do a 'repmat' (ie: find the 'one' along the array, and replace all zeroes by one) along the RL axis, to make sure this does not happen. In fact, for the centerline detection project, I suggest we do a new QC on SCT, showing the sagittal and coronal images, with a repmat of the centerline along RL and AP respectively.
Do you use JPEG? I always use PNG instead of JPEG to avoid compression errors. Anyway, to confirm this issue, you can simply zoom heavily on this region, and re-do your GIF.
Nope, I use PNGs itself! But, can confirm that the zooming and creating the gifs again has solved this issue. Thank you for the suggestion! I would not have realized it otherwise.
I created another gif (I promise, this is the last 😅) according to the suggestion in this comment. Seems better for visualization than the overlaying that I was doing earlier. But I guess overlaying is better we want to understand what's false positive and false negative compared to the GT.
zoomed gif
Lastly, if it is not too much to ask, could you please suggest this comment in a new issue in this repository? Also, I don't think I understand
I suggest we do a new QC on SCT, showing the sagittal and coronal images, with a repmat of the centerline along RL and AP respectively.
Maybe @valosekj gets it? I will discuss with him more on this
But I guess overlaying is better we want to understand what's false positive and false negative compared to the GT.
Yes, good point
Lastly, if it is not too much to ask, could you please suggest #15 (comment) in a new issue in this repository?
(p.s. I unassigned myself as I won't be actively working on this, but i'm obviously happy to review work, follow progress, give feedback, etc.)
Heads up: during the 2023-02-14 CanProCo meeting, we decided to hold on to this project for a while. We will try to obtain SC seg directly without needing the centerline. Relevant project: contrast-agnostic-softseg.