def main(): dir_path = os.path.dirname(os.path.realpath(__file__)) MODEL_DIR = dir_path + "/models/" COCO_MODEL_PATH = "/Data/mask_rcnn_coco.h5" #os.path.normpath(dir_path + "/Data/mask_rcnn_coco.h5") DATA_PATH = "E:\\Data\\" #dir_path + "/Data/data_set.obj" DATASET_FILE = "/Data/Data/dataset.obj" MODEL_PATH = "./models/mask_rcnn_moles.h5" ITERATION = 0 SHOW_SAMPLES = False config = MolesConfig() if not os.path.exists(DATASET_FILE): print('No preprocessed version found') dataset_train, dataset_val = serialize_dataset(DATASET_FILE, DATA_PATH, config) else: dataset_train, dataset_val = deserialize_dataset(DATASET_FILE) # Show some random images to verify that everything is ok if SHOW_SAMPLES: print(dataset_train.image_ids) image_ids = np.random.choice(dataset_train.image_ids, 3) for image_id in image_ids: image = dataset_train.load_image(image_id) mask, class_ids = dataset_train.load_mask(image_id) visualize.display_top_masks(image, mask, class_ids, dataset_train.class_names) # Create the MaskRCNN model print('creating model...') model = modellib.MaskRCNN(mode="training", config=config, model_dir=MODEL_DIR) # Use the coco model as start point print('loading weights...') model.load_weights(COCO_MODEL_PATH, by_name=True, exclude=["mrcnn_class_logits", "mrcnn_bbox_fc", "mrcnn_bbox", "mrcnn_mask"]) # Train the model on the train dataset # First only the header layers if ITERATION < 30: print("training the heads") model.train(dataset_train, dataset_val, initial_epoch=ITERATION, learning_rate=config.LEARNING_RATE, epochs=30, layers='heads') # After all the layers if ITERATION < 90: print("training the layers") model.train(dataset_train, dataset_val, initial_epoch=ITERATION-30 if ITERATION > 30 else 0, learning_rate=config.LEARNING_RATE/10, epochs=90-(30-ITERATION if ITERATION > 30 else 0), layers="all") print("Trained finished!")
dataset_train = MoleDataset() dataset_train.load_shapes(train_data, config.IMAGE_SHAPE[0], config.IMAGE_SHAPE[1]) dataset_train.prepare() dataset_val = MoleDataset() dataset_val.load_shapes(val_data, config.IMAGE_SHAPE[0], config.IMAGE_SHAPE[1]) dataset_val.prepare() del train_data del val_data # Show some random images to verify that everything is ok image_ids = np.random.choice(dataset_train.image_ids, 3) for image_id in image_ids: image = dataset_train.load_image(image_id) mask, class_ids = dataset_train.load_mask(image_id) visualize.display_top_masks(image, mask, class_ids, dataset_train.class_names) # Create the MaskRCNN model model = modellib.MaskRCNN(mode="training", config=config, model_dir=MODEL_DIR) # Use as start point the coco model model.load_weights(COCO_MODEL_PATH, by_name=True, exclude=[ "mrcnn_class_logits", "mrcnn_bbox_fc", "mrcnn_bbox", "mrcnn_mask" ]) # Train the model on the train dataset # First only the header layers model.train(dataset_train,