Toward Better Generalization: Shape Feature-Enhanced Fastener Defect Detection With Diffusion Model

Fastener defect detection is important for making sure that railroads operate reliably. A robust fastener detector should perform well in most railway track circumstances, especially in an unseen environment. However, with a limited fastener dataset, this requirement remains challenging due to a large domain shift between the available dataset and unseen fasteners. This article aims to build a robust fastener defect detection architecture by instance segmentation, where only a single domain of data is available for training. To this end, we propose a novel end-to-end fastener detection model that can be divided into a base detector and two key components: a mask prediction head (MPH) and a mask refinement head (MRH). MPH generates the initial masks based on the predicted boxes supplied from the base detector. By combining the shape priors of the fastener, a simple yet effective shape enhancement module is introduced to encourage MPH to better capture the shape feature. MRH is presented to improve the initial masks with an iteratively refined process, which is achieved by a denoising diffusion model (DDM). Meanwhile, a shape guidance module (SGM) is designed to enhance the recovery capability of DDM. Experiments on the track fastener dataset revealed that the proposed heads, i.e., MPH and MRH, can be equipped with various detection frameworks and have shown consistent generalization improvements on many base detectors. Moreover, our approach significantly boosts the detection performance on different unseen sites.