YOLOv5屏蔽区域/选择区域 训练(重点区域训练)
- 前期准备
- labelme选择mask区域
- 代码改动
前期准备
思路就是通过使用一个mask掩膜,对我们想要屏蔽或者选择的区域进行定向遮挡或者定向选择,在训练的时候v5需要加载数据即图片和标注信息,我们在数据加载的过程中,将掩膜覆盖到输入图像上,即完成了屏蔽区域/选择区域训练。
labelme选择mask区域
标注好我们想检测的区域之后,我这里将标注区域命名成了mask,然后我们点击保存导出json文件:
然后我们再打开终端,cd到json所在的文件夹路径下,输入如下命令:
labelme_json_to_dataset xxx.json
这里有的兄弟会遇到一些报错,主要原因还是环境的版本问题,可以在网上查到,替换一下对应版本的库就行了。
我们就会在这个路径下获得一个文件夹:
打开文件夹之后会看到里面有这么五个文件,我们需要的就是label.png,其他的不用管:
到这里,我们的前期准备工作就完成了。
代码改动
打开utils/datasets.py文件,在代码的最面前,首先读取mask图像,然后将其转成二值图像:
然后再下方设置好超参数,方便在训练和检测时的设置:
要想对训练时的图像加载的图像进行操作,显然我们首先要到train.py中寻找数据加载的函数,我们可以定位到train函数下调用的create_dataloader()函数:
摁住ctrl+左键进入该函数:
来到了utils/datasets.py下的函数:
这里我们注意到该函数返回的dataset是通过LoadImagesAndLabels()函数来读取的,进入该函数之后,我们往下拉,找到加载图片的代码,可以在getitem()函数中看到:
查看这一段代码,我们发现当设置mosaic的时候,是一种加载方式,当未设置mosaic的时候就是直接加载的方式,v5默认的是mosaic之后再训练,因此我们直接进入load_mosaic()函数中:
进入到load_mosaic()函数之后,我们发现代码先调用load_image()函数来读取图像,然后通过顺序判断来进行mosaic组合,我们这里继续进入load_image()函数中,感觉胜利即将在前方了:
果然在进入load_image之后我们看到了加载图像的那句代码!这里的红框就是我自己加的代码,通过bool型来判断是否需要进行mask以及mask的作用是屏蔽还是选择,而这些bool型变量已经在这篇博客的开始被我们定义好了,可以根据需求来排列组合。这样就完成了我们的选择区域/屏蔽区域训练的需求。