U-Net segmentation of nerve histologies

The problem

Vagus nerve stimulation is used to treat epilepsy, rheumatoid arthritis and other serious conditions, but current stimulation is not selective enough and causes side effects. Nerve fibres are organised into bundles called fascicles, and knowing their layout helps target stimulation more precisely. That layout comes from segmenting stained histological sections — which, done by hand, does not scale to large datasets.

Approach

We trained a U-Net to perform semantic segmentation of the fascicles, replacing the classical skip connections with Xception-style residual connections. With only 72 annotated images available, the work focused on making the most of little data and on strongly imbalanced classes (fascicles cover a small fraction of each image):

• data augmentation with random affine transformations (rotation, translation, shear, scaling) and colour/brightness shifts;
• a comparison of loss functions — binary cross-entropy, focal loss, and their combination;
• post-processing: removal of implausibly small predictions, watershed separation of merged fascicles, and flagging of suspicious shapes.

Outcome

Combining binary cross-entropy with focal loss generalised best, raising the validation IoU from 0.60 (cross-entropy alone) to 0.69. The post-processing brought the distribution of predicted fascicle areas and shapes much closer to the ground truth.