文章目录[隐藏]
前言引入
想必大家在做yolo目标检测的时候很容易会遇上这样的问题,图片中的某个部分其实我并不想去检测里面的物品,它无关紧要…
正如这个转角,我的本意只想检测前方通道的目标,这个转角应该交给另外的摄像头去识别,所以我不希望因为这里出现了一小部分物体,导致整个检测结果受到影响,因此,想要让yolo不检测这一部分区域…
所以为了实现这个功能,我们可以用一个mask覆盖掉这个区域,让yolo在检测的时候,假装“看不到这个区域”,就能成功实现yolo检测忽略这一块内容的目的了;
先演示一下效果,为了更加直观,我用了yolov5给的测试样例图片进行测试,下面是正常情况下,检测zidane.jpg
的结果图(这里我使用了yolov5s6.pt
模型进行检验):
可以看到是非常标准的一个识别结果,然后我加上了一个mask掩模图,此图如下:
红色区域就是我想要让yolo去识别的部分,黑色区域则不让yolo进行识别,然后我们来看一下检测结果:
可以看到非常明显,右边的人不再被识别到,而且左边的人的手也因为区域的限制没有被完全选中,证明了我们这个方法的可行性;如果这个演示示例符合你的预期功能,那可以继续往下看hhh
mask图的制作
所以,其实最重要的就是把mask图给制作出来,这里也非常容易,我们只需要安装一个labelme
软件即可,在anaconda的环境下执行命令:
#labelme需要pyqt5的支持
pip install pyqt5
pip install labelme==3.16.2
这里指定了labelme的版本,对于之后的大数据处理有用处,但是本文暂时不会介绍,可以百度查看
安装完成之后,我们打开anaconda prompt
(如果是linux系统那就打开命令窗口即可,前提是已经安装anaconda)
直接输入命令
然后就会弹出一个窗口
这样就证明我们已经安装完毕,然后我们先做一点准备工作,选择一个文件夹进入,然后在里面新建4个文件夹
然后我们将我们要进行标记掩膜的图片放入到pic中,在labelme选择Open Dir
打开,然后选择Create Polygons进行多边形标记
标记完成之后会让你输入标签值,这里我们只是做掩模图而不是数据集,我们直接写上mask
即可
然后选择旁边的Save
,路径选择到我们刚刚创建的json文件夹
中,然后就会发现我们的文件夹中生成了一个json
文件
然后就是至关重要的一步,我们要用这个json
文件去生成mask图
同样打开命令窗口,然后cd到我们的json文件夹下,然后利用labelme的转换生成mask图:
#windows下(输入后半部分命令即可)
(base) C:\Users\11942> D:
(base) D:\> cd D:\python\labelme_json\json
(base) D:\python\labelme_json\json> labelme_json_to_dataset zidane.json
#linux下
cd python/json
labelme_json_to_dataset zidane.json
然后我们到json文件夹下,就能看到我们生成了一个zidane.json
文件夹,打开之后会有五个文件,其中label.png
就是我们需要的mask图了
yolo基于mask的检测
制作好mask图之后,我们要如何把它添加到yolo识别中呢?
我们可以看到mask图中除了红色区域,其他区域的颜色均为黑色,同时我们也知道,黑色的RGB值为0,所以我们可以利用这一点,将我们想要进行检测的图片,根据mask图进行区域变黑,这样yolo就只会“看到”需要检测的区域和黑色区域,就不会对区域外的内容进行检测了
所以现在的目标就是,我们要如何根据mask图进行区域的选择,这里我们可以利用opencv来实现这一点
我们新建一个py文件,放到json文件夹目录下,然后用下面的代码进行转换:
import cv2
image = cv2.imread('label.png')
yuantu = cv2.imread('zidane.jpg')
img2gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) #将mask转换为灰度图
ret,mask = cv2.threshold(img2gray,0,255,cv2.THRESH_BINARY) #大于0的部分变为白,小于等于0的部分变黑
img1_bg = cv2.bitwise_and(yuantu,yuantu, mask=mask) #按位与,0直接变黑,原图贴上白色部分不变
cv2.imshow("111", img1_bg)
cv2.waitKey(0)
然后运行代码,可以看到我们输出的图片为:
效果非常明显,除了目标区域其他部分都变黑了,那么如果我们把这张图放到yolo去检测,就可以完美地解决区域检测问题啦!
那么现在就剩下最后一个问题,如何将这部分代码替换到yolov5的检测代码里面去?
(温馨提示:这里的替换方法只针对yolov5的loadimage函数,如果是yolo的其他函数或版本,虽然同理,但是需要自己找到替换的位置)
现在我假设你已经下载好yolov5的检测代码了,如果没有可以访问官网:yolov5_github
然后我们进入到utils/datasets.py
这个py文件里面,在开头添加上一段制作mask的代码
###################################################################
#制作mask
image_label = cv2.imread('/home/jr410/yolov5/utils/label/label.png') #根据自己路径修改mask位置
img2gray = cv2.cvtColor(image_label, cv2.COLOR_BGR2GRAY)
ret,mask = cv2.threshold(img2gray,0,255,cv2.THRESH_BINARY)
###################################################################
然后锁定到大概154行的位置(由于yolo代码在持续迭代,行数可能有略微的出入),找到LoadImages
这个类
然后找到def __next__(self)
这个函数的位置,我们要处理的部分主要在此
然后我们锁定到这一行(大概在220行左右)
很明显这个时候他已经要对图片进行处理了,我们只需要在他处理之前给他套上mask图,就能够实现我们的检测目的啦!
把上面截图那段代码替换为:
##########################################################
#设置是否添加mask检测,flag=1即为maks检测, 0则不然 #
flag = 1 #
##########################################################
if (flag):
img = cv2.bitwise_and(img0, img0, mask=mask)
img = letterbox(img, self.img_size, stride=self.stride)[0]
else:
# Padded resize
# 对图片进行resize+pad
img = letterbox(img0, self.img_size, stride=self.stride)[0]
这样修改完成之后,我们回退到yolo的上一个目录文件夹下,同时运行detect.py
文件(记得运行之前要先去官网下载yolov5s.pt权重文件)
python detect.py --source ./data/images/zidane.jpg --weights weights/yolov5s.pt
然后我们就可以得到下面的检测结果
跟官方直接检测的结果进行对比(将flag设置为0即可取消mask检测):
非常明显地看出右侧的人已经不在我们的识别区域内了!
至此,我们的yolo-mask目标检测就完成啦!
版权声明:本文为CSDN博主「wksgogogo」的原创文章,遵循CC 4.0 BY-SA版权协议,转载请附上原文出处链接及本声明。
原文链接:https://blog.csdn.net/qq_45940870/article/details/119647522
暂无评论