Colour augmentation for improved semi-supervised semantic segmentation

The key idea behind consistency regularization based semi-supervised classification is clearly illustrated in the $\pi$-model of Laine et al.  [Laine and Aila, 2017], in which a network is trained by minimizing both supervised and unsupervised loss terms. The supervised loss term applies traditional cross-entropy loss to supervised samples with ground truth annotations. Unsupervised samples are stochastically augmented twice and the unsupervised loss term encourages the network to predict consistent labels under augmentation.

The Mean Teacher model of Tarvainen et al.  [Tarvainen and Valpola, 2017] uses two networks; a teacher and a student. The weights of the teacher are an exponential moving average of those of the student. The student is trained using gradient descent as normal. The teacher network is used to generate pseudo-targets for unsupervised samples that the student is trained to match under stochastic augmentation.

The UDA approach of [Xie et al., 2019] adopted RandAugment [Cubuk et al., 2020]; a rich image augmentation scheme that chooses 2 or 3 image operations to apply from a menu of 14. While only one network is used instead of two, we note an important similarity with Mean Teacher; just as the teacher network is used to predict a pseudo target, UDA predicts a pseudo-target for an un-augmented image that is used as a training target for the same iamge with RandAugment applied.

The FixMatch approach of [Sohn et al., 2020] refines this approach further. They separate their augmentation scheme into weak – consisting of simple translations and horizontal flips – and strong that uses RandAugment. They predict hard pseudo-labels for weakly augmented unsupervised samples that are used as training targets for strongly augmented variants of the same samples.

[Hung et al., 2018] and [Mittal et al., 2019] adopt GAN-based adversarial learning, using a discriminator network that distinguishes real from predicted segmentation maps to guide learning.

[Perone and Cohen-Adad, 2018] and [Li et al., 2018] two early applications of consistency regularisation to semantic segmentation that we are aware of. Both come from the medical imaging community, tackling MRI volume segmentation and skin lesion segmentation respectively. Both approaches use standard augmentation to provide perturbation, as in the $\pi$-model [Laine and Aila, 2017] and Mean Teacher [Tarvainen and Valpola, 2017]. [Ji et al., 2019] developed a semi-supervised over-clustering approach that can be applied to natural photographic images, where the list of ground truth classes is highly constrained.

[French et al., 2020] analysed the problem of semantic segmentation, finding that it has a challenging data distribution to which the cluster assumption does not apply. They offer this as an explanation as to why consistency regularization had not been successfully applied to semantic segmentation of photographic images. They present an approach that drives the Mean Teacher [Tarvainen and Valpola, 2017] algorithm using an augmentation scheme based on Cutmix [Yun et al., 2019], achieving state of the art results.

Approaches based on contrastive learning [Henaff, 2020, He et al., 2020, Chen et al., 2020b, Chen et al., 2020a] train a residual network [He et al., 2016] using only unlabelled input images. Afterwards the network backbone (consisting of convolutional layers) is frozen and a linear model is trained in a supervised fashion using it’s feature representations as inputs and ground truth labels as targets. The resulting image classifiers – in which only the last linear layer was trained using ground truth labels – are able to achieve ImageNet results that are competitive with those obtained by traditional supervised learning in which the whole network is trained.

In contrast to prior work [Henaff, 2020] the MoCo model of He et al.  [He et al., 2020] simplified contrastive learning using standard augmentation to generate stochastically augmented variants of batches of unlabelled images. The network is encouraged to predict embeddings that are more similar for augmented variants of the same input image than for different images. The augmentation scheme used is very similar to the standard scheme used to train residual networks [He et al., 2016] and by Mean Teacher [Tarvainen and Valpola, 2017] for their ImageNet results. Chen et al.  [Chen et al., 2020a] conducted a rigorous ablation study of the augmentations used for contrastive learning, assessing the effectiveness of each augmentation operation. They found that strong colour augmentation is essential for good performance, as without it the network is able cheat by using image colour statisticts as a short-cut to discriminate between images, rather than having to focus on image content. Strong colour augmentation masks this signal, forcing the network to focus on the image content, extracting features suitable for accurate image classification and other downstream tasks.

We note the similarities between recent contrastive learning approaches and the Information Invariant Clustering approach of Ji et al.  [Ji et al., 2019], who also encourages consistency under stochastic augmentation.

The recent work of [Liu et al., 2021] adapt contrastive learning – typically used for classification – for semantic segmentation, achieving impressive results with very few labelled imgaes.

Our choice to apply colour augmentation to the unsupervised loss term in a semi-supervised semantic segmentation setting was inspired by the thorough ablation study performed by Chen et al.  [Chen et al., 2020a] in which they explore the effects of augmentation in the similar setting of self-supervised learning. As stated in Section 2.3, they found that colour augmentation prevents the network from learning to use image colour statistics as a short-cut for image instance discrimination. Colour augmentation modifies the colour statistics sufficiently to prevent the network from using them to match image instances with one another trivially, forcing the network to focus on image content.

While [French et al., 2020] offer the challenging data distribution present in semantic segmentation problems as an explanation as to why consistency regularization driven by standard augmentation had yielded few successes when applied to photographic image datasets such as Pascal VOC 2012 [Everingham et al., 2012], we offer colour statistics as an alternative explanation.

The consistency loss term in equation 2 offers the opportunity for the network to minimize $L_{cons}$ using colour statistics. This is further illustrated in Figure 1, in which the yellow arrows follow a single pixel from the input image $x$ through both the student and teacher sides of the consistency loss term. Given the care taken to maintain the input-to-target pixel correspondence as stated in Section 3, most input pixels in $x$ (geometric augmentation can move some parts of an image outside the bounds of the image, hence correspondence for some pixels will be missing) will have corresponding pixels in the same locations in the prediction maps from both the student side $f_{\theta}(t_{\alpha}(x))$ and the teacher side $t_{\alpha}(g_{\phi}(x))$. Given that the consistency loss term $L_{cons}$ penalises the network for giving inconsistent class predictions for each pixel, a simple way to minimize $L_{cons}$ is to predict the class of a pixel in the output segmentation maps using only the corresponding pixel in the input image, ignoring surrounding context. Thus, the network effectively clusters the colour of individual input pixels, rather than using surrounding context to identify the type of object that the pixel lies within.

Following [Chen et al., 2020a] we propose using colour augmentation to prevent the network from utilizing this short-cut. We acknowledge that [Ji et al., 2019] applied colour augmentation in an unsupervised semantic segmentation setting. While their codebase uses the same colour augmentation approach as [He et al., 2020] and [Chen et al., 2020a] they describe it simply as ‘photometric augmentation’ in their paper, giving little hint that it is in fact key to the success of consistency regularization based techniques in this problem domain.

Our implementation is directly based on that of [French et al., 2020]. We add colour augmentation to their implementation of standard augmentation, ICT [Verma et al., 2019], VAT [Miyato et al., 2017], Cutout [DeVries and Taylor, 2017] and Cutmix based regularizers. This allows us to assess its effect on a variety of regularizers across three datasets; Cityscapes, Pascal VOC 2012 and the ISIC Skin Lesion segmentation dataset [Codella et al., 2018]. Our colour augmentation randomly modifies the brightness, contrast, saturation and hue of an input image and is implemented using the ColorJitter transformation from the torchvision [Chintala et al., 2017] package.

Cityscapes is a photograpic image dataset of urban scenery captured from the perspective of a car. Its’ training set consists of 2975 images.

Our Cityscapes results are presented in Table 1 as mean intersection-over-union (mIoU) percentages, where higher is better. The addition of colour augmentation results in a slight improvement to the CutOut and CutMix results across the board.

Pascal VOC 2012[Everingham et al., 2012] is a photographic image dataset consisting of various indoor and outdoor scenes. It consists of only 1464 training images, and thus we follow the lead of [Hung et al., 2018] and augment it using Semantic Boundaries[Hariharan et al., 2011], resulting in 10582 training images.

Our Pascal VOC 2012 experiments evaluate regularizers based on standard augmentation, ICT [Verma et al., 2019] and VAT [Miyato et al., 2017], Cutout and Cutmix as in [French et al., 2020].

Our results are presented in Table 2.

The ISIC skin lesion segmentation dataset [Codella et al., 2018] consists of dermoscopy images focused on lesions set against skin. It has 2000 images in its training set and is a two-class (skin and lesion) segmentation problem, featuring far less variation than Cityscapes and Pascal. Our results are presented in Table 3.

While colour augmentation improved the performance of all regularizers on the Pascal dataset when using the DeepLab v2 architecture, the results for ISIC 2017 are less clear cut. It harms the performance of VAT and ICT, although we note that we increased the consistency loss weight of ICT to match the value used for Pascal. It yields a noticeable improvement when using standard augmentation and Cutout. Colour augmentation increases the variance of the accuracy when using CutMix, making it slightly less reliable. We hypothesized the the hue jittering component of the colour augmentation may harm performance in this benchmark as colour is a useful queue in lesion segmentation, so we tried disabling it when using ICT and VAT. This did not however improve colour augmentation results.

While we have demonstrated that colour augmentation can improve semi-supervised segmentation performance when using a simple consistency regularization based approach, we acknowledge that our results do not match those of the recent Classmix [Olsson et al., 2021], DMT [Feng et al., 2021] and ReCo [Liu et al., 2021] approaches that use more recent semi-supervised regularizers.

We also note that [Liu et al., 2021] focused on situations in which a very small number of labelled samples were used. As their work did not feature experiments with a comparable number of labelled samples to our own, we were unable to directly compare their results with ours in Tables 1 and  2.