Add a new segments / labels to a pre-trained model?

[ross-mitchell]

Hello:

Is there a (simple) method to add additional labels to a pre-trained model?

I would like to leverage the learning in a pre-trained model (swin_unetr_btcv_segmentation_v0.3.0) but use it to segment additional structures in my data (CT volumes of the abdomen). I am looking for suggestions about how to do this. Any insight appreciated.

The model currently is trained to segment multiple organs (liver, stomach, gall bladder, kidneys, spleen, aorta, etc). I am interested in segmenting these organs, along with organs the model is not yet aware of: the prostate, the small and large intestines, and the colon.

I imagine this could be done by adding additional channels to the tensors stored in the pre-trained model '.pt' file, (create a new model with additional channels, then copy existing weights from the model .pt file)

Alternatively, perhaps, I could load the pre-trained model into MonaiLabel, use it to segment the organs it has been trained on, then segment additional organs assigned to additional labels, then finally train a new model from scratch with the enhanced set of labels.

Thank you in advance for your suggestions.

[tangy5]

Hi @ross-mitchell , thanks for posting this. I love your question. You are talking about some "incremental learning" technique for medical segmentation. But unfortunately, no, now most segmentation models do not support adding a channel/token as output and continue fine-tune the model. The only way is to train a new model with different number of channels, but you can use pre-trained weights with encoder part or entire pretrained weights except the last segmentation head, but this might not be a optimal option regarding performance. This is a really good topic, and under discussion in the medical segmentation community, I think there should be some research exploration.

Hi @dongyang0122 @finalelement, just for a discussion interest, can we explore some "incremental learning" to achieve training with additional segment class with a pretrained segmentation model? I know there are several papers in recent year, but haven't seen anything from our team. But this might be useful for application. Especially for MONAI Label users.

[pieper]

@ross-mitchell we have been working on a similar problem, adding lymph nodes together with a set of other segments. This repository has the info on where we are now. It's starting to work, but still needs improvement. Basically we used the first model to segment a bunch of classes (bones, organs, lungs...) and then added our lymph node segmentations into those labels and trained a new model based on the resulting labelmaps. The source data, segmentations, and monai app are alll here:

https://github.com/LymphNodeQuantification/Monailabel-LNQ

Any ideas how to improve this process would be appreciated. As @tangy5 says it's an important class of problems.

[lassoan]

We have a similar plan (have not started to work on it yet), because the Total Segmentator is a good enough model for segmenting "background" structures (that are not the main target of the treatment, but they give some anatomical context, can be used for initial registration, nice for visualization, etc), but not accurate enough for surgical planning of the target organ.

[pieper]

Yes, in our case having Total Segmentator provide a lot of 'true negative' sementation appears to help provide context for our 'true positive' lymph nodes. We also hope it will allow the model to learn how to identify lymph nodes that are not segmented in our 'true positive' set, but that's yet to be seen.

[diazandr3s]

This is a great question.
We're working with @danieltudosiu to easily do this sort of "model surgery" in MONAI Label. The idea is to have a big model trained on huge datasets that you can use to retrain/fine-tune on any number of segments. More updates on this soon :)
In the meantime, my suggestion is as @ross-mitchell said: create a MONAI Label model from scratch.
For this, you could use either total segmentator (for CT images) or a pretrained model in MONAI Label to first segment the organs it has been trained on, and then segment additional organs assigned to additional labels.

Hope this makes sense

[ross-mitchell]

Thanks all for your input. Also for pointing out TotalSegmentator! That works quite well on our data. My plan is to use TotalSegmentator to generate 'original' labels for MONAI Label, then have human experts refine these.

[danieltudosiu]

Thanks a lot @diazandr3s for the shoutout <3

@ross-mitchell you could try some sort of "model surgery" through transfer learning have a look at the following two papers:

Tang, Y., Yang, D., Li, W., Roth, H.R., Landman, B., Xu, D., Nath, V. and Hatamizadeh, A., 2022. Self-supervised pre-training of swin transformers for 3d medical image analysis. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 20730-20740).

Liu, Q., Xu, Z., Jiao, Y. and Niethammer, M., 2022. iSegFormer: Interactive Segmentation via Transformers with Application to 3D Knee MR Images. In International Conference on Medical Image Computing and Computer-Assisted Intervention (pp. 464-474). Springer, Cham.

Hope they help,

Dan

[pieper]

@ross-mitchell Yes, your approach sounds good. We have had good luck training MONAI Label using a combination of TotalSegmentator output with manual annotations of lymph nodes. The MONAI Label result almost exactly matches the TotalSegmentator result when run on scans that were not part of the training set and is starting to learn the lymph nodes well too.

[rajellumen]

@diazandr3s @ross-mitchell I could use a little bit of help. We are trying to do something similar. We used lungmask to create initial segments of the lung lobes and used monai label and these segments to train a model from scratch. Now i have a model that is trained on 5 segments and it is doing a good job of providing initial segments on new studies. However, I would like to add new labels and use this 5 segement trained model as the pretrained model. Running into issues as the outchannels is now more than 5. Any suggestion on how i can make this work? thank you very much for your time.