AxonSynth

Pretrained three-class 3D U-Net for axon segmentation in light-sheet microscopy, trained on domain-randomized synthetic images.

Code, environment, and documentation: github.com/lincbrain/axonsynth

Model

• Architecture: MONAI 3D U-Net, approximately 4.8 million parameters.
• Classes: background, axon shell, and axon interior.
• Training patches: 128 × 128 × 128 voxels at 0.8 µm isotropic resolution.
• Checkpoint: best_model.pt, the main three-class model selected at epoch 136.

Binary axon foreground probability is calculated by summing the shell and interior probabilities.

Download

Using the Hugging Face CLI:

hf download edziocodes/axonsynth best_model.pt --local-dir weights

Run inference

Clone the code repository and create its environment:

conda env create -f environment.yml
conda activate axonsynth

From the code repository root, run inference on a volume prepared at 0.8 µm isotropic resolution in model axis order. See the repository README for preparation instructions.

python inference/infer_lsm.py \
    --input prepared_0p8um.nii.gz \
    --checkpoint weights/best_model.pt \
    --output-dir predictions \
    --segmentation-mode three_class_shell_interior \
    --stitch-device cpu \
    --output-profile binary_focused \
    --threshold 0.5 \
    --require-cuda

This command saves the normalized input, foreground probability, binary mask, and inference settings. The threshold is configurable.

Evaluation

To compare a prediction against a dense ground-truth volume of the same shape, run from the code repository root:

python tools/evaluate_lsm_case.py \
    --predictions-dir predictions \
    --prefix prepared_0p8um \
    --gt-path ground_truth.nii.gz \
    --threshold 0.5

This writes metrics.json, summary.txt, and true-positive/false-positive/false-negative overlay images to evaluation/. The underlying reusable metric functions are in paper_analysis/metrics.py.

License

This model checkpoint is released under the GNU Affero General Public License v3.0 (AGPL-3.0). See LICENSE for the full terms. The source code is available in the AxonSynth GitHub repository under the same license.