def load_network(self): config = Config() config.NAME = 'predict' config.NUM_CLASSES = 1 + 1 config.IMAGES_PER_GPU = 1 config.GPU_COUNT = 1 additional_info = json.load(open(self.config_path)) for i,j in additional_info.items(): try: setattr(config,i,eval(j)) except: setattr(config,i,j) config.__init__() from mrcnn import model as modellib self.model = modellib.MaskRCNN(mode="inference", model_dir='./',config=config)
router_train = DatasetXML()
router_train.load_dataset(dataset_dir, label, train_imgs)
router_train.prepare()
print('Train: %d' % len(router_train.image_ids))
router_test = DatasetXML()
router_test.load_dataset(dataset_dir, label, test_imgs)
router_test.prepare()
print('Test: %d' % len(router_test.image_ids))
else: # pkl format
file_name = 'data_val2017-laptop'
pkldata = pkl.load(open(os.path.join(dataset_dir, file_name+'.pkl'), 'rb'))
# =============== train models ===============
# prepare config
config = Config()
config.NAME = label + "_cfg" # Give the configuration a recognizable name
config.STEPS_PER_EPOCH = n_steps # Number of training steps per epoch
config.NUM_CLASSES = 1 + 1 # Number of classes (background + router)
model = MaskRCNN(mode='training', model_dir=model_root_dir, config=config)
model.load_weights(model_dir, by_name=True,
exclude=["mrcnn_class_logits", "mrcnn_bbox_fc", "mrcnn_bbox", "mrcnn_mask"])
model.train(router_train, router_test, learning_rate=config.LEARNING_RATE, epochs=n_epochs, layers='heads')