CINFormer: Transformer network with multi-stage CNN feature injection for surface defect segmentation

Traditional machine visions utilize low-level texture features to detect defects. These methods can not effectively defect complex defects with low contrast appearance and small size. Recently, with the rapid development of deep learning, some CNN-based methods have been widely applied to surface defect detection. For instance, Tabernik et al. [14] proposed a two-stage network, which introduces a segmentation subnetwork in the classification network. Cui et al. [29] integrate fine-grained details of preceding layers at each deep layer, which facilitates the detection of small objects.

However, it is hard for these methods to identify defects in complex scenes. Benefiting from the attention mechanism that highlights important features, some methods introduce it into the model to strengthen the feature representation. Based on the human visual gain mechanism, Wei et al. [18] constructed a Faster VG-RCNN that integrates attention-related visual gain mechanism to better detect small defects. Su et al. [30] integrate the improved multi-head nonlocal self-attention into FPN to highlight defect features. Yang et al.[31] introduce the spatial attention module into each residual unit of the backbone to make the network adaptively focus on defect regions. Similarly, Jiang et al. [32] use joint channel-spatial attention weights to selectively fuse high-level and low-level features. Xiang ea al. [33] introduce the attention module after high-level and low-level feature fusion to strengthen feature representation. However, these CNN-based methods are still challenged by defects in complex scenes because CNN features are vulnerable to background interference.

Recently, Dosovitskiy et al. [21] proposed a vision transformer, which shows the strong ability to capture global features benefiting from the self-attention mechanism. To deal with variations in the scale of objects, Liu et al. [26] proposed a hierarchical transformer based on window self-attention. Wang et al. [34] proposed a CrossFormer to improve Transformer’s ability to build cross-scale interactions. Xie et al. [35] constructed a hierarchical transformer as an encoder to generate multi-level features, and fused these features through MLP layers in the decoder. These methods focus on drawing global relationships, which are robust against background interference. However, these methods are prone to lose some detailed information, Because it does not pay attention to local information.

To deal with the above problem, Liu et al. [26] proposed a local window mechanism to perform self-attention operations within the window. Yang et al. [36] proposed a Focal Self-Attention (FSA), which pays fine-grained attention to the area around the current token and coarse-grained attention to the area far away from the current token. In this way, it can more effectively capture local and global attention. Xia et al. [37] proposed a deformable self-attention module in which key and value in self-attention are sampled to flexibly learn informative features. Fan et al. [38] proposed a multiscale Transformer, which creates a multiscale feature pyramid that learns low-level information at early layers and semantic information at deep layers. Ren et al. [39] proposed a Shunted Self-Attention (SSA), which can capture multi-scale features through multi-scale token aggregation. However, these methods still do not solve the problem that self-attention is low-pass filtering, making it difficult for them to deal with weak or minor defects.

Due to the local properties of CNN, it is prone to bring in noises when dealing with defects in complex scenes. In contrast, the transformer model can draw long-distance dependencies and suppress noise, while it is weak in capturing detailed information about objects. Therefore, some methods improve the networks’ performance by combining CNN and transformer networks. For example, Peng et al. [23] designed a hybrid network structure termed Conformer, which enhances representation learning by combining local and global features in an interactive way. Wu et al. [25] use convolutional projection instead of linear projection in each self-attention block, which can make the transformer model capture local details. Xie et al. [40] proposed a novel one-stage framework called Pyramid Grafting Network (PGNet), which utilizes global information to strengthen CNN’s feature representation in the decoder. Li et al. [24] designed a novel Unified transformer (UniFormer) to learn both local and global token affinity by integrating 3D convolution and spatiotemporal self-attention. Although these methods inherit the local properties of CNNs while maintaining the global merits of transformers, they still confront challenges posed by weak defects and complex background.

Different from the above methods, the proposed CINformer presents a one-way multi-stage feature injection method from the CNN to the transformer network, which can better preserve the local details of the CNN. Furthermore, a Top-K self-attention module is proposed to make the model focus on defect information by selecting important tokens or channels to calculate self-attention.

As depicted in Fig. 2, the proposed network is constructed based on a U-shaped architecture. A multi-level transformer network is adopted as the encoder in which each level is injected by CNN features. Specifically, we adopt the ResNet-18 [5] and the Swin-T [26] as the CNN and transformer backbones, respectively. Based on the ResNet-18, a CNN stem is constructed to generate hierarchical CNN features, denoted as {${R_{i}|i=1,2,3,4}$}. It is noted that the transformer network takes the convolutional feature $R_{1}$ instead of images as the input. The other convolutional features $R_{2}$, $R_{3}$, and $R_{4}$ are injected into features of different stages of the transformer network, which can provide the transformer with detailed defect features. The detailed CNN features are beneficial for segmenting those small defects.

The decoder consists of four decoding stages, each of which fuses high-level features and the corresponding low-level features in the encoder. Besides, a Top-K self-attention module is presented to alleviate the impact of redundant background information. The proposed Top-K self-attention module selects important tokens or channels to calculate self-attention, which is used in the last two stages of the transformer network to help the model attend to important defect regions.

The illustration of the transformer with CNN injection is shown in Fig. 2. In CNN, features at deep layers tend to contain rich semantic information about defects, while features at shallow layers contain abundant detailed information about defects. So we construct a CNN stem through FPN [41] to generate multi-level features. The FPN fuses high-level and low-level features from top to down. In FPN, high-level features are upsampled and concatenated with the low-level features, and fed into 3$\times$3 convolutions for further fusion. Compared with the multi-level features of the CNN backbone, the multi-level features generated by FPN contain richer semantic information and detailed information. So we use the multi-level features {${R_{i}|i=1,2,3,4}$} in the top-down path as the features injected into Transformer.

The convolutional feature $R_{1}$ is taken as the input of the transformer network, generating transformer features {${S_{i}|i=1,2,3,4}$} of different stages. The other features are injected into the subsequent transformer stages, respectively. Each transformer stage takes the fused features of transformer features and convolution features as input. Specifically, the convolutional feature $R_{i}$ is applied a $1\times 1$ convolution to adjust the number of channels. The transformed convolutional feature is reshaped from $\mathbb{R}^{C\times H\times W}$ into $\mathbb{R}^{C\times N}$, where $N=H\times W$. The reshaped feature is fused with transformer feature $S_{i-1}$ and fed into transformer stage $i$. Mathematically, the injection process can be described as follows:

where $Reshape(.)$ represents reshape operation. $Concat(.)$ represents concatenation operation. $Conv(.)$ denotes the convolution operation. And the $Linear(.)$ denotes linear projection.

Traditional self-attention computes the feature similarity between any two spatial locations. As a result, this brings redundant information, which is likely to drown out some detailed defect information. This is because the defects are usually weak and small, which are easily confused by the background. Traditional self-attention utilizes all tokens to produce the attention map, which brings in redundant background information. To remedy this problem, we present a simple but effective Top-K self-attention module, which is used in the last two transformer stages. The proposed self-attention module selects the top $k$ important tokens and channels to compute the attention. As a result, it can highlight important defect features and filter out some redundant background information. At the same time, it brings lower computational complexity than the self-attention module. The details of the proposed self-attention module are shown in Fig. 3.

We will present the proof by visualizing the feature maps of the original self-attention and Top-K self-attention module in Fig. 7 in our paper. The feature map of the Top-K self-attention can focus on defect information while suppressing redundant background information compared with that of the original self-attention.

In our module, given input feature $X\in\mathbb{R}^{N\times C}$, it is linearly mapped into queries Q, keys K, and values V. Mathematically, we have:

where $W_{q}$, $W_{k}$ and $W_{v}$ represent three matrices of $\mathbb{R}^{C\times C}$.

With K, Q, and V obtained, we use the Top-K self-attention mechanism to select the top $k$ most important tokens and channels to keep the most useful information and suppress the redundant background information. Specifically, for token selection, we compute the variance statistics of tokens along the channel dimension. If the variance of a token is higher, it tends to contain more useful information for defects. Therefore, we select the top $k$ important tokens by ranking their variances and generating the indexes of these selected tokens. Similarly, for channel selection, we first compute the variance statistics of channels along the spatial dimension and then select the top $k$ important channels by sorting the channel variances, generating the indexes of these selected channels.

In the Top-K self-attention module, the Top-K token indexes and channel indexes are first computed on the Q. Then they are used to select the most important tokens and channels from queries Q, generating Q^′. Meanwhile, they are used to select the most important channels from keys K and the most important tokens from values V, generating K^′ and V^′, respectively. Notably, the Top-K token indexes and channel indexes are only computed on the Q once because the self-attention module first calculates similarities of Q and K and then performs matrix multiplication with V. If the indexes are re-computed on K and V, the corresponding tokens and channels at the same position in the selected Q, K, and V do not match. With the obtained K^′,Q^′and V^′, the Top-K self-attention is calculated as follows:

where the $Softmax(.)$ denotes the Softmax function. Compared with the original K, Q, and V, the K^′, Q^′and V^′ have fewer dimensions. So the proposed self-attention can also reduce the computational cost of the self-attention module.

In addition, we introduce a constraint vector to counteract the expansion of unselected regions due to the softmax operation. Because in the Top-K self-attention, each unselected token in the key K also computes similarity with selected tokens in the query Q, generating a score in the attention map. The constraint vectors (6) and (7) are introduced to suppress the scores of these unselected tokens. Mathematically, we have:

where $mean(.)$ is the mean operation, and $Sig(.)$ is the Sigmoid function. $\gamma$ is a learnable parameter.

As stated in  [42], the relations of different tokens and semantic information at shallow layers are weak. Therefore, the proposed Top-K self-attention mechanism is only used in the last two stages of the network. In the ablation experiments, we also observe a performance degradation when using the Top-K self-attention module in the first two stages.

In the experiment, NEU [43], DAGM 2007 [44], and Magnetic tile [45] (MT) defect datasets are selected to demonstrate the effectiveness of the proposed method. The details of the three defect datasets are illustrated as follows.

NEU is a steel strip surface defect dataset, which contains three types of defects: patch, inclusion, and scratch. There are 300 images for each type of defect, and each image is provided with pixel-level ground truth. The intra-class difference and inter-class similarity of defects bring great challenges to defect segmentation.

DAGM 2007 is artificially generated but the samples are very similar to those in the real world, which contains 10 types of defects generated from various texture and defect models. The defect region in each defect image is roughly marked by an ellipse. Most of the defects are low-contrast and affected by complex background interference, so it is hard to detect them accurately.

MT is a magnetic tile defect dataset, which contains 1344 images. There are 392 defect images, including five categories of defects: uneven, fray, crack, blowhole, and break. These defects exhibit complicated appearances with multi-scale and low contrast, which bring challenges to defect segmentation.

Image Size: The images in the three defect datasets vary from 200$\times$200 to 512$\times$512 in size. It is a common practice to resize the image to a fixed resolution in the defect segmentation task [17, 46].

We implement the proposed method based on Pytorch [47]. All experiments are conducted on the Tesla V100 platform. The CNN stem adopts the ResNet-18 pre-trained on the ImageNet [48]. We use AdamW [49] optimizer with a learning rate of 0.00075 and weight decay of 0.005 to train the model. Meantime, a cosine decay learning rate scheduler is adopted during the training stage. For each dataset, 70% of the images in each category are selected as the training set and the remaining images are used as the test set. Each image is resized to 224$\times$224 during both the training and test stages by following the common practice [17, 46]. With a batch size of 4, the model is trained for 150 epochs, 100 epochs, and 100 epochs on NEU, DAGM 2007, and MT defect datasets, respectively, where the first 20 epochs are used for warming up. In the experiment, the mean Intersection of Union (mIoU) is adopted to evaluate the performance of different models. We utilize the cross-entropy loss as the supervision to train the network.

The proposed CINFormer is compared with 13 state-of-the-art models on DAGM 2007, NEU, and MT defect datasets. These methods are classified as CNN-based methods including UNet [50], ResNet [5], SPNet [51], GCAPNet [52], and OCNet [53], vision transformer based methods including PVT [28], Segformer [35], Swin-T [26], and DAT [37], and methods of combining CNN and transformer including Uniformer [24], Conformer [23], VST [54], PGNet [40].

The comparison results of various methods are given in Table I. It is seen that CNN-based methods are superior to transformer-based methods on MT. This is because there are many weak defects on the MT dataset and CNN is better at capturing defect details than the transformer. On the contrary, transformer-based methods outperform CNN-based methods on NEU. This is because the most of defects have relatively complex background disturbances. This indicates that the transformer models can better suppress the background interference in complex scenes than CNN. On the DAGM dataset, the transformer-based methods and the CNN-based methods perform similarly. As to the existing hybrid methods of combining CNN and transformer, they perform differently on three datasets. They generally perform worse than CNN on MT, perform almost the same as the transformer on NEU, and perform a little better on DAGM. This means the existing hybrid methods do not effectively exploit their respective merits on the surface defect datasets. The proposed CINFormer outperforms the above methods on three defect datasets. It is noted that the CINFormer can only fine-tune the Transformer part while keeping the weights of the ResNet-18 frozen during training. Even so, CINFormer outperforms Conformer on the defect datasets. When trained in an end-to-end way, CINFormer obtains the best result on three defect datasets. In general, the experimental results demonstrate the effectiveness of the proposed one-way CNN injection.

Fig. 4 presents some segmentation results of various methods. It is found that some methods fail to detect small or low-contrast defects, such as ResNet, UNet. Some methods such as Segformer and Swin even misclassify the defects, manifested by inconsistent colors between predictions and ground truth (GT) on defect regions in Fig. 4. On the contrary, the proposed CINFormer can accurately detect those defect regions. In addition, Fig. 5 presents the feature visualization results of various models on three defect samples. It is found that our method can better focus on defect regions while suppressing background interference.

To further demonstrate the effectiveness of the proposed CNN injection strategy and the Top-K self-attention module, we conduct the following ablation experiments.

Ablation on the Transformer with CNN Injection: To demonstrate the priority of the proposed CNN feature injection manner, we conduct several experiments based on the models with and without the Top-K self-attention module. We first compare the Swin-T model with the CINFormer model without the Top-K self-attention module denoted as CINFormer (w/o) and then compare the two models having the Top-K self-attention module. The Swin-T model with the Top-K self-attention module is denoted as Swin-T (Top-K). The experimental results in Table II show that CINFormer outperforms the Swin-T in both conditions. CINFormer (w/o) surpasses Swin-T by 6.5 points and CINFormer surpasses Swin-T (Top-K) by 7.0 points in terms of mIoU on DAGM. This indicates that the injected multi-level CNN features are beneficial for the transformer to detect defects.

Furthermore, the proposed multi-stage CNN injection manner is compared with other hybrid structures of CNN and transformer, as shown in Fig. 6 (a) and (b). As shown in Table III, the proposed method obtains improvements of 1.0 and 1.8 points over the other two structures, respectively. It is noted that structure (a) is a bidirectional feature injection pattern with features of CNN and transformer interacting with each other, which is similar to Conformer [23]. Though this kind of bidirectional interaction can promote the full merging of CNN features and transformer features, it impairs the ability of the CNN branch to represent the detailed features. As a result, our proposed one-way CNN feature injection method outperforms it on the defect dataset.

Ablation on the Top-K Self-attention Module: To demonstrate the effectiveness of the proposed Top-K self-attention module in the transformer network, we use it to replace the self-attention module of different transformer models. Specifically, we only replace the self-attention module in the last two stages of the transformer models. The experimental results in Table IV show that the introduction of Top-K self-attention module improves mIoU by 0.8, 0.8, 0.7, and 1.3 points compared with Swin-T, Conformer, Uniformer, and CINFormer (w/o) on the DAGM defect dataset, respectively. This indicates that the proposed Top-K self-attention module is a plug-and-play module, which can improve the performance of various transformer models in defect detection. Meanwhile, the proposed Top-K self-attention module reduces the computational complexity of the transformer model because it only selects important tokens and channels to calculate attention. Furthermore, we present the feature visualization of CINFormer and CINFormer (w/o), as shown in Fig. 7. The introduction of the Top-K self-attention module can make CINFormer focus on defect regions and suppress background interference.

Configuration of the Top-K Self-attention Module: First, we compare the performance of models using Top-K self-attention at different stages. To be specific, experiments are conducted based on the models with and without the CNN injection module, as shown in Table V. The Swin-T model contains four stages, denoted as 1, 2, 3, and 4, respectively. We gradually add the proposed Top-K self-attention in the transformer network starting from the last stage. The proposed CINFormer achieves the best performance when the Top-K self-attention is used in the last two stages. This is because a single token in the earlier stage is not highly correlated with its surrounding tokens [42], which makes it hard to reconstruct the abandoned tokens. Therefore, using the Top-K self-attention mechanism in the early stage will make the model lose some important detailed information during feature extraction, thereby reducing the accuracy of recognition. In contrast, a single token in the latter stage has a strong correlation with its surrounding tokens, the defect (target) tokens are easily affected by background noise. To sum up, when Top-K self-attention modules are used in the last two stages, the proposed CINFormer can achieve the best balance between accuracy and cost.

Secondly, we further investigate the effect of the hyperparameter K on model performance. The experiments are conducted on the model under the CNN-injected module and the non-CNN-injected module. The K is set to 14, 21, 28, 35, and 42, respectively. As shown in Table VI, the proposed CINFormer achieves the best performance when K is set to 28. Because the K with small values will discard a lot of tokens, including some useful information. And the K with large values will overwhelm the detailed information, which is the same as the original self-attention mechanism.