图片分割之_训练模型和预测
图片分割之 _ 训练模型和预测
1. 说明
本篇使用 Mask R-CNN 算法,以及十几张从网络上下载的香蕉图片,训练一个模型。用于识别图像中的香蕉,不同于苹果,桔子,香蕉从不同的角度看差异很大,尤其是三五根香蕉放连在一起,或者整把香蕉的形态和单根香蕉差异很大。可以算是一种识别起来相对困难的水平。
下图是用训练好的模型识别出的香蕉图片,可以看到,基本识别正确。
操作步骤可分为:安装工具,标注图片,修改源码,模型训练和模型预测。我的工作环境是 Ubuntu,硬件有 GPU 支持,操作过程中使用了 python,图片标注工具,以及 shell 脚本。
2. 安装工具
(1) 下载程序源码
1 | $ git clone https://github.com/matterport/Mask_RCNN.git # (大概200多M) |
(2) 下载相关软件
1 | $ sudo pip install opencv-python |
3. 标注图片
(1) 收集图片
香蕉图片可以从网上下载,也可用手机拍照,图片分辨率不用太高,1000x1000 以下即可,如果分辨率太高,可用 linux 中的 convert 命令缩放。我使用的 15 张图片如下图所示:
需要注意的是,图片需要包括香蕉的各个角度,以及常见的多根组合的几种形态。
(2) 用软件标注图片
1 | $ labelme 图片文件名.jpg |
labelme 为一个图形化的标注工具,使用左侧面板中的 create polygons,将图片中所需识别的香蕉圈出来,如果某个点画错了,用 Backspace 可删除最后设置的点(用法类似于 photoshop 中的多边形套锁工具),标注完加入填入 label 名,这个名字后面在程序中会到,标注完注意保存文件,文件名默认为:图片名.json。锚点的细密程度请参考下图:
标注无需粒度过细,Labelme 工具比较智能,只要位置相近,就能把锚点自动贴近边界。好的工具让标注事半功倍,一般情况下,十几张图片半个多小时即可标注完成,另外,一个图中也可标注多个区域,label 名都设置为 banana 即可。
(3) 解析和拆分标注文件
使用 labelme 自带的 labelme_json_to_dataset 命令工具,可将 json 文件拆分成目录,目录中数据如下:
一条命令可以转换一个图片,当图片多时,建议使用 shell 脚本处理,shell 脚本示例如下,请根据环境调整。
1 | for file in `ls *.json` |
(4) mask 文件转码
由于不同版本的 labelme 生成的文件格式不同,有的 mask 是 24 位色,有的是 8 位色,用以下 python 程序看一下图片格式:
1 | from PIL import Image |
如果 image.mode 是 P,即 8 位彩色图像,直接使用即可,如果是其它格式,使用以下程序将其转换成 8 位图片:
1 | Img_8 = img.convert("P") |
将转换后的图片复制到另一文件夹即可,复制方法请参考以下 shell 脚本
1 | mkdir ../cv2_mask |
(5) 调整目录结构
把上述的原图放在 pic 目录中, 标注文件放在 json 目录中, 拆分后的标注文件放在 labelme_json 目录中,掩码 mask 放在 cv2_mask 目录中,调整之后的目录结构如下图所示:
其中的 mine 本例中所有程序和数据,数据放在 data 目录中,训练好的模型放在 models 目录中。
4. 训练和预测
(1) 训练模型
源码
1 | import os |
运行程序
1 | $ python train.py |
在程序运行过程中,如果因为 tensorflow 版本与 mask_rcnn 不匹配,引起找不到 keepdims 问题,需要修改 Mask_RCNN/mrcnn/model.py,将其中的 keepdims 改为 keep_dims 即可。
我的机器训练完不到 15 分钟,如果把两次训练的迭代次数分别设成 1 和 2 则 2 分钟完成训练。
(2) 预测模型
源码
1 | # -*- coding: utf-8 -*- |
运行程序
1 | $ python test.py |
5. 分析总结
- 自动标注:当图片数量很多时,可以先训练少量图片,生成模型,让模型自动标注,人为检查标注是否正确,对于不正确的人工重新标注。
- 建议使用 GPU:相比 GPU,我用 4 核的 CPU 计算,速度目测差了 50 倍左右。个人觉得没有 GPU,训练速度几乎是无法接受的。
- 迭代次数:迭代次数可以调整,如果同一个图,用网上下载的基础模型完全识别不出。而用第 1 次迭代和第 30 次迭代结果差不太多,以后再训练就可以减少迭代次数,以节约时间。
- 生成模型:由于迭代训练了 30 次,models 目录下产生了 30 个模型文件,占空间比较大,不用的可以删除掉。
- 更多例程请参考原代码中的 Mask_RCNN/samples/ 目录。
6. 问题及解决方法
- 问题: 在 CPU 上运行时可能报错 SVD did not converge,
分析及解决:该问题发生成 resize 图片时,代码 mrcnn/utils.py 计算 resize 的 scale 里用两个 int 型相除,结果 scale 变成 0,导致 resize 出错,解决访问是添加:scale = float(max_dim) / image_max 强制类型转型即可。
7. 参考
- Mask RCNN 训练自己的数据集
https://blog.csdn.net/l297969586/article/details/79140840
- mask rcnn 训练自己的数据集