Geo-UNet: A Geometrically Constrained Neural Framework for Clinical-Grade Lumen Segmentation in Intravascular Ultrasound

TL;DR

This work tackles the challenge of precise lumen segmentation in venous IVUS (v-IVUS) for accurate stent sizing, addressing the limitations of standard UNet-based approaches under limited data and the radial geometry of IVUS. It introduces Geo-UNet, a geometry-informed, two-task network that operates in polar space to predict both a single lumen contour and a dense pixel-wise lumen mask, using a novel CDFeLU activation to fuse contour information into pixel predictions. The model optimizes a unified loss that combines area-, distance-, and contour-based penalties and adds an inference-time Geo-UNet++ step to reduce wrap-around artifacts at the 0/2π boundary. On a venous IVUS dataset, Geo-UNet and Geo-UNet++ outperform several baselines, achieving high Dice scores and clinically relevant lumen-diameter accuracy, with particularly notable improvements in contour contiguity and minor-diameter estimation. The approach suggests a practical pathway to clinical-grade lumen segmentation in radial imaging settings and offers avenues for extension to other radially structured modalities and 3D contexts.

Abstract

Precisely estimating lumen boundaries in intravascular ultrasound (IVUS) is needed for sizing interventional stents to treat deep vein thrombosis (DVT). Unfortunately, current segmentation networks like the UNet lack the precision needed for clinical adoption in IVUS workflows. This arises due to the difficulty of automatically learning accurate lumen contour from limited training data while accounting for the radial geometry of IVUS imaging. We propose the Geo-UNet framework to address these issues via a design informed by the geometry of the lumen contour segmentation task. We first convert the input data and segmentation targets from Cartesian to polar coordinates. Starting from a convUNet feature extractor, we propose a two-task setup, one for conventional pixel-wise labeling and the other for single boundary lumen-contour localization. We directly combine the two predictions by passing the predicted lumen contour through a new activation (named CDFeLU) to filter out spurious pixel-wise predictions. Our unified loss function carefully balances area-based, distance-based, and contour-based penalties to provide near clinical-grade generalization in unseen patient data. We also introduce a lightweight, inference-time technique to enhance segmentation smoothness. The efficacy of our framework on a venous IVUS dataset is shown against state-of-the-art models.

Figures (5)

• Figure 1: Geo-UNet Architecture for Lumen Segmentation: The feature extractor is a fully convolutional UNet module with inputs of polar 2D IVUS frames. The top branch produces a probability map for the lumen contour ($\mathbf{P_{\text{c}}}$) via a row-wise softmax, which is converted to a single contour segmentation ($\mathbf{S_{\text{c}}}$) via a row-wise expectation function. The bottom branch produces a per-pixel probability map ($\mathbf{P_{\text{pix}}}$) via a channel-wise softmax. $\text{CDFeLU}(\cdot)$ allows the top branch to inform the bottom, refining the pixel-wise probabilities to give the segmentation ($\mathbf{S_{\text{pix}}}$) that is compared against the (polar) ground-truth lumen mask. The loss functions are highlighted in grey.
• Figure 2: Example of lumen segmentation performance. (Green-Predicted, Blue-Truth)
• Figure 3: Geo-UNet++: Inference-time Segmentation Enhancement: The bottom middle image shows the performance of Geo-UNet when given a polar input image. The green is the prediction, and the blue is the ground truth. In the bottom left image, note the sudden jumps and misalignment in the green prediction at the top and the bottom of the image, corresponding to $0$ and $2\pi$, respectively, yielding discontinuity at $0$/$2\pi$ in the Cartesian representation. Starting from the top left input image and to the right, we illustrate the ideas behind Geo-UNet++. Exploiting convolution's lack of dependence on input dimensions, we perform inference using the same trained Geo-UNet model, on an input wrap-padded with a repetition of the top of the original input, as highlighted by the orange braces. The padding provides additional context near 0. To recover the segmentation, we take the middle portion from $-\frac{\pi}{3}$ to $\frac{5\pi}{3}$ which typically avoids border discontinuities at the top and bottom of the padded image. This is essentially segmentation prediction on the Cartesian input image rotated counterclockwise by $\frac{\pi}{3}$ and does not affect the clinical objective of diameter estimation from the segmentation mask. On the lower right, we see that the Geo-UNet++ result is smoother and nearly perfectly aligned with the ground truth.
• Figure 4: Variations in appearance that are all considered as N1 frames with normal anatomy.
• Figure 5: Additional segmentation result comparison across Geo-UNet++, Geo-UNet, and baselines for both N1 and N2 frames.