在目录Mask_RCNN-master\samples\shapes中,根据train_shapes.ipynb创建自己的训练文件
训练代码如下,模型文件会保存在该文件的上级目录的logs文件夹中,每个epoch结束保存一次
# -*- coding: utf-8 -*-
"""
Created on Wed Nov 6 19:06:46 2019
@author: Lenovo
"""
import os
import sys
import random
import math
import re
import time
import numpy as np
import cv2
import matplotlib
import matplotlib.pyplot as plt
from PIL import Image
import yaml
import skimage.io
# 添加文件跟目录到搜索路径
ROOT_DIR = os.path.abspath("../../")
# 引入MaskRCNN
sys.path.append(ROOT_DIR)
from mrcnn.config import Config
from mrcnn import utils
import mrcnn.model as modellib
from mrcnn import visualize
from mrcnn.model import log
################################保存模型以及下载预训权重###############################################
# 保存日志和训练模型的文件夹
MODEL_DIR = os.path.join(ROOT_DIR, "logs")
# 预训练权重文件的路径,没有则自动下载
COCO_MODEL_PATH = os.path.join(ROOT_DIR, "mask_rcnn_coco.h5")
if not os.path.exists(COCO_MODEL_PATH):
utils.download_trained_weights(COCO_MODEL_PATH)
################################自定义配置文件################################################
class HooksConfig(Config):
"""
在shapes数据集上进行配置修改,即创建继承于基类Config的子类并且修改对应的属性值即可
"""
NAME = "hooks"
GPU_COUNT = 1
IMAGES_PER_GPU = 2 #每个GPU处理的图片数量
NUM_CLASSES = 1 + 1 # background + hooks
# 设置图片resize时的最大最小尺寸
IMAGE_MIN_DIM = 128*8
IMAGE_MAX_DIM = 128*8
# 使用的anchor尺寸,根据resize信息进行设置,8*8指的是正方形anchor的边长,也就是此时anchor的面积为(8*8)^2
RPN_ANCHOR_SCALES = (8*8, 16*6, 32*8, 64*8, 128*8)
# 每张图片中ROIs的数量,适当减少可以加速训练,目的在于允许ROI采样1/3的正类ROIs
TRAIN_ROIS_PER_IMAGE = 32
# 当数据比较简单时,可以使用较小的epoch
STEPS_PER_EPOCH = 100
# epoch较小时,也使用较小的验证步长
VALIDATION_STEPS = 50
# 生成config对象并打印
config = HooksConfig()
config.display()
##########################用于可视化######################################################
def get_ax(rows=1, cols=1