Drosophilla_melanogaster_Bouton_3d_segmentation

MicroSAM Large: Bouton Segmentation in the Drosophila Mushroom Body Calyx

<table> <tr> <td align="center" width="50%"> <img src="gifs/ground_truths.gif" width="300"/><br/> <em>Ground truth</em> </td> <td align="center" width="50%"> <img src="gifs/predicitons.gif" width="300"/><br/> <em>Predicted segmentation</em> </td> </tr> </table> <p align="center"> <img src="gifs/original.gif" width="300"/><br/> <em>Original raw volume</em> </p>

Model Description

In this repository, the weights of a fine tuned MicroSAM model are provided. The model was developed for the automated instance segmentation of projection neuron boutons in confocal fluorescence microscopy volumes of the Drosophila melanogaster mushroom body calyx. The starting point of this work is the vit_l_lm variant of MicroSAM, referred to throughout the underlying thesis as MicroSAM Large, which is a generalist light microscopy model that builds on the Segment Anything Model and that was fine tuned by its original authors on a curated collection of over 17,000 light microscopy images. In the work underlying this repository, that generalist encoder was further fine tuned on a custom annotated bouton dataset using the Automatic Instance Segmentation mode of MicroSAM, in which an additional decoder is trained jointly with the encoder rather than relying on the interactive point or box prompting mode of the original Segment Anything Model.

This model was developed as part of the master thesis "Image Segmentation and Quantification of Synaptic Markers in the Mushroom Body Calyx of Drosophila melanogaster," carried out at the Software Engineering Group, RWTH Aachen University.

A second, smaller variant referred to in the thesis as MicroSAM Base is also discussed below, since it was trained for direct comparison and offers a favourable trade off when compute or GPU memory are limited.

Architecture

The model retains the encoder decoder structure of the Segment Anything Model, in which a Vision Transformer encoder produces a dense image embedding that can be reused across queries. In the Automatic Instance Segmentation configuration used here, an additional lightweight decoder is trained jointly with the fine tuned encoder, and this decoder predicts three outputs at each pixel, namely a foreground probability map, a centroid distance map, and a boundary distance map. These three outputs are combined through a watershed step to produce instance labels, which allows the model to operate without manual prompts at inference time.

Because the Segment Anything framework expects three channel input and the imaging data consist of two fluorescence channels, each two dimensional slice was arranged as a three channel array in which the first channel contains the BRP shortcherry signal, the second channel contains the Kenyon cell signal, and the third channel was set to zero throughout. Slices were normalized using percentile based clipping at the 0.00001th and 99.99999th percentiles of the full volume, rescaled to the zero to one range, and subsequently rescaled to the zero to 255 range expected by the SAM image encoder.

Training Data

Training data were prepared as two dimensional axial slices extracted from preprocessed three dimensional confocal volumes of the mushroom body calyx. The complete dataset comprises 13 annotated volumetric image stacks acquired from seven distinct Drosophila brain preparations, using two different acquisition systems, namely Airyscan super resolution confocal microscopy and standard confocal laser scanning microscopy. At the volume level, these stacks are split into nine training volumes, two validation volumes, and two test volumes, with both acquisition modalities represented in each of the three partitions. Because the split is made at the level of individual brain preparations rather than at the level of slices, the training, validation, and test partitions each draw on different brain specimens, so that the two held out test volumes discussed in the evaluation section below come from brains that were not seen during either training or validation.

Only slices containing at least one annotated foreground instance were retained for training, since slices without labelled signal contribute no positive supervisory signal to the instance segmentation decoder. For the standard confocal volumes, several preprocessing variants were generated, including the original unenhanced volumes, Difference of Gaussians sharpened volumes, and Richardson Lucy deconvolved volumes, and the checkpoint provided in this repository was trained on all of these preprocessing variants jointly.

Training Procedure

Fine tuning was performed with the AIS decoder enabled, using the AdamW optimizer with an initial learning rate of one times ten to the negative fifth power and a plateau based learning rate scheduler. Training ran for a maximum of 200 epochs with early stopping applied at a patience of 40 epochs, using a batch size of one and training patches of 1024 by 1024 pixels, with patches required to contain at least 0.1 percent foreground voxels. Geometric augmentations included random horizontal and vertical flipping, random rotation by multiples of 90 degrees, and random elastic deformation, while intensity augmentations included additive Gaussian noise and random brightness scaling applied independently to the two fluorescence channels.

At inference time, each axial slice was processed independently using the AIS decoder in a tiled configuration with an 896 by 896 pixel non overlapping core and a 64 pixel halo on each side, producing a 1024 by 1024 context window per tile. The resulting per slice instance maps were assembled into a three dimensional volume and postprocessed with three dimensional label stitching based on 26 connected component analysis, together with minimum and maximum volume filtering and axial span filtering.

Evaluation

The model was evaluated against three baseline approaches, namely CellposeSAM, Swin UNETR, and nnU-Net, as well as against the existing semi manual Imaris based workflow used by the collaborating Tavosanis laboratory. The test set comprised two brain specimens contributing a combined total of 61 ground truth bouton instances once all preprocessing variants of the second specimen were included. Because missed boutons represent a permanent counting error while spurious detections can in principle be filtered downstream, the evaluation in the thesis prioritizes recall, matched instance mean intersection over union, and a recall weighted panoptic quality score, while also reporting false positive counts and volume error as practical measures of correction burden.

Under this evaluation framework, the MicroSAM Large checkpoint provided in this repository recovered all 61 ground truth boutons across the full evaluation set and achieved the highest aggregate matched instance mean intersection over union among all evaluated models, at 0.758. In a matched two image comparison restricted to the deconvolved preprocessing condition, it achieved a recall weighted panoptic quality of 0.787 together with complete recall, which was the strongest result among the evaluated configurations. Its main limitation is a comparatively high false positive count, with 30 false positives in the full evaluation aggregate and 16 in the matched two image comparison, some of which likely correspond to large boutons being split into multiple predicted components rather than to spurious detections in empty regions.

The smaller MicroSAM Base variant, which uses the vit_b_lm encoder, achieved nearly comparable performance, also recovering all 61 boutons with a mean intersection over union of 0.755 and a recall weighted panoptic quality of 0.771, while substantially reducing the false positive count to 12 in the full aggregate and 6 in the matched two image comparison. For deployments in which compute or GPU memory are limiting factors, MicroSAM Base is recommended as the practical substitute for the MicroSAM Large checkpoint provided here.

Benchmark Comparison Table

The Config column in the table below refers to preprocessing configurations, not to separate biological datasets. Three stored input representations were produced for the standard confocal LSM volumes used in this thesis. The Original representation is the unenhanced reconstructed volume. The DoG representation applies a Difference of Gaussians sharpening filter, which is a scale selective band pass operation that emphasizes structures within a chosen spatial scale range. The PSF representation, short for point spread function deconvolution, applies Richardson Lucy deconvolution, which is an iterative algorithm that inverts the blurring introduced by the microscope's own point spread function using an analytically derived approximation of that function. Airyscan volumes, which make up the other acquisition modality in the dataset, were left in their original reconstructed form throughout, since the DoG and Richardson Lucy steps were only applied to the standard confocal LSM data.

A checkpoint labelled All, as in the MicroSAM_L_all and MicroSAM_B_all checkpoints provided in this repository, was trained jointly on all three of these representations, meaning that during training the model saw the Original, DoG sharpened, and Richardson Lucy deconvolved versions of the same LSM training volumes together with the unmodified Airyscan volumes, rather than being trained separately on any single representation. Because the All preprocessing checkpoint is trained this way, it can also be applied at inference time to any of the three representations of a held out test volume, whereas a checkpoint trained on only one representation, such as the Original or PSF only checkpoints, can only be sensibly evaluated on the matching representation of the test data.

The All-PSF label in the Config column reflects which inference time representation of the held out LSM test volume was selected to pair with the All preprocessing checkpoint for fair comparison. In the underlying thesis, the All preprocessing checkpoint for each model family was evaluated separately on the raw, DoG sharpened, and Richardson Lucy deconvolved versions of the held out LSM test specimen, and the representation giving the best recall and matched instance mean intersection over union was kept as that model's representative result. For MicroSAM Large, MicroSAM Base, SwinUNETR, and nnU-Net, the Richardson Lucy deconvolved representation performed best, so their rows are labelled All-PSF, meaning the All preprocessing checkpoint applied to Richardson Lucy deconvolved input. For CellposeSAM XYZ 2D, the DoG sharpened representation performed best instead, so its row is labelled All-DoG. In every case, the Airyscan held out volume is included in its original unmodified form alongside whichever LSM representation was selected, since Airyscan data was never deconvolved or sharpened.

The table below shows each model's best, all preprocessing configuration on this matched two image test set, together with the Original only checkpoints for MicroSAM Large and MicroSAM Base, for a combined ground truth count of 33 boutons per row. mIoU is mean intersection over union over matched pairs, and PQrec is mIoU multiplied by recall.

The All-PSF rows correspond to the MicroSAM_L_all and MicroSAM_B_all checkpoints in this repository, trained on all preprocessing variants jointly, while the Original rows correspond to the MicroSAM_L_Original and MicroSAM_B_Original checkpoints, trained only on the unenhanced volumes. MicroSAM Large All-PSF gives the highest mIoU and PQrec of any configuration, at the cost of the most false positives. MicroSAM Large Original reaches nearly the same quality with half as many false positives, and MicroSAM Base Original is the leanest checkpoint that still achieves complete recall. The full per preprocessing breakdown, including the DoG and PSF only rows for each model, is reported in Table 4.6 of the accompanying thesis.

Input Format Requirements

In order to obtain the segmentation quality reported in the accompanying thesis, input images must be arranged in the same three channel order that was used during fine tuning, since the model has not been validated under alternative channel orderings and its predictions may degrade substantially if this convention is not followed. The first channel must contain the BRP shortcherry signal, which is the presynaptic active zone marker that defines the bouton instances the model was trained to detect. The second channel must contain the Kenyon cell signal, which served as anatomical context during fine tuning rather than as the primary segmentation target, even though it was not discarded from the input. The third channel must be set to zero, which is the convention adopted during fine tuning in place of the channel averaging strategy that the original MicroSAM authors found to perform somewhat better in their own experiments. In addition, input slices should be lateral axial sections, since the model was fine tuned exclusively on this orientation and was not exposed to coronal or sagittal reslicing during training. Departing from this format, for instance by reversing the channel order or by filling the third channel with non zero values, is likely to produce degraded segmentation quality, since the fine tuned decoder has learned to associate spatial patterns in the first channel specifically with bouton structures.

Available Checkpoints

This repository provides four fine tuned checkpoints, corresponding to the two encoder sizes discussed above and to two different training data compositions. Each checkpoint folder contains the weights in both the original best.pt format and a best.safetensors format.

The evaluation numbers reported above correspond to the models/MicroSAM_L_all and models/MicroSAM_B_all checkpoints, since the all preprocessing training condition is the one analyzed in the final model selection section of the thesis.

Usage

from micro_sam.automatic_segmentation import automatic_instance_segmentation

result = automatic_instance_segmentation(
    input_path="path/to/slice.tif",
    checkpoint_path="models/MicroSAM_L_all/best.pt",
    model_type="vit_l_lm",
)

To use the smaller, faster substitute checkpoint instead, point checkpoint_path to models/MicroSAM_B_all/best.pt and set model_type to vit_b_lm. Both best.pt and best.safetensors contain the same weights, so either file can be used depending on which format your inference environment expects. Refer to the MicroSAM library documentation for the exact arguments expected by your installed version, since this repository contains only the fine tuned weights and not the full preprocessing or postprocessing pipeline described in the thesis.

Limitations

The model was fine tuned on a comparatively small annotated dataset drawn from two brain specimens, so performance on data acquired with substantially different imaging settings, labelling strategies, or biological preparations has not been verified beyond the evaluation set described above. Because individual boutons appear as clusters of discrete fluorescent puncta rather than as homogeneous filled objects, the model produces a non negligible number of false positives, several of which likely correspond to large boutons being split into multiple predicted components in regions with weak connecting signal between active zone puncta. Predictions should therefore be reviewed before being used for unsupervised quantitative analysis, particularly in dense or low signal regions.

Citation

If you use these weights in your work, please cite the accompanying thesis.

@mastersthesis{rawat2026microsam,
  author  = {Aditya Rawat},
  title   = {Image Segmentation and Quantification of Synaptic Markers in the Mushroom Body Calyx of Drosophila melanogaster},
  school  = {RWTH Aachen University, Software Engineering Group},
  year    = {2026}
}

License

This model is released under the MIT License. Please refer to the LICENSE file in this repository for full terms.

Acknowledgments

This work was carried out at the Software Engineering Group, RWTH Aachen University, under the supervision of Prof. Dr. Abigail Morrison, with Prof. Dr. Ing. Johannes Stegmaier serving as second examiner. The Imaris reference segmentations used for comparison were provided by collaborating biologists in the Tavosanis laboratory.