Self-Supervised Edge Detection Reconstruction for Topology-Informed 3D Axon Segmentation and Centerline Detection

Abstract

Many machine learning-based axon tracing methods rely on image datasets with segmentation labels. This requires manual annotation from domain experts, which is labor-intensive and not practical for large-scale brain mapping on hemisphere or whole brain tissue at cellular or sub-cellular resolution. Additionally, preserving axon structure topology is crucial to understanding neural connections and brain function. Self-supervised learning (SSL) is a machine learning framework that allows models to learn an auxiliary task on unannotated data to aid performance on a supervised target task. In this work, we propose a novel SSL auxiliary task of reconstructing an edge detector for the target task of topology-oriented axon segmentation and centerline detection. We pretrained 3D U-Nets on three different SSL tasks using a mouse brain dataset: our proposed task, predicting the order of permuted slices, and playing a Rubik's cube. We then evaluated these U-Nets and a baseline model on a different mouse brain dataset. Across all experiments, the U-Net pretrained on our proposed task improved the baseline's segmentation, topology-preservation, and centerline detection by up to 5.03%, 4.65%, and 5.41%, respectively. In contrast, there was no consistent improvement over the baseline observed with the slice-permutation and Rubik's cube pretrained U-Nets.