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)