ε-Seg: Sparsely Supervised Semantic Segmentation of Microscopy Data

TL;DR

The paper tackles the challenge of semantic segmentation in electron microscopy with extremely sparse labels. It introduces ε-Seg, a sparsely supervised framework that combines a hierarchical VAE with a Gaussian Mixture Model prior, center-region inpainting, and contrastive learning, plus a direct latent-space segmentation head to avoid clustering. It demonstrates competitive performance on BetaSeg and liver FIBSEM datasets and shows cross-domain applicability to fluorescence microscopy, while maintaining a compact model size suitable for practical use. Limitations include its 2D scope and opportunities for more adaptive loss strategies and a 3D extension to handle volumetric EM data.

Abstract

Semantic segmentation of electron microscopy (EM) images of biological samples remains a challenge in the life sciences. EM data captures details of biological structures, sometimes with such complexity that even human observers can find it overwhelming. We introduce ε-Seg, a method based on hierarchical variational autoencoders (HVAEs), employing center-region masking, sparse label contrastive learning (CL), a Gaussian mixture model (GMM) prior, and clustering-free label prediction. Center-region masking and the inpainting loss encourage the model to learn robust and representative embeddings to distinguish the desired classes, even if training labels are sparse (0.05% of the total image data or less). For optimal performance, we employ CL and a GMM prior to shape the latent space of the HVAE such that encoded input patches tend to cluster wrt. the semantic classes we wish to distinguish. Finally, instead of clustering latent embeddings for semantic segmentation, we propose a MLP semantic segmentation head to directly predict class labels from latent embeddings. We show empirical results of ε-Seg and baseline methods on 2 dense EM datasets of biological tissues and demonstrate the applicability of our method also on fluorescence microscopy data. Our results show that ε-Seg is capable of achieving competitive sparsely-supervised segmentation results on complex biological image data, even if only limited amounts of training labels are available.

Figures (8)

• Figure 1: The overall pipeline of $\bm\epsilon$-Seg which is trained on an inpainting task (of center-region masked inputs). $\phi$ and $\theta$ are encoder and decoder of the network, respectively. Dotted arrows show sampling from a distribution (Gumbel-Softmax (Categorical-like distribution) for segmentation head and Normal distribution for conditional posterior). $h$ is an intermediate feature embedding of input $x$ coming from the encoder $\phi$. $f(h)$ is a logit vector and $|f(h)|=C$ with $C$ being the number of different classes/GMM prior components (equal to 4 for "BetaSeg" Muller2021-mg). $\beta$ and $\gamma$ are feature-wise linear modulation (FiLM's film) parameters (shifting and scaling factors) of features $h$. $h'$ are the posterior distribution's parameters and are divided into two chunks shown as $\bm\mu_{\bm{L}}(\bm{x})$ and $\bm\sigma_{\bm{L}}(\bm{x})$ by c being the corresponding label of the masked center region of each input patch $x$ in the batch. $z_L$ is a sample from $\mathcal{N}(\bm\mu_{\bm{L}}(\bm{x}),\bm\sigma_{\bm{L}}^2(\bm{x}))$. $\bm{y}'$ is a differentiable sample from a Gumbel-Softmax gumbel distribution. Green arrow shows positive pair of patches having similar labels, and red arrows show negative pairs of patches having dissimilar labels. $\mathcal{L}_{CL}$ is then computed on $\bm\mu s$ (further explanation can be found in section \ref{['mini:cl']}). For $\mathcal{L}_{I}$ inpainting loss, $\mathcal{L}_{CL}$ contrastive loss, $\mathcal{L}_{CE}$ cross-entropy loss and $\mathcal{L}_{KL}$ refer to Equations \ref{['eq:mse_inpainting_loss']}, \ref{['eq:cl']}, \ref{['eq:ce']} and \ref{['eq:kl']} respectively.
• Figure 2: Qualitative segmentation result on part of the test image stack (here we show section $627$ of $high\_c4$ of the "BetaSeg" dataset Muller2021-mg).
• Figure 3: Qualitative segmentation result on two crops of the whole 3D volume. (a) and (b) are section $80$ and $26$ of crop$00$ and crop$10$ in "liver FIBSEM" dataset respectively. The U-Net is sparsely-supervised (for the fully-supervised U-Net result, see Figure \ref{['fig:seg_models_zerial_full']}).
• Figure 4: Qualitative results on a representative 2-channel image from the overlapping subset of the "Aitslab-bioimaging1" and "Aitslab-bioimaging2" datasets. The first two panels show the fluorescence microscopy channels: EGFP-Galectin-3-labeled cytoplasm (left) and Hoechst 33342-stained nuclei (center-left). The center-right panel (GT) displays the ground truth semantic segmentation with nuclei (cyan) and cytoplasm (magenta). The rightmost panel ($\bm\epsilon$-Seg) shows the prediction from our method.
• Figure S1: The overall pipeline of Vanilla HVAE in Table \ref{['tab:baseline']} (first row in Table \ref{['tab:loss']}), which is trained on an inpainting task (of the center-region masked inputs). $\phi$ and $\theta$ are encoder and decoder of the network, respectively. Dotted arrows show sampling from a distribution. $h$ is an intermediate feature embedding of input $\bm{x}$ coming from the encoder $\phi$ and it is posterior distribution's parameters which is divided into two chunks shown as $\bm{\mu}_{L}$ and $\bm\sigma_{L}$. $\bm{z}_L$ is a sample from $\mathcal{N}(\bm\mu_{L}(\bm{x}),\bm\sigma_{L}^2(\bm{x)})$. For $\mathcal{L}_{I}$ inpainting loss and $\mathcal{L}_{KL}$ refer to Equations \ref{['eq:mse_inpainting_loss']} and \ref{['eq:hvae_kl']} respectively.