axondeepseg / sam_myelin_seg_tem Goto Github PK

Pretty straightforward. This would make this pipeline fully automatic.

From my understanding, it would be possible to get accurate results because the SOTAs for object detection are pretty powerful. Also, the data is already there. For other datasets than TEM, some minor preprocessing would be required to get the axon centroids.

The TEM dataset was the biggest of our private datasets, but we will also need to move on to other modalities to compare performance with ivadomed and nnunetv2.

The ultimate goal of using SAM was to train a single "foundation" model for every type of contrasts/resolution, so we will eventually want to train SAM on an aggregation of all our private datasets. For comparison, we will need to train models on every dataset individually before aggregating. This necessary step will also allow us to fix bugs with these datasets before we use them all at once.

It would be important to keep the model that gives the best validation metrics. Currently, this is not the case and we use the final checkpoint for testing, but this is not optimal.

The exact way to process batches is described at the end of this notebook: https://github.com/facebookresearch/segment-anything/blob/main/notebooks/predictor_example.ipynb

Both the images and the prompts can be batched.

We would like to know how easily SAM can be fine-tuned for fully automatic segmentation (i.e. prompting with the whole image as a bbox instead of prompting with a ROI of interest).

The perfect pretext to try this is axon segmentation.

1. If it works well, this axon segmentation could then be used to generate the bbox prompts used subsequently for the myelin segmentation
2. There is no overlap between instances in the axon class. The myelin class would not be well suited because there are many overlaps (or rather touching myelin sheaths). A big advantage of segmenting the myelin with localized bbox prompts is that we get a clean and reliable instance segmentation. The way AxonDeepSeg currently works it that the semantic segmentation is "subdivided" (semantic to instance) with a watershed algorithm. Although it works fine in a lot of cases, sometimes this process deteriorates the segmentation. See example below, and look for small axons touching big axons. In the instance segmentation, a "leak" artifact occurs, where the myelin of the small axon is wrongly attributed to the big axon.
   

Currently, bounding boxes are directly loaded for training. They are generated by extracting the thightest bounding box around the annotation. The exact coordinates of the bounding box should have a random component to avoid overfit. Similar to what was done in MedSAM:

https://github.com/bowang-lab/MedSAM/blob/8432244ac07be6baba120dcb786e8a694c188eb9/train_one_gpu.py#L104-L107

It would be interesting to see if training the image encoder as well could help. MedSAM trains the encoder as well and they specify that all the weights in the image encoder were updated.

Currently, the axon and myelin segmentation models are trained separately and independently. The 2 cascaded models should be trained at once. This would also allow parameter sharing, like having a common image encoder for both models.

This month, the myelin segmentation should get good enough to move on to this "cascaded" training. I already expect a lot of autograd problems...

However, after this, the model should be ready for a public release.

This is a general overview of what needs to be done in this project before moving on to other datasets. Currently, both the axon and myelin seg models outperform ivadomed, but the overall pipeline is not efficient and nnUNet still beats SAM.

• integrate "patch-based" training: similar to how we trained U-Nets, this would allow bigger batch sizes and would eventually allow for data augmentation. Ideally, implement this with the MONAI dataloader for easier dataAug integration
• merge axon and myelin image encoders (halves overall model size, allows parameter sharing, more efficient training pipeline); see #10. Eventually, all datasets would be aggregated and the image encoder would learn to process all modalities.
• implement multi-GPU training to further increase the batch size and be able to train longer

I just realized that I forgot to re-compute the image embeddings for ViT_H training. Not sure I understand why the training could still be completed... This needs to be fixed for proper ViT_H results.

I would like to add some regularization to the loss function for robustness to discourage the model to produce "glitchy" segmentations. For a perfect illustration, see the image below, taken from the validation set of https://github.com/brainhack-school2023/collin_project/tree/main (first iteration of this project).

1. Pink axon has discontinuities.
2. Brown axon is not complete

I am not yet entirely sure how to regularize the myelin prediction, but will update this issue later.

(Maybe we should keep the training scripts separate from the rest)
This would be pretty straightforward to integrate. The model checkpoints will need a release, and we can use the inference script as a reference. The only time-consuming part would be to add tests.

Order in which samples are loaded should be shuffled.