[深度学习][原创]unet训练自己数据集整个流程和问题探索

项目地址：https://github.com/zhixuhao/unet

测试环境：

ubuntu18.04

cuda10.0+cudnn7.6.5

anaconda3+python3.7

keras==2.2.5

tensorflow-gpu==1.14.0

第一步：准备数据集

使用labelme标注数据，然后使用labelme_json-to_dataset批量转换为5个文件，类似下面这样

关于怎么转化问题具体参考这个博客详细的不得了。https://blog.csdn.net/qq_29462849/article/details/81037343

为了偷懒，我自己使用的VOC数据集然后写了一个脚本把xml格式转成labelme对应格式，这样我瞬间得到很多labelme数据集,VOC数据集很多比如VOC2012，VOC2007等，得到json文件，再然后用labelme_json-to_dataset。然而很不幸的是unet还不能直接使用labelme直接转化的图片。其中有2个重要的事情要做。

第一，比如上图的img.png必须全部转成8位图片，经过查看发现labelme出来是24位，我训练后测试是全灰色的情况，这个要转化，代码网上有，这里我放上了

import os
import glob
import cv2


def togrey(img, outdir):
    src = cv2.imread(img)
    try:
        dst = cv2.cvtColor(src, cv2.COLOR_BGR2GRAY)
        cv2.imwrite(os.path.join(outdir, os.path.basename(img)), dst)
    except Exception as e:
        print(e)


for file in glob.glob('../data/mydata/train/pic/*.png'):
    togrey(file, '../data/mydata/train/pic/')

for file in glob.glob('../data/mydata/test/pic/*.png'):
    togrey(file, '../data/mydata/test/pic/')

第二，掩码文件需要转成黑白的，就是二值化，不然后面训练测试模型是全黑的，这个很重要。经过测试黑背景或者白背景都行，我用的ostu全局阈值批量二值化。

 

 

上面第一个图是labelme转化的结果，第二个就是训练用的掩码文件。经过反复测试证明，如果我把第二个图片反色，弄成白背景的化训练测试结果也是反色，没有影响。

最后将图片和掩码文件放在对应路径改改unet源码路径就可以训练起来。

最后给一个训练测试结果

 

白色一团就是车辆大体位置，由于图片属于公司内部东西，这里就不放上来了。我把labelme转化放上来对比看看

 

还有一个num_class大家需要注意，如果只有一个类，unet代码基本不用改了，num_class=2，为什么是2不是1？其实很简单，背景算一个类。此外测试时候图片必须是0.png,1.png,...,类似这样，比如我是10张，名字依次0-9.png，然后测试代码num_image=10，还有predict_generator第2个参数也要改成10，如果coding能力可以也可以自己改改不要数字的。下面是修改的代码

testGene = testGenerator("data/mydata/test/pic",num_image=10)
model = unet()
model.load_weights("./model/model.h5")
results = model.predict_generator(testGene, 10, verbose=1)
saveResult("data/mydata", results,num_class=2)

。关于多分类的我暂时还没研究，等我研究出来，再发一篇博客吧。

总体感觉：

文件代码很少，但是问题很多，网上都是预测全灰、全白、全黑情况。我查了都没解决问题，最后我发现label文件问题，好了本次测试总结到此结束！

第二次更新：上面方法可行，但是后面深入研究发现预测全黑很可能是对的情况。因为有label并不是labelme直接生成的，而且按照类别索引建立的，比如像素都是0,1,2,3这样的，当然是全黑的！你可以生成预测图片并写个脚本去输出像素看看，目前我发现很多数据集都是按照索引构建的标签数据集！