OAK和 yolox完美结合 ,实现对足球,篮球,排球的识别和测量

OAK和 yolox完美结合 ,实现对足球,篮球,排球的识别和测量

背景

全国大学生电子设计竞赛需要用机器视觉的方法识别足球,篮球,排球。并且能够测量足球、篮球、排球的大小和距离。这需要深度视觉、神经网络的共同配合,目前的openMV和K210等相机都无法很好的实现以上功能。同时还需要有足球、篮球、排球的大量数据集,才能训练出这样的球类分类的神经网络,实现起来有一定难度。

目前的解决方案

目前没有现成的球类分类识别模型,必须自己去标注制作三种球类的数据集。训练集目前常用的开源目标检测数据标注工具如labelimg完全依赖于人力标注:标注人员需要同时标注目标的位置和目标的类别。在图片数量较多、目标较密集的情况下这种标注的工作量巨大,标注员在长时间进行标注后或标注较小的目标时会出现标注精度下降的情况,且在待选类别过多时标注也颇费时间。

我们创新做法是提供了一种人机协同的图像目标检测数据半自动标注方法:目标检测模型辅助用户进行标注,降低标注负担、提高标注容错率;用户对标注结果进行监督,对少数错误结果进行校正。生成三种球类的数据集后,我们采用最新的yolox神经网络进行训练生成模型后,将模型部署到OAK相机里,就可以同时实现球类分类和球类大小和距离测量了。

实现我们做法的总体流程介绍

实现我们方法的详细步骤:

1 找到 yolox 的球类识别模型

首先在yolox网站上下载预训练模型,搭建yolox训练环境 yolox手把手教学网址

使用 yolox COCO 预训练模型,使用其中的 sports ball 类别

2 分别对三种球类的视频、图像进行识别,依据结果完成数据标注,依照 Pascal VOC 格式生成数据集

PASCAL VOC数据集目录结构:

   VOC2021
  ├── Annotations
  ├── ImageSets
  │   └── Main
  └── JPEGImages
  • Annotations目录下放的是标注文件,xml格式,这是用于目标检测任务的一个总的标注集合,里面存放有大而全的xml文件
  • JPEGImages目录下是所有的jpg图片
  • ImageSets/Main目录中是训练集、验证集、测试集这里面全是train.txt、val.txt、trainval.txt、test.txt。这些文件里的每一行都是个文件名,其有2层含义,一方面其对应的图片来自JPEGImages里的同名文件,另一方面其对应的标注结果来自Annotations里的同名文件。

3 将数据集送入yolox重新训练

4 保存训练好的yolox模型

5 将模型转成onnx模型 export_onnx.py

dummy_input = torch.randn(1, 3, exp.test_size[0], exp.test_size[1])
torch.onnx._export(
   model,
   dummy_input,
   args.output_name,
   input_names=[args.input],
   output_names=[args.output],
   opset_version=args.opset,
)

6 将模型转成 openvino 模型 onnx -> OpenVino

mo.py --input_model yolox.onnx --data_type FP16 --reverse_input_channels --scale_values "[58.395, 57.12 , 57.375]" --mean_values "[123.675, 116.28 , 103.53]" 

7 将模型转成blob文件 OpenVino -> OAK

$MYRIAD_COMPILE -m yolox.xml -ip U8 -VPU_NUMBER_OF_SHAVES 6 -VPU_NUMBER_OF_CMX_SLICES 6 -o yolox_openvino_2021.4_6shave.blob

8 DepthAI部署测试

Depthai 的部署可以参照 OAKChina/yolox

  • 修改 label_classes.py 文件 ,修改 COCO_CLASSES 内容匹配所训练类型与顺序
  • 修改 main.py 文件,修改 blob 文件路径和修改模型输入大小

检验我们的结果