Try fully automatic axon segmentation

[hermancollin]

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.

[hermancollin]

In a now deleted notebook on the MedSAM repo, this was used to get the full bbox:

B,_, H, W = gt.shape
boxes = torch.from_numpy(np.array([[0,0,W,H]]*B)).float().to(device)

Taken from https://github.com/bowang-lab/MedSAM/blob/d9298cdf04914b0934c6c0188dea874f3a7f2405/finetune_and_inference_tutorial_auto_seg.ipynb

[hermancollin]

Preliminary results from the ac/add_axon_segmentation_training branch trained for 100 epochs. Not very impressive.

GT: prediction:

I can see a trend in the 8 validation images where the right part of the image is very poorly segmented compared to the left part (as can be seen above). I suspect there's a problem with the prompt.

[hermancollin]

This was the case. The validation bboxes were rotated 90 degrees. Fixed in a62a7ecf4424840989167eab46e9fa9752eb5466

Now it's slightly better