PDSAM: Prompt-Driven SAM for Track Defect Detection

Abstract

The track defect detection is critical for ensuring the safety and reliability of railway systems. Existing machine vision-based approaches are hindered by three key issues: high-time complexity stemming from end-to-end network training, limited availability of training data (with only a few hundred labeled images), and suboptimal prediction precision. To address these challenges, this article introduces the prompt-driven segment anything model (PDSAM), a novel image semantic segmentation framework that introduces a paradigm shift in problem formulation. The core contribution lies in reformulating the segmentation task as a prompt generation problem, which offers two correlated advantages. First, a simplified prompt generation network reduces both training time and data requirements compared with standalone segmentation networks. Second, an upscaling and visual prompting technique restores spatial resolution and mitigates the risk of local optima in feature optimization, enabling more precise and fine-grained segmentation outputs. Experimental evaluations on benchmark datasets demonstrate that PDSAM outperforms state-of-the-art methods in both prediction accuracy and computational efficiency for railway track defect detection. The proposed framework’s source code and pretrained models (PTMs) are publicly available to facilitate reproducibility and further research, accessible at: https://github.com/FreddyDylan/PDSAM/