基于改进YOLOv5_OBB的中华绒螯蟹旋转目标检测

袁红春, 白宝来, 陶 磊

(上海海洋大学信息学院,上海 201306)

中华绒螯蟹(Eriocheirsinernsis),俗称河蟹、毛蟹或大闸蟹,是我国特有的水产珍品,具有极高的经济及医用价值[1]。由于缺乏规范化、标准化管理措施,目前我国中华绒螯蟹养殖的自动化程度低[2]。通过目标检测方法精确地检测与定位中华绒螯蟹,可以为精确投放饵料、质量估计、形态学分析和行为分析等提供指导方案,实现自动化养殖。Faster RCNN[3]、YOLO系列[4-7]和FCOS[8]检测方法是通过水平检测框来定位目标,检测精度高,但不适用于中华绒螯蟹目标检测。由于中华绒螯蟹的活动方向是任意的,水平检测框会造成目标范围不准确(包含大量背景区域),且中华绒螯蟹在摄食、避险等情况下水平检测框之间会出现强烈重叠,从而降低检测精度,无法获得目标运动方向的信息[9]。使用具有角度信息的旋转目标检测方法可以克服水平检测框定位目标检测方法的局限性,例如:Yang et al[10]在Faster RCNN基础上添加一个角度预测分支构成旋转区域卷积神经网络(rotational region CNN,R2CNN)模型,但该模型存在检测速度慢、角度回归不准确等问题;Xu et al[11]采用一种新的旋转框表示方法建立旋转目标检测(gliding vertex )模型,但这一模型存在检测框预测不准确、检测速度慢等问题;Yang et al[12]采用从粗粒度到细粒度逐步回归的方法(R3Det)来快速、准确地检测目标,但该方法存在训练速度慢等问题;Zhou et al[13]通过在模型中使用极坐标系建立一种无锚框的旋转检测方法(rotation-equivariant detector, P-RSDet),但检测精度不高。由此可知,现有的旋转目标检测方法难以同时兼顾检测精度、训练速度和检测速度。

YOLOv5_OBB是在YOLOv5[14]基础上使用环形平滑标签(circular smooth label, CSL)技术[15]完成旋转目标的检测的,具有轻量化和实时性等优点,但检测准确率不够高。针对中华绒螯蟹数据集存在小目标和密集目标的情况,本研究拟对YOLOv5_OBB模型进行改进,基于中华绒螯蟹数据集验证各种改进措施的检测精度,实现对中华绒螯蟹的快速检测和准确定位。

1 材料与方法

1.1 图像采集与处理

中华绒螯蟹数据集图像拍摄于2022年1月11日至15日,在宽38 cm、长56 cm的收纳筐中放入8～23只中华绒螯蟹、3～7个石头和1～4块水草,以模拟中华绒螯蟹天然养殖环境。图像尺寸为3 248×2 160,共1万余张。通过筛选去除相似图片、水体浑浊图片和模糊图片,保留5 000张图片。

使用rolabelimg工具对数据集进行标注,将每张图片上的所有目标,包括中华绒螯蟹、石头和水草,用最小任意方向的矩形框标注。为了缩短图像标注的时间,降低成本,设计一种半自动数据标注方法:(1)使用旋转标记工具(rolabelimg)工具标注少量数据集,记为已标注数据集A;(2)使用数据集A训练模型M,选择部分未标注的数据集B,通过模型M检测数据集B中的目标,保存检测框信息;(3)在rolabelimg工具中手动调整检测框位置,并将数据集B放入已标注的数据集A中。如果数据集U中所有图像均标注完毕,即A=U,则结束半自动标注;否则,重复上述(2)、(3)步骤,直至标注完毕。

1.2 YOLOv5_OBB网络结构

YOLOv5_OBB旋转目标检测算法以YOLOv5目标检测算法为基础,而YOLOv5目前共有5个版本,占用的计算资源和检测准确率有所不同。将旋转目标检测框的角度回归任务转换为角度分类任务,完成多方向目标检测。网络模型由输入端、主干网络(backbone)、颈部模块(neck)和输出端4个部分组成。其中,输入端、主干网络和Neck模块依然沿用YOLOv5的网络结构,在输出端添加角度预测信息,包括损失计算和目标检测框筛选。在YOLOv5_OBB模型中使用一种没有边界问题的角度预测方法,即将角度预测作为分类问题。采用CSL技术完成角度分类任务。CSL的表达式为:

(1)

式中:g(x)为窗口函数;r为窗口半径;θ为当前检测框的角度。

1.3 改进策略

由于中华绒螯蟹个体较小,在拍摄距离较大时数据集出现大量小目标。在YOLOv5_OBB中,随着数据流在网络中不断加深,目标特征表现逐渐减弱,会造成小目标漏检。本研究通过在空间金字塔池化(spatial pyramid pooling -fast,SPPF)模块之前添加通道关注(efficient channel attention, ECA)[16]模块,增强网络对重要通道的学习能力,提高模型对小目标的识别率。此外,本研究采用加权双向特征金字塔网络(bidirectional feature pyramid network, BiFPN)多尺度特征融合,通过双向跨尺度连接和加权特征融合得到包含高层和底层语义信息的全局特征,加强特征融合,提高检测速度。

YOLOv5_OBB是一种基于锚框的目标检测算法,在不同尺度的特征图上每个空间点安排3个先验框,从而产生大量负样本,正负样本之间不平衡,在密集目标检测中会出现漏检或误检,故以变焦损失(varifocal loss, VFL)[17]代替焦点损失(focal loss)[18]。

(2)

式中:p∈[0,1],为正样本预测概率;pγ为负样本调制因子;q为目标分数;对于正样本,q为边界框与真实值(ground truth, GT)之间的交并比(intersection over union, IOU)值。为了平衡正样本和负样本间的损失,在负样本项中添加了一个可调整的比例因子α。

通过系统试验,对不同的改进策略进行验证,评估这些策略对模型性能的提升效果,论证这些策略在提高模型性能方面的有效性。

1.4 试验设置

操作系统的配置为Ubuntu16.04.4 LTS系统,深度学习框架为Pytorch 1.11.0和CUDA 11.2。计算机硬件配置: AMD 5950X CPU,GeForce RTX 3090 24 GB GPU,16 GB内存。共拍摄5 000张图像,按照8∶2划分训练集和测试集。超参数设置:批大小(batch size)为16,训练周期(epochs)为200,初始学习率为0.01。中华绒螯蟹具有一定的长宽比,图像中特定角度明显,使用CSL时窗口函数为高斯函数,窗口半径设置为6。

采用精确率P(precision)、召回率R(recall)和平均精度均值(mean average precision, mAP)来评价模型的准确性。

2 结果与分析

2.1 基准模型的选择

YOLOv5_OBB是基于YOLOv5改建的,而YOLOv5共有5个版本。添加通道注意力(efficient channel attention, ECA)模块后,以不同版本YOLOv5为基准模型,采用BiFPN和变焦损失计算精确率、召回率、平均精度均值、权重,并进行对比,结果如表1所示。

表1 YOLOv5基准模型的对比结果Table 1 Comparison of YOLOv5 benchmark model

由表1可知,以YOLOv5n为基准模型的召回率与最优结果仅相差0.2%,而精确率、平均精度均值、模型大小和计算量在不同模型中的表现均好,因此,选择YOLOv5n作为YOLOv5_OBB的基准模型。

2.2 消融试验

通过消融试验验证各个改进策略对模型的优化效果。基于自建的中华绒螯蟹数据集,以精确率、召回率和平均精度均值作为评价指标,验证改进效果,结果如表2所示。

表2 消融试验结果1) Table 2 Results of ablation experiment

由表2可知,本研究采用的改进策略可提高中华绒螯蟹多方向目标检测精度,效果最佳。

2.3 改进前后模型检测精度的对比

通过检测结果和可视化结果分析改进后YOLOv5_OBB模型对目标检测精度的提升效果。在相同配置下将改进后YOLOv5_OBB模型和原始YOLOv5_OBB模型训练200轮,得到精确率、召回率和平均精度均值的对比曲线(图1)。由图1可知:前25轮两种模型都迅速收敛;75轮之后模型趋于稳定;两种模型的训练效果均没有出现过拟合和欠拟合现象,且改进后YOLOv5_OBB模型的各项评价指标均有所改善。

图1 不同评价指标的对比曲线Fig.1 Comparison curves of different evaluation indicators

小目标检测结果如图2A、2B所示,原YOLOv5_OBB对于右上角小目标存在漏检现象,而改进后的YOLOv5_OBB成功检测出此小目标,表明改进后模型的检测精度更高。密集目标检测结果如图2C、2D所示,原YOLOv5_OBB存在3处误检(图2C中的黑色标注),而改进后的YOLOv5_OBB没有出现误检,表明改进后模型的检测准确率提高。

图2 可视化结果Fig.2 Visualized results

基于自建中华绒螯蟹数据集,原YOLOv5_OBB模型和改进的YOLOv5_OBB模型在精确率、召回率、平均精度均值、权重和运算量方面的对比结果如表3所示。从表3可知,改进后的模型在权重增加0.1 MB、10亿次浮点运算量增加0.1 s的情况下,精确率提高了1.2%,召回率提高了1.9%,平均精度均值提高了1.3%。可视化结果(图2)和数据分析结果(表3)均验证了改进策略的可行性。

表3 原始模型和改进后模型的对比Table 3 Comparison between the original and improved model

2.4 不同模型的对比试验

使用Rotated FasterRCNN、RoI Transformer[19]、ReDet[20]和Oriented R-CNN[21]多方向目标检测模型与本文改进的YOLOv5_OBB模型进行对比,使用自建中华绒螯蟹数据集进行训练和测试,得到各模型在精确率、召回率、平均精度均值和权重方面的对比结果(表4)。

表4 各模型对比结果Table 4 Comparison of each model

由表4可知,基于中华绒螯蟹数据集改进的YOLOv5_OBB模型的精确率、平均精度均值和模型权重均优于其他多方向目标检测算法,其召回率略低于RoI Transformer,但平均精度均值比RoI Transformer提高11.1%,模型大小仅4.4 MB。这表明改进的YOLOv5_OBB模型的综合性能更佳。

3 小结

本研究基于自建的中华绒螯蟹数据集,提出了一种基于改进 YOLOv5_OBB的旋转目标检测方法,首先引入了高效通道注意模块和BiFPN网络结构,采用VariFocal Loss进行训练,在YOLOv5_OBB的主干网络中进行特征融合,实现双向跨尺度连接和加权特征融合。以精确度、召回率和平均精度均值作为评价指标进行检测,结果表明:改进的YOLOv5_OBB模型的检测精度比原模型有所提高,克服了原模型存在的小目标和密集目标漏检或误检等问题;改进的YOLOv5_OBB模型具有轻量化、检测精度高等优点。