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
# Create model
if args.command == "train":
model = MaskRCNN(mode="training", config=config, model_dir=args.logs)
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=[
test_set.load_dataset('dataset', is_train=False)
test_set.prepare()
print('Test: %d' % len(test_set.image_ids))
# prepare config
config = PlaneConfig()
config.display()
# define the model
model = MaskRCNN(mode='training', model_dir='./logs/', config=config)
# load weights (mscoco) and exclude the output layers
# Select weights file to load
if args.model.lower() == "coco":
model_path = 'mask_rcnn_plane.h5'
elif args.model.lower() == "last":
# Find last trained weights
model_path = model.find_last()
else:
model_path = args.model
print("Training on: " + model_path)
model.load_weights(model_path,
by_name=True,
exclude=[
"mrcnn_class_logits", "mrcnn_bbox_fc", "mrcnn_bbox",
"mrcnn_mask"
])
# train weights (output layers or 'heads')
model.train(train_set,
test_set,
learning_rate=config.LEARNING_RATE,
epochs=100,
data_train.load_dataset(train_ids, TRAIN_IMAGES_DIR)
data_train.prepare()
data_val = SteelDataset(df_defects)
data_val.load_dataset(val_ids,
TRAIN_IMAGES_DIR) # val images are in the train dir
data_val.prepare()
steel_config = SteelConfig()
model = MaskRCNN(mode="training",
config=steel_config,
model_dir=str(MODEL_DIR))
if args.model.lower() == "last":
model_path = model.find_last() # Find last trained weights
elif args.model.lower() == "default":
model_path = str(COCO_MODEL_PATH)
# Exclude the last layers because they require a matching number of classes
# ^ Original comment at:
# https://github.com/matterport/Mask_RCNN/blob/master/samples/balloon/balloon.py
model.load_weights(filepath=model_path,
by_name=True,
exclude=[
"mrcnn_class_logits", "mrcnn_bbox_fc", "mrcnn_bbox",
"mrcnn_mask"
])
model.train(train_dataset=data_train,
val_dataset=data_val,
