def load_model(weights, dataset, logs='../../'):
class InferenceConfig(ProfileConfig):
# 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
config = InferenceConfig()
# config.display()
model = modellib.MaskRCNN(mode="inference", config=config, model_dir=logs)
weights_path = 'mask_rcnn_profile_0010.h5'
model.load_weights(weights_path, by_name=True)
return model
def main(*args):
# 配置参数
config = MyTrainConfig()
# 图片尺寸统一处理为1024，可以根据实际情况再进一步调小
config.IMAGE_MIN_DIM = FLAGS.image_dim
config.IMAGE_MAX_DIM = FLAGS.image_dim
# 每轮训练的step数量
config.STEPS_PER_EPOCH = FLAGS.epochs_step
# 每轮验证的step数量
config.VALIDATION_STEPS = FLAGS.validation_step
config.display()
# 生成训练集
dataset_train = CocoDataset()
dataset_train.load_coco(COCO_DIR, "train") # 加载训练数据集
dataset_train.prepare()
# 生成验证集
dataset_val = CocoDataset()
dataset_val.load_coco(COCO_DIR, "val") # 加载验证数据集
dataset_val.prepare()
# 模型实例训练
model = modellib.MaskRCNN(mode="training", config=config, model_dir="log")
# 训练权重
model.load_weights(COCO_MODEL_PATH, by_name=True)
# 模型训练
model.train(dataset_train,
dataset_val,
learning_rate=config.LEARNING_RATE,
epochs=FLAGS.max_epochs,
layers='all')
# 保存模型训练权重
save_wrapper = os.path.join(local_output_path, "train_mask_rcnn.h5")
model.keras_model.save_weights(save_wrapper)
# 复制保存到桶
file.copy_parallel(local_output_path, FLAGS.train_url)
config = CustomConfig()
else:
class InferenceConfig(CustomConfig):
# 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
config = InferenceConfig()
config.display()
# Create model
if args.command == "train":
model = modellib.MaskRCNN(mode="training",
config=config,
model_dir=args.logs)
else:
model = modellib.MaskRCNN(mode="inference",
config=config,
model_dir=args.logs)
# Select weights file to load
if args.weights.lower() == "coco":
weights_path = COCO_WEIGHTS_PATH
# Download weights file
if not os.path.exists(weights_path):
utils.download_trained_weights(weights_path)
elif args.weights.lower() == "last":
# Find last trained weights
weights_path = model.find_last()[1]