Computer Science
Computer Science-2021.gz:143169
Computer Science
58e666f49eedebae1127c3f436166a7c525ce1c3
BI-GCN: Boundary-Aware Input-Dependent Graph Convolution Network for Biomedical Image Segmentation

Segmentation is an essential operation of image processing. The convolution operation suffers from a limited receptive field, while global modelling is fundamental to segmentation tasks. In this paper, we apply graph convolution into the segmentation task and propose an improved \textit{Laplacian}. Different from existing methods, our \textit{Laplacian} is data-dependent, and we introduce two attention diagonal matrices to learn a better vertex relationship. In addition, it takes advantage of both region and boundary information when performing graph-based information propagation. Specifically, we model and reason about the boundary-aware region-wise correlations of different classes through learning graph representations, which is capable of manipulating long range semantic reasoning across various regions with the spatial enhancement along the object's boundary. Our model is well-suited to obtain global semantic region information while also accommodates local spatial boundary characteristics simultaneously. Experiments on two types of challenging datasets demonstrate that our method outperforms the state-of-the-art approaches on the segmentation of polyps in colonoscopy images and of the optic disc and optic cup in colour fundus images.


Introduction
Accurate examination of anatomical structures in biomedical images helps clinicians determine the likelihood of many medical conditions and diseases. For example, shapes of the optic disc (OD) and optic cup (OC) of the retina are essential in the diagnosis of glaucoma, a chronic degenerative disease that causes irreversible loss of vision [32]. Likewise, colorectal polyps are linked with an increased risk of having colorectal cancer, the third most common cancer worldwide [29]. However, it is impractical to manually annotate these structures in clinics due to the time-consuming, labour-intensive, and error-prone process. Automated and accurate image segmentation is increasingly needed.
Deep learning-based pixel-wise biomedical image segmentation methods, such as [8,10,28,35,39,42], are popular because of Convolutional Neural Network's (CNN's) excellent ability to extract and reason about semantic features, and correlations among different classes. Fan et al. [8] and Zhang et al. [39] studied the correlations between the foreground and background classes using a similar foreground erasing mechanism, which achieved state-of-the-art performance in the colonoscopy polyps segmentation tasks. Due to limited receptive fields of CNN, they incorporated dense dilation [36] convolutions to enlarge the receptive regions for long range context reasoning. Along the same line, Fu et al. [10] proposed multiscale input and side-output mechanism with deep supervision, which achieves multi-level receptive field fusion for long range relationship aggregation. However, this is inefficient since stacking local cues cannot always precisely handle long range context relationships. Especially for pixel-level classification problems, such as segmentation, performing long range interactions is an important factor for reasoning in complex scenarios [4]. For example, it is prone to assign visually similar pixels in a local region into the same category. Meanwhile, pixels of the same object but distributed in a distance are difficult to construct dependencies. Recent graph propagation-based reasoning approaches [5,16,21,38] that use Laplacian smoothing-based graph convolution [18], provide specific benefits in the sense of global long range information reasoning. They estimated the initial graph structure from a data-independent Laplacian matrix defined by randomly initialized adjacency matrix [5,38] or hand-crafted adjacency matrix [18,21]. However, one may make a model to learn a specific long range context pattern [19], which is less related to the input features. Differently, as seen in previous works, the graph structure can be estimated with the similarity matrix from the input data [20], we estimate the initial adjacency matrix in a data-dependent way. Specifically, we implement two diagonal matrices to perform channel-wise attention on the dot product distance of input vertex embeddings, and measure spatial-wise weighted relations among different vertices. In this way, our model is capable of learning an input-related long range context pattern, which improves the model segmentation performance, please read Ablation Study (Section 3.4) for more details.
Different to the dense pixel classification, recent polygon-based methods [22,23] explored the spatial features of an object's boundary. They regressed the predefined vertex position along the object borders then linked the predicted vertices to form a polygon to indicate the object's boundary. For example, Meng et al. proposed a CNN and Graph Convolution Network (GCN) aggregated network [22] to directly regress the vertices' coordinates of the OC and OD boundaries. It is well established that boundary knowledge is essential in obtaining spatial features in segmentation tasks. With regard to biomedical image segmentation, the boundary accuracy is often more critical than that of the regional pixel-wise coverage [23]. In this work, we explicitly exploit the boundary information to fuse the boundary features into the proposed graph convolution, proposing a boundary-aware graph convolution.  Our experimental results prove that the proposed boundary-aware graph convolution can improve the region segmentation performance. Along the same line, recent methods, such as [9,24,26,34,40], considered object's boundary segmentation while tackling object region segmentation tasks. However, they treated it as a multi-task learning problem, leaving the model to learn a robust backbone. Differently, we explicitly fuse the boundary information to improve the classic graph convolution's performance on the segmentation task.
This work proposes a novel GCN that reasons the boundary-aware region-wise correlations of different classes through learning graph representations. Specifically, with an input image, multi-level features were learned through a Backbone Module. Further, we constructed two feature aggregation modules upon shallow-and deep-level backbone features respectively for boundary and region characteristics learning, namely Boundary Segmentation Module (BSM) and Region Segmentation Module (RSM). The outputs of BSM and RSM were fed into the proposed graph reasoning module GRM for long range feature reasoning. In detail, the region-wise output (R S ) from RSM and the boundary-wise output (B S ) from BSM were embedded into graph representation in terms of vertex embeddings. Within the GRM, we proposed a new Boundary-aware Input-dependent Graph Convolution (BI-GConv). Specifically, the input-dependent adjacency matrix was estimated from the input features of vertex embeddings that were refined by spatial and channel attentions. Additionally, the object's boundary features were fused into the constructed Laplacian matrixL, enhancing the characteristics along the boundaries to emphasize the boundary-aware correlations across different regions. Our experimental results show that the proposed framework makes a large improvement over state-of-the-art approaches.

Method
There are four major components in the proposed model: a backbone module, a region segmentation module (RSM), a boundary segmentation module (BSM), and a graph reasoning module (GRM). Fig. 1 provides an overview of the model architecture and Fig. 2 highlights the proposed GRM and BI-GConv. Given an input image, we extract its multi-level features through a backbone module. Following the previous PraNet [8], we adopt the truncated Res2Net [12] as backbone module due to its superior ability to extract discriminative features in segmentation tasks. Feature Aggregation. For the RSM and BSM, we applied two feature aggregation modules to different levels of features of the backbone to extract full spectrum of semantic and spatial characteristics. As observed by previous studies [37,42], deep-and shallow-layer features from different levels complement each other, where the deep-layer features contain rich semantic region information while the shallow-layer features preserve the sufficient spatial boundary information. Thus, we applied feature aggregation module on relative deep-level and low-level features for BSM and RSM respectively. In detail, the feature aggregation module contains several feature aggregation blocks to fuse low-to-high level features hierarchically and iteratively. Notably, in the feature aggregation block, two 1 × 1 convolutions are applied on two resolution's features respectively to reduce and maintain the same channel size. After feature aggregation module, the output B s and R s have sizes of 64 × 64 × 1 and 32 × 32 × 2, respectively. We downsampled B s by 2 and fed it into the GRM together with R s . Downsample is for maintaining the same feature map size and reducing the computational overhead. Classic Graph Convolution. We first revisit the classic graph convolution. Given a graph where I is the identity matrix, A is the adjacency matrix, and D is a diagonal matrix that represents the degree of each vertex in V , such that D ii = ∑ j A i, j . The Laplacian of the graph is a symmetric and positive semi-definite matrix, so L can be diagonalized by the Fourier basis U ∈ R N×N , such that L = UΛU T . Thus, the spectral graph convolution of i and j can be defined as i * j = U((U T i) (U T j)) in the Fourier space. The columns of U are the orthogonal eigenvectors U = [u 1 , ..., u n ], and Λ = diag([λ 1 , ..., λ n ]) ∈ R N×N is a diagonal matrix with non-negative eigenvalues. Since U is not a sparse matrix, this operation is computationally expensive. Defferrard et al. [7] proposed that the convolution operation on a graph can be defined by formulating spectral filtering with a kernel g θ using a recursive Chebyshev polynomial in Fourier space. The filter g θ is parameterized as a Chebyshev polynomial expansion of order K, such that g θ (L) = ∑ k θ k T k (L), where θ ∈ R K is a vector of Chebyshev coefficients, andL = 2L/λ max − I N represents the rescaled Laplacian. T k ∈ R N×N is the Chebyshev polynomial of order K. In [18], Kipf et al. further simplified the graph convolution as and θ is the only Chebyshev coefficient left. The corresponding graph Laplacian adjacency matrixÂ is hand-crafted, which leads the model to learn a specific long range context pattern rather than the input-related one [19]. Thus, we refer to the classic graph convolution as input-independent graph convolution. Boundary-Aware Input-dependent Graph Convolution (BI-GConv). Given the input region features R s ∈ R N×C , where C is the channel size; N = H × W is the number of spatial locations of input feature, which is referred to as the number of vertexes. Firstly, we construct the input-dependent adjacency matrix. In detail, we implement two diagonal matrices (Λ c andΛ s ) to perform channel-wise attention on the dot-product distance of input  vertex embeddings, and measure spatial-wise weighted relations among different vertices, respectively. For example,Λ c (R s ) ∈ R C×C is the diagonal matrix that contains channel-wise attention on the dot-product distance of the input vertex embeddings;Λ s (R s ) ∈ R N×N is the spatial-wise weighted matrix, measuring the relationships among different vertices.
where Pool c (·) denotes the global max pooling for each vertex embedding; MLP(·) is a multi-layer perceptron with one hidden layer. On the other hand, where Pool s (·) represents the global max pooling for each position in the vertex embedding along the channel axis; Conv(·) is a 1 × 1 convolution layer. The data-dependent adjacency matrixĀ is given by spatial and channel attention-enhanced input vertex embeddings. We initialize the input-dependent adjacency matrixĀ as: where · represents matrix product; denotes Hadamard product; ψ(R s ,W ψ ) ∈ R N×C and φ (R s ,W φ ) ∈ R N×C are both linear embeddings (1 × 1 convolution); W ψ and W φ are learnable parameters. Secondly, we fuse the boundary information through integrating the predicted boundary function map Bs ∈ R N×1 into the built Laplacian matrixL, which allows us to leverage the spatial hint of the object's boundary. Specifically, we fuse it into the spatial-wise weighted matrixΛ s (R s ). The boundary-aware spatial weighted matrixΛ s b (R s , B s ) is given as follows: where is the broadcasting Hadamard product along channels; Conv(·) is a 1 × 1 convolution layer. In this way, the features of boundary pixels are emphasized by assigning larger weights. As this is the case, the proposed boundary-aware graph convolution can take the spatial characteristics of the boundary into account when reasoning the correlations between different regions. Then, the boundary-aware input-dependent adjacency matrixÃ will be given as:Ã where ζ (R s ,W ζ ) ∈ R N×C is 1 × 1 convolution; W ζ is learnable parameter. Fig. 2 shows a detailed demonstration ofÃ. With the constructedÃ, the normalized Laplacian matrix is given asL where I is the identity matrix,D is a diagonal matrix that represents the degree of each vertex, such thatD ii = ∑ jÃi, j . Inspired by [19], in order to reduce the complexity (O(N 2 )) of calculating Laplacian matrixL, we calculate the degree matrixD as:D where 1 denotes all-one vector in R N , and the calculation in the inner bracket is executed first. By doing this, we can override the quadratic order of computation, thus reduce the computational overhead. Given computedL, with R s as the input vertex embeddings, we formulate the single-layer BI-GConv as : where W G ∈ R C×C denotes the trainable weights of the BI-GConv; σ is the ReLu activation function; Y is the output vertex features. Moreover, we add a residual connection to reserve the features of input vertices. Graph Reasoning Module. We build the GRM with the chained BI-GConv layers in order to focus on the long-term dependency between chained graph layers. Fig. 2 shows an overview of the proposed GRM with three BI-GConv layers of Gated Recurrent Unit (GRU) [2] connection, which gains the best performance. Please note, these layers can also be connected by Residual [13] connections; comprehensive comparison results of these two connection manners can be found in the supplementary material.

Datasets
We use a retinal and a colorectal dataset, each of them is merged from five different datasets of its kind. The details of boundary ground truth generation can be found in the supplementary material. Fundus OD and OC images: 2,068 images from five datasets (Refuge [27], Drishti-GS [30], ORIGA [41], RIGA [1], RIM-ONE [11]) are pooled together. 613 images are randomly selected for test and the remaining 1,455 are for training and validation. We follow [22] to crop a subimage of 256 × 256 pixels centered at the OD center for the subsequent analysis. The training and test data come from the same data source (seen data), OD and OC segmentation is used for evaluate the model's learning ability.

Implement Details and Evaluation Metrics
Implement Details To augment the dataset, we randomly rotate and horizontally flip the training dataset with the probability of 0.3. The rotation ranges from −30 to 30 degree. 10% of the training dataset are randomly selected as the validation dataset. The network is trained end-to-end by Adam optimizer [17] for 50 thousands iterations, with a learning rate of 6e-3 and decayed to 0.36e-3 after 30 thousand iterations. The batch size is set at 48. All the training processes are performed on a server with four TESLA V100, and all the test experiments are conducted on a local workstation with Intel(R) Xeon(R) W-2104 CPU and Geforce RTX 2080Ti GPU. Code is available at: https://github.com/ smallmax00/BI-GConv Loss Function. The overall loss function is defined as: L total = L R + α · L B , where α balances the losses between region and boundary predictions; a detailed ablation study on α is presented in the supplementary material. We adopt standard Dice Loss [25] for the region (L R ) and boundary (L B ) segmentation supervision. Evaluation Metrics. Typical region segmentation metrics such as Dice similarity score (Dice) and balanced accuracy (B-Acc) are applied, where B-Acc is the mean value of Sensitivity and Specificity. Boundary Intersection-over-Union (BIoU) [6] is used as the boundary segmentation measurement. Details of the utilization of BIoU can be found in the supplementary material. To show statistical significance, 95% confidence interval is generated by using 2000 sample bootstrapping for all the evaluation metrics.

Results
We show qualitative (Fig. 3)    database comprised of 15 short colonoscopy sequences. ETIS [29] is an early established dataset for colorectal cancer early diagnosis. EndoScene-CVC300 [33] is a re-annotated benchmark with associated polyp and background (mucosa and lumen). In Tab. 1, our model achieves 72.7 % Dice score, which outperforms the cutting-edge polyps segmentation methods PraNet [8] and ACSNet [39]

Ablation Study
We conducted extensive ablation studies, and all the results consolidated our model's effectiveness. For example, the ablation results on different graph convolutions and model components are shown in Tab. 3 and Tab. 4. In the supplementary material, we report the ablation study experimental results related to the GRM, such as number (N) of BI-GConv layers, different multi-layers connection manners (GRU [2] or Residual [13]). Graph Convolution. In this section, we evaluate the effectiveness of the proposed graph convolution operation. Firstly, we employ the classic graph convolution [18] to reason the correlations between regions. Then we explore the input-dependent graph convolutions in terms of channel attention (channel) and spatial attention (spatial) mechanism, respectively and simultaneously (channel and spatial). Additionally, we adopt three more potent GRMs to show the superiority of our proposed BI-GConv. In detail, we apply the SGR [21], Dual-GCN [38], and GloRe module [5] respectively, where the SGR module exploits knowledge graph mechanism; DualGCN explores the coordinate space and feature space graph convolution; and GloRe utilizes projection and re-projection mechanism to reason the semantics between different regions. These methods achieved state-of-the-art performance on different segmentation tasks, however, they all belong to the data-independent Laplacian based graph convolution. Tab. 3 shows that our model achieves more accurate and reliable results than [18] and outperforms the SGR, DualGCN, and GloRe by 1.9 %, 1.5 % and 1.3 % in terms of mean Dice on the two test datasets. Model Components. We evaluate the effectiveness of each component of our proposed model. Firstly, we remove the BSM and the GRM; we refer to the rest of the structure as the Baseline model. Then we add BSM (w/ BSM) and single layer BI-GConv (w/ BI-GConv, N = 1) to evaluate each component's effectiveness gradually. Note that, in order to eliminate the model size difference and maintain a similar number of parameters (e.g. 32.8 million for our model), we replace the GRMs with several feed-forward CNN blocks in other ablation models (e.g., Baseline, w/ BSM, etc.) with standard convolution layer of kernel size 3 × 3, padding 1, followed by a Batch Normalization [14] layer. Tab. 4 shows that BSM and single layer BI-GConv can improve from Baseline model by 1.0 % and 3.4 % via the mean Dice on the two test datasets.

Conclusion
In this paper, we present a graph-propagation based framework to tackle biomedical image segmentation problems. We proposed a boundary-aware input-dependent graph convolution (BI-GConv) to reason about the boundary-enhanced long range correlations between regions in biomedical image segmentation tasks. Our experiment results show that the proposed GRM can efficiently reason about the semantic region features while explicitly considering spatial boundary features on the segmentation tasks of the colonic polyps and OD and OC. We will extend the proposed model to tackle 3D biomedical image segmentation tasks to learn the graph representation of volumes and object surfaces in the future.