用maskrcnn训练自己的数据集

https://blog.csdn.net/Xiongchao99/article/details/79106588文章中提到如何构建自己的数据集

自己的方法就是改写demo和train_shapes,数据格式仍然采用CocoDataset，自己标定好直接替换coco数据
1.训练配置参数
2.制作训练集和验证集
3.模型生成及训练
4.模型预测

1.训练配置参数
MydataConfig类，主要改动类别数

class MydataConfig(coco.CocoConfig):
    # 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
        # 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 = 128
   # IMAGE_MAX_DIM = 128

    # Use smaller anchors because our image and objects are small
   # RPN_ANCHOR_SCALES = (8, 16, 32, 64, 128)  # 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 = 32

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

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

config = MydataConfig()
config.display()

2.数据集构建
用coco的cocodataset,只需要改变他所需要的文件名字就好了。

注意：COCO_DIR目录下，存放val2017和train2017的img

# Load my_train_COCO dataset    训练集json
COCO_DIR='H:\\pyprogram\\mask_rcnn\\labelme-master'
dataset_train = coco.CocoDataset()
dataset_train.load_coco(COCO_DIR, "train",2017)
dataset_train.prepare()
print(dataset_train.class_names)
class_names=dataset_train.class_names

# Load my_val_COCO dataset    验证集json
COCO_DIR='H:\\pyprogram\\mask_rcnn\\labelme-master'
dataset_val = coco.CocoDataset()
dataset_val.load_coco(COCO_DIR, "val",2017)
dataset_val.prepare()
print(dataset_val.class_names)
class_val_names=dataset_val.class_names

3.训练模型

生成“training”模式的模型，加载coco权重（只需要改变COCO_MODEL_PATH改变权重文件来接着训练），调用训练命令

# 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
    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)

# 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=2, 
            layers="all")

4.模型测试
再次生成一个配置类和模型，调用1中定义的Mydataconfig类，模型是“inference”模式。
通过加载权重，选择预测图片，预测，可视化。

class InferenceConfig(MydataConfig):
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1
inference_config = InferenceConfig()

# Recreate the model in inference mode
model_test = modellib.MaskRCNN(mode="inference", model_dir=MODEL_DIR, config=config)

# Load weights trained on MS-COCO
COCO_MODEL_TEST_PATH='H:\\pyprogram\\mask_rcnn\\maskRCNN\\mydata_logs\\coco20180410T2304\\mask_rcnn_coco_0002.h5'
model_test.load_weights(COCO_MODEL_TEST_PATH, by_name=True)

# Load a random image from the images folder
file_names = next(os.walk(IMAGE_DIR))[2]
image = skimage.io.imread(os.path.join(IMAGE_DIR, random.choice(file_names)))

# Run detection
results = model_test.detect([image], verbose=1)

# Visualize results
r = results[0]
visualize.display_instances(image, r['rois'], r['masks'], r['class_ids'], 
                            class_names, r['scores'])


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