Mask_RCNN訓練自己的資料

看過之前文章的童鞋，應該已經能夠使用VIA完成資料的标注工作：

Mask_RCNN訓練自己的資料，制作類似于COCO資料集中所需要的Json訓練集

現在來具體說一下如何使用标注生成的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.0001
           

4.執行訓練，并顯示結果

在配置完以上步驟之後，确認你的訓練集和測試集以及train_mask.py的檔案組織方式如下：

    ├── train

    ├── val

    ├── train_mask.py

    └── model.py...

執行train_mask即可

5.用訓練好的模型預測

具體看之前的博文：

Mask_RCNN：使用自己訓練好的模型進行預測

結果圖：

最後，夜王鎮樓：