SAMBA: A Trainable Segmentation Web-App with Smart Labelling

Segmentation is the assigning of a semantic class to every pixel in an image and is a prerequisite for various statistical analysis tasks in materials science, like phase quantification, physics simulations or morphological characterization. The wide range of length scales, imaging techniques and materials studied in materials science means any segmentation algorithm must generalise to unseen data and support abstract, user-defined semantic classes. Trainable segmentation is a popular interactive segmentation paradigm where a classifier is trained to map from image features to user drawn labels. SAMBA is a trainable segmentation tool that uses Meta's Segment Anything Model (SAM) for fast, high-quality label suggestions and a random forest classifier for robust, generalizable segmentations. It is accessible in the browser (https://www.sambasegment.com/) without the need to download any external dependencies. The segmentation backend is run in the cloud, so does not require the user to have powerful hardware.


Statement of need
Trainable segmentation tools like ilastik (Berg et al., 2019) and Trainable Weka Segmentation (Arganda-Carreras et al., 2017) are popular options for image segmentation in biology and materials science due to their flexibility.They apply common image processing operations like blurs, edge detection, and texture filters to create a feature stack, then train a classifier (usually a random forest ensemble) to map from the stack to user drawn labels on a pixelwise basis.ilastix offers fast segmentations and a responsive live-view, whereas Trainable Weka Segmentation is part of and interoperates with the FIJI ImageJ library (Schindelin et al., 2012).Other trainable segmenters are based on the napari (Ahlers et al., 2023) image viewer, including napari-feature-classifier (Luethi & Hess, n.d.) and the GPU-accelerated napari-apoc (Haase et al., 2023).
Accurate, representative labels are important for segmentation quality, but tooling is limited to polygons and/or pixel brushes which can make labelling boundaries time-consuming.SAMBA (Segment Anything Model Based App) uses the recent object-detecting foundation model Segment Anything Model (Kirillov et al., 2023) as a 'Smart Labelling' tool, which offers suggestions for associated objects as the cursor is moved which can then be added as labels.This 'Smart Labelling' can quickly pick out portions of the same phase and is resilient against noise, varying exposure, and certain types of artefacts etc. SAM has been applied to cell microscopy annotation and segmentation (Archit et al., 2023), but not to multi-phase segmentation.
SAMBA is also unique in being fully web-based, meaning it can be used without downloading a large binary, installing libraries or requiring significant local computational power.An associated gallery page allows users to contribute to and view an open dataset of micrographs, labels, segmentations and experimental metadata.It is designed to be usable by researchers regardless of prior programming or image analysis experience.an embedding of an image generated by a pretrained autoencoder (Dehghani et al., 2023) through a lightweight decoder network to produce a segmentation for the prompt (i.e, the object at the mouse cursor).The decoder is lightweight enough that it can be run in the browser in real-time when exported via the ONNX framework (Bai et al., 2019), though the embedding must still be generated by a large autoencoder model running on a computer with PyTorch.Despite its impressive zero-shot performance, SAM is primarily an object-detection model and cannot be used for the semantic phase segmentation desired in materials science.However, its understanding of shape, texture and other lower-level features apply well to instances of certain phases.This provides the motivation for and structure of SAMBA: a Python web server supplies the image embedding (and later performs random forest segmentation) for the user's image, then browser-based labelling is performed by client-side object segmentations as label suggestions.Labels are confirmed with a left click, and a right click switches length scales of the smart labelling -this was achieved by exposing SAM's hidden multi-mask output.Normally, SAM produces three masks per prompt -each representing a different length scale and with an associated internal quality estimate.Giving the user the choice of length scales allows greater flexibility, which is useful given SAM is not designed for micrographs, and therefore its internal estimate of quality may not always match the user's.

Design and usage
Standard drawing tools like a variable width brush, polygon, and eraser allow the user to create labels when SAM struggles, for example with very high frequency spatial features.Previous labels are not overwritten (unless erased), so new labels can be added on top for boundary fine-tuning or rapid whole-image labelling.These labels can then be downloaded for use in offline workflows, like training task-specific neural networks.
After labelling is complete, the Python web server trains a random forest ensemble to map the feature stack (computed in the background) to the labels.The feature stack creation involves a Python reimplementation of most of the image processing operations outlined in Trainable Weka Segmentation.This is achieved predominantly using the scikit-image library (Walt et al., 2014).Once completed, the resulting segmentation of the full image is returned to the user as a variable opacity overlay, allowing for additional labels to refine the output.An optional overlay highlights pixels the classifier is least certain about, which could be useful starting points for new user labels.Improvements in this 'Human-In-The-Loop' (HITL) feedback (i.e, feature-weighted uncertainty highlights) are interesting avenues for future work.The trained classifier is a scikit-learn object (Pedregosa et al., 2011) -it can be downloaded and applied to future data either on SAMBA or in other Python workflows.

Figure 1 :
Figure 1: (a) screenshot of the SAMBA website, displaying the different labelling options including SAM powered 'Smart Labelling'.(b) shows how changing the Smart Label region sizes affects the suggested label at the same mouse position (red), giving the user the flexibility to focus on different length scales.(c) an example output segmentation of the tool, which can be saved as .tifffor later analysis.