图片分割之 _ 训练模型和预测

1. 说明

本篇使用 Mask R-CNN 算法，以及十几张从网络上下载的香蕉图片，训练一个模型。用于识别图像中的香蕉，不同于苹果，桔子，香蕉从不同的角度看差异很大，尤其是三五根香蕉放连在一起，或者整把香蕉的形态和单根香蕉差异很大。可以算是一种识别起来相对困难的水平。

下图是用训练好的模型识别出的香蕉图片，可以看到，基本识别正确。

操作步骤可分为：安装工具，标注图片，修改源码，模型训练和模型预测。我的工作环境是 Ubuntu，硬件有 GPU 支持，操作过程中使用了 python，图片标注工具，以及 shell 脚本。

2. 安装工具

(1) 下载程序源码

1	$ git clone https://github.com/matterport/Mask_RCNN.git # (大概200多M)

(2) 下载相关软件

$ sudo pip install opencv-python
$ sudo pip install tensorflow
$ sudo pip install scikit-image
$ sudo pip install keras==2.0.8
$ sudo pip install labelme # 标注工具

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 脚本示例如下，请根据环境调整。

for file in `ls *.json`
do
    echo labelme_json_to_dataset $file
    labelme_json_to_dataset $file
done
mkdir ../labelme_json/
mv *_json ../labelme_json/

(4) mask 文件转码

由于不同版本的 labelme 生成的文件格式不同，有的 mask 是 24 位色，有的是 8 位色，用以下 python 程序看一下图片格式：

1
2
3

from PIL import Image
img = Image.open('label.png')
print(img.mode)

如果 image.mode 是 P，即 8 位彩色图像，直接使用即可，如果是其它格式，使用以下程序将其转换成 8 位图片：

1 2	Img_8 = img.convert("P") Img_8.save('xxx.png')

将转换后的图片复制到另一文件夹即可，复制方法请参考以下 shell 脚本

mkdir ../cv2_mask
cd ../cv2_mask
for file in `ls ../labelme_json`
do
    echo 'cp ../labelme_json/'$file'/label.png '$file.png 
    cp '../ labelme_json /'$file'/label.png' $file.png 
done

(5) 调整目录结构

把上述的原图放在 pic 目录中, 标注文件放在 json 目录中, 拆分后的标注文件放在 labelme_json 目录中，掩码 mask 放在 cv2_mask 目录中，调整之后的目录结构如下图所示：

其中的 mine 本例中所有程序和数据，数据放在 data 目录中，训练好的模型放在 models 目录中。

4. 训练和预测

(1) 训练模型

源码

import os
import sys
sys.path.append(xxxx) # 加入Mask_RCNN源码所在目录
import random
import math
import re
import time
import numpy as np
import cv2
import matplotlib
import matplotlib.pyplot as plt
import tensorflow as tf
from mrcnn.config import Config
from mrcnn import model as modellib,utils
from mrcnn import visualize
import yaml
from mrcnn.model import log
from PIL import Image

ROOT_DIR = os.getcwd()
MODEL_DIR = os.path.join(ROOT_DIR, "models")
iter_num=0
COCO_MODEL_PATH = os.path.join(ROOT_DIR, "mask_rcnn_coco.h5")

# 从网上下载训练好的基础模型
if not os.path.exists(COCO_MODEL_PATH):
    utils.download_trained_weights(COCO_MODEL_PATH)

# 配置
class ShapesConfig(Config):
    NAME = "shapes" # 命名
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1
    NUM_CLASSES = 1 + 1  # 背景一类，香蕉一类，共两类
    IMAGE_MIN_DIM = 320
    IMAGE_MAX_DIM = 384
    RPN_ANCHOR_SCALES = (8 * 6, 16 * 6, 32 * 6, 64 * 6, 128 * 6)
    TRAIN_ROIS_PER_IMAGE = 100 # Aim to allow ROI sampling to pick 33% positive ROIs
    STEPS_PER_EPOCH = 100
    VALIDATION_STEPS = 50

config = ShapesConfig()
config.display()

# 重写数据集
class DrugDataset(utils.Dataset):
    def get_obj_index(self, image):
        n = np.max(image)
        return n

    # 获取标签
    def from_yaml_get_class(self, image_id):
        info = self.image_info[image_id]
        with open(info['yaml_path']) as f:
            temp = yaml.load(f.read())
            labels = temp['label_names']
            del labels[0]
        return labels

    # 填充mask
    def draw_mask(self, num_obj, mask, image,image_id):
        info = self.image_info[image_id]
        for index in range(num_obj):
            for i in range(info['width']):
                for j in range(info['height']):
                    at_pixel = image.getpixel((i, j))
                    if at_pixel == index + 1:
                        mask[j, i, index] = 1
        return mask
	
    # 读入训练图片及其配置文件
    def load_shapes(self, count, img_floder, mask_floder, imglist, dataset_root_path):
        self.add_class("shapes", 1, "banana") # 自定义标签 
        for i in range(count):
            filestr = imglist[i].split(".")[0]
            mask_path = mask_floder + "/" + filestr + "_json.png"
            yaml_path = dataset_root_path + "labelme_json/" + filestr + "_json/info.yaml"
            cv_img = cv2.imread(dataset_root_path + "labelme_json/" + filestr + "_json/img.png")
            self.add_image("shapes", image_id=i, path=img_floder + "/" + imglist[i],
                           width=cv_img.shape[1], height=cv_img.shape[0], mask_path=mask_path, yaml_path=yaml_path)

	# 读取标签和配置 
    def load_mask(self, image_id):
        global iter_num
        print("image_id",image_id)
        info = self.image_info[image_id]
        count = 1  # number of object
        img = Image.open(info['mask_path'])
        num_obj = self.get_obj_index(img)
        mask = np.zeros([info['height'], info['width'], num_obj], dtype=np.uint8)
        mask = self.draw_mask(num_obj, mask, img,image_id)
        occlusion = np.logical_not(mask[:, :, -1]).astype(np.uint8)
        for i in range(count - 2, -1, -1):
            mask[:, :, i] = mask[:, :, i] * occlusion
            occlusion = np.logical_and(occlusion, np.logical_not(mask[:, :, i]))
        labels = []
        labels = self.from_yaml_get_class(image_id)
        labels_form = []
        for i in range(len(labels)):
            if labels[i].find("banana") != -1: # 自定义标签
                labels_form.append("banana")
        class_ids = np.array([self.class_names.index(s) for s in labels_form])
        return mask, class_ids.astype(np.int32)

#基础设置
dataset_root_path="data/"
img_floder = dataset_root_path + "pic" # 基本图片目录
mask_floder = dataset_root_path + "cv2_mask" # mask图片目录
imglist = os.listdir(img_floder)
count = len(imglist)

# 构造训练集
dataset_train = DrugDataset()
dataset_train.load_shapes(count, img_floder, mask_floder, imglist, dataset_root_path)
dataset_train.prepare()

# 构造验证集
dataset_val = DrugDataset()
dataset_val.load_shapes(7, img_floder, mask_floder, imglist, dataset_root_path)
dataset_val.prepare()

# 建立模型
model = modellib.MaskRCNN(mode="training", config=config,
                          model_dir=MODEL_DIR)

# 定义模式
init_with = "coco"  # imagenet, coco, or last

if init_with == "imagenet":
    model.load_weights(model.get_imagenet_weights(), by_name=True)
elif init_with == "coco":
    model.load_weights(COCO_MODEL_PATH, by_name=True,
                       exclude=["mrcnn_class_logits", "mrcnn_bbox_fc",
                                "mrcnn_bbox", "mrcnn_mask"])
elif init_with == "last":
    model.load_weights(model.find_last()[1], by_name=True)

model.train(dataset_train, dataset_val,
            learning_rate=config.LEARNING_RATE,
            epochs=10,
            layers='heads')

model.train(dataset_train, dataset_val,
            learning_rate=config.LEARNING_RATE / 10,
            epochs=30,
            layers="all")

运行程序

1	$ python train.py

在程序运行过程中，如果因为 tensorflow 版本与 mask_rcnn 不匹配，引起找不到 keepdims 问题，需要修改 Mask_RCNN/mrcnn/model.py，将其中的 keepdims 改为 keep_dims 即可。

我的机器训练完不到 15 分钟，如果把两次训练的迭代次数分别设成 1 和 2 则 2 分钟完成训练。

(2) 预测模型

源码

# -*- coding: utf-8 -*-

import os
import sys
sys.path.append(os.path.dirname(os.getcwd())) # 注意：加mask_rcnn目录
import skimage.io
from mrcnn.config import Config
from datetime import datetime 
import mrcnn.model as modellib
from mrcnn import visualize

ROOT_DIR = os.getcwd()
sys.path.append(ROOT_DIR)
MODEL_DIR = os.path.join(ROOT_DIR, "models")

# 配置，同train
class ShapesConfig(Config):
    NAME = "shapes"
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1
    NUM_CLASSES = 1 + 1
    IMAGE_MIN_DIM = 320
    IMAGE_MAX_DIM = 384
    RPN_ANCHOR_SCALES = (8 * 6, 16 * 6, 32 * 6, 64 * 6, 128 * 6)
    TRAIN_ROIS_PER_IMAGE =100
    STEPS_PER_EPOCH = 100
    VALIDATION_STEPS = 50

class InferenceConfig(ShapesConfig):
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1

config = InferenceConfig()
model = modellib.MaskRCNN(mode="inference", model_dir=MODEL_DIR, config=config)
model.load_weights('models/shapes20190117T1428/mask_rcnn_shapes_0001.h5', by_name=True) # 注意换成你模型的路径
#model.load_weights('models/shapes20190117T1428/mask_rcnn_shapes_0030.h5', by_name=True) # 注意换成你模型的路径
#model.load_weights('mask_rcnn_coco.h5', by_name=True) # 注意换成你模型的路径

class_names = ['BG', 'banana']
image = skimage.io.imread('/tmp/banana.jpg') # 注意事换成你要识别的图片

a=datetime.now() 
results = model.detect([image], verbose=1)
b=datetime.now() 
print("@@ detect duration",(b-a).seconds, 'second')
r = results[0]
# 画图
visualize.display_instances(image, r['rois'], r['masks'], r['class_ids'], 
                            class_names, r['scores'])

运行程序

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 训练自己的数据集

https://blog.csdn.net/qq_29462849/article/details/81037343