tensorflow1.4 用自己的数据训练mask r_cnn

详情 我的GitHubhttps://github.com/Monkeone/maskr-_cnn-
1.数据预处理
win10 pip安装labelme，运行会失败
使用labelme标注图片
labelme的安装
先安装Anaconda，2、进入Anaconda port
3、输入conda create --name=labelme python=3.7
4、输入activate labelme
然后建立的labelme环境就会被激活
5、在此激活的环境下输入pip install pyqt5安装pyqt5
6、pyqt5安装完成后输入pip install labelme安装labelme
7、在labelme的激活环境下输入labelme即可打开labelme使用
8、如果要退出labelme环境的话，直接输入deactivate即可
在anaconda prompt中输入命令
activate labelme
labelme
将图片和原图放在同一文件夹
2.开始训练

import os
import sys
import numpy as np
import json
import cv2
import pickle
# 工程的根目录
ROOT_DIR = os.path.abspath("resources/")
# 导入Mask RCNN库
sys.path.append(ROOT_DIR)  # To find local version of the library
from mrcnn.config import Config
from mrcnn import utils
import mrcnn.model as modellib
# 定义保存文件夹，用来训练日志和训练好的权重文件
MODEL_DIR = os.path.join(ROOT_DIR, "logs")
# 定义COCO数据集预训练权重文件的路径
COCO_MODEL_PATH = os.path.join(ROOT_DIR, "mask_rcnn_coco.h5")
# 如果此路径不存在，则报错
assert os.path.exists(COCO_MODEL_PATH), 'you need read resources/readme.md guide file and download mask_rcnn_coco.h5'

# 获取文件夹中的文件路径，返回。。。.jpg
def get_filePathList(dirPath, partOfFileName=''):
    allFileName_list = next(os.walk(dirPath))[2]
    fileName_list = [k for k in allFileName_list if partOfFileName in k]
    filePath_list = [os.path.join(dirPath, k) for k in fileName_list]
    #print(filePath_list)
    return filePath_list

# 根据配置json文件路径解析出字典
def get_jsonDict(config_jsonFilePath):
    with open(config_jsonFilePath, 'r', encoding='utf8') as file:
        fileContent = file.read()
    json_dict = json.loads(fileContent)
    className_list = json_dict['className_list']
    className_list = [k.strip() for k in className_list]
    className_list = sorted(className_list, reverse=False)
    json_dict['className_list'] = className_list
    #print(json_dict)
    return json_dict

# 模型配置类
class ModelConfig(Config):
    def __init__(self, json_dict):
        super(ModelConfig, self).__init__()
        self.NAME = json_dict['source']
        self.BACKBONE = json_dict['backbone']
        self.GPU_COUNT = 1
        self.IMAGES_PER_GPU = json_dict['batch_size']
        self.NUM_CLASSES = 1 + len(json_dict['className_list'])
        self.IMAGE_MIN_DIM = min(json_dict['image_width'], json_dict['image_height'])
        self.IMAGE_MAX_DIM = max(json_dict['image_width'], json_dict['image_height'])
        self.IMAGE_SHAPE = np.array([json_dict['image_height'], json_dict['image_width'], 3])
        self.IMAGE_META_SIZE = 15
        self.RPN_ANCHOR_SCALES = (8, 16, 32, 64, 128)
        self.TRAIN_ROIS_PER_IMAGE = 32
        self.STEPS_PER_EPOCH = json_dict['steps_per_epoch']
        self.VALIDATION_STEPS = 20
        self.LEARNING_RATE = 0.001

# 数据集类
class ShapesDataset(utils.Dataset):
    # 类初始化方法
    def __init__(self, imageFilePath_list, json_dict):
        super(ShapesDataset, self).__init__()
        self.className_list = json_dict['className_list']
        # 添加物体种类
        for i, className in enumerate(self.className_list, 1):
            self.add_class(source=json_dict['source'], class_id=i, class_name=className)
        # 添加图片
        for i, imageFilePath in enumerate(imageFilePath_list):
            jsonFilePath = os.path.splitext(imageFilePath)[0] + '.json'
            mask_pickleFilePath = os.path.splitext(imageFilePath)[0] + '_mask.pickle'
            self.add_image(source=json_dict['source'], image_id=i,
                           path=imageFilePath, jsonFilePath=jsonFilePath,
                           canLoadPickle=False, mask_pickleFilePath=mask_pickleFilePath)

    # 加载图片
    def load_image(self, image_id):
        info = self.image_info[image_id]
        image = cv2.imread(info['path'])[:,:,::-1]
        return image

    # 加载掩码
    def load_mask(self, image_id):
        info = self.image_info[image_id]
        # 每张图片样本第一次调用load_mask时需要重新生成掩码，并把掩码和种类Id列表保存为pickle文件
        # 加快训练速度
        if info['canLoadPickle'] and os.path.exists(info['mask_pickleFilePath']):
            with open(info['mask_pickleFilePath'], 'rb') as file:
                result_tuple = pickle.load(file)
        else:
            jsonFilePath = info['jsonFilePath']
            with open(jsonFilePath, 'r', encoding='utf8') as file:
                fileContent = file.read()
            json_dict = json.loads(fileContent)
            shapes = json_dict['shapes']
            #print(shapes)
            #获取label的个数
            shape_number = len(shapes)
            #print(shape_number)
            image_ndarray = cv2.imread(info['path'])
            height, width, _ = image_ndarray.shape
            mask_ndarray = np.zeros((height, width, shape_number), np.uint8)
            label_list = []
            for i, shape in enumerate(shapes):
                self.draw_mask(mask_ndarray, i, shape, label_list)
            # 解决不同实例的掩码重叠的问题，次序在后的掩码在最上层，次序在前的掩码的重叠部分被遮挡，即赋值为0
            canDrawRegion_ndarray = np.logical_not(mask_ndarray[:, :, -1]).astype(np.uint8)
            for i in range(len(shapes) - 2, -1, -1):
                mask_ndarray[:, :, i] = mask_ndarray[:, :, i] * canDrawRegion_ndarray
                canDrawRegion_ndarray = np.logical_and(canDrawRegion_ndarray, np.logical_not(mask_ndarray[:, :, i]))
            # 把物体种类名称转换为物体种类Id
            classId_list = [self.class_names.index(k) for k in label_list]
            result_tuple = (mask_ndarray.astype(np.bool), np.array(classId_list))
            info['canLoadPickle'] = True
            with open(info['mask_pickleFilePath'], 'wb') as file:
                pickle.dump(result_tuple, file)
        return result_tuple

    # 利用json文件中的信息绘制掩码
    def draw_mask(self, mask_ndarray, i, shape, label_list):
        if 'shape_type' not in shape:
            shapeType = 'polygon'
        else:
            shapeType = shape['shape_type']
        shapeType_list = ['polygon', 'rectangle', 'circle', 'line', 'point', 'linestrip']
        label = shape['label']
        label_list.append(label)
        point_list = shape['points']
        if shapeType not in shapeType_list:
            print('shape have wrong shape_type! please check json file.')
            return
        elif shapeType == 'polygon' or shapeType == 'line' or shapeType == 'linestrip':
            point_ndarray = np.array(point_list)[np.newaxis, ...]
            mask_ndarray[:, :, i:i+1] = self.draw_fillPoly(mask_ndarray[:, :, i:i+1].copy(), point_ndarray, 128)
        elif shapeType == 'rectangle':
            leftTop_point, rightDown_point = point_list
            mask_ndarray[:, :, i:i+1] = self.draw_rectangle(mask_ndarray[:, :, i:i+1].copy(), leftTop_point, rightDown_point, 128, -1)
        elif shapeType == 'circle':
            center_point, contour_point = point_list
            center_x ,center_y = center_point
            contour_x, contour_y = contour_point
            radius = int(((center_x - contour_x) ** 2 + (center_y - contour_y) ** 2) ** 0.5)
            mask_ndarray[:, :, i:i+1] = self.draw_circle(mask_ndarray[:, :, i:i+1].copy(), tuple(center_point), radius, 128, -1)
        elif shape_type == 'point':
            center_point = point_list[0]
            radius = 3
            mask_ndarray[:, :, i:i+1] = self.draw_circle(mask_ndarray[:, :, i:i+1].copy(), tuple(center_point), radius, 128, -1)

    # 画多边形
    def draw_fillPoly(self, mask, point_ndarray, color):
        cv2.fillPoly(mask, point_ndarray, 128)
        return mask

    # 画矩形
    def draw_rectangle(self, mask, leftTop_point, rightDown_point, color, thickness):
        cv2.rectangle(mask, leftTop_point, rightDown_point, color, thickness)
        return mask

    # 画圆形
    def draw_circle(self, mask, center_point, radius, color, thickness):
        cv2.circle(mask, center_point, radius, color, thickness)
        return mask

# 随机划分为3个数据集：训练集，验证集，测试集
def get_train_val_test(dirPath, image_suffix, json_dict):
    #获取文件路径
    imageFilePath_list = get_filePathList(dirPath, image_suffix)
    np.random.seed(2019)
    np.random.shuffle(imageFilePath_list)
    val_percent, test_percent = 0.05, 0.10
    sample_number = len(imageFilePath_list)
    val_number = max(int(sample_number * val_percent), 1)
    test_number = max(int(sample_number * test_percent), 1)
    train_number = sample_number - val_number - test_number
    print('训练集样本数量:%d，验证集样本数量:%d，测试集样本数量:%d' %(train_number, val_number, test_number))
    train_imageFilePath_list = imageFilePath_list[-train_number : ]
    val_imageFilePath_list = imageFilePath_list[: val_number]
    test_imageFilePath_list = imageFilePath_list[val_number : val_number+test_number]
    train_dataset = ShapesDataset(train_imageFilePath_list, json_dict)
    val_dataset = ShapesDataset(val_imageFilePath_list, json_dict)
    test_dataset = ShapesDataset(test_imageFilePath_list, json_dict)
    train_dataset.prepare()
    val_dataset.prepare()
    test_dataset.prepare()
    return train_dataset, val_dataset, test_dataset

# 解析调用代码文件时传入的参数
import argparse
def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('-i', '--in_dirPath',
        type=str,
        help='输入图片文件夹路径',
        default='resources/resized_images_640x640')
    parser.add_argument('--image_suffix',
        type=str,
        help='图片文件的后缀',
        default='.jpg')
    parser.add_argument('-c', '--config',
        type=str,
        help='模型配置json文件路径',
        default='resources/model_config.json')
    argument_namespace = parser.parse_args()
    return argument_namespace

if __name__ == '__main__':
    # 解析出调用代码文件时传入的参数
    argument_namespace = parse_args()
    #文件路径
    in_dirPath = argument_namespace.in_dirPath.strip()
    #文件后缀
    image_suffix = argument_namespace.image_suffix.strip()
    #json配置文件
    config_jsonFilePath = argument_namespace.config.strip()
    # 根据配置文件和配置类ModelConfig实例化模型配置对象
    #解析成字典
    json_dict = get_jsonDict(config_jsonFilePath)
    #print(json_dict)
    #实例化，获取参数
    modelConfig = ModelConfig(json_dict)
    modelConfig.display()
    # 随机划分训练集、验证集、测试集，在模型训练中，测试集没有被用到

    train_dataset, val_dataset, _ = get_train_val_test(in_dirPath, image_suffix, json_dict)
    # 实例化模型对象
    model = modellib.MaskRCNN(mode="training",
        config=modelConfig,
        model_dir=MODEL_DIR)
    # 模型加载COCO数据集的预训练权重
    model.load_weights(COCO_MODEL_PATH,
        by_name=True,
        exclude=["mrcnn_class_logits",
            "mrcnn_bbox_fc",
            "mrcnn_bbox",
            "mrcnn_mask"])
    # 模型fine-tune第1步：训练卷积核以外的层
    print('模型fine-tune第1步：训练卷积核以外的层')
    model.train(train_dataset,
        val_dataset,
        learning_rate=modelConfig.LEARNING_RATE,
        epochs=1,
        layers='heads')
    # 模型fine-tune第2步：训练全部层
    print('模型fine-tune第2步：训练全部层')
    model.train(train_dataset,
        val_dataset,
        learning_rate=modelConfig.LEARNING_RATE / 10,
        epochs=20,
        layers="all")

测试一张图片

import os
import sys
import cv2
import time
import numpy as np
import json
# 工程的根目录
ROOT_DIR = os.path.abspath("../resources/")
# 导入Mask RCNN库
sys.path.append(ROOT_DIR) 
from mrcnn.config import Config
import mrcnn.model as modellib
from mrcnn import visualize

# 获取文件夹中的文件路径
def get_filePathList(dirPath, partOfFileName=''):
    allFileName_list = next(os.walk(dirPath))[2]
    fileName_list = [k for k in allFileName_list if partOfFileName in k]
    filePath_list = [os.path.join(dirPath, k) for k in fileName_list]
    return filePath_list 
    
# 根据配置json文件路径解析出字典
def get_jsonDict(config_jsonFilePath):
    with open(config_jsonFilePath, 'r', encoding='utf8') as file:
        fileContent = file.read()
    json_dict = json.loads(fileContent)
    className_list = json_dict['className_list']
    className_list = [k.strip() for k in className_list]
    className_list = sorted(className_list, reverse=False)
    json_dict['className_list'] = className_list
    return json_dict

# 模型测试配置类
class InferenceConfig(Config):
    def __init__(self, config_dict):
        super(InferenceConfig, self).__init__()
        self.NAME = config_dict['source']
        self.BACKBONE = config_dict['backbone']
        self.GPU_COUNT = 1
        self.IMAGES_PER_GPU = 1
        self.BATCH_SIZE =1
        self.NUM_CLASSES = 1 + len(config_dict['className_list'])
        self.IMAGE_MIN_DIM = min(config_dict['image_width'], config_dict['image_height'])
        self.IMAGE_MAX_DIM = max(config_dict['image_width'], config_dict['image_height'])
        self.IMAGE_SHAPE = np.array([config_dict['image_height'], config_dict['image_width'], 3])
        self.IMAGE_META_SIZE = 15
        self.RPN_ANCHOR_SCALES = (8, 16, 32, 64, 128)
        self.TRAIN_ROIS_PER_IMAGE = 32
        self.STEPS_PER_EPOCH = config_dict['steps_per_epoch']
        self.VALIDATION_STEPS = 20
        self.LEARNING_RATE = 0.001

# 获取模型对象    
def get_model(model_dirPath, config_dict):
    model = modellib.MaskRCNN(mode="inference", 
        config=InferenceConfig(config_dict),
        model_dir=model_dirPath)
    weights_path = model.find_last()
    print('模型加载权重文件，权重文件的路径：%s' %weights_path)
    model.load_weights(weights_path, by_name=True)
    return model

# 定义检测函数，并将检测结果用matplotlib库绘制出来    
def detect_image(model, imageFilePath, config_dict):
    image_ndarray = cv2.imread(imageFilePath)[:, :, ::-1]
    results = model.detect([image_ndarray], verbose=0)
    r = results[0]
    className_list = ['BG'] + config_dict['className_list']
    visualize.display_instances(image_ndarray, r['rois'], r['masks'], r['class_ids'], 
                            className_list, r['scores'], figsize=(8, 8))

# 解析调用代码文件时传入的参数    
import argparse
def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('-m', '--model_dirPath',
        type=str,
        help='模型权重文件夹路径',
        default='resources/logs')
    parser.add_argument('-i', '--image_dirPath',
        type=str,
        help='图片文件夹路径',
        default='resources/n01440764')
    parser.add_argument('--image_suffix',
        type=str,
        help='图片文件的后缀',
        default='.jpg')
    parser.add_argument('-c', '--config',
        type=str,
        help='模型配置json文件路径',
        default='resources/model_config.json')
    argument_namespace = parser.parse_args()
    return argument_namespace 
    
if __name__ == '__main__':
    # 解析传入的参数
    argument_namespace = parse_args()
    model_dirPath = argument_namespace.model_dirPath.strip()
    image_dirPath = argument_namespace.image_dirPath.strip()
    image_suffix = argument_namespace.image_suffix.strip()
    config_jsonFilePath = argument_namespace.config.strip()
    # 获取模型配置字典，并实例化模型对象 
    config_dict = get_jsonDict(config_jsonFilePath)
    model = get_model(model_dirPath, config_dict)
    # 获取图片文件路径
    imageFilePath_list = get_filePathList(image_dirPath, image_suffix)
    assert len(imageFilePath_list), 'no image in image directory path, please check your input parameters: image_dirPath , image_suffix'
    imageFilePath = np.random.choice(imageFilePath_list, 1)[0]
    print('对此文件路径的图片做检测：%s'%imageFilePath)
    # 对单张图片做检测
    detect_image(model, imageFilePath, config_dict)