Citation: Li, H.; Li, H.; Hou, Z.; Song,
H.; Liu, J.; Dai, P. Feature
Augmentation Based on Pixel-Wise
Attention for Rail Defect Detection.
Appl. Sci. 2022, 12, 8006. https://
doi.org/10.3390/app12168006
Academic Editor: Andrés Márquez
Received: 29 June 2022
Accepted: 9 August 2022
Published: 10 August 2022
Publisher’s Note: MDPI stays neutral
with regard to jurisdictional claims in
published maps and institutional affil-
iations.
Copyright: © 2022 by the authors.
Licensee MDPI, Basel, Switzerland.
This article is an open access article
distributed under the terms and
conditions of the Creative Commons
Attribution (CC BY) license (https://
creativecommons.org/licenses/by/
4.0/).
Article
Feature Augmentation Based on Pixel-Wise Attention for Rail
Defect Detection
Hongjue Li
1
, Hailang Li
2
, Zhixiong Hou
2
, Haoran Song
2
, Junbo Liu
2
and Peng Dai
2,
*
1
School of Astronautics, Beihang University, Beijing 100191, China
2
Infrastructure Inspection Research Institute, China Academy of Railway Science Corporation Limited,
Beijing 100081, China
* Correspondence: daipeng_iic@protonmail.com
Abstract:
Image-based rail defect detection could be conceptually defined as an object detection
task in computer vision. However, unlike academic object detection tasks, this practical industrial
application suffers from two unique challenges, including object ambiguity and insufficient annota-
tions. To overcome these challenges, we introduce the pixel-wise attention mechanism to fully exploit
features of annotated defects, and develop a feature augmentation framework to tackle the defect
detection problem. The pixel-wise attention is conducted through a learnable pixel-level similarity
between input and support features to obtain augmented features. These augmented features contain
co-existing information from input images and multi-class support defects. The final output features
are augmented and refined by support features, thus endowing the model to distinguish between
ambiguous defect patterns based on insufficient annotated samples. Experiments on the rail defect
dataset demonstrate that feature augmentation can help balance the sensitivity and robustness of the
model. On our collected dataset with eight defected classes, our algorithm achieves 11.32% higher
mAP@.5 compared with original YOLOv5 and 4.27% higher mAP@.5 compared with Faster R-CNN.
Keywords: object detection; pixel-wise attention; feature augmentation; rail defect
1. Introduction
Discovering defects on rail is the first step for rail health maintenance and is vital
for the safe operation of high speed trains. Recent progress in high-speed photography
technology offers the possibility of capturing real-time rail images from a running train
and further paved the way to solving this practical industrial problem from the perspective
of object detection using computer vision approaches. In computer vision, object detection
approaches based on a deep convolutional neural network (CNN) have achieved great
progress in both accuracy and efficiency [
1
–
3
]. Current CNN-based object detection is
mainly built on two alternatives: two-stage methods [
4
] and one-stage methods [
5
]. Two-
stage methods achieve high detection accuracy by separately conducting region proposal
and detection process. As a comparison, one-stage detection methods are less accurate but
faster in achieving real-time detection.
In academic research, both two-stage and one-stage methods are usually trained on
large-scale benchmarks such as MS COCO [
6
] and ImageNet [
7
], and have been successfully
applied to various tasks, such as defect detection [
8
,
9
], medical detection [
10
,
11
], etc.
Nevertheless, detecting a defect on the rail image is not as easy as detecting a normal
object on natural images due to the following two reasons. First, the concerned defects
are usually tiny and ambiguous because the images are captured from a running train at
very high speed. Complex illumination environment makes defects look similar to other
non-defect patterns such as dirt or gap. Second, there is a lack of benchmark annotations of
these railway defects and some defects are difficult to distinguish from 2D images. The
insufficient and ambiguous defect images that are substantially different from natural
photos may advise against such pretrain-finetune knowledge transfer.
Appl. Sci. 2022, 12, 8006. https://doi.org/10.3390/app12168006 https://www.mdpi.com/journal/applsci
