Multi-Task Perception Algorithm for Rail Transit Scenarios Based on Triplet Attention

Aiming at the challenges of insufficient object detection accuracy and low detection speeds, and the pursuit of an accuracy-speed balance in environmental perception within rail transit scenarios, this paper proposes a multi-task perception model that features simultaneous detection and segmentation. Firstly, to address the current lacking of datasets in rail transit scenarios, a railway environmental perception dataset was created using sample images captured from track scenes for model training and evaluation. Secondly, to overcome the limited data size and insufficient diversity in the dataset, Differential RandAugment was applied to the dataset for data augmentation. Moreover, by introducing differentiable parameters to optimize the data augmentation strategy, more diverse training samples were generated. Finally, a triplet attention mechanism was incorporated to address the original YOLOP algorithm's deficiency of losing key information during feature extraction. By weighting attention in three directions of the feature map, the original model's capacities to extract key information were improved, thereby enhancing its performance in small object detection tasks. Through validation using datasets from actual scenes, the model exhibited significant improvements in small object detection, detection speed, and environmental adaptability. For the self-created datasets, the detection accuracy of objects across all categories reached 0.945, the recall rate was 0.853, and the average accuracy was 0.900. In terms of small object detection, the average accuracy reached 0.837, and the detection speed reached 23.5 frames per second. These study outcomes provide an effective solution for object detection research in the field of rail transit and highlight important prospects for practical applications.