Ensemble model for rail surface defects detection
Hailang Li, Fan Wang, Junbo Liu, Haoran Song, Zhixiong Hou, Peng Dai
Abstract
The detection of rail surface defects is vital for high-speed rail maintenance and management. The CNN-based computer vision approach has been proved to be a strong detection tool widely used in various industrial scenarios. However, the CNN-based detection models are diverse from each other in performance, and most of them require sufficient training samples to achieve high detection performance. Selecting an appropriate model and tuning it with insufficient annotated rail defect images is time-consuming and tedious. To overcome this challenge, motivated by ensemble learning that uses multiple learning algorithms to obtain better predictive performance, we develop an ensemble framework for industrialized rail defect detection. We apply multiple backbone networks individually to obtain features, and mix them in a binary format to obtain better and more diverse sub-networks. Image augmentation and feature augmentation operations are randomly applied to further make the model more diverse. A shared feature pyramid network is adopted to reduce model parameters as well as computation cost. Experimental results substantiate that the approach outperforms single detecting architecture in our specified rail defect task. On the collected dataset with 8 defect classes, our algorithm achieves 7.4% higher mAP.5 compared with YOLOv5 and 2.8% higher mAP.5 compared with Faster R-CNN.