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