一种基于预训练大模型的轨道缺陷检测方法及装置

专利主权项内容

1.一种基于预训练大模型的轨道缺陷检测方法，其特征在于，所述方法包括：S1、构建基于Transformer的初始掩码自编码模型；基于图像随机掩码重建的预训练方法，根据预先获取的通用图像样本对所述初始掩码自编码模型进行第一阶段预训练，得到第一中间训练模型；S2、基于图像随机掩码重建的预训练方法，根据预先获取的大规模无标注轨道图像样本对所述第一中间训练模型进行第二阶段预训练，得到第二中间训练模型；其中，图像随机掩码重建是通过借鉴预训练语言模型BERT采用的自掩码预训练机制，对需要进行预训练的图像输入进行随机的部分遮盖，要求网络对原始的图像进行重建，从而实现对于图像数据的预训练；S3、将第二中间训练模型的编码器部分与分割头进行组合，得到待训练部件检测模型，基于ViTDet的训练方法，使用轨道重要部件标注图像，对所述待训练部件检测模型进行训练，得到轨道重要部件位置检测模型；S4、根据所述轨道重要部件位置检测模型以及预先获取的轨道图像样本，得到轨道图像样本对应的重要部件位置，基于所述重要部件位置选取掩码进行图像重建，对所述第二中间训练模型进行训练，得到第三中间训练模型；S5、基于所述第三中间训练模型对预先获取的真实无病害轨道图像样本进行图像重构，生成无病害样本；S6、基于所述第一中间训练模型对所述真实无病害轨道图像样本进行图像重构，生成第一部分有病害样本；S7、对所述真实无病害轨道图像样本进行真实病害特征模拟，生成第二部分有病害样本；S8、根据所述真实无病害轨道图像样本、真实有病害轨道图像样本、生成的无病害样本、第一部分有病害样本以及第二部分有病害样本，对所述第三中间训练模型进行微调，得到训练好的轨道缺陷检测模型；S9、将待检测的轨道图像输入到训练好的轨道缺陷检测模型，得到所述待检测的轨道图像对应的轨道缺陷检测结果；其中，所述初始掩码自编码模型包括Transformer编码器和Transformer解码器；所述S1的基于图像随机掩码重建的预训练方法，根据预先获取的通用图像样本对所述初始掩码自编码模型进行第一阶段预训练，得到第一中间训练模型，包括：S11、对预先获取的通用图像样本进行切片，得到多个分片，随机选择部分分片进行掩码操作，通过1×1卷积核对剩余分片中的每个分片进行编码，在第一个分片前面加入特殊令牌作为起始分片，并给所有分片加上对应的位置编码，将编码后的分片送入所述Transformer编码器；S12、将所述Transformer编码器的输出进行归一化，通过线性层将编码器的分片的维度特征调整到所述Transformer解码器的维度特征，将调整后的分片重新放回到对应的位置，在原本被掩码的位置放入掩码令牌，对处理后的分片加上位置编码，将经过位置编码后的分片输入到所述Transformer解码器；S13、将所述Transformer解码器的输出进行归一化处理以及还原处理，得到还原图像；S14、采用平均平方误差损失函数，根据所述还原图像以及通用图像样本，对所述初始掩码自编码模型进行第一阶段预训练，得到第一中间训练模型；其中，所述S5的基于所述第三中间训练模型对预先获取的真实无病害轨道图像样本进行图像重构，生成无病害样本，包括：S51、将预先获取的真实无病害轨道图像样本切分为分片序列；S52、将所述分片序列随机排列后平均切分为L个部分，对所有分片使用1×1卷积核进行编码，在每个部分的头部分别添加特殊令牌作为起始分片，并整体叠加位置编码，得到L次重构的输入；S53、将所述L次重构的输入分别输入到所述第三中间训练模型，得到L个高质量重构图像；S54、对所述L个高质量重构图像计算平均值，得到无病害样本；其中，所述S6的基于所述第一中间训练模型对所述真实无病害轨道图像样本进行图像重构，生成第一部分有病害样本，包括：S61、将所述真实无病害轨道图像样本切分成分片序列；S62、将所述分片序列随机排列后平均切分为L个部分，对所有分片使用1×1卷积核进行编码，在每个部分头部分别添加特殊令牌作为起始分片，并整体叠加位置编码，得到L次重构的输入；S63、将所述L次重构的输入分别输入到所述第一中间训练模型，得到L个低质量重构图像；S64、对所述L个低质量重构图像计算平均值，得到第一类有病害样本；S65、将所述第一类有病害样本随机配对，将所述第一类有病害样本的像素颜色值两两进行加和计算，并计算和值的平均值，通过第三中间训练模型进行一次重构，得到第二类有病害样本；其中，所述S7的对所述真实无病害轨道图像样本进行真实病害特征模拟，生成第二部分有病害样本，包括：S71、对所述真实无病害轨道图像样本添加泊松噪声，得到轨道擦伤类有病害样本；S72、在所述真实无病害轨道图像样本的扣件位置扣除部分图像，替换为噪声图像，并将扣件图像翻转并粘贴到所述真实无病害轨道图像样本的非轨道区域或扣件区域，得到扣件丢失、异物类有病害样本；S73、在所述真实无病害轨道图像样本的轨道非扣件位置生成多条模拟裂纹曲线，得到道床类有病害样本。。 （来自 马克数据网）