Efficient Lightweight Railway Track Segmentation Network for Resource-Constrained Platforms with TensorRT

Abstract

Accurate and rapid railway track segmentation is the fundamental for foreign object intrusion detection, inspection, online monitoring, and non-destructive assessment of transportation infrastructure. Recently, vision-based track segmentation algorithms have demonstrated strong performance. However, most existing models struggle to meet the real-time requirements on resource-constrained edge devices. Considering this challenge, we propose an edge-enabled real-time railway track segmentation algorithm, which is optimized to be suitable for edge applications by optimizing the network structure and quantizing the model after training. Initially, Ghost convolution is introduced to reduce the complexity of the backbone, thereby achieving the extraction of key information of the interested region at a lower cost. To further reduce the model complexity and calculation, a new lightweight detection head is proposed to achieve the best balance between accuracy and efficiency. Subsequently, we introduce quantization techniques to map the model’s floating-point weights and activation values into lower bit-width fixed-point representations, reducing computational demands and memory footprint, ultimately accelerating the model’s inference. Finally, we draw inspiration from GPU parallel programming principles to expedite the pre-processing and post-processing stages of the algorithm by doing parallel processing. The approach is evaluated with public and challenging dataset RailSem19 and tested on Jetson Nano. Experimental results demonstrate that our enhanced algorithm achieves an accuracy level of 83.3% alongside with 25 FPS inference speed when the input size is 480 × 480. The code can be found at: https://github.com/ccl-1/light-yolov8-seg-quantization-tensorrt.