Given a graph $\mathcal{G}$ representing our 3D volumetric data, at the inference time, we aim to obtain a refined segmentation $Y^{*}$ as the results of our GCN model $\Gamma$,

where the graph $\mathcal{G}$ is constructed from the set of volumes $S=\{\mathbb{E},\mathbb{U},V,Y\}$ (see section 2.1 and Fig. 2) and $\phi$ represents the GCN’s parameters.

Since most of the voxels in the volume are irrelevant for the refinement process and given that graphs are not restricted to the rectangular structured representation of data, we define an ROI tailored to our target anatomy. We define our working region as $\hbox{ROI}(x)=\hbox{dilation}(U_{b}(x))\cup\mathbb{E}_{b}(x)$ with $\mathbb{E}_{b}$ the expectation binarized by a threshold of $0.5$. Since the entropy is usually high in boundary regions, including the dilated $U_{b}$ ensures that the ROI is bigger enough to contain the organ. Also, this allows us to include high confidence background predictions ($Y=0$) for training the GCN. Including the expectation in the ROI give us high confidence foreground predictions for training the model. This ROI reduces the number of nodes of the graph and, in consequence, the memory requirements. The voxels $x\in\hbox{ROI}$ define the nodes for $\mathcal{G}$. Each node is represented by a feature vector containing intensity $V(x)$, expectation $\mathbb{E}(x)$, and entropy $\mathbb{U}(x)$. Finally, we labeled each node in the graph according to its uncertainty level using the next rule:

The most straightforward connectivity option is to consider the connectivity with adjacent voxels (n6 or n26 adjacent voxel neighborhood). However, this simple nearest neighborhood scheme may not be adequate for our problem for two reasons; First, with this scheme, every single voxel is connected with its local neighborhood but lacks global information. Considering the original volumetric representation of the data, this means that the main source of information is coming only from adjacent voxels, while information from the global context (not-adjacent voxels) is mainly ignored.

Second, voxels with high uncertainty tend to shape contiguous clusters. Because of this, a voxel with high uncertainty will be most probably surrounded by other high uncertainty points, reducing the connection with the low uncertainty points. Hence, a simple n6 or n26 connectivity might limit the propagation of information from the confidence to the uncertain regions.

A fully-connected graph can take advantage of the relationships between local and long-range connections, and propagate information from both certain and uncertain regions during training and inference time. The main disadvantage of a fully-connected graph for a GCN model is the prohibitive memory requirements.

In our work, we evaluate an intermediate solution. For a particular node (or voxel) $x$, we create connections with its six perpendicular immediate neighbors in the volume coordinate system (N6) to consider local information. Additionally, we randomly select $k$ additional nodes in the graph and create a connection between these random elements and $x$. This defines a sparse representation that considers local and long-range connections between high and low uncertainty elements. The $k$ random nodes can be taken from any part inside the ROI used to define the nodes of the graph. We empirically found that $k=16$ offers a balance in performance and graph size, and kept this value during our experiments.

To define the weights for the edges, we use a function based on Gaussian kernels considering the intensity $V(x)$ and the 3-D position $x\in\mathbb{R}^{3}$ associated with the node:

where $\lambda$ and $\beta$ are balancing factors, $\hbox{div}(\cdot)$ is given by the diversity between the nodes (Zhou et al., 2017b), defined as $\hbox{div}(x_{i},x_{j})=\sum_{c=1}^{M}(P^{c}(x_{i})-P^{c}(x_{j}))\log{\frac{P^{c}(x_{i})}{P^{c}(x_{j})}}$ with $M=2$, $P^{1}(x_{i})=\mathbb{E}(x_{i})$ and $P^{2}(x_{i})=1-\mathbb{E}(x_{i})$ for our binary case. We opt for an additive weighting, instead of a multiplicative one, because the GCN can take advantage of connections with both similar and dissimilar nodes (in intensity and space) in the learning process, and using a multiplicative weighting could completely cut these connections. We found out that the diversity can indirectly bring information about the similarity of two nodes, in terms of class probability.

Since the diversity does not have an upper bound, it is possible to apply a non-linear transformation in order to normalize its value to the range $(0,1]$, leading to the following version of the diversity:

We can integrate this into eq.(6) replacing the regular diversity:

Similarly, if we only take the exponential part of eq.(7), we get a similarity metric between the expectation of the nodes $x_{i}$ and $x_{j}$. We use this version of the diversity as a third variation of our weighting function, leading to the following expression:

where the function $\hbox{inv\_div}(x_{i},x_{j})$ is given by

It is worth mentioning that if we set $\lambda=0$ and keep the same value of $\beta$ all the weighting functions become the same.

In addition to the uncertainty threshold analysis presented in Soberanis-Mukul et al. (2020), we extend this analysis to the different variations of the weighting functions introduced, exploring their effects in the refinement scheme. The corresponding analysis is presented in section 3.4.4.

At this point, we have reformulated the refinement task as a semi-supervised graph learning tasks. As mentioned previously, we use the proposal from Kipf and Welling (2017) to train the GCN in a semi-supervised way. A convolutional $H$ layer is defined as:

The variable $X\in\mathbb{R}^{N\times K}$ represents the input feature matrix of the layer with $N$ the number of nodes and $K$ the number of input features per node. $W\in\mathbb{R}^{K\times K^{\prime}}$ is the weight matrix of the current layer with $K^{\prime}$ the number of output features per node. $\hat{A}$ follows the renormalization proposed by Kipf and Welling (2017) with $\tilde{A}=A+I_{N}$, $A$ the adjacency matrix of $\mathcal{G}$, and the diagonal matrix $\tilde{D}=\Sigma_{\hbox{cols}}\tilde{A}$ that sums across the columns of $\tilde{A}$. In general, we keep the same GCN architecture, but employing a sigmoid activation function at the output layer, for binary classification. Representing our graph $\mathcal{G}$ by its adjacency matrix $A$ and feature matrix $X$, this leads to the following GCN model:

Finally, it is worth to mention that, even if we validated our proposal using this GCN learning architecture, given the modular nature of our method, it is possible to employ a different semi-supervised graph learning strategy.

We chose a 2D U-Net to be our CNN model (Ronneberger et al., 2015). We included dropout layers at the end of every convolutional block of the U-Net, as indicated by the MCDO method. The U-Net was trained considering a binary segmentation problem. Since we are employing a 2D model, we trained the models using axial slices. The model was trained with the dice loss function, using the Adam optimizer. We train the model for around 200 epochs keeping the overall best performing model during the entire training procedure. At inference time, we predicted every slice separately and then we stacked all the predictions together to obtain a volumetric segmentation (a similar strategy was used to perform the uncertainty analysis). As a post-processing step, we compute the largest connected component in the prediction, to reduce the number of false positives. At this point, it is worth mentioning that the U-Net was used only for testing purposes and different architectures can be used instead. This is mainly because our refinement method uses the model-independent MCDO analysis.

We utilized MCDO to compute the expectation and entropy using a dropout rate of $0.3$ and a total of $T=20$ stochastic passes. To obtain volumetric uncertainty from a 2D model, we performed the uncertainty analysis on every individual slice of the input volume and then we stacked all the results together to obtain the volumetric expectation and entropy. We tested different values for the uncertainty threshold $\tau$ (see section 3.4.2).

The GCN model is a two-layered network with 32 feature maps in the hidden layer and a single output neuron for binary node-classification. The graphical network is trained for 200 epochs with a learning rate of $1e-2$, binary entropy loss, and the Adam optimizer. We kept these same settings for the refinement of both segmentation tasks. After the refinement process, we can replace only the uncertain voxels with the GCN prediction, or we can replace the entire CNN prediction with the GCN output. We use the second approach since we found it producing better results.

Given that with small sets, the statistical significance tests can fail (Szucs and Ioannidis, 2017; Biau et al., 2008), we generate the dice score in slice-wise to increase the sample size (up with 278-1700 slices for the spleen, and pancreas respectively). Then we have run the non-parametric statistical significance test, namely Kolmogorov–Smirnov test. We perform a statistical significance analysis between the results of the GCN refinement and the initial results of the CNN. A single start (*) indicates a p-value $<0.05$, while a double-start (**) indicates a p-value $<0.01$ with respect to the original CNN prediction.

We evaluate the performance of the GCN refinement when the base CNN is trained with a small number of samples. For this, we randomly selected 10 out of the 45 training samples of the NIH dataset. For spleen, we selected nine. Results are presented in Table 2.

Note the increment in the standard deviation of all the models. A reason for this can be that the CNN does not generalize adequately to the testing set, due to the small number of training examples. Similar to the previous results, the increment in the dice score for the GCN refinement is about 2.4% with respect to the CNN base model for the pancreas, and improvement of 2.3% for spleen, compared with the base CNN.

In our experiments, we evaluate the influence of different values for $\tau$. We tested the method with values of $\tau\in\{0.001,0.3,0.5,0.8,0.999\}$. In this way, we covered a wide range of conditions that define a voxel as “uncertain”. After training the GCN, we replaced all the CNN predictions with the GCN output. Table 3 compares the CNN output with the GCN refinement at different values of $\tau$ for the tasks of the pancreas and spleen segmentation.

The parameter $\tau$ controls the minimum requirement to consider a voxel as uncertain. Lower values lead to a higher number of uncertain elements. This has a direct relationship with the number of high certainty nodes in the graph representation, and hence, in the number of training examples for the GCN. This also influences the quality of the training voxels for the GCN, since a high threshold relaxes the amount of uncertainty necessary to rely on the prediction of the CNN.

However, from the results of Table 3, except for pancreas-10 and spleen-9, there is no significant impact on the choice of this parameter. One reason can be that there is a clear separation between high and low uncertainty points. Therefore, changing $\tau$ may add (remove) a few number of nodes that are insignificant for the learning process of the GCN.

For the pancreas-10 model, we notice a progressive decrease in the dice score. Since this model uses fewer training examples, it is expected to have low confidence in their predictions (in contrast with the model trained with 45 volumes). In this scenario, a higher uncertainty threshold increases the chance to include high-uncertainty nodes as ground truth for training the GCN. A lower $\tau$ includes fewer points but with higher confidence. This appears to be beneficial in the pancreas segmentation model trained with fewer examples.

The opposite occurs with spleen-9, where higher $\tau$ are beneficial. This might indicate a dependency on the characteristics of the anatomies since the pancreas presents more inter-patient and inter-slice variability (see Fig. 5).

In general, our results suggest that $\tau$ parameter should be selected based on the target anatomy. Further, $\tau$ appears to have more influence in conditions of high uncertainty, e.g. when the model is trained with fewer examples. In the cases where $\tau$ has no significant impact, intermediate values are preferred, since they lead to a lower number of nodes, and in consequence to lower memory requirements.

We compare the refinement performance of our method when a classical n26 neighborhood is employed. We repeated the experiments for different uncertainty thresholds, and the CNN model trained with different numbers of examples on the pancreas and spleen datasets. Results are presented in Tables 4 and 5 for each organ, respectively. In all tables, the weighting employed is eq. 6 with $\lambda=0.5$ and $\beta=1$.

The results show a better refinement when including the random long-range connections, especially when working with the CNN trained with limited data, for both tasks. It is worth to mention that when comparing with the original CNN output presented in Tables 1 and 2, our refinement method using the n26 connectivity still having a slight improvement over the original CNN prediction.

Finally, we analyze the influence of different variations of the components of our weighting function. We perform our refinement strategy using the weighting functions given by equations 6, 8, and 9, presented in section 2.2.2. In all the experiments of this section, we keep a fixed value of $\tau=0.5$. We will use the notation $w_{i,(\lambda,\beta)}$, $i\in\{1,2,3\}$ to indicate the parameters employed by the corresponding function, for example, the notation $w_{1,(0.5,1)}$ holds for weighting function $w_{1}$ with $\lambda=0.5$ and $\beta=1$. We explore the situations when: a) either only diversity or only the Gaussian kernels are employed (sec. 3.4.5); b) using a diversity normalized to the range [0, 1] (sec. 3.4.6); c) using similarity in expectation given by the inverse of the diversity (sec. 3.4.7); and d) we discuss deep insights regarding the three different variations of the diversity employed in the previous sections (sec. 3.4.8).

We can divide the weighting function into two metrics, diversity and Gaussian similarity kernels (in intensity and position). In this experiment, we keep one of these components at a time. This is done by setting the values of $\lambda$ and $\beta$ to zero, accordingly. The results are presented in Table 6. Note that $w_{1,(0.5,1)}$ corresponds to the weighting function used so far.

Different weighting functions and variations still outperform the initial CNN prediction. Now, we will focus on the small differences between these results. As we mentioned, our weighting scheme considers diversity together with intensity and position similarity. The two later Gaussian components are commonly used in the literature and follow the intuition that two components that are similar in intensity and close to each other are likely to belong to the same class. The diversity is an additional component that allows us to include the results from the MCDO analysis in the edge weighting. The diversity has a lower bound of zero but in contrast with the Gaussian kernels, in an ideal form, it does not have an upper bound. When two nodes have a similar expectation, the diversity between those elements will be small, and the weighting function will only rely on the Gaussian similarities. When the nodes have important differences in expectation (e.g 0 vs. 1), the unbounded nature of the diversity will ignore the much smaller contribution of the Gaussian kernels (the diversity can reach values of $30$, while each Gaussian kernel has a maximum at $1.0$), and bias the weight to the value of the diversity.

One of the questions, we had, during this work whether the normalization of the diversity would have a positive/negative impact on the behavior of the GCN. In fact, $w_{2}$ presents a negative impact on the learning process. As we mentioned, the unbounded nature of the diversity makes $w_{1}$ to prefer connections with opposed expectations (and ignore the Gaussian similarity in those cases). On the other hand, the Gaussian similarity is only considered when the expectation of both nodes is basically the same, and even in those cases, the weight is small compared with the values of the diversity. In this sense, the GCN can learn mostly from examples that are different in expectation or (with a considerably lower contribution ) from examples similar in intensity and position.

Using $w_{2}$, which normalizes the diversity into a value between 0 and 1, makes the contribution of the Gaussian kernels more representative. The normalized diversity of $w_{2}$ will no longer ignore the Gaussian similarity, and it will assign the highest weights to connections between nodes that are similar in intensity and position, but at the same time different in expectation, which is counter-intuitive. Table 7 shows the drop in the performance using $w_{2}$ as weighting function. Given that the diversity is in range between 0 and 1, we use $\lambda=1$ for $w_{2}$.

If we take only the exponential part of eq. 7, we will obtain values in the range [0,1] that gives high weights to similar expectations, presenting a better agreement with the Gaussian kernels. This weighting scheme is given by $w_{3}$. The results, compared with the U-Net and the initial $w_{1,(0.5,1)}$ weighting function are presented in Table 8. Results, again show an improvement for the GCN refinement, supporting our discussion about $w_{2}$.

In this section, as final analysis for the weighting function, we compare $w_{1}$, $w_{2}$, and $w_{3}$ when $\beta=0$. The results are presented in Table 9. Two points are worth mentioning. Comparing the results of $w_{2,(1,0)}$ in Table 9 with the results of $w_{2,(1,1)}$ in Table 7, it is clear that the only-diversity version of $w_{2}$ has better performance, supporting the idea of the inconsistent connections of $w_{2,(1,1)}$.

The second interesting fact comes when comparing the diversity of $w_{1}$ with the normalized diversity of $w_{2}$, both in Table 9. Even though these two functions are expected to behave similarly, we notice a difference in their performance. To understand the possible causes of these differences, we take a closer look at the assigned weights of these functions to a subset of nodes. Let’s consider a central node with expectation close to zero, and four neighbors nodes with expectations of approximately 0.0, 1.0, 0.999, and 0.368, respectively (see Fig. 6). As expected, $w_{1,(0.5,0)}$ will assign high values to opposed expectations (36.2 and 19.9 in Figure 6.a), while the weight will be low for similar expectations, and close to zero to identical expectations (weights of 3.7 and 0 in Fig. 6.a). However, for the same node structure, $w_{2,(1,0)}$ will weight with a high value (close to 1) almost all the connections, except for the neighbors with the same expectation (see Fig, 6.b). Even though the difference between the expectations values of 1 and 0.368 is not high enough, $w_{2,(1,0)}$ will give an importance that is similar to the nodes with completely opposed expectations, leading to an inconsistent weighting scheme. In contrast, $w_{3,(1,0)}$ (which acts as a kind of similarity in expectation) assigns weights close to zero to all the nodes that do not have the same expectation, even to the pair (0, 0.368), see Fig. 6.c. This can suggest a better consistency for the inverse diversity, compared with diversity normalized to [0, 1].

We can also see these conclusions if we plot the values for these three variations of the diversity as a function of the difference between $p1=0$ and $p2\in[0,1]$ (see Fig. 7). We can see how the diversity grows exponentially when the difference between $p1$ and $p2$ is close to 1 (Fig. 7 Diversity). In a similar way, the inverse diversity assigns a weight of 1 to differences close to zero and decreases exponentially at the moment this difference starts to increment (Fig. 7 Inv. Diversity). Finally, the normalized diversity starts assigning a weight of zero to equal expectations, however, the curve grows exponentially until reaching a weight of 1 when the difference in expectations still close to 0.2 (Fig. 7 Norm. Diversity), leading to the inconsistencies mentioned before.

In general, a weighting function like $w_{2}$ appears to be not recommendable for the refinement with GCNs. The inverse diversity of $w_{3}$ can be a better option, however, it is also possible that under certain conditions, this weighting could introduce noisy connections. For example, when the expectation has values close to 0.5 for both nodes, the function will assign high weights, however, this might not heavily contribute to the GCN given the fact that these nodes might have high uncertainty as well. This suggests that, under high uncertainty environments, it is better to weigh based on dissimilarity in expectation. On the other hand, the inverse of diversity can take advantage of well-separated expectations.