看过之前文章的童鞋，应该已经能够使用VIA完成数据的标注工作：
现在来具体说一下如何使用标注生成的json进行自定义模型的训练~
1.训练函数主干：
其中自带的train_shapes.ipynb里面基本包含了训练所需的参数，但是jupyter notebook的调试起来太烦了，因此我新建一个名为train_mask.py的文件，该文件用于调用训练集，验证集的加载以及训练过程中超参数的配置等等。
1.1.导入所需要的库
import os
import sys
import visualize
import model as modellib
import balloon这里面的visualize,model,balloon是该版本mask_rcnn自带的库函数，也是后面我们需要针对自己的数据集进行修改的！如果导入的时候带有红色的下划线，像这样：
不用担心，只需要保证这3个函数文件和刚才新建的train_mask.py在同一级目录下即可
1.2.获得训练根路径并加载coco预训练权重
#该文件用于训练模型
# 获取该project的根目录
ROOT_DIR = os.path.abspath('')
# 存放训练模型的logs路径
MODEL_DIR = os.path.join(ROOT_DIR, "logs")
# 加载预训练的COCO权重
COCO_MODEL_PATH = os.path.join(ROOT_DIR, "mask_rcnn_coco.h5")
sys.path.append(ROOT_DIR)
# 从balloon文件中引入配置
config = balloon.BalloonConfig()
#datasets文件夹下是训练集和验证集
BALLOON_DIR = os.path.join(ROOT_DIR, "datasets")1.3.加载验证集和训练集并可视化
#加载训练集
dataset_train = balloon.BalloonDataset()
dataset_train.load_balloon(BALLOON_DIR, "train")
dataset_train.prepare()
#打印出训练集信息
print("Train Image Count: {}".format(len(dataset_train.image_ids)))
print("Class Count: {}".format(dataset_train.num_classes))
for i, info in enumerate(dataset_train.class_info):
print("{:3}. {:50}".format(i, info['name']))
#从训练集中选出4张用与显示训练集是否正确
image_ids = dataset_train.image_ids[:4]
for image_id in image_ids:
print(image_id)
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)
#加载验证集
dataset_val = balloon.BalloonDataset()
dataset_val.load_balloon(BALLOON_DIR, "val")
dataset_val.prepare()
#打印出验证集信息
print("Val Image Count: {}".format(len(dataset_val.image_ids)))
print("Class Count: {}".format(dataset_val.num_classes))
for i, info in enumerate(dataset_val.class_info):
print("{:3}. {:50}".format(i, info['name']))
#从验证集中选出4张用与显示验证集是否正确
image_ids = np.random.choice(dataset_val.image_ids, 4)
for image_id in image_ids:
image = dataset_val.load_image(image_id)
mask, class_ids = dataset_val.load_mask(image_id)
visualize.display_top_masks(image, mask, class_ids, dataset_train.class_names)这时你会看到，训练集和测试集的相关信息显示出来，并且标记的图框也会跟着图片显示：
1.3.建立模型并设定权重初始化方式以及超参数
#创建一个训练模型
model = modellib.MaskRCNN(mode="training", config=config,
model_dir=MODEL_DIR)
#设定权重的初始化方式，有imagenet，coco,last三种
init_with = "coco"
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)
#设定训练参数，如学习率，epoch等
model.train(dataset_train, dataset_val, learning_rate=config.LEARNING_RATE, epochs=50, layers="all")这个时候我们已经完成了train_mask的全部任务了，接下来针对via格式的标注文件，完成载入方式的适配！
2. 载入VIA所标记的标记文件并生成mask
在这一步，主要是为via生成的标注格式量身定制一套数据加载方法，主要是根据balloon.py进行修改，这一步应该算是核心了。
首先，打开程序自带的balloon.py,将balloon数据集改为自己的数据集的名字（其实不用改也行，只是一个label而已，因为这里只演示一个类，我就直接用balloon啦~），例如，你可以随意将Something换为任何类别名称：
class SomethingDataset(utils.Dataset):在该类之下，添加之前文章所写的必须包含的几个函数：
def load_balloon(self, dataset_dir, subset):
"""定义数据集有哪些类别.
dataset_dir: 数据集根路径.
subset: train or val
"""
# 添加类别的名称和ID号，在这里，为了简便，我只添加一类：‘car’，如果是多类，挨着添加即可，例如我还有‘cat’,‘dog’....
self.add_class("balloon", 1, "balloon")
self.add_class("balloon", 2, "cat")
self.add_class("balloon", 3, "dog")
# 加载训练集还是测试集?
assert subset in ["train", "val"]
dataset_dir = os.path.join(dataset_dir, subset)
# 加载数据集
# VIA标记工具的格式如下:
# { 'filename': '28503151_5b5b7ec140_b.jpg',
# 'regions': {
# '0': {
# 'region_attributes': {},
# 'shape_attributes': {
# 'all_points_x': [...],
# 'all_points_y': [...],
# 'name': 'polygon'}},
# ... more regions ...
# },
# 'size': 100202
# }
# 读入json标注数据集
annotations = json.load(open(os.path.join(dataset_dir, "via_region_data.json")))
# 获得字典的键值
annotations = list(annotations.values())
# 拿到每一个regions的信息并组成一个列表
annotations = [a for a in annotations if a['regions']]
# 加载图片
for a in annotations:
# 获取组成每个物体实例轮廓的多边形点的x、y坐标。
# 这些坐标保存在r['shape_attributes'中，（参见上面的json格式）
polygons = [r['shape_attributes'] for r in a['regions'].values()]
image_path = os.path.join(dataset_dir, a['filename'])
image = skimage.io.imread(image_path)
height, width = image.shape[:2]
self.add_image(
"balloon",
image_id=a['filename'],
path=image_path,
width=width, height=height,
polygons=polygons)
def load_mask(self, image_id):
"""为图像生成实例mask.
Returns:
masks: 一个bool数组，每个实例一个mask,
其中每个mask为：[高, 宽, 数目]
class_ids: 每个mask对应的类别.
"""
image_info = self.image_info[image_id]
if image_info["source"] != "balloon":
return super(self.__class__, self).load_mask(image_id)
# 将多边形转换为一个二值mask,[高, 宽, 数目]
info = self.image_info[image_id]
# print("这是P",info["height"], info["width"], len(info["polygons"]))
mask = np.zeros([info["height"], info["width"], len(info["polygons"])],
dtype=np.uint8)
for i, p in enumerate(info["polygons"]):
# 获取多边形内的像素索引并将其设置为1
rr, cc = skimage.draw.rectangle((p['y'], p['x']), extent=(p['height'], p['width']))
mask[rr, cc, i] = 1
# print("info['class_ids']", info['class_ids'])
# 返回mask和类别的信息
return mask, np.ones([mask.shape[-1]], dtype=np.int32)
def image_reference(self, image_id):
"""Return the path of the image."""
info = self.image_info[image_id]
if info["source"] == "balloon":
return info["path"]
else:
super(self.__class__, self).image_reference(image_id)3.修改配置函数，适配自己的数据集
同时，你还要针对自己的数据集，新建一个配置函数，名字应该和上面对应：
class SomethinConfig(Config):由于现在你只添加了一类物品，因此需要将NUM_CLASSES改为1+1：
# 类别的数目（包括背景）
NUM_CLASSES = 1 + 1并且根据自己的图片设置数据维度，例如我的图片大小为1024*800，因此我设置为：
IMAGE_MIN_DIM = 800
IMAGE_MAX_DIM = 1024因为数据集较小，可以把STEPS_PER_EPOCH改小一点，这样训练函数的loss下降曲线会非常平滑~~：
# 每一个epoch迭代多少次
STEPS_PER_EPOCH = 100我的显卡有12G的显存，我就填的2，如果你的显卡不太好的话，你可以把IMAGES_PER_GPU设为1,如果这个时候还是出现OOM问题的话，可以将backbone设置为稍小一点的resnet50：
# TITAN X有12gb显存可以一次装两张，因此设置为2
IMAGES_PER_GPU = 2
# Backbone架构，用于提取特征，包括两种: resnet50, resnet101
BACKBONE = "resnet50"接着根据数据的特性设置学习率和权重衰减因子，
# 学习率和动量，不同的优化器有不同的配置，而且不同框架之间的实现也不一样，因此需要自己尝试
LEARNING_RATE = 0.001
LEARNING_MOMENTUM = 0.9
# 权重衰减因子
WEIGHT_DECAY = 0.00014.执行训练，并显示结果
在配置完以上步骤之后，确认你的训练集和测试集以及train_mask.py的文件组织方式如下：
├── train
├── val
├── train_mask.py
└── model.py...
执行train_mask即可
5.用训练好的模型预测
具体看之前的博文：
结果图：
最后，习大大镇楼：
