Mask R-CNN tensorflow 训练自己的数据【从标注数据到最终训练和测试】超全教程，呕血踩坑，Ubuntu 16.04 完美运行

前言

楼主由于工作需要，必须使用Mask-Rcnn 去训练自己的数据，写这篇博客之前楼主从百度搜索了各种训练的方法，但是搜索到的博文 写的都比较有歧义，最终我通过各种尝试 终于让训练跑了起来，也特在此处 写上这篇博文 送给大家。

本教程 适用于 Ubuntu 系统的用户、Windows 系统的用户

我使用的Mask RCNN-->https://github.com/matterport/Mask_RCNN

可以直接使用以下命令 克隆到自己的电脑上：

git clone https://github.com/matterport/Mask_RCNN

使用的主要框架如下：

• Tensorflow-gpu 版本 1.4.0

• keras 2.1.0

• python 3.5.2

推荐大家使用python3 去运行项目，我下面的教程也是基于python3的（如果大家后面运行的时候 出现一些意想不到的错误，基本上都是因为这几个框架的版本问题了。对着这几个版本安装 肯定是没问题的）

 

制作Mask Rcnn的数据源

 

1.数据标注

我使用 labelme 进行数据的标注， 关于labelme的安装我主要参考了https://blog.csdn.net/u010205128/article/details/80975818 这篇博文（上面说要用Annaconda ，我之前是没用Annaconda，但是无奈最后自己试了其他方法，发现还是需要装上 Annaconda 的！）

启动 labelme

打开一张图片进行标注

保存之后，可以看到 labelme给我们生成了一个json 文件 ，之后 来到我们保存json 文件的目录下

输入以下命令，将json 文件进行转化(注意 你使用的 命令行工具 是需要激活 labelme 环境的～具体的可以参考前面提到的labelme安装的博文中)

labelme_json_to_dataset 002.json(替换成你生成的json文件名)

转化完成后 可以看见生成的文件目录

打开之后可以看见生成的文件

截止到目前位置，我们都还很顺利～   查看别的博文，很多博主写的文章都有说到  ---------》 要将“labelme生成的掩码标签 label.png为16位存储”转化为 “opencv默认读取8位” 其实只要你生成的label.png文件是彩色的他就是默认的八位的编码的了！主要是labelme 版本升级了，截止到我写博文为止 labelme 已经升级，升级后的版本 生成的 label.png文件就已经是 八位编码的！！！！也就是说，我们不需要在进行复杂的编码转换！

 

下面我们进行 数据源的制作，其实制作步骤也是很简单。

首先 创建 四个文件夹 分别是 cv2_mask、json、labelme_json、pic

cv2_mask----》放labelme 生成的 label.png 文件  注意 这里我们把 label.phg 重命名为 002.png

json 文件夹----》 存放我们使用labelme 标注后生成的那个 json文件， 文件名注意要对上 上一步我们重命名的 002.png ----》 那么你的json 文件名称 就应该为 002.json

labelme_json 文件夹 ---》 存放我们 前面转化生成的 002_json文件夹（这里的命名规则想必大家也都懂了，我们要使用和上面 图片文件对应的命名格式 xxx_json 这种格式！）

pic文件夹  ----》 存放我们的原始图片  注意！名称要和前面几个步骤的名称对上 002.png

完成了上面的全部 步骤 我们的数据准备工作就已经完成了！ 下面可以开始正式的准备训练了！！！！！

 

首先，我们在 项目的根目录下创建 一个 train,py Python文件，接着把下面的代码 考进去，这里需要大家注意的是，我们创建的训练文件 必须在根目录下，否则会出现 找不到包的文件。这个很重要，真的

下面 贴出训练文件的代码，然后我再告诉大家需要修改哪些地方（ 这里的代码也是我从别处偷过来的～～～ 但是已经忘记是从哪里偷的了）

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

import os
import sys
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
# import utils
from mrcnn import model as modellib, utils
from mrcnn import visualize
import yaml
from mrcnn.model import log
from PIL import Image

# os.environ["CUDA_VISIBLE_DEVICES"] = "0"
# Root directory of the project
ROOT_DIR = os.getcwd()

# ROOT_DIR = os.path.abspath("../")
# Directory to save logs and trained model
MODEL_DIR = os.path.join(ROOT_DIR, "logs")

iter_num = 0

# Local path to trained weights file
COCO_MODEL_PATH = os.path.join(ROOT_DIR, "mask_rcnn_coco.h5")
# Download COCO trained weights from Releases if needed
if not os.path.exists(COCO_MODEL_PATH):
    utils.download_trained_weights(COCO_MODEL_PATH)


class ShapesConfig(Config):
    """Configuration for training on the toy shapes dataset.
    Derives from the base Config class and overrides values specific
    to the toy shapes dataset.
    """
    # Give the configuration a recognizable name
    NAME = "shapes"

    # Train on 1 GPU and 8 images per GPU. We can put multiple images on each
    # GPU because the images are small. Batch size is 8 (GPUs * images/GPU).
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1

    # Number of classes (including background)
    NUM_CLASSES = 1 + 1  # background + 1 shapes

    # Use small images for faster training. Set the limits of the small side
    # the large side, and that determines the image shape.
    IMAGE_MIN_DIM = 320
    IMAGE_MAX_DIM = 384

    # Use smaller anchors because our image and objects are small
    RPN_ANCHOR_SCALES = (8 * 6, 16 * 6, 32 * 6, 64 * 6, 128 * 6)  # anchor side in pixels

    # Reduce training ROIs per image because the images are small and have
    # few objects. Aim to allow ROI sampling to pick 33% positive ROIs.
    TRAIN_ROIS_PER_IMAGE = 100

    # Use a small epoch since the data is simple
    STEPS_PER_EPOCH = 100

    # use small validation steps since the epoch is small
    VALIDATION_STEPS = 50


config = ShapesConfig()
config.display()


class DrugDataset(utils.Dataset):
    # 得到该图中有多少个实例（物体）
    def get_obj_index(self, image):
        n = np.max(image)
        return n

    # 解析labelme中得到的yaml文件，从而得到mask每一层对应的实例标签
    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

    # 重新写draw_mask
    def draw_mask(self, num_obj, mask, image, image_id):
        # print("draw_mask-->",image_id)
        # print("self.image_info",self.image_info)
        info = self.image_info[image_id]
        # print("info-->",info)
        # print("info[width]----->",info['width'],"-info[height]--->",info['height'])
        for index in range(num_obj):
            for i in range(info['width']):
                for j in range(info['height']):
                    # print("image_id-->",image_id,"-i--->",i,"-j--->",j)
                    # print("info[width]----->",info['width'],"-info[height]--->",info['height'])
                    at_pixel = image.getpixel((i, j))
                    if at_pixel == index + 1:
                        mask[j, i, index] = 1
        return mask

    # 重新写load_shapes，里面包含自己的自己的类别
    # 并在self.image_info信息中添加了path、mask_path 、yaml_path
    # yaml_pathdataset_root_path = "/tongue_dateset/"
    # img_floder = dataset_root_path + "rgb"
    # mask_floder = dataset_root_path + "mask"
    # dataset_root_path = "/tongue_dateset/"
    def load_shapes(self, count, img_floder, mask_floder, imglist, dataset_root_path):
        """Generate the requested number of synthetic images.
        count: number of images to generate.
        height, width: the size of the generated images.
        """
        # Add classes
        self.add_class("shapes", 1, "person")

        for i in range(count):
            # 获取图片宽和高
            print(i)
            filestr = imglist[i].split(".")[0]
            # print(imglist[i],"-->",cv_img.shape[1],"--->",cv_img.shape[0])
            # print("id-->", i, " imglist[", i, "]-->", imglist[i],"filestr-->",filestr)
            # filestr = filestr.split("_")[1]
            mask_path = mask_floder + "/" + filestr + ".png"
            yaml_path = dataset_root_path + "labelme_json/" + filestr + "_json/info.yaml"
            print(dataset_root_path + "labelme_json/" + filestr + "_json/img.png")
            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)

    # 重写load_mask
    def load_mask(self, image_id):
        """Generate instance masks for shapes of the given 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("person") != -1:
                # print "car"
                labels_form.append("person")
            elif labels[i].find("leg") != -1:
                # print "leg"
                labels_form.append("leg")
            elif labels[i].find("well") != -1:
                # print "well"
                labels_form.append("well")
        class_ids = np.array([self.class_names.index(s) for s in labels_form])
        return mask, class_ids.astype(np.int32)


def get_ax(rows=1, cols=1, size=8):
    """Return a Matplotlib Axes array to be used in
    all visualizations in the notebook. Provide a
    central point to control graph sizes.

    Change the default size attribute to control the size
    of rendered images
    """
    _, ax = plt.subplots(rows, cols, figsize=(size * cols, size * rows))
    return ax


# 基础设置
dataset_root_path = "samples/trinmy/myinfo/"
img_floder = dataset_root_path + "pic"
mask_floder = dataset_root_path + "cv2_mask"
# yaml_floder = dataset_root_path
imglist = os.listdir(img_floder)
count = len(imglist)

# train与val数据集准备
dataset_train = DrugDataset()
dataset_train.load_shapes(count, img_floder, mask_floder, imglist, dataset_root_path)
dataset_train.prepare()

# print("dataset_train-->",dataset_train._image_ids)

dataset_val = DrugDataset()
dataset_val.load_shapes(count, img_floder, mask_floder, imglist, dataset_root_path)
dataset_val.prepare()

# print("dataset_val-->",dataset_val._image_ids)

# Load and display random samples
# image_ids = np.random.choice(dataset_train.image_ids, 4)
# for image_id in image_ids:
#    image = dataset_train.load_image(image_id)
#    mask, class_ids = dataset_train.load_mask(image_id)
#    visualize.display_top_masks(image, mask, class_ids, dataset_train.class_names)

# Create model in training mode
model = modellib.MaskRCNN(mode="training", config=config,
                          model_dir=MODEL_DIR)

# Which weights to start with?
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":
    # Load weights trained on MS COCO, but skip layers that
    # are different due to the different number of classes
    # See README for instructions to download the COCO weights
    # print(COCO_MODEL_PATH)
    model.load_weights(COCO_MODEL_PATH, by_name=True,
                       exclude=["mrcnn_class_logits", "mrcnn_bbox_fc",
                                "mrcnn_bbox", "mrcnn_mask"])
elif init_with == "last":
    # Load the last model you trained and continue training
    model.load_weights(model.find_last()[1], by_name=True)

# Train the head branches
# Passing layers="heads" freezes all layers except the head
# layers. You can also pass a regular expression to select
# which layers to train by name pattern.
model.train(dataset_train, dataset_val,
            learning_rate=config.LEARNING_RATE,
            epochs=10,
            layers='heads')

# Fine tune all layers
# Passing layers="all" trains all layers. You can also
# pass a regular expression to select which layers to
# train by name pattern.
model.train(dataset_train, dataset_val,
            learning_rate=config.LEARNING_RATE / 10,
            epochs=10,
            layers="all")

 

1.首先 第一个需要大家注意的地方，在代码的33行

我们需要使用预训练的 mask_rcnn_coco.h5 这个训练好的模型作为我们参数的初始化，（能否不用预训练的模型？----》说实话我也不知道行不行，但是大家都是这么做的，那我也只好跟着这么做了。呵呵哈哈哈，也许你可以自己尝试一下不用预训练的模型 ）

关于 mask_rcnn_coco.h5  这个文件 怎么获取，大家可以直接 复制，百度去搜索基本都可以找的到(不过好像都是csdn 要积分的下载)，或者在大家有需要的也可以给我留言，我直接邮箱转发给你们。（Ubuntu用户也免去了找百度云的那些鬼软件）

这里我需要说明一下，下载了 mask_rcnn_coco.h5  这个文件之后，你直接把这个文件丢到 根目录下就可以的了，不需要再去做别的配置（别的博文说需要配置？我其实一脸蒙，这个怎么配置他们也没说，就说了配置！！ 不得不吐槽一下）。如下图所示：

 

2. 第二个需要修改的地方 在 第53行，类别数的修改！这一个很重要，我最开始就是漏掉了这里，导致后面运行的时候报错

    NUM_CLASSES = 1 + 1  # background + 1 shapes

这里修改为 背景 + 训练数据的类别数， 我这里只有一个类别 就是 人脸， 所以 NUM_CLASSES = 1+ 1

3. 第三个需要修改的地方 在186行，需要修改刚刚我们制作数据源的路径。指向我们制作数据的路径就可以了

 

下面贴出 myinfo 这个文件夹下的文件，也就是我们刚刚制作数据源的文件夹

4.第四个地方需要修改的 代码 122 行 和 代码160 - 169

这里是添加标注的类别名称， 注意要把 第三个参数（原来是 person） 改成你自己标注数据时写的那个类别名称（我前面标注写的是face）所以修改成 face 。

如果你还有更多的分类，直接拷贝 然后修改 第二个和第三个参数即可。

代码160 - 169：

这里 labels[i].find("face")  ！= -1  对应前面122行写的 face，你需要改成你自己的类别名称，然后162行labels_form.append("face") 也要改成你自己的类别名称。后面两个elif 你可以把他删除掉。

如果你有更多的类别，只需要在后面加 elif 然后 labels_form.append(”你的类别名称“)即可

=========================================================================================================

其他还有一些配置大家都可以到代码的45-71行去修改，无非也就是 训练次数、图片大小、诸如此类的配置。在这里我就不在去修改这些配置了。

 

完成了上述的全部步骤，可以正式开始训练了！我用python3 去运行 train.py这个文件，用python2 会出现一下让我匪夷所思的错误！啧啧，最终我用python3 成功的跑通了。

python3 train.py 

如果大家运行出现 错误的话，首先检查 是否已经正确修改了我上面提到的那几个地方。

其次如果错误是提示你 Shape dimension 之类的错误 基本上都是因为你的类别数目和代码53行中 NUM_CLASSES = 1+ 1 对不上，如果提示没有包，或者其他的依赖，你直接使用 sudo pip3 install 包名 这种方式安装（我们用的是python3  所以要用pip3 安装）

关于常见的错误 大家可以到 下面的这篇博文查看，里面详细的描述了各种可能的错误

 https://blog.csdn.net/lovebyz/article/details/80138261

如果顺利的话，一次就能运行成功

成功运行后 你看到的应该是这样的：

运行了一段时间后（前面默认设置的是 训练100 步保存一次模型），可以到 logs 文件夹下找到你训练好的模型文件

这里可以看到生成的很多日志文件，是我前面有训练其他的模型 产生的，你只需要找到当前文件下 最后的那个文件夹就可以了，也就是说你的模型是保存在排在最后的那个文件夹中，打开那个文件夹你可以看到：

这个mask_rcnn_shapes_0001.h5 文件就是生成的模型文件了

 

测试

下面我们就要使用 训练好的模型就行测试了！

首先，在你的项目根目录下 创建 fortest.py 文件（与train.py文件同在一个目录即可），然后把下面的代码帖进去

# -*- coding: utf-8 -*-
import os
import sys
import random
import math
import numpy as np
import skimage.io
import matplotlib
import matplotlib.pyplot as plt
import cv2
import time
from mrcnn.config import Config
from datetime import datetime
# Root directory of the project
ROOT_DIR = os.getcwd()

# Import Mask RCNN
sys.path.append(ROOT_DIR)  # To find local version of the library
from mrcnn import utils
import mrcnn.model as modellib
from mrcnn import visualize
# Import COCO config
# sys.path.append(os.path.join(ROOT_DIR, "samples/coco/"))  # To find local version
# from samples.coco import coco


# Directory to save logs and trained model
MODEL_DIR = os.path.join(ROOT_DIR, "logs")

# Local path to trained weights file
COCO_MODEL_PATH = os.path.join(MODEL_DIR ,"mask_rcnn_shapes_0010.h5")
# Download COCO trained weights from Releases if needed
if not os.path.exists(COCO_MODEL_PATH):
    utils.download_trained_weights(COCO_MODEL_PATH)
    print("cuiwei***********************")

# Directory of images to run detection on
IMAGE_DIR = os.path.join(ROOT_DIR, "images")

class ShapesConfig(Config):
    """Configuration for training on the toy shapes dataset.
    Derives from the base Config class and overrides values specific
    to the toy shapes dataset.
    """
    # Give the configuration a recognizable name
    NAME = "shapes"

    # Train on 1 GPU and 8 images per GPU. We can put multiple images on each
    # GPU because the images are small. Batch size is 8 (GPUs * images/GPU).
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1

    # Number of classes (including background)
    NUM_CLASSES = 1 + 1  # background + 3 shapes

    # Use small images for faster training. Set the limits of the small side
    # the large side, and that determines the image shape.
    IMAGE_MIN_DIM = 320
    IMAGE_MAX_DIM = 384

    # Use smaller anchors because our image and objects are small
    RPN_ANCHOR_SCALES = (8 * 6, 16 * 6, 32 * 6, 64 * 6, 128 * 6)  # anchor side in pixels

    # Reduce training ROIs per image because the images are small and have
    # few objects. Aim to allow ROI sampling to pick 33% positive ROIs.
    TRAIN_ROIS_PER_IMAGE =100

    # Use a small epoch since the data is simple
    STEPS_PER_EPOCH = 100

    # use small validation steps since the epoch is small
    VALIDATION_STEPS = 50

#import train_tongue
#class InferenceConfig(coco.CocoConfig):
class InferenceConfig(ShapesConfig):
    # Set batch size to 1 since we'll be running inference on
    # one image at a time. Batch size = GPU_COUNT * IMAGES_PER_GPU
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1

config = InferenceConfig()

model = modellib.MaskRCNN(mode="inference", model_dir=MODEL_DIR, config=config)


# Create model object in inference mode.
model = modellib.MaskRCNN(mode="inference", model_dir=MODEL_DIR, config=config)

# Load weights trained on MS-COCO
model.load_weights(COCO_MODEL_PATH, by_name=True)

# COCO Class names
# Index of the class in the list is its ID. For example, to get ID of
# the teddy bear class, use: class_names.index('teddy bear')
class_names = ['BG', 'person']
# Load a random image from the images folder
file_names = next(os.walk(IMAGE_DIR))[2]
image = skimage.io.imread("./images/5951960966_d4e1cda5d0_z.jpg")

a=datetime.now()
# Run detection
results = model.detect([image], verbose=1)
b=datetime.now()
# Visualize results
print("shijian",(b-a).seconds)
r = results[0]
visualize.display_instances(image, r['rois'], r['masks'], r['class_ids'],
                            class_names, r['scores'])

运行之前，我们需要修改几个地方：

代码31行，修改模型的路径，把路径替换成 刚刚我们训练好的模型路径

COCO_MODEL_PATH = os.path.join(MODEL_DIR ,"shapes20190221T1101/mask_rcnn_shapes_0010.h5")

不要直接复制上面的代码，因为你的模型路径可能和我的不一样，自己根据你的模型路径文件填写。

第二个需要修改的地方，代码 99行，把其中的路径 改为你需要测试图片的路径。

image = skimage.io.imread("./images/002.png")

完成上述步骤 可以进行测试了，命令如下：

python3 fortest.py

然后可以看到

过了一会

 

至此，Mask Rcnn 模型的训练与测试我们就完成了，当然这是很简单的去跑一下，其中还有很多的超参数我们没有去修改，上面 训练 和 测试文件的代码大家在跑通后 好好的去读一下和理解。如果有其他问题大家也可以给我留言，初次写博文，多有不当的地方还望大家指教！