def main():
FLAGS = parser.parse_args()
# Calculate the predictions for all the images in the input_img_dir.
for img_name in os.listdir(FLAGS.input_img_dir):
if img_name.endswith('.png'):
original_img = load_img(
FLAGS.input_img_dir + '/' +
img_name) # load_img function exists in file utils.py
# Resizing image because of the small memory size
input_img = tl.prepro.imresize(original_img,
[FLAGS.img_size, FLAGS.img_size])
input_img = np.reshape(input_img,
[1, FLAGS.img_size, FLAGS.img_size, 3])
unet = UNet(FLAGS.img_size)
# The output is an array of the size(img_size * img_size, 1)
prediction = unet.predict(input_img, FLAGS.model_save_dir)
# Saving the image given the probabilities
# save_img fnuction exists in file utils.py
save_img(
prediction, original_img, FLAGS.img_size, FLAGS.input_img_dir +
'/' + img_name.split('.')[0] + '_pred.png')
def main(argv):
data = get_data(FLAGS.train_data, FLAGS.num_classes)
train_data = data.batch(16, drop_remainder=True)
model = UNet(num_classes=FLAGS.num_classes)
model.compile(optimizer=tf.keras.optimizers.Adam(),
loss='binary_crossentropy')
model.fit(train_data, epochs=25)
for index, (image, label) in enumerate(data.batch(1).take(5)):
prediction = model.predict(image)
plot_result(f'results/{index}.png', image, label, prediction)
def train_model(name, experiment, image_size, training_data_list, training_mask_list,
model_spec=[16, 32, 64, 128, 256], preprocess_list=None,
preprocess_stretch=False, preprocess_mask=None,
preprocess_fisher=False, keep_image=True,
load_model=False, epochs=15):
# make copies of the input array before shuffling
training_data_list = list(training_data_list)
training_mask_list = list(training_mask_list)
random.Random(experiment*42).shuffle(training_data_list)
random.Random(experiment*42).shuffle(training_mask_list)
# we're augmenting data -- expand the list of training data
train_input_img_paths = training_data_list[:-(test_samples
+ val_samples)] * random_factor
train_target_img_paths = training_mask_list[:-(test_samples
+ val_samples)] * random_factor
val_input_img_paths = training_data_list[-(
test_samples + val_samples):-val_samples]
val_target_img_paths = training_mask_list[-(
test_samples + val_samples):-val_samples]
test_input_img_paths = training_data_list[-test_samples:]
test_target_img_paths = training_mask_list[-test_samples:]
pp = None
# Chain of preprocessing functions, first one added is performed first
if preprocess_list is not None:
# Instantiate data Sequences for each split
if not preprocess_stretch:
pp = ImagePreprocessGradient(preprocess_list, keep_image, pp)
else:
pp = ImagePreprocessStretchedGradient(preprocess_list, pp)
if preprocess_mask is not None:
# Apply mask after gradients - masking first only gets overwritten
pp = ImagePreprocessMask(preprocess_mask, pp)
if preprocess_fisher is True:
pp = ImagePreprocessFisherize(pp)
if pp is not None:
# Instantiate pre-processed data sequences for each split
train_gen = RoadSeq(batch_size, image_size,
train_input_img_paths, train_target_img_paths, augment_data=True,
preprocess_fn=pp.preprocess())
val_gen = RoadSeq(batch_size, image_size,
val_input_img_paths, val_target_img_paths, augment_data=False,
preprocess_fn=pp.preprocess())
test_gen = RoadSeq(len(test_input_img_paths), image_size,
test_input_img_paths, test_target_img_paths, augment_data=False,
preprocess_fn=pp.preprocess())
else:
# use the images as they are
train_gen = RoadSeq(batch_size, image_size,
train_input_img_paths, train_target_img_paths, augment_data=True)
val_gen = RoadSeq(batch_size, image_size,
val_input_img_paths, val_target_img_paths, augment_data=False)
test_gen = RoadSeq(len(test_input_img_paths), image_size,
test_input_img_paths, test_target_img_paths, augment_data=False)
model_name = name+'.'+str(experiment)+'.h5'
model = UNet(image_size, model_spec)
model.compile(optimizer="adam",
loss="binary_crossentropy", metrics=["acc"])
if load_model:
model.load_weights(model_name)
model.summary()
callbacks = [
keras.callbacks.ModelCheckpoint(
model_name, save_best_only=True)
]
model.fit(train_gen, epochs=epochs, verbose=1,
validation_data=val_gen, callbacks=callbacks)
x, y = test_gen.__getitem__(0)
start = timer()
results = model.predict(x)
end = timer()
prediction_time = (end - start) / len(results)
results = np.array(results > 0.5).astype(np.uint8)
return calculate_error(results, test_target_img_paths) + (prediction_time,)
scheduler = CyclicLR(optimizer, lr, lr_max, cycle_momentum=False)
for epoch in range(25):
if train:
# Training loop
model.train()
total_training_loss = 0.0
for i, (data, target, class_ids) in enumerate(train_loader):
data, target = data, target
if use_gpu:
data = data.cuda()
target = target.cuda()
optimizer.zero_grad()
output = model.predict(data, use_gpu, class_ids)
loss = criterion(output, target)
loss.backward()
optimizer.step()
scheduler.step()
total_training_loss += loss.item()
img = (data[0].transpose(0, 1).transpose(1,
2).detach().cpu().numpy())
output_mask = np.abs(
create_boolean_mask(output[0][0].cpu().detach().numpy()) *
(-1))
target_mask = target[0][0].cpu().numpy().astype(np.bool)
def run_demo(name,
experiment,
image_size,
training_data_list,
training_mask_list,
model_spec=[16, 32, 64, 128, 256],
preprocess_list=None,
preprocess_stretch=False,
preprocess_mask=None,
keep_image=True):
# make copies of the input array before shuffling
training_data_list = list(training_data_list)
random.Random(experiment * 42).shuffle(training_data_list)
test_input_img_paths = training_data_list[-test_samples:]
if training_mask_list is not None:
training_mask_list = list(training_mask_list)
random.Random(experiment * 42).shuffle(training_mask_list)
test_target_img_paths = training_mask_list[-test_samples:]
else:
test_target_img_paths = None
pp = None
# Chain of preprocessing functions, first one added is performed first
if preprocess_list is not None:
# Instantiate data Sequences for each split
if not preprocess_stretch:
pp = ImagePreprocessGradient(preprocess_list, keep_image, pp)
else:
pp = ImagePreprocessStretchedGradient(preprocess_list, pp)
if preprocess_mask is not None:
# Apply mask after gradients - masking first only gets overwritten
pp = ImagePreprocessMask(preprocess_mask, pp)
if pp is not None:
# Instantiate pre-processed data sequences for each split
test_gen = RoadSeq(len(test_input_img_paths),
image_size,
test_input_img_paths,
test_target_img_paths,
augment_data=False,
preprocess_fn=pp.preprocess())
else:
# use the images as they are
test_gen = RoadSeq(len(test_input_img_paths),
image_size,
test_input_img_paths,
test_target_img_paths,
augment_data=False)
model_name = name + '.' + str(experiment) + '.h5'
model = UNet(image_size, model_spec)
model.compile(optimizer="adam",
loss="binary_crossentropy",
metrics=["acc"])
model.load_weights(model_name)
x, y = test_gen.__getitem__(0)
results = model.predict(x)
results = np.array(results > 0.5).astype(np.uint8)
display_results(test_input_img_paths, test_target_img_paths, [results])
return results
from model import UNet
from config import UnetConfig
import cv2
config=UnetConfig()
unet=UNet(config=config)
unet.loadWeight("./logs")
imgdir=os.listdir("./dataset-postdam/test")
index=0
for name in imgdir:
img=cv2.imread("./dataset-postdam/test/"+name)
img=cv2.resize(img,(config.IMAGE_H,config.IMAGE_W))
result=unet.predict(img)
unet.visual(result,"./test-result/"+str(index)+".png")
index=index+1
