SemiCurv: Semi-Supervised Curvilinear Structure Segmentation

We briefly review previous work on curvilinear structure segmentation in the context of road segmentation from satellite images, medical image segmentation and crack detection in pavements. Early work in road segmentation [17] used fully connected networks, but were soon followed by CNNs [18]. Recent work has built on the CNN approach, focusing on designing deeper and wider networks and considering information at multiple scales [19, 20]. In the medical imaging community, efforts have focused on further improving the popular UNet backbone  [21, 22]. These fully supervised methods have shown to be competitive in delineating linear structures like membrane segmentation [23]. Progress in the area of crack segmentation proceeded in a similar trajectory, where recent progress is due to improved backbone network design [24, 1]. Orthogonal to the above, in this work, we address the semi-supervised learning setting to reduce the amount of labelled data needed to train high performance models by exploiting unlabelled data.

To model spatial correlation and connectivity in curvilinear structures, a few works have incorporated topological constraints when training the segmentation model. One study utilized intermediate features from a VGG network to capture topological structure [4], while another work used concepts from persistent homology to provide more principled topological features; this idea was implemented by defining a loss between ground-truth and predicted persistent diagrams [3], which assumes access to ground-truth labels. In this work, we adopt a simpler approach by using positional encodings to capture spatial correlation.

We briefly review the two genres of semi-supervised learning (SSL) methods, namely the consistency-based SSL and adversarial training based SSL; for a more complete review, please refer to [5]. Consistency-based SSL methods [25, 26, 27] utilize the unlabelled data to constrain learning by enforcing the model’s predictions on the unlabelled data to be consistent with a specified target. Consistency regularization was firstly introduced in the context of deep learning by [28]. To better exploit the power of ensemble learning [25] used temporal ensembling of models’ predictions as a target for consistency. [26] further proposed to ensemble the model parameters over time to obtain more stable consistency targets. Different from the previously mentioned approaches, the VAT method [27] proposed to enforce consistency of a model’s predictions with that on an adversarially perturbed input [29]. Orthogonal to the above, the MixMatch method [30] proposed to interpolate between labelled and unlabelled samples to diversify the training set. Due to the high accuracy of temporal ensembling and low memory footprint, we adopted Mean Teacher [26] as the backbone framework. Another line of works adapt generative adversarial training to semi-supervised learning. Specifically, in the seminal work [31], a discriminator was introduced to distinguish generated samples from real ones and at the same time classify real samples into respective categories.

Semi-Supervised learning has been applied to semantic image segmentation to alleviate the expensive image annotation task. Following the adversarial training approach [31], [32] first employed a discriminator to distinguish generated fake images from real ones to improve semi-supervised learning. A conditional generative adversarial network (GAN) was further employed by ensuring that a discriminator cannot distinguish predicted segmentation masks on unlabelled images from ground-truth segmentation masks (on the labelled images) [6]; [7] further imposes cycle consistency (building on CycleGAN) to constrain the learning on unlabelled data. However, under low labeled data training a GAN is prone to instability and may not generalize well to unlabeled data. Following the consistency-based SSL approaches, [9] proposed to enforce consistency between the segmentation predictions of images subject to a MixUp augmentation, i.e. random crops of the two images are superimposed and consistency is enforced on the segmentation predictions of respective patches. In the context of medical image segmentation, a framework similar to Mean Teacher has been applied to skin lesion and CT-scan segmentation [33]. All these existing methods only consider data augmentations that are either pixel-wise perturbations [9] or rotations that are multiples of $90^{\circ}$ [33] with pairwise consistency; such augmentations are not strong enough to fully exploit the power of consistency regularization. These works also do not address the issue where the MSE loss used to enforce consistency is prone to collapse.

We first provide an overview of our SemiCurv approach, then describe the individual components in greater detail. Our SemiCurv approach, illustrated in Fig. 2, is built upon the consistency-based SSL framework (Sect. III-B). Both labelled and unlabelled images are first augmented by a differentiable geometric transformation (Sect. III-C) and then fed into student and teacher networks. The teacher network is a parameter-wise temporal ensemble (exponential moving average) of the student network and generates a target (pseudo-label) to encourage consistency on the unlabelled data. Sinusoid positional encoding is appended to each convolution layer in the encoder to explicitly capture spatial correlations (Sect. III-E). The predicted segmentation maps (posterior probabilities) from both networks are transformed back to the original pose via the differentiable inverse transform, and the dice loss (Sect. III-F) and N-pair loss (Sect. III-D) are used on the labelled and unlabelled data respectively for training; in particular, we show how the N-pair loss resolves a “collapsing” issue associated with the MSE loss commonly used for consistency on the unlabelled data.

We denote a training set of $N_{l}$ labelled images and their ground-truth segmentation maps as $\mathcal{D}_{l}=\{(\mathbf{X}^{l}_{i},\mathbf{Y}^{l}_{i})\}_{i=1}^{N_{l}}$, and denote the set of $N_{u}$ unlabelled images as $\mathcal{D}_{u}=\{\mathbf{X}^{u}_{j}\}_{j=1}^{N_{u}}$. The loss function optimized by semi-supervised learning methods can be generally written as Eq. (1), where $\mathit{l}_{l}$, $\mathit{l}_{u}$, and $\gamma$ are supervised loss on labelled images, unsupervised loss on unlabelled images and balancing hyperparameter, respectively.

We note that the loss functions of multiple SSL methods [26, 25, 27] can be written in this form. Most existing works adopt the cross-entropy loss or its variants as the supervised loss $l_{l}$ for classification tasks. In this work, we demonstrate in Sect. III-F that for segmentation with severe class imbalance, Dice loss is more appropriate. For the unsupervised loss, the $\Pi$ model [25], mean teacher model (MT) [26], VAT [27], MixMatch [30] and variants, all apply mean square error (MSE) consistency between the prediction posteriors of an unlabelled image under two different augmentations $t_{1}(\cdot)$ and $t_{2}(\cdot)$, and two different segmentation networks $f_{1}(\cdot)$ and $f_{2}(\cdot)$ as in Eq. 2. We show in Sect. III-D that under severe class imbalance the N-pair loss is more effective than MSE in avoiding trivial solutions.

SemiCurv is based on the mean teacher (MT) method [26] for semi-supervised classification, which has also been previously applied to the segmentation task. Briefly, the MT method maintains two networks: one fully trainable network called the student network, and another non-trainable network called the teacher network. The teacher network provides pseudo-labels to supervised the student network to learn. To improve stability of training, the parameters of teacher are the exponential moving average (EMA) of the student network’s parameters. In the rest of the paper, we denote the student network as $f(\cdot)$ and teacher network as $\hat{f}(\cdot)$.

Strong data augmentation is key to the success of consistency-based semi-supervised learning [34]. As discussed, existing semi-supervised segmentation methods often consider a limited set of augmentations. To further increase the variation of poses to improve feature learning we introduce affine transformations for stronger data augmentation. We denote the transformation applied to input image $\mathbf{X}$ as $t(\mathbf{X})$. To enable computing pixel-wise consistency loss, we further apply the inverse transformation to the output of both student and teacher networks as $t^{-1}(f(t(\mathbf{X})))$. With the inverse transformation, predictions on two arbitrarily augmented images are directly comparable at the pixel level.

Concretely, an affine transformation involves an arbitrary combination of scaling, rotation, shearing and translation, and can be formulated as a matrix multiply with the following transformation matrix,

For every pixel at $(\hat{w},\hat{h})$ in the image after affine transformation, we find its coordinate in the original image by $[w,h,1]^{\top}=\mathbf{H}^{-1}[\hat{w},\hat{h},1]^{\top}$. We employ bilinear interpolation to produce the pixel intensity at every pixel after transformation. Concretely, we can find the 4 neighbouring pixels before the transform as $\{x_{\left\lfloor w\right\rfloor,\left\lfloor h\right\rfloor},x_{\left\lceil w\right\rceil,\left\lfloor h\right\rfloor},x_{\left\lfloor w\right\rfloor,\left\lceil h\right\rceil},x_{\left\lceil w\right\rceil,\left\lceil h\right\rceil}\}$. The pixel intensity value after transformation $\hat{x}_{\hat{w},\hat{h}}$ can be determined by bilinear interpolation as,

Because of the bilinear transformation $\hat{x}_{\hat{w},\hat{h}}$ is obviously differentiable w.r.t. $\{x_{w,h}|w\in\{0,W-1\},h\in\{0,H-1\}$ where $H$ and $W$ are the height and width of image respectively, i.e. $t(\mathbf{X})$ is differentiable w.r.t. $\mathbf{X}$. For the inverse transformation, we simply apply $\mathbf{H}^{-1}$ for affine transformation and thus $t^{-1}(f(t(\mathbf{X})))$ is still differentiable w.r.t. $\mathbf{X}$.

We also notice that to generate diverse and realistic augmented image samples, we extrapolate images by mirroring the image over the edges. To give an example, the result of a 1D sequence “$abcd$” after mirror flipping is “$dcb|abcd|cba$”. This is achieved by clipping the pixel location before transformation with $x=(x//W\%2)*(W-x\%W)+(1-x//W\%2)*(x\%W)$ where $//$ and $\%$ are floor and modulo divisions respectively. An illustration of augmentation for image segmentation is given in Fig. 3 (a). With the differentiable transformation, we are able to randomly generate two augmented images $t_{1}(\mathbf{X})$ and $t_{2}(\mathbf{X})$ that are respectively fed into the student and teacher networks $f$ and $\hat{f}$. The two prediction posteriors can be then aligned via the differentiable inverse transformation $t_{1}^{-1}(f(t_{1}(\mathbf{X})))$ and $t_{2}^{-1}(\hat{f}(t_{2}(\mathbf{X})))$. For the rest of the paper, we denote prediction posteriors after alignment as $g_{i1}=t_{1}^{-1}(f(t_{1}(\mathbf{X}_{i})))$ (student) and $\hat{g}_{j2}=t_{2}^{-1}(\hat{f}(t_{2}(\mathbf{X}_{j})))$ (teacher).

Consistency-based SSL methods commonly use mean square error (MSE) for the consistency loss on unlabelled data. In this section, we first point out that MSE loss allows collapse of model predictions on unlabelled data to the majority class as this is a trivial solution of the MSE consistency loss. We then introduce the N-pair loss as a way to mitigate this issue thus enabling the unlabelled data to better regularize the model.

Without loss of generality, suppose we have two scalar predictions (pixel-wise predictions) $g_{1}$ and $g_{2}$. We visualize the loss surface for MSE pairwise consistency in Fig. 3 (b), and observe that the MSE loss is flat along the diagonal because MSE loss is minimized when the predictions from the two networks match. When the class distribution is highly imbalanced as in the case of curvilinear segmentation, where there are far more background pixels than foreground pixels, the two networks can achieve zero MSE consistency loss just by assigning every single pixel to the majority class. In this case every pixel will be predicted as background. We term this all majority class prediction as a collapsed prediction. An instance of this behaviour is shown in Fig. 8 (a) where training IoU gets closer to $1$, yet validation IoU collapses to $0$. Avoiding this often requires very careful selection of the EMA hyperparameter and consistency weights, which requires expensive tuning runs.

In this work, inspired by the recent progress in contrastive learning [14], we propose to use an N-pair loss on unlabelled data [13] to avoid collapsed predictions. The N-pair loss simultaneously exploits all unlabelled samples in a training mini-batch to construct one positive pair to encourage similarity and multiple negative pairs to encourage diversity in predictions. By enforcing diversity with the negative pairs, the N-pair loss effectively penalizes models that give collapsed predictions on all unlabelled data. We illustrate the construction of the N-pair loss for curvilinear structure segmentation in Fig. 4: in a mini-batch of $N_{B}$ unlabelled images the positive pair is chosen as the predictions from student and teacher on an anchor image $\mathbf{X}_{i}$, $\{(g_{i1},\hat{g}_{i2})|i=1\cdots N_{B}\}$, while negative pairs are chosen as the predictions between the student’s predictions on the same anchor image and teacher’s predictions on other images in the mini-batch, $\{(g_{i1},\hat{g}_{j2})|i\neq j;i,j=1\cdots N_{B}\}$. Formally, the N-pair loss is given by,

We now formally show how N-pair loss prevents collapsed predictions that can occur when MSE is used. Specifically, we show that the N-pair loss will be lower for correct predictions than collapsed predictions, unlike MSE loss. First consider the case when predictions all collapse to a single class, often the background resulting in all zero predictions, i.e. $g_{1}=\mathbf{0}$ and $\hat{g}_{2}=\mathbf{0}$. Then, the cosine similarity for all positive and negative pairs are 1; the same result holds when predictions are all foreground, i.e. $g_{1}=\mathbf{1}$ and $\hat{g}_{2}=\mathbf{1}$. Then, the N-pair loss for this collapsed prediction is

We note that in this case, the MSE loss is $\tilde{l}_{mse}=0$. When the predictions are all correct, the positive pair similarity is $1$, and for negative pairs this is approximately $\epsilon<<1$, because the segmentation masks for arbitrary two images are unlikely to significantly match, and the N-pair loss becomes

while the MSE loss is still $l_{mse}=0$. Since $\exp\frac{\epsilon-1}{\tau}<1$, it is easy to verify that $l_{\textrm{N-pair}}<\tilde{l}_{\textrm{N-pair}}$ when $N_{B}$ is larger than 1 and the ratio $l_{\textrm{N-pair}}/\tilde{l}_{\textrm{N-pair}}$ is getting smaller with increased batchsize $N_{B}$. This suggests the gap between N-pair loss under collapsed prediction and correct prediction is more significant with larger batchsize. We therefore conclude that the N-pair loss can easily distinguish good predictions from collapsed ones as opposed to MSE loss, for which the loss is $0$ in both cases.

The nature of curvilinear structure segmentation requires the network to encode spatial correlation. For example, cracks, roads, and blood vessels are spatially continuous thus pixels adjacent to other positive (foreground) pixels are likely to be positive as well. Though it is difficult to handcraft features that capture this correlation, encoding spatial locations has been shown to be effective [16]. One straightforward way to encode spatial locations is by appending linear coordinates, normalized to between 0 and 1, to the intermediate feature layers as additional feature channels. However, learning directly on this absolute positional encoding risks overfitting to specific locations, especially when training with very few labelled samples. Inspired by the positional encoding of [35], we apply a sinusoid coordinate encoding as in Eq. (9) where $S_{x},S_{y}$ are the linear positional encoding normalized to between 0 and 1 and $K$ is a period parameter. Such periodic positional encoding allows the network to be aware of relative location while reducing the risk of overfitting to absolute locations. The positional encodings are channel-wise concatenated to each intermediate activation map of the encoder network as shown in Fig. 5.

It is well-known that class imbalance can affect the performance of machine learning models [36, 37, 38]. The cross-entropy (CE) loss commonly used for supervised segmentation tasks averages the pixel-wise CE loss over all pixels in the image; CE loss biases learning towards the majority class when there is an imbalanced class distribution. In curvilinear structure segmentation tasks, this imbalance is particularly severe due to a low positive class ratio (see Table I). As a result, models trained with CE may have high per-pixel accuracy but low intersection over union (IoU), which is the most common evaluation metric for segmentation tasks. In order to mitigate this issue, Dice loss [39] (Eq. 10) was proposed as an alternative for segmentation; it focuses on the intersection of predictions with ground truth only for the positive class.

The final loss function is the weighted combination of supervised loss and unlabelled loss: $l_{total}=l_{dice}+\gamma l_{\textrm{N-pair}}$. As commonly done in consistency-based SSL, we do not apply the unlabelled loss $l_{\textrm{N-pair}}$ at the beginning of training but instead gradually increase the strength of $\gamma$ following a sigmoid function $\gamma=\exp(-10(\frac{min(T,T_{rp})}{T_{rp}}-1)^{2})$ where $T$ and $T_{rp}$ are epoch number and hyperparameter, respectively. We set the mean teacher EMA hyperparameter $\alpha$ to 0.999. For all experiments, we use the SGD optimizer, fixing the total number of training epochs to 1000 and $T_{rp}=500$, the initial learning rate at $0.001$ and continuously decay by half every 500 epochs. Each epoch is defined as cycling all labeled data once. The sinusoid spatial encoding parameter was set as $K=4$. The overall algorithm is summarized in Algorithm 1.

To evaluate the quality of segmentation predictions, we compute Intersection over Union (IoU) using a threshold of 0.5 on prediction posteriors for all test images and report the mean IoU. We also evaluate the F1 measure of precision and recall at threshold 0.5. This metric treats segmentation as a binary classification task. For all 6 datasets, we set the background pixel label as 0 and foreground as 1.

We extensively compare SemiCurv with existing fully supervised methods (with different backbone networks) and state-of-the-art generic semi-supervised learning methods, as described in the following.

We present the comparison of both fully supervised and semi-supervised methods with $5\%$ and $1\%$ labelled data in Table III; for MITRoad and EM128 we consider different labeling budgets due to the size of dataset. In the fully supervised block, UNet ($*95\%$) and UNet ($*90\%$) indicate the relative performance $95\%*m$ and $90\%*m$ where $m$ is the performance achieved by the fully supervised method trained with $100\%$ of the labelled data. From the extensive comparison, we make the following observations. First, the adapted UNet is a very strong backbone network for curvilinear structure segmentation. It outperforms the state-of-the-art semantic segmentation backbone DeepLabV3+, edge detection network HED, and network specifically designed for crack segmentation FPHBN on all datasets. Moreover, under the semi-supervised setting, with only $5\%$ and $1\%$ labelled data SemiCurv can match and even outperform the $95\%$ and $90\%$ relative performance of fully supervised counterparts in terms of IoU and F1, respectively. By comparing to the state-of-the-art semi-supervised methods for segmentation, we still observe very significant advantages for SemiCurv. In particular, the improvement from baseline is more significant in the lower label regime. For example, on CrackForest IoU improved $17\%$ and $2.2\%$ from the UNet baseline trained using $1\%$ and $5\%$ labelled data respectively. The improvement on all other datasets also exhibits similar patterns with large improvements over the UNet baseline. We also observe CutMix to be a very competitive method, in particular on the MIT Road dataset. We speculate that this is because the MIT Road dataset resembles generic semantic segmentation problems, such that it benefits from cutmix style augmentation. Finally, TCSM produces much worse results compared to both MT and SemiCurv. We speculate that this is due to weak augmentation adopted in TCSM hampering semi-supervised performance.

Here we carry out ablation studies to investigate the effectiveness of each component and present the results in Table V.

For alternative positional encodings, we evaluate two additional methods on CrackForest segmentation task. The results are presented in Tab. IV. We observe a small drop of performance when swapping out our proposed sinusoid encoding with learned position embedding [35] and rotary position embedding [47]. Overall, the proposed sinusoid position encoding is still the better option.

The quantitative comparison between using cosine similarity and L2 distance induced similarity is shown in Table V under the similarity metric (Sim.) column. For both similarity metrics, we use the SemiCurv framework and keep the training protocols unchanged. We observe from the comparison that cosine similarity outperforms L2 distance consistently on all datasets. Also, stability of training suffers with L2 distance compared to SemiCurv with cosine distance. This can be attributed to the normalizing effect of cosine similarity, while L2 distance is affected by the absolute number of positive pixel predictions.

We further present qualitative comparisons of UNet, CutMix, MT, and SemiCurv on 4 datasets in Fig. 7. We make the following observations: first, for both MT and SemiCurv, the visual quality of segmentation results are more appealing and better mimicks the ground-truth compared to the fully supervised baseline UNet. In particular, our method is able to generate more well-connected predictions for $037$ (1st row) thanks to the contribution of appending positional encoding. We further notice that SemiCurv produces, in general, cleaner outputs for $037$ (1st row), $022$ (2nd row), $066\_1$ (3rd row). We also observe, in $0192\_541\_1$ (4th row), the proposed approach captures very subtle cracks on the top. The segmentation results on DRIVE dataset (5th/6th rows) also suggest the superiority of SemiCurv. It produces relatively high fidelity and well-delineated blood vessel segmentation maps. In comparison, UNet, CutMix, and MT tend to mingle different blood vessels together which is less common in SemiCurv’s predictions. For satellite image segmentation, the visualization covers industrial areas (row 7) and rural areas (rows 8-9). We observe consistent improvements of SemiCurv compared to both UNet and state-of-the-art semi-supervised image segmentation CutMix. Interestingly, the SemiCurv predictions sometimes contain valid road branches that are missing in the ground-truth annotation (last row).