Topology-Preserving Automatic Labeling of Coronary Arteries via Anatomy-aware Connection Classifier

We start by extracting centerlines from the vessel segmentation annotations in a CCTA image, using a traditional 3D thinning algorithm [9]. Next, we use the minimum spanning tree algorithm [5] to construct two coronary artery trees, one for the left domain (LD) and one for the right domain (RD). Following the branch bifurcation rules in [22] (details can be found in supplementary materials), the coronary artery trees are split into several segments, denoted by $\mathcal{S}=\{S_{i}\}_{i=1}^{N}$, where $N$ is the number of segments, $S_{i}\in\mathbb{R}^{L_{i}\times 3}$ denotes the $i$-th segment comprised of the 3D positions of $L_{i}$ centerline points. Our model takes as input the vessel segments $\mathcal{S}$ and their connections $\mathcal{C}=\{(i_{1},j_{1}),(i_{2},j_{2}),...,(i_{N_{c}},j_{N_{c}})\}$, where $(i_{k},j_{k})$ indicates that the $i_{k}$-th segment is connected with the $j_{k}$-th segment. The objective is to predict the classes of the vessel segments.

As illustrated in Figure 2, we design a novel framework named TopoLab for the automatic labeling of coronary arteries. In Sec. 3.2, we will introduce the hierarchical feature extraction module comprised of intra-segment feature aggregation and inter-segment feature interaction. In Sec. 3.3, we will elaborate on the details of AC-Classifier which exploits the category topological trees effectively.

We first feed the CCTA image $X\in\mathbb{R}^{H\times W\times D}$ into an image encoder (e.g. U-Net [2]) to obtain the downscaled feature map $F\in\mathbb{R}^{H/4\times W/4\times D/4\times C}$, where $C$ is the channel size.

For each segment $S_{i}$, trilinear interpolation in the downsampled feature map $F$ is adopted for the centerline point $v\in S_{i}$ to obtain the corresponding point features $f(v)=\text{Tri}(v,F)\in\mathbb{R}^{C}.$ The segment features are denoted as $\{f(v)|v\in S_{i}\}$, simplified as $E_{i}\in\mathbb{R}^{L_{i}\times C}$.

We introduce Transformer [17] as our intra-segment feature aggregator for its strong capability to model the relationships among the sequences with varying lengths. Concretely, we set a learnable embedding $q\in\mathbb{R}^{C}$ called segment query to aggregate intra-segment point features. The segment query is concatenated with the segment features $E_{i}$ to obtain the new tensor $\tilde{E}_{i}\in\mathbb{R}^{(L_{i}+1)\times C}$. Following [4], we feed $\tilde{E}_{i}$ augmented by the learnable 3D positional encodings into a Transformer encoder containing several standard sub-blocks. Each standard sub-block has the same architecture as in [17], which consists of multi-head attention, feed-forward network, layer normalization and ReLU activation. The state of the segment query $q$ at the output of the Transformer encoder serves as the aggregated segment representation $\hat{E}_{i}\in\mathbb{R}^{C}$.

Specifically, let $\hat{E}\in\mathbb{R}^{N\times C}$ denotes the aggregated segment features where the $i$-th item of $\hat{E}$ is $\hat{E}_{i}$, and $A\in\mathbb{R}^{N\times N}$ is the adjacency matrix for the vessel segment graph derived from the segment connections $\mathcal{C}$. The process of GCN layers is as follows:

where $\sigma$ is the ReLU activation, $W_{l}\in\mathbb{R}^{C\times C}$ is the learnable parameters for the $l$-th GCN layer, the input for the first layer is $\hat{E}_{0}=\hat{E}$.

Finally, we fuse the input segment features $\hat{E}$ and the output of the final GCN layer $\hat{E}_{f}$ to obtain $\overline{E}=[\hat{E}_{f},\hat{E}]W_{f}\in\mathbb{R}^{N\times C}$ with the parameters $W_{f}\in\mathbb{R}^{(C+C)\times C}$. The enhanced segment features $\{\overline{E}_{1},\overline{E}_{2},...,\overline{E}_{N}\}$ are forwarded to the classifier for segment labeling, where $\overline{E}_{i}\in\mathbb{R}^{C}$ is the $i$-th item of $\overline{E}$.

The direct approach for labeling the coronary arteries is to use a linear layer to classify each segment independently as in previous methods. To incorporate the category topological trees into the classifier design, we propose to conduct the classification task for every connected segment pair.

We begin by defining the ground truth segment connections which are composed of the topology-conforming connections derived from the category topological trees (like LM$\rightarrow$LAD, LCX$\rightarrow$OM, etc.). Note that the self-connections (e.g. RCA$\rightarrow$RCA) are also considered as the ground truth connections. The ground truth segment connections have the corresponding template embeddings for classification, which are represented by $G\in\mathbb{R}^{N_{g}\times 2C}$, where $N_{g}$ is the number of ground truth segment connections. Denote $g_{i}=\text{Concate}(\text{Enc}(x),\text{Enc}(y))\in\mathbb{R}^{2C}$ as the $i$-th item of $G$, where $x$ and $y$ stand for the segment classes indexed by the $i$-th ground truth connection, Enc denotes the sinusoidal encoding as in [17]. The connection templates $G$ are used to enable classification for the connected segment pairs.

Then, given the segment features $\{\overline{E}_{i}\}_{i=1}^{N}$, and all of the connected segment pairs $\mathcal{C}=\{(i_{1},j_{1}),(i_{2},j_{2}),...,(i_{N_{c}},j_{N_{c}})\}$, the connection features $P\in\mathbb{R}^{N_{c}\times 2C}$ are obtained by rearrangement of the segment features, where the $k$-th item of $P$ is $\text{Concate}(\overline{E}_{i_{k}},\overline{E}_{j_{k}})\in\mathbb{R}^{2C}$. We use an MLP layer to further fuse the features $\hat{P}=\text{MLP}(P)\in\mathbb{R}^{N_{c}\times 2C}$.

where $y_{i}$ is the ground truth of the $i$-th segment connection, $\hat{p}_{i}$ denotes the $i$-th item of $\hat{P}$. sim in Eq. 3 stand for the cosine similarity, $\text{sim}(x,y)=\frac{x^{T}y}{||x||_{2}||y||_{2}}$, and the temperature $\tau$ is a hyperparameter which is set as $0.05$ by default.

Datasets. We train and evaluate our method on two datasets. The orCaScore [19] MICCAI 2014 Challenge contains 72 contrast-enhanced CTA images and non-contrast enhanced CT scans. As the original dataset only contains the labels of calcifications, we have annotated vessel segmentation and anatomical categories for coronary arteries with experienced radiologists. Considering the small amount of data, we randomly split the dataset into five folds to perform cross-validation. The mean values of cross-validation results are reported. We also collect an In-house Dataset containing 1200 CTA scans which have been annotated by at least two experts. The dataset is collected in compliance with the terms of the licensing agreement and ethical certification. We randomly split the dataset into train, validation and test set with 800, 200 and 200 scans respectively. The annotations for both datasets adhere to the same standard, which includes 14 classes of coronary artery segments (see Figure 1(b)). It is a challenging task, surpassing the scope of studies like TreeLab-Net[20] and CPR-GCN[22], which only consider 10 and 11 categories, respectively.

3D ResUNet [25] is employed as the image encoder for feature extraction with channel dimension $C=64$. The transformer encoder has 3 standard blocks and the number of graph convolution layers is set to 4. We train the network using AdamW optimizer [10] with a base learning rate of 5e-4 and the cosine learning rate schedule during the training stage. The batch size is set to 4. All networks are implemented by Pytorch [12] and trained on four NVIDIA GeForce RTX 3090 GPUs. We train our model on the orCaScore dataset for 3.5k iterations and in-house dataset for 12.5k iterations.

In this subsection, we explore the effectiveness of different components in TopoLab on orCaScore dataset, as shown in Table 3.