基于U-net和YOLOv4-tiny的锁孔中心定位算法*

谢孟添 刁云峰 程文明 唐 鑫 肖施睿

1西南交通大学机械学院 成都 610031 2轨道交通运维技术与装备四川省重点实验室 成都 610031

0 引言

近年来，随着我国经济的飞速发展，铁路集装箱货运需求量大大增加。为了应对运量的上升，起重机的自动化是至关重要的一步。目前，在集装箱装卸工作中，起重机操作人员需要在地面员工的辅助下进行集装箱的装卸，大大降低了集装箱的装卸效率。更重要的是，长时间的手工工作可能导致工作事故。调查显示[1]，2015年～2020年，人为因素造成的事故占到集装箱码头所有事故的一半以上。因此，目前迫切需要一种高效、安全的自动化方法来替代人工操作。

随着计算机视觉领域的发展，基于图像视觉的识别和定位技术逐渐运用在了集装箱领域[2-4]。Yoon H J等[5]利用双目视觉系统实现了集装箱高度距离的检测，但该方法计算复杂度高，对环境光照敏感，且无法定位集装箱的横向与纵向位置，故难以应用于实际的集装箱货场。李岩等[6]提出了一种基于机器学习的集装箱锁孔检测和定位方法，该方法运用改进的HOG算法提取锁孔特征，并利用SVM分类器实现了锁孔的实时检测和定位，但检测准确率并未达到实际应用的水平。文献[7]利用传统的边缘检测算法来提取锁孔的边缘，并在理想的实验条件下通过微缩集装箱模型进行了测试，得到了较好的对位精度。但实际条件下的集装箱通常会受到环境腐蚀，锁孔周围的锈迹会大大降低边缘检测算法的精度。Li Y等[8]利用目标检测的方法对集装箱锁孔进行检测，通过检测框的中心位置来进行锁孔定位。虽然该方法取得了较好的检测精度，但仅通过检测框中心难以计算出准确的锁孔中心。

为了解决上述问题，本文提出了一种基于图像分割的锁孔中心定位方法。该方法分为2个阶段：第1阶段利用基于深度学习的目标检测方法对锁孔进行检测，提取锁孔所在的矩形区域；第2阶段采用图像分割方法对矩形框进行预测，从而精确计算出锁孔区域及其中心位置。同时，在铁路集装箱货场进行数据采集，并在该数据集下进行了实验评估，证实了本文提出方法的有效性。

1 算法

1.1 算法整体框架

本文采用两段式的方法对锁孔中心进行定位。因为基于图像分割的方法在处理大分辨率的图像时会产生较大计算量，同时集装箱货场环境复杂，难以保证分割精度。因此，在算法的第1步中，采用基于深度学习的目标检测算法YOLOv4-tiny[9]来对原始图片进行检测，利用矩形框将锁孔从原始图像中提取出来。第2步，利用基于U-net[10]的图像分割的方法对提取的检测框进行预测，从而准确地得出锁孔区域，并根据图像处理的方法计算出区域中心的准确位置。其算法流程如图1所示。

图1 算法流程图

1.2 基于YOLOv4-tiny的锁孔检测

本文选用了YOLOv4-tiny算法作为第1阶段的目标检测算法。因为与YOLOv4相比，YOLOv4-tiny结构更轻量化，运算速度更快。此外，相比于two-stage目标检测算法和one-stage目标检测算法，YOLOv4-tiny在保持检测准确率的基础上还能达到较快的运算速度，更利于集装箱货场这类需求实时检测的应用环境。其网络结构如图2所示。

图2 YOLOv4-tiny网络结构

其网络结构主要由2部分组成，骨干网络和主干网络。其中，骨干网络为CSPDarknet53-tiny。图片输入后，骨干网络会对其进行特征提取和融合，以生成不同深度的特征。随后，主干网络会对尺寸为13×13的深层特征282进行上采样操作，并与尺寸为26×26的浅层特征17进行特征融合，丰富浅层特征的语义信息，提升网络检测能力。融合的特征在头部网络中进行预测，其中尺寸为26×26的特征用于小目标的预测，13×13的特征用于大目标的预测。根据生成预测目标的置信度和边界框位置，通过置信度阈值和非极大值抑制来筛选出准确的目标，得到的锁孔区域检测结果如图3所示。

1.3 基于U-net的锁孔中心计算

在深度学习领域主要有2类图像分割算法：1）基于卷积神经网络（CNN）的方法，该方法通过卷积运算进行图像特征的提取，但在卷积过程中容易损失图像信息。且由于其类别预测采用的全连接层，故难以做到像素级的精确分割。2）利用反卷积层替代全连接层，采用全卷积神经网络（FCN）对输入图片进行处理。输出与原图尺寸相同的分割图，实现对图像的精确分割。

本文选用了U-net模型作为第2阶段的图像分割模型，因为U-net模型结构简单，同时能够在较小的数据集下取得较好的分割效果，大大降低了实际应用时数据筹备的时间。该网络是经典的全卷积神经网络（见图4），网络采用3×3的卷积核对输入图片进行一系列的卷积和池化，得到4个不同尺寸的特征，然后对小尺寸的深层特征进行上采样操作，同时采用级联的方式与同尺寸的浅层特征进行拼接，使得深层特征和浅层特征得到了充分的利用，大大提高了图像分割效果。

图4 U-net模型框架

将第1阶段提取的锁孔矩形框输入U-net图像分割算法，可得到精确分割后的锁孔区域灰度图。由于锁孔区域是规则的几何形状，故通过拟合该区域的形心即可精确计算出锁孔的实际中心。

不同行业中的企业对人才需求的类别、特质不同，同一行业中的企业对人才需求的时间、能力也有所不同，如何抽取企业需求，分解能力指标，区分能力度，形成可初初始化的需求库，是构建共享平台的关键问题之一，这需要对行业和企业进行深度调研，有针对性的构建人才需求平台。

本文采用基于图像处理的方法对分割后的灰度图进行二值化处理，精确计算出锁孔内所有白色像素点的坐标，并通过式（1）计算出锁孔中心在图像中的像素坐标。

式中：xa、ya分别为锁孔中心的横、纵坐标，n为白色像素点的总数,xi、yi为第i个白色像素点的横、纵坐标。

图5 U-net锁孔中心计算效果图

2 实验

2.1 数据集及标注

由于目前在集装箱锁孔检测领域并没有公共的数据集，为了保证算法检测的准确度和有效性，本文在铁路集装箱货场进行了数据采集，同时为了增加数据环境的多样性，提高模型的泛化能力，在晴天、多云以及雨天的环境下进行数据采集，共采集1 409张集装箱图片和5 237锁孔图片。

对于集装箱图片，本文采用LabelImg标注软件将其转换成VOC2007[11]数据集格式，并按照7:2:1的比例将其划分为训练集、验证集和测试集，以此构成第1阶段的数据集。为了保证算法在锁孔检测时矩形框能完整的提取锁孔，在标注过程中，标注框略应大于锁孔边缘。

对于锁孔图片，本文采用二阶贝塞尔曲线[12]对锁孔边缘进行精确的标注，根据获取的曲线坐标重新绘制锁孔边缘，并分别用白色像素和黑色像素填充锁孔边缘的内部和外部。最终将其作为第2阶段的数据集。完整的数据集处理过程如图6所示。图中黄色边缘线为二阶贝塞尔曲线标注的锁孔边缘。

图6 U-net标注的锁孔数据集

2.2 实验环境和实验参数

本实验基于Pytorch深度学习平台进行，操作系统为Ubuntu18.04，CPU为AMD Ryzen 5 3500X，GPU为NVIDIA GTX 1060SUPER 6G。第1阶段的模型训练采用ADAM优化器，训练批尺寸为16，初始学习率设置为0.001，并采用学习率衰减的方式进行训练，一共迭代100次。第2阶段的模型训练同样采用Adam优化器，网络分辨率为256×256，训练批尺寸为2，学习率初始化为10-5，一共循环100次。

2.3 评价指标

为了评价算法能否准确地从集装箱图片中检测锁孔，本文采用准确率acc作为锁孔检测的评价指标。该指标代表所有预测为锁孔的对象N中正确预测T的比例，计算公式为

2.3.2 锁孔中心检测评价指标

本文选用了实际中心与预测中心像素坐标的欧氏距离dis作为第2阶段锁孔中心检测的评价指标。当距离越小时，代表预测中心与实际中心越接近，算法精度越高。

式中：xt、yt为像素坐标系下锁孔实际中心的横坐标与纵坐标，xp、yp为像素坐标系下锁孔预测中心的横坐标与纵坐标。

2.4 实验结果与分析

2.4.1 锁孔检测实验结果

为了验证在第一阶段的锁孔检测操作中，YOLOv4-tiny算法能够准确快速地提取锁孔，本文将YOLOv4-tiny与Faster R-CNN[13]和RetinaNet[14]算法进行对比。在相同数据集下的检测结果如表1所示。

表1 锁孔检测算法对比

由表1可看出，RetinaNet的锁孔检测平均成功率和YOLOv4-tiny较为接近，但本文选用的YOLOv4-tiny算法在检测速度上拥有明显的优势，完全能满足第一阶段锁孔检测的实时性需求。

2.4.2 锁孔中心定位检测结果

本文模型在相同的数据集和评价标准下与基于目标检测的方法和DIAO等[7]提出的方法做了精度对比。其中，基于目标检测的方法选用了YOLOv4-tiny作为对比，因为该方法是目前实时检测算法中检测精度较高的算法，能保证本次对比实验的公平性。实验对比的可视化结果如图7所示。其中红圈和黄圈中心分别为实际锁孔中心和预测锁孔中心，其半径均为5像素值。当两圈重合面积较大时代表预测中心与实际中心更接近。

图7 不同算法效果展示

从对比结果中可以看出DIAO等提出的方法在锁孔锈蚀较为严重时检测精度很差，完全偏离了实际的锁孔中心。即使在检测表面状况良好的锁孔时，其检测精度也无法达到实际应用需求。此外，由于YOLOv4-tiny算法以检测框中心近似替代实际锁孔中心，故中心定位精度较差。而本文方法所拟合的锁孔中心几乎与实际中心重叠。

同时，在相同数据集下对不同方法的中心误差进行计算，统计了3种方法在不同误差内的图片比例。对比结果如表2所示。

表2 不同算法结果对比

由表2可以看出，本文提出的方法在中心定位精度方面，远超另外两类算法，且在3像素误差内的检测成功率可达94%，完全满足实际应用的需求。

3 结论

快速准确的锁孔中心定位算法是实现集装箱装卸自动化的关键技术之一，本文提出了一种基于深度学习的解决方案。该方案包括：1）利用YOLOv4-tiny算法对集装箱图片进行检测，获取锁孔区域的矩形框；2）采用基于U-net的图像分割模型对锁孔矩形框进行分割，实现锁孔区域的提取；3）应用图像处理的方法对分割的锁孔区域进行计算，从而拟合出锁孔区域的中心。结果表明：本方案在实际的集装箱数据集上的识别精度远高于其他锁孔中心定位算法，能够符合集装箱锁孔中心实时定位的需求。