基于注意力YOLOv5模型的自动水果识别

Automatic Fruit Recognition Based on Attention YOLOv5 Model

近年来,人工智能在各个领域有着广泛的应用.针对超市及菜市场人工称重操作耗时、计价流程繁杂的问题,本文提出一种基于注意力YOLOv5模型的水果自动识别算法.首先,为了提升仅有局部特征不同,全局特征相似水果的识别准确率,本文在YOLOv5的SPP(spatial pyramid pooling)层后增加SENet(squeeze-and-excitation networks),采用注意力机制自动学习每个特征通道的重要程度,进而按照重要程度强化对水果识别任务有用的特征并抑制没有用的特征;其次,针对水果识别预测框与目标框重叠时,GIOU不能准确表达边框重合关系问题,本文将原有的边框回归损失函数GIOU替换为CIOU,同时考虑目标框与预测框的高宽比和中心点之间的关系,从而使水果预测框更加接近真实框,提升预测精度.实验结果表明,改进后的模型在常见场景下水果识别能力有明显提升,平均精度mAP达99.10%,识别速度FPS达到82,能够满足实际应用需要.

In recent years,artificial intelligence has been widely used in various fields.To address time-consuming manual weighing and complicated pricing procedures in supermarkets and vegetable markets,this study proposes an automatic fruit recognition model based on attention YOLOv5.First,to improve the recognition accuracy of fruits with different local features but similar global features,the study adds squeeze-and-excitation networks(SENet)after the spatial pyramid pooling(SPP)layer of YOLOv5 and uses the attention mechanism to automatically learn the importance of each feature channel.Further,the useful features for fruit recognition tasks according to the importance are strengthened and those useless are suppressed.Second,when the fruit recognition prediction frame overlaps the target frame,GIOU cannot accurately express the overlapping relationship of the frames.In response,this study replaces the original frame regression loss function GIOU with CIOU and considers the relationships of aspect ratio and center point between the target frame and the prediction frame.In this way,the fruit prediction frame is closer to the real frame,and thereby the prediction accuracy is improved.Experimental results show that the improved model has significantly improved fruit recognition ability in common scenarios with a mean average precision(mAP)of 99.10%and a recognition speed of 82 FPS,which can meet the needs of practical applications.