class Segment(object): def __init__(self, backend, input_size, nb_classes): """ Model Factory that fetches the corresponding model based on the backend that has been defined and initiates the training process :param backend: define the backbone architecture for the training :param input_size: the size of input image :param nb_classes: the number of classes """ self.input_size = input_size self.nb_classes = nb_classes if backend == "UNET": self.feature_extractor = UNET(self.input_size, self.nb_classes).build() else: raise ValueError( 'No such arch!... Please check the backend in config file') def train(self, train_configs): """ Train the model based on the training configurations :param train_configs: configurations for the training """ optimizer = adam(learning_rate=train_configs['learning_rate']) train_times = train_configs['train_times'] # Data sequence for training sequence = DataSequence( train_configs['data_directory'] + '/data_road/training', train_configs['batch_size'], self.input_size) steps_per_epoch = len( sequence) * train_times # increase steps_per_epoch # configure the model for training self.feature_extractor.compile(optimizer=optimizer, loss='categorical_crossentropy', metrics=['accuracy']) # define the callbacks for training tb = TensorBoard(log_dir=train_configs["logs_dir"], write_graph=True) mc = ModelCheckpoint(mode='max', filepath=train_configs["save_model_name"], monitor='accuracy', save_best_only='True', save_weights_only='False', verbose=1) es = EarlyStopping(mode='max', monitor='accuracy', patience=6, verbose=1) model_reducelr = callbacks.ReduceLROnPlateau( mode='min', monitor='loss', factor=0.2, patience=5, verbose=1, min_lr=0.05 * train_configs['learning_rate']) callback = [tb, mc, es, model_reducelr] # Train the model on data generated batch-by-batch by the DataSequence generator self.feature_extractor.fit_generator(sequence, steps_per_epoch=steps_per_epoch, epochs=train_configs['nb_epochs'], verbose=1, shuffle=True, callbacks=callback, workers=3, max_queue_size=8)
class Segment(object): def __init__(self, backend, input_size, nb_classes): """ Model Factory that fetches the corresponding model based on the backend that has been defined and initiates the training process :param backend: define the backbone architecture for the training :param input_size: the size of the input image :param nb_classes: the number of classes """ self.input_size = input_size self.nb_classes = nb_classes if backend == "ENET": self.feature_extractor = ENET(self.input_size, self.nb_classes).build() elif backend == "VGG": self.feature_extractor = VGG(self.input_size, self.nb_classes).build() elif backend == "UNET": self.feature_extractor = UNET(self.input_size, self.nb_classes).build() else: raise ValueError( 'No such arch!... Please check the backend in config file') def train(self, train_configs, model_configs): """ Train the model based on the training configurations :param train_configs: Configuration for the training """ optimizer = Adam(train_configs["learning_rate"]) # Data sequence for training train_gen = DataSequence(train_configs["train_images"], train_configs["train_annotations"], train_configs["train_batch_size"], model_configs['classes'], model_configs['im_height'], model_configs['im_width'], model_configs['out_height'], model_configs['out_width'], do_augment=True) # Data sequence for validation val_gen = DataSequence(train_configs["val_images"], train_configs["val_annotations"], train_configs["val_batch_size"], model_configs['classes'], model_configs['im_height'], model_configs['im_width'], model_configs['out_height'], model_configs['out_width'], do_augment=True) # Configure the model for training # https://www.depends-on-the-definition.com/unet-keras-segmenting-images/ # if model_configs['classes'] == 2: # loss_function = 'binary_crossentropy' # else: # loss_function = 'categorical_crossentropy' # Segmentation models losses can be combined together by '+' and scaled by integer or float factor dice_loss = sm.losses.DiceLoss() focal_loss = sm.losses.BinaryFocalLoss( ) if model_configs['classes'] == 2 else sm.losses.CategoricalFocalLoss( ) total_loss = dice_loss + (1 * focal_loss) iou = sm.metrics.IOUScore(threshold=0.5) fscore = sm.metrics.FScore(threshold=0.5) metrics = [iou, fscore] self.feature_extractor.compile(optimizer=optimizer, loss=total_loss, metrics=metrics) # Load pretrained model if train_configs["load_pretrained_weights"]: print("Loading pretrained weights: " + train_configs["pretrained_weights_path"]) self.feature_extractor.load_weights( train_configs["pretrained_weights_path"]) # Define the callbacks for training tb = TensorBoard(log_dir=train_configs["logs_dir"], write_graph=True) mc = ModelCheckpoint( mode="max", filepath=str(train_configs["save_model_name"]).replace(".h5", "") + ".{epoch:03d}.h5", monitor="f1-score", save_best_only=True, save_weights_only=False, verbose=2) callback = [tb, mc] # Train the model on data generated batch-by-batch by the DataSequence generator self.feature_extractor.fit_generator(train_gen, steps_per_epoch=len(train_gen), validation_data=val_gen, validation_steps=len(val_gen), epochs=train_configs["nb_epochs"], verbose=1, shuffle=True, callbacks=callback, workers=6, max_queue_size=24, use_multiprocessing=True)