YOLOv5屏蔽区域检测以及选择区域检测
- 前期准备
- labelme选择mask区域
- 代码改动
前期准备
思路就是通过一个mask掩膜,对我们想要屏蔽或者选择的区域进行遮挡处理,在推理的时候,将有mask掩膜的图像输入,将最后的结果显示在原始图像上,即完成了屏蔽区域检测。
labelme选择mask区域
标注好我们想检测的区域之后,我这里将标注区域命名成了mask,然后我们点击保存导出json文件:
然后我们再打开终端,cd到json所在的文件夹路径下,输入如下命令:
labelme_json_to_dataset xxx.json
这里有的兄弟会遇到一些报错,主要原因还是环境的版本问题,可以在网上查到,替换一下对应版本的库就行了。
我们就会在这个路径下获得一个文件夹:
打开文件夹之后会看到里面有这么五个文件,我们需要的就是label.png,其他的不用管:
到这里,我们的前期准备工作就完成了。
代码改动
打开yolov5的项目,进入utils/datasets.py中,在代码的开头加上读取我们导出的label.png的代码:
读入我们导出的mask图片文件,然后将其转成灰度图,然后使用opencv的阈值函数,将大于0的像素值全部变成255,其余的都为0,得到一张mask的二值图。
然后我们再看到yolov5自带的推理代码中,即detect.py文件。看到代码中的数据加载部分:
我们这里是对图像进行推理,因此我们加载的是图像,不进入webcam中,因此我们加载推理数据的方法是LoadImages这个类,cltr+左键进入这个类中,在这个位置下加入以下代码:
这段代码主要实现了两个功能,其中flag参数决定了我们是否使用mask掩膜对目标进行遮挡检测,reverse参数决定了我们是进行选择区域检测还是屏蔽区域检测(其实这两者是一个取反的操作)。
这里我的参数是flag=1,reverse=0,因此此时如果我运行程序的话,我们会看到的结果是只对我们刚才框定的mask区域的目标进行检测,而对其他地方并不会检测:
可以看到,和我们预想的一样。
我们再将参数改成flag=1, reverse=1,看看结果:
此时我们检测的区域是除了mask区域的以外区域,没问题。
那么最后我们将两个参数都设置为0:
没有任何问题!
后面我也会使用openvino和tensorRT在c++进行推理部署,实现选择区域或屏蔽区域检测的功能,欢迎大家关注。