Mask R-CNN(https://arxiv.org/abs/1703.06870v3/)是基于Faster R-CNN算法改进的一个目标检测与实例分割算法,本文并不着重探讨算法本身,而就Mask R-CNN的一个Python实现(https://github.com/matterport/Mask_RCNN/)的编程实践进行探讨。作者给出关于训练自己的数据集的样例可以在(https://github.com/matterport/Mask_RCNN/blob/master/samples/shapes/train_shapes.ipynb/)上找到。

Config

config.py中主要实现了Config类,即训练参数类。

Base configuration class. For custom configurations, create a
sub-class that inherits from this one and override properties
that need to be changed.

作者建议直接继承Config类,并重写需要更改的参数,这里对几个重要参数进行说明。

GPU_COUNT

GPU数量,默认为1

注意,当仅使用CPU时这个值应当被设置为1。

IMAGES_PER_GPU

每个GPU上所处理的图像数量,默认为2

Number of images to train with on each GPU. A 12GB GPU can typically handle 2 images of 1024x1024px.

作者建议,一个有12G显存的GPU可以处理两张1024*1024px的图像。

STEPS_PER_EPOCH

每一次迭代的步骤数,默认为1000

由于每一次迭代末尾都需要进行Tensorboard的更新和验证集的验证,这个值不应该设置的过小。

NUM_CLASSES

需要识别的物体类的个数(包括背景),默认为1

实际上这里只需要设置参数为1+需要识别的物体的个数即可。

BACKBONE

骨干网络结构,默认为resnet101

在这个实现中还有resnet50可以使用,如果需要可以重写一个自定义的骨干网络结构。

RPN_ANCHOR_SCALES

RPN(Region Proposal Network)锚点标度,默认为(32, 64, 128, 256, 512)

rpn.png

引用一张来自Faster R-CNN论文(https://arxiv.org/abs/1506.01497v3)中的一张图片,这个参数实际上就是对应图右侧所示候选区域的大小。默认参数给出了五种不同的大小,我们需要根据实际需要去调整最合适的参数(和图片像素大小对应)。

IMAGE_RESIZE_MODE & IMAGE_MIN_DIM & IMAGE_MAX_DIM

在这个实现中,为了能够让多组图像在同一批进行训练,作者在实现中对图像进行了重置大小操作。默认模式(IMAGE_RESIZE_MODE)为square,即重置为一个方形。具体实现是使用的双线性插值(Bilinear Interpolation),可见utils.py:388

IMAGE_MIN_DIM与IMAGE_MAX_DIM即为图像最小最大大小。若图像大小小于最小值则会被放大,若大于最大值则会被缩小。

这个参数当考虑硬件实际情况酌情设置,过大的大小会拖慢训练速度,过小的大小会使训练结果变差。

注意,这里的IMAGE_MAX_DIM大小要设置为2的倍数以防止在降尺度与升尺度时出现分数。(至少要能被2整除6次)

TRAIN_ROIS_PER_IMAGE

每张图片的RoI(Region of Interest)数量,默认为200

简单的来说就是重要的区域的个数,详见上述论文。由于RPN生成的候选区的情况不一,可以根据实际情况动态调整。

LEARNING_RATE

学习率,默认为0.001

作者提到,在Mask R-CNN论文中使用的学习率为0.02,但这个值放在TensorFlow实现下会导致权重爆炸,故在此设置一个较低值。


在重写完配置并实例化后,我们可以调用display方法查看我们的训练配置。

Utils

utils.py中主要实现了一些工具,包括计算IOU,数据集类(Dataset)等,这里就如何重写数据集类进行说明。

首先,继承Dataset类,并至少实现以下操作:

  • 添加需要识别的物体
  • 添加图像
  • load_mask()

下面给出一个具体的示例:

具体示例

class ELDataset(utils.Dataset):
    def load_el(self, count, img_folder, mask_folder, img_list, dataset_root_path):
          # 向数据集中加入需要识别的物体
        self.add_class("shapes", 1, "el")

        for i in range(count):
            file_name = img_list[i].split(".")[0]
            mask_path = mask_folder + file_name + ".png"
            mask_img = np.array(PIL.Image.open(mask_path).convert('L'))
            
            # 向数据集中加入图像
            self.add_image("shapes", image_id=i, path=img_folder + img_list[i], 
                           mask_path=mask_path, width=mask_img.shape[1], 
                           height=mask_img.shape[0])
            
    def load_mask(self, image_id):
        info = self.image_info[image_id]
        # Mask可以存在多层,在这个示例里仅使用了一层
        mask = np.zeros([info['height'], info['width'], 1], dtype=np.uint8)
        mask[:, :, 0] =  np.array(PIL.Image.open(info['mask_path']).convert('L'))
        
        return mask.astype(np.bool), np.array([1]).astype(np.int32)
      
# 路径设定
train_dataset_root_path = "data_train/"
train_img_folder = train_dataset_root_path + "img/"
train_mask_folder = train_dataset_root_path + "mask/"
train_img_list = os.listdir(train_img_folder)
train_count = len(train_img_list)

val_dataset_root_path = "data_val/"
val_img_folder = val_dataset_root_path + "img/"
val_mask_folder = val_dataset_root_path + "mask/"
val_img_list = os.listdir(val_img_folder)
val_count = len(val_img_list)

# 测试数据集
dataset_train = ELDataset()
dataset_train.load_el(train_count, train_img_folder, train_mask_folder, train_img_list, 
                          train_dataset_root_path)
dataset_train.prepare()

# 验证数据集
dataset_val = ELDataset()
dataset_val.load_el(1, val_img_folder, val_mask_folder, val_img_list, 
                        val_dataset_root_path)
dataset_val.prepare()

Model

model.py中主要实现了算法模型。

我们可以通过以下语句创建训练模型:

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

在训练模式下,我们可以对模型进行训练。模型的权重设置为以下三种:

  • imagenet : 使用ImageNet预训练模型。
  • coco : 使用COCO预训练模型。
  • last : 使用上一次训练的模型。

最后,我们可以通过以下语句开始训练:

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

Train in two stages:

  1. Only the heads. Here we're freezing all the backbone layers and training only the randomly initialized layers (i.e. the ones that we didn't use pre-trained weights from MS COCO). To train only the head layers, pass layers='heads' to the train() function.
  2. Fine-tune all layers. For this simple example it's not necessary, but we're including it to show the process. Simply pass layers="all to train all layers.

我们可以更改layers参数为headsall选择我们需要训练的网络层,更改epoch参数设定训练迭代次数。

Visualize

visualize.py中主要实现了对数据及结果进行可视化的方法。

例如,如下方法可以实现对识别结果的可视化:

visualize.display_instances(original_image, r['rois'], r['masks'], r['class_ids'], 
                            dataset_val.class_names, r['scores'], ax=get_ax())

visualize.png


引用目录


本文实际完成于2020.4.27,月更博主这个月没有鸽!

Last modification:April 27, 2020
博客维护不易,如果你觉得我的文章有用,请随意赞赏