def main(): # 1936 x 1216 input_size = (320, 480, 3) classes = 20 train_dataset_x = '../seg_train_images/seg_train_images' train_dataset_y = '../seg_train_annotations/seg_train_annotations' test_size = 0.2 batch_size = 8 datasets_paths = get_data_paths(train_dataset_x, train_dataset_y) train_data, test_data = train_test_split(datasets_paths, test_size=test_size) net = Unet(input_size, classes) #net = SegNet(input_size, classes) net.summary() train_gen = DataGenerator(train_data, input_size, classes, batch_size) val_gen = DataGenerator(test_data, input_size, classes, batch_size) callbacks = [ ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=3, verbose=1), EarlyStopping(monitor='val_loss', min_delta=0, patience=9, verbose=1), ModelCheckpoint('checkpoint/ep{epoch:03d}-loss{loss:.5f}-val_loss{val_loss:.5f}.h5', monitor='val_loss', verbose=1, save_best_only=True, mode='min') ] net.compile(optimizer=Adam(1e-3),loss=categorical_crossentropy) history = net.fit_generator( train_gen, steps_per_epoch=train_gen.num_batches_per_epoch, validation_data=val_gen, validation_steps=val_gen.num_batches_per_epoch, initial_epoch=0, epochs=50, callbacks=callbacks ) net.save_weights('checkpoint/first_stage.h5') train_data, test_data = train_test_split(datasets_paths, test_size=test_size) train_gen = DataGenerator(train_data, input_size, classes, batch_size) val_gen = DataGenerator(test_data, input_size, classes, batch_size) net.compile(optimizer=Adam(1e-4),loss=categorical_crossentropy) history = net.fit_generator( train_gen, steps_per_epoch=train_gen.num_batches_per_epoch, validation_data=val_gen, validation_steps=val_gen.num_batches_per_epoch, initial_epoch=50, epochs=100, callbacks=callbacks ) net.save_weights('checkpoint/final_stage.h5')
batch_size=batch_size, seed=1) train_generator = zip(image_generator, mask_generator) for (img, mask) in train_generator: img = img / 255. mask = mask / 255. mask[mask > 0.5] = 1 mask[mask <= 0.5] = 0 yield (img, mask) model = Unet(1, image_size=512) trainset = DataGenerator("dataset/train", batch_size=2) model.fit_generator(trainset, steps_per_epoch=100, epochs=1) model.save_weights("model.h5") testSet = DataGenerator("dataset/test", batch_size=1) alpha = 0.3 model.load_weights("model.h5") if not os.path.exists("./results"): os.mkdir("./results") for idx, (img, mask) in enumerate(testSet): oring_img = img[0] # 开始用模型进行预测 pred_mask = model.predict(img)[0] pred_mask[pred_mask > 0.5] = 1 pred_mask[pred_mask <= 0.5] = 0 # 如果这里展示的预测结果一片黑,请调整lr,同时注意图片的深度是否为8 # cv2.imshow('pred_mask', pred_mask)