VesselMorph: Domain-Generalized Retinal Vessel Segmentation via Shape-Aware Representation

Unlike ML models, our visual interpretation of vessels is rarely affected by data distribution shifts. Mimicking the human vessel recognition can thus help address the DG problem. In addition to intensity values, human perception of vessels also depends on the local contrast and the correlation in a neighborhood, which is often well described by the local Hessian. Inspired by the use of DTI to depict the white matter tracts, we create a Hessian-based bipolar tensor field to represent the morphology of vessels. Given a 2D input image $\mathbf{x}\in\mathbb{R}^{h\times w}$ and scale $\sigma$, the classical Frangi vesselness $\mathcal{V}(\sigma)$ [7] is defined as:

Here, $\lambda_{1},\lambda_{2}$ are the sorted eigenvalues of the Hessian $\mathcal{H}$, $\mathcal{R}_{B}=\lambda_{1}/\lambda_{2}$, $S$ is the Frobenius norm of the Hessian ($\|\mathcal{H}\|_{F}$), $\beta=0.5$ and $c=0.5$. Note that we assume vessels are brighter than the background; fundus images are negated to comply. To represent vessels of different sizes, we leverage the multiscale vesselness filter that uses the optimal scale $\sigma^{*}$ for the Hessian $\mathcal{H}(\mathbf{x}_{ij},\sigma)$ at each pixel $(i,j)$. This is achieved by grid search in the range $[\sigma_{min},\sigma_{max}]$ to maximize the vesselness $\mathcal{V}(\sigma)$, i.e., $\sigma^{*}=\operatorname*{argmax}_{\sigma_{min}\leq\sigma\leq\sigma_{max}}\mathcal{V}(\sigma)$. Then the optimized Hessian is represented by a $2\times 2$ matrix:

where $G(\mathbf{x}_{ij},\sigma^{*})$ is a 2D Gaussian kernel with standard deviation $\sigma^{*}$. Then we apply the eigen decomposition to obtain the eigenvalues $\lambda_{1}$, $\lambda_{2}$ ($|\lambda_{1}|\leq|\lambda_{2}|$) and the corresponding eigenvectors $\mathbf{v}_{1}$, $\mathbf{v}_{2}$ at the optimal $\sigma^{*}$.

Instead of solely analyzing the signs and magnitudes of the Hessian eigenvalues as in the traditional Frangi filter, we propose to leverage the eigenvectors along with custom-designed magnitudes to create our tensor field as shown in Fig. 3(Left). The core idea of Frangi filter is to enhance the tubular structure by matching the vessel diameter with the distance between the two zero crossings in the second order derivative of Gaussian ( $2\sqrt{2}\sigma^{*}$). However, the solution is not guaranteed to land in range $[\sigma_{min},\sigma_{max}]$, especially for small vessels. Consequently, we observe that the inaccurate estimation of $\sigma^{*}$ results in a blurring effect at the vessel boundary, which is problematic for segmentation. As an example in Fig. 3(Left), the direction of $\mathbf{v}_{1}$ at $p_{2}$ aligns with that at $p_{1}$, even though $p_{1}$ is inside the vessel while $p_{2}$ is in the background but close to the boundary. This makes it difficult for the vector orientations alone to differentiate points inside and outside the vessel. To tackle this, we introduce the idea of a bipolar tensor by assigning a large magnitude to the orthogonal eigenvector $\mathbf{v}_{2}$ to points in the background, as shown in the blue dashed ellipse. Specifically, we define the magnitudes $\alpha_{1}$ and $\alpha_{2}$ associated with the eigenvectors $\mathbf{v}_{1}$ and $\mathbf{v}_{2}$ as:

where $P(\mathbf{x}>\mathbf{x}_{ij})$ is the probability that the intensity of a random pixel $x$ in the image is greater than $\mathbf{x}_{ij}$. This is equivalent to normalizing the histogram by the factor $hw$ and computing the cumulative distribution function at $\mathbf{x}_{ij}$. This term thus provides a normalized brightness function in the range $[0,1]$. The exponential term represents how vessel-like the voxel is by using a normalized eigenvalue, and is in the $[0,1]$ range as well. $\epsilon$ is a constant that controls the sensitivity, which is empirically set to $0.5$. With the custom magnitudes $\alpha_{1}$ and $\alpha_{2}$, the two poles can better differentiate vessels from the background. Fig. 3(Right) is an example of BTF on an OCTA image. In practice, we stack the two vectors as the input to the structural encoding network, i.e., $\Psi(\mathbf{x}_{ij})=\left[\alpha_{1}\mathbf{v}_{1}^{\top},\alpha_{2}\mathbf{v}_{2}^{\top}\right]^{\top}\in\mathbb{R}^{4\times 1}$.

Preserving the spatial resolution for the bottleneck of models with U-Net backbone is a common strategy to emphasize the structural features in unsupervised segmentation [17, 11] and representation disentanglement [3, 21]. We employ a network that has a full-resolution ($h\times w$ pixels) latent space as the feature extraction model. We propose to extract vessel structure from both the intensity image $\mathbf{x}\in\mathbb{R}^{h\times w}$ and its corresponding BTF, $\Psi(\mathbf{x})\in\mathbb{R}^{4\times h\times w}$. Therefore, in Fig. 2, the intensity $D(E^{I}(\cdot))$ and structure $D(E^{S}(\cdot))$ encoding pathways share the decoder D, and the latent images $\mathbf{z}^{I},\mathbf{z}^{S}\in\mathbb{R}^{h\times w}$. To distribute more workload on the encoder, $D$ has a shallower architecture and will be discarded in testing. For the intensity encoding, the model is optimized by minimizing the segmentation loss function defined as the combination of cross-entropy and Dice loss:

where $N=h\times w$ is the total number of pixels in the image, $\mathbf{y}$ is the ground truth and $\hat{\mathbf{y}}^{I}$ is the prediction from the training-only decoder $D$. Although there is no explicit constraint on the latent image $E^{I}(\mathbf{x})=\mathbf{z}^{I}$, we note that the segmentation-based supervision encourages it to include the vessels while most other irrelevant features are filtered out. Hence, we can view the latent feature as a vessel representation.

Our approach is slightly different for the structure encoding as we observe that it is hard for the feature extraction network to generate a stable latent image that is free of artifacts when the number of input channels is greater than 1. Thus, it is necessary to use $E^{I}$ as a teacher model that provides direct supervision on the vessel representation. In other words, we first train the intensity encoding path to get $E^{I}$ and $D$, then train the $E^{S}$ by leveraging both the segmentation loss in Eq. 4 and a similarity loss defined as:

which is a weighted sum of $L_{1}$ norm and structural similarity loss SSIM [10]. SSIM is defined as $\mathrm{SSIM}(A,B)=2\frac{(2\mu_{A}\mu_{B}+c_{1})(2\sigma_{AB}+c_{2})}{(\mu_{A}^{2}+\mu_{B}^{2}+c_{1})(\sigma_{A}^{2}+\sigma_{B}^{2}+c_{2})},$ where $\mu$ and $\sigma$ represent the mean and standard deviation of the image, and we set $c_{1}=0.01$ and $c_{2}=0.03$. The overall loss function for the structural encoding is thus $\mathcal{L}(\Psi(\mathbf{x}),\mathbf{y})=\omega_{1}\mathcal{L}_{seg}(\hat{\mathbf{y}}^{S},\mathbf{y})+\omega_{2}\mathcal{L}_{sim}(\mathbf{z}^{S},\mathbf{z}^{I})$, with empirically determined weights $\omega_{1}=1$, $\omega_{2}=5$. Experimentally, we found that the $\mathbf{z}^{I}$ is good at preserving small vessels, while $\mathbf{z}^{S}$ works better on larger ones.

Given the two synthesized vessel representations $\mathbf{z}^{I}$ and $\mathbf{z}^{S}$, we need to introduce a fusion method to take advantage of both intensity and structure features. Naively stacking these two channels as input to the segmentation network is prone to inducing bias: if $\mathbf{z}^{I}$ is consistently better for images from the source domain, then the downstream task model $D^{T}$ would learn to downplay the contribution of $\mathbf{z}^{S}$ due to this biased training data. As a result, despite its potential to improve performance, $\mathbf{z}^{S}$ would be hindered from making a significant contribution to the target domain during testing. To circumvent this issue, we propose a simple weight-balancing trick. As illustrated in Fig. 2, we randomly swap some patches between the two latent images so that $D^{T}$ does not exclusively consider the feature from a single channel, even for biased training data. This trick is feasible because $\mathbf{z}^{S}$ and $\mathbf{z}^{I}$ are in the same intensity range, due to the similarity constraints applied in Eq. 5. Thus the input to $D^{T}$ is $\tilde{\mathbf{x}}=\Gamma(\mathbf{z}^{I},\mathbf{z}^{S})$, where $\tilde{\mathbf{x}}\in\mathbb{R}^{2\times h\times w}$. The loss function leveraged for $D^{T}$ is the same as Eq. 4.

The complete algorithm for training of VesselMorph is shown in Algorithm.1. Briefly, we first train the intensity encoder $E^{I}$ as it is easier to generate a stable vessel representation $\mathbf{z}^{I}$. Then a structure encoder $E^{S}$ is trained with the supervision of the ground truth and teacher model $E^{I}$ so that an auxiliary representation $\mathbf{z}^{S}$ is extracted from the structural descriptor BTF. The last step is to train a segmentation network $D^{T}$ with the fusion of the two vessel maps $\Gamma(\mathbf{z}^{I},\mathbf{z}^{S})$. During testing, the patch-swapping is no longer needed, so we simply concatenate $E^{I}(\mathbf{x})$ and $E^{S}(\Psi(\mathbf{x}))$ as the input to $D^{T}$.

Fig. 4 shows a qualitative ablation study: it illustrates that the intensity representation $\mathbf{z}^{I}$ may miss large vessels in the very high-resolution HRF images, while $\mathbf{z}^{S}$ remains robust. In contrast, $\mathbf{z}^{I}$ provides sharper delineation for very thin vessels in ROSE. The fusion of both pathways outperforms either pathway for most scenarios. These observations are further supported by the quantitative ablation study in Fig.6. We note that $\mathbf{z}^{S}$ and $\mathbf{z}^{I}$ can be used as synthetic angiograms that provide both enhanced vessel visualization and model interpretability.

Fig. 5 shows the t-SNE plots [20] of the datasets. The distribution gaps between datasets are greatly reduced for the two latent vessel representations.

Table 2 compares all methods on the target domain $\mathcal{T}$. For the diseased ARIA data, all methods show comparable performance and are not significantly different from the baseline. VesselMorph has the best OOD outcome for both cross-modality (dark gray) and cross-resolution (light gray) scenarios, except the OCTA500 dataset where VFT, MASF and VesselMorph perform similarly. The results of VFT and VesselMorph prove the value of the shape information.