class MaskR:
def __init__(self, model_dir, init_with='coco'):
self.model_dir = model_dir
self.init_with = init_with
self.test_model = None
def train(self, dataset_train, dataset_val, epochs):
'''
'''
t0 = time.time()
print ('GPU available:', tf.test.is_gpu_available())
self.mode = 'training'
self.config = C.TrainingConfig()
if not os.path.exists(self.model_dir):
os.mkdir(self.model_dir)
self.model = MaskRCNN(self.mode, self.config, self.model_dir)
# Which weights to start with?
# imagenet, coco, or last
# Local path to trained weights file
COCO_MODEL_PATH = os.path.join(self.model_dir, "mask_rcnn_coco.h5")
# Download COCO trained weights from Releases if needed
if not os.path.exists(COCO_MODEL_PATH):
utils.download_trained_weights(COCO_MODEL_PATH)
if self.init_with == "imagenet":
self.model.load_weights(self.model.get_imagenet_weights(), by_name=True)
elif self.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
self.model.load_weights(COCO_MODEL_PATH, by_name=True,
exclude=["mrcnn_class_logits", "mrcnn_bbox_fc",
"mrcnn_bbox", "mrcnn_mask"])
elif self.init_with == "last":
# Load the last model you trained and continue training
self.model.load_weights(model.find_last(), by_name=True)
print ('MaskRCNN Setup complete after', time.time()-t0, 'seconds')
t0 = time.time()
history = History()
self.model.train(dataset_train, dataset_val, custom_callbacks=[history],
learning_rate=self.config.LEARNING_RATE,
epochs=epochs,
layers='heads')
p.dump(history.history, open(os.path.join(self.model_dir, "history.p"), "wb"))
print ('MaskRCNN Training complete after', time.time()-t0, 'seconds')
return history
def predict(self, images, verbose=False):
'''
'''
if not self.test_model:
model = MaskRCNN(mode="inference",
config=C.TestingConfig(),
model_dir=self.model_dir)
weights = model.find_last()
model.load_weights(weights, by_name=True)
self.test_model = model
results = []
for image in images:
results.append(self.test_model.detect([image])[0])
if verbose:
r = results[0]
visualize.display_instances(images[0], r['rois'], r['masks'], r['class_ids'],
["",""], r['scores'],figsize=(10,10))
return results
else:
model = MaskRCNN(mode="inference", config=config, model_dir=args.logs)
# Select weights file to load
COCO_WEIGHTS_PATH = "./mask_rcnn_coco.h5"
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()
elif args.weights.lower() == "imagenet":
# Start from ImageNet trained weights
weights_path = model.get_imagenet_weights()
else:
weights_path = args.weights
# Load weights
print("Loading weights ", weights_path)
if args.weights.lower() == "coco":
# Exclude the last layers because they require a matching
# number of classes
model.load_weights(
weights_path,
by_name=True,
exclude=[
"mrcnn_class_logits",
"mrcnn_bbox_fc",
"mrcnn_bbox",
# define a configuration for the model
class PlotImageConfig(Config):
# define the name of the configuration
NAME = "plotimage_cfg"
# number of classes (background + kangaroo)
NUM_CLASSES = 1 + 1
# number of training steps per epoch
STEPS_PER_EPOCH = 131
# prepare train set
train_set = PlotImageDataset()
train_set.load_dataset(input_annotations_dir, is_train=True)
train_set.prepare()
print('Train: %d' % len(train_set.image_ids))
# prepare test/val set
test_set = PlotImageDataset()
test_set.load_dataset(input_annotations_dir, is_train=False)
test_set.prepare()
print('Test: %d' % len(test_set.image_ids))
# prepare config
config = PlotImageConfig()
config.display()
# define the model
model = MaskRCNN(mode='training', model_dir=output_model_dir, config=config)
# load weights (mscoco) and exclude the output layers
# model.load_weights('mask_rcnn_coco.h5', by_name=True, exclude=["mrcnn_class_logits", "mrcnn_bbox_fc", "mrcnn_bbox", "mrcnn_mask"])
model.load_weights(model.get_imagenet_weights(), by_name=True)
# train weights (output layers or 'heads' or 'all')
model.train(train_set, test_set, learning_rate=config.LEARNING_RATE, epochs=5, layers='all')
model.keras_model.save_weights(output_model_path)
