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

1. 说明

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

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

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

2. 安装工具

(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) 用软件标注图片

$ 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程序看一下图片格式：

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

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

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")

运行程序

$ 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'])

运行程序

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