StyleSeg V2: Towards Robust One-shot Segmentation of Brain Tissue via Optimization-free Registration Error Perception

TL;DR

StyleSeg V2 tackles the bottleneck of one-shot brain tissue segmentation by introducing optimization-free registration error perception through mirror-based analysis, enabling reliable use of unlabeled data and more faithful style transformations. It refines atlas-based training with Weighted Image-Aligned Style Transformation (WIST) and a Confidence-Guided Dice loss (L_cgd), achieving higher segmentation and registration performance than prior methods. Across OASIS, CANDIShare, and MM-WHS 2017, StyleSeg V2 yields consistent Dice gains of approximately 2% over StyleSeg and outperforms other state-of-the-art approaches, with ablations confirming the efficacy of the proposed components. This work enhances practical applicability of one-shot segmentation in neuroimaging by leveraging symmetry-inspired error perception to balance diversity, fidelity, and supervision from unlabeled data.

Abstract

One-shot segmentation of brain tissue requires training registration-segmentation (reg-seg) dual-model iteratively, where reg-model aims to provide pseudo masks of unlabeled images for seg-model by warping a carefully-labeled atlas. However, the imperfect reg-model induces image-mask misalignment, poisoning the seg-model subsequently. Recent StyleSeg bypasses this bottleneck by replacing the unlabeled images with their warped copies of atlas, but needs to borrow the diverse image patterns via style transformation. Here, we present StyleSeg V2, inherited from StyleSeg but granted the ability of perceiving the registration errors. The motivation is that good registration behaves in a mirrored fashion for mirrored images. Therefore, almost at no cost, StyleSeg V2 can have reg-model itself "speak out" incorrectly-aligned regions by simply mirroring (symmetrically flipping the brain) its input, and the registration errors are symmetric inconsistencies between the outputs of original and mirrored inputs. Consequently, StyleSeg V2 allows the seg-model to make use of correctly-aligned regions of unlabeled images and also enhances the fidelity of style-transformed warped atlas image by weighting the local transformation strength according to registration errors. The experimental results on three public datasets demonstrate that our proposed StyleSeg V2 outperforms other state-of-the-arts by considerable margins, and exceeds StyleSeg by increasing the average Dice by at least 2.4%.

Figures (11)

• Figure 1: The schematic diagram of (a) StyleSeg and (b) StyleSeg V2.
• Figure 2: The overview of iterative training in StyleSeg V2. (a) We first frozen reg-model to predict a pseudo mask (warped atlas mask) and a registration confidence map for supervising seg-model on both warped atlas image and unlabeled one. Weighted image-aligned style transformation style-transformed (WIST) increases both diversity and fidelity of the warped atlas image, and confidence guided Dice loss $\mathcal{L}_{cgd}$ constrains the seg-model using the aligned regions of the unlabeled image. (b) We then fix the seg-model to produce a refined prediction of unlabeled image, which can be further exploited to weakly supervise the reg-model, triggering a new round of iteration
• Figure 3: The calculation of the registration error map formulated in Eq. \ref{['eq:2']}. Registration path (a): direct registraion, to produce deformation flow $\phi$; Registration path (b): flipping-registration-flipping, to produce composed flow $\Phi$. The registration error is the L2 distance between $\phi$ and $\Phi$.
• Figure 4: (a) The detail of original IST, which swaps the Fourier amplitude component of the warped atlas with that of the unlabeled images to transfer the style. (b)The detail of weighted IST (WIST), which divides the warped atlas $I_{\tilde{a}}$ and unlabeled $I_{u}$ images into multiple sub-images with different confidence intervals, and applies lesser style-transformation strength on the sub-images with lower confidence, and vice versa. The style-transformed sub-images are added together to give the final WIST-transformed warped atlas image $\overline {I_{\tilde{a}}}$.
• Figure 5: The flaw of IST: artifacts appear if the images are mis-aligned due to imperfect registration, and become severer as the transformation strength goes higher.