def run():
images_path = '/data/buildings/test'
input_size = (320, 320)
batch_size = 2600
weights = 'weights_train/weights.unet-experiment-{}.h5'
generator, _ = create_generator(images_path,
input_size,
batch_size,
1,
rescale_masks=False,
binary=True)
images, masks = next(generator)
for i, options in enumerate(experiments):
print('Running prediction for experiment {}'.format(str(i)))
dropout = options['dropout']
model = build_unet(input_size, nb_classes=1, dropout=dropout)
model.compile('sgd', 'mse')
model.load_weights(weights.format(i))
probs = model.predict(images, verbose=1)
iou = batch_general_jaccard(masks, probs, binary=True)
f1 = K.eval(f1_score(K.variable(masks), K.variable(probs)))
print('mean IOU for {}: {}'.format(i, np.mean(iou)))
print('F1 score for {}: {}'.format(i, f1))
K.clear_session()
def lrtest_densenet(data_dir, logdir, weights_dir, weights_name, input_size,
nb_classes, batch_size, config, initial_epoch,
pre_trained_weight, augment):
session_config()
model = build_densenet(input_size, nb_classes, config=config)
gpus = get_number_of_gpus()
print('Found {} gpus'.format(gpus))
if gpus > 1:
model = ModelMGPU(model, gpus)
binary = nb_classes == 1
if binary:
loss = binary_soft_jaccard_loss
else:
loss = soft_jaccard_loss
model.compile(optimizer=Adam(lr=1e-3),
loss=loss,
metrics=['acc', binary_jaccard_distance_rounded])
train_generator, num_samples = create_generator(os.path.join(
data_dir, 'train'),
input_size,
batch_size,
nb_classes,
rescale_masks=False,
binary=binary,
augment=augment)
steps_per_epoch = num_samples // batch_size
if augment:
steps_per_epoch = steps_per_epoch * 4
steps_per_epoch = 1000
clr = CyclicLR(base_lr=0,
max_lr=1e-1,
step_size=10 * steps_per_epoch,
mode='triangular')
model.fit_generator(generator=train_generator,
steps_per_epoch=steps_per_epoch,
epochs=10,
verbose=True,
workers=8,
callbacks=[clr])
h = clr.history
lr = h['lr']
acc = h['acc']
print(lr)
print()
print(acc)
def lrtest_unet(data_dir, logdir, weights_dir, weights_name, input_size,
nb_classes, batch_size, initial_epoch, pre_trained_weight,
learning_rate, augment):
session_config()
model = build_unet(input_size, nb_classes)
binary = nb_classes == 1
if binary:
loss = binary_soft_jaccard_loss
else:
loss = soft_jaccard_loss
model.compile(optimizer=Adam(lr=learning_rate),
loss=loss,
metrics=['acc', binary_jaccard_distance_rounded])
train_generator, num_samples = create_generator(os.path.join(
data_dir, 'train'),
input_size,
batch_size,
nb_classes,
rescale_masks=False,
binary=binary,
augment=augment)
steps_per_epoch = num_samples // batch_size
if augment:
steps_per_epoch = steps_per_epoch * 4
clr = CyclicLR(base_lr=0,
max_lr=1e-1,
step_size=5 * steps_per_epoch,
mode='triangular')
model.fit_generator(generator=train_generator,
steps_per_epoch=steps_per_epoch,
epochs=5,
verbose=True,
workers=8,
callbacks=[clr],
initial_epoch=initial_epoch)
h = clr.history
lr = h['lr']
acc = h['acc']
print(lr)
print()
print(acc)
def run():
parser = argparse.ArgumentParser()
parser.add_argument("--weights-path", type=str)
parser.add_argument("--epoch-number", type=int, default=5)
parser.add_argument("--test-images", type=str, default="")
parser.add_argument("--output-path", type=str, default="")
parser.add_argument("--input-size", type=int, default=713)
parser.add_argument("--batch-size", type=int, default=713)
parser.add_argument("--model-name", type=str, default="standard")
parser.add_argument("--save-imgs", type=bool, default=True)
args = parser.parse_args()
model_name = args.model_name
images_path = args.test_images
input_size = args.input_size
batch_size = args.batch_size
save_imgs = args.save_imgs
model_choices = {
'densenet': build_densenet,
'unet': build_unet,
'unet-old': build_unet_old
}
model_choice = model_choices[model_name]
model = model_choice((input_size, input_size), 1)
model.compile(optimizer=Adam(lr=1e-4),
loss=binary_soft_jaccard_loss,
metrics=['acc', binary_jaccard_distance_rounded])
model.load_weights(args.weights_path)
generator, _ = create_generator(
images_path, (input_size, input_size),
batch_size,
1,
rescale_masks=False,
with_file_names=True,
binary=True,
mean=np.array([[[0.36654497, 0.35386439, 0.30782658]]]),
std=np.array([[[0.19212837, 0.19031791, 0.18903286]]]))
images, masks, file_names = next(generator)
probs = model.predict(images, verbose=1)
probs = np.round(probs)
iou = batch_general_jaccard(masks, probs)
f1 = batch_classwise_f1_score(masks, probs)
print('mean IOU: {}'.format(iou))
print('F1 score: {}'.format(f1))
if not save_imgs:
return
# wow such hack
from keras_utils.prediction import get_real_image, get_geo_frame, geo_reference_raster
for i, prob in enumerate(probs):
# mask_result = np.argmax(masks[i], axis=2)
# img = get_real_image(images_path, file_names[i])
mask = masks[i]
raster = get_real_image(images_path, file_names[i], use_gdal=True)
R = raster.GetRasterBand(1).ReadAsArray()
G = raster.GetRasterBand(2).ReadAsArray()
B = raster.GetRasterBand(3).ReadAsArray()
img = np.zeros((512, 512, 3))
img[:, :, 0] = B
img[:, :, 1] = G
img[:, :, 2] = R
prob = np.round(prob)
prob = (prob[:, :, 0] * 255.).astype(np.uint8)
mask = (mask[:, :, 0] * 255.).astype(np.uint8)
pred_name = "pred-{}.tif".format(i)
pred_save_path = "{}/{}".format(args.output_path, pred_name)
cv2.imwrite(pred_save_path, prob)
cv2.imwrite("{}/image-{}.tif".format(args.output_path, i), img)
cv2.imwrite("{}/mask-{}.tif".format(args.output_path, i), mask)
try:
# Get coordinates for corresponding image
ulx, scalex, skewx, uly, skewy, scaley = get_geo_frame(raster)
# Geo reference newly created raster
geo_reference_raster(pred_save_path,
[ulx, scalex, skewx, uly, skewy, scaley])
except ValueError as e:
print("Was not able to reference image at path: {}".format(
pred_save_path))
def train_unet(data_dir, logdir, weights_dir, weights_name, input_size,
nb_classes, batch_size, initial_epoch, pre_trained_weight,
learning_rate, augment):
session_config()
model = build_unet(input_size, nb_classes)
model.summary()
binary = nb_classes == 1
if binary:
loss = binary_soft_jaccard_loss
else:
loss = soft_jaccard_loss
model.compile(optimizer=Adam(lr=learning_rate),
loss=loss,
metrics=['acc', binary_jaccard_distance_rounded])
train_generator, num_samples = create_generator(
os.path.join(data_dir, 'train'),
input_size,
batch_size,
nb_classes,
rescale_masks=True,
binary=binary,
augment=augment,
mean=np.array([[[0.01279744, 0.01279744, 0.01279744]]]),
std=np.array([[[0.11312577, 0.11312577, 0.11312577]]]))
val_generator, val_samples = create_generator(
os.path.join(data_dir, 'val'),
input_size,
batch_size,
nb_classes,
rescale_masks=True,
binary=binary,
augment=augment,
mean=np.array([[[0.01279744, 0.01279744, 0.01279744]]]),
std=np.array([[[0.11312577, 0.11312577, 0.11312577]]]))
if pre_trained_weight:
print('Loading weights: {}'.format(pre_trained_weight))
model.load_weights(pre_trained_weight)
steps_per_epoch = num_samples // batch_size
if augment:
steps_per_epoch = steps_per_epoch * 4
cb = [ValidationCallback(val_samples // batch_size, val_generator)
] + callbacks(logdir,
filename=weights_name,
weightsdir=weights_dir,
monitor_val='mIOU',
base_lr=0.0002,
max_lr=0.002,
steps_per_epoch=steps_per_epoch)
model.fit_generator(generator=train_generator,
steps_per_epoch=steps_per_epoch,
epochs=10000,
verbose=True,
workers=8,
callbacks=cb,
initial_epoch=initial_epoch)
def run():
np.random.seed(2)
tf.set_random_seed(2)
data_dir = '/data/{}/'
weights_dir = 'weights_train'
for i, run in enumerate(runs):
base_lr = run['base_lr']
max_lr = run['max_lr']
input_size = (run['input_size'], run['input_size'])
weights_name = run['name']
logs_dir = 'logs/{}'.format(run['name'])
batch_size = run['batch_size']
print("Running for config {}".format(run))
for j, dataset in enumerate(run['datasets']):
binary = True if dataset != 'multiclass' else False
nb_classes = 1 if binary else 5
print('Running training for {}'.format(dataset))
train_generator, num_samples = create_generator(
os.path.join(data_dir.format(dataset), 'train'),
input_size,
batch_size,
nb_classes=nb_classes,
rescale_masks=run['rescale_masks'],
binary=binary,
augment=False,
mean=np.array([[[0.36654497, 0.35386439, 0.30782658]]]),
std=np.array([[[0.19212837, 0.19031791, 0.18903286]]]))
val_generator, val_samples = create_generator(
os.path.join(data_dir.format(dataset), 'val'),
input_size,
batch_size,
nb_classes=nb_classes,
rescale_masks=run['rescale_masks'],
binary=binary,
augment=False,
mean=np.array([[[0.36654497, 0.35386439, 0.30782658]]]),
std=np.array([[[0.19212837, 0.19031791, 0.18903286]]]))
if run['network'] == 'unet':
model = build_unet(input_size, nb_classes=nb_classes)
else:
model = build_densenet(input_size, nb_classes, 67)
model.summary()
gpus = get_number_of_gpus()
print('Fund {} gpus'.format(gpus))
if gpus > 1:
model = ModelMGPU(model, gpus)
if binary:
loss = binary_soft_jaccard_loss
else:
loss = soft_jaccard_loss
model.compile(optimizer=Adam(),
loss=loss,
metrics=['acc', binary_jaccard_distance_rounded])
if run['pre_weights_name']:
pre_weights_name = run['pre_weights_name'].format(dataset)
weight = 'weights_train/weights.{}.h5'.format(pre_weights_name)
print('Loading weights: {}'.format(weight))
model.load_weights(weight)
steps_per_epoch = num_samples // batch_size
cyclic = 'triangular2'
cb = [
ValidationCallback(
val_samples // batch_size, val_generator, binary=binary)
] + callbacks(logs_dir.format(dataset),
filename=weights_name.format(dataset),
weightsdir=weights_dir,
monitor_val='mIOU',
base_lr=base_lr,
max_lr=max_lr,
steps_per_epoch=steps_per_epoch,
cyclic=cyclic)
model.fit_generator(generator=train_generator,
steps_per_epoch=steps_per_epoch,
epochs=100,
verbose=True,
workers=8,
callbacks=cb)
K.clear_session()
def pred():
image_path = '/data/{}/test'
for dataset in datasets:
im_path = image_path.format(dataset['name'])
for model in models:
input_size = model['input_size']
generator, _ = create_generator(
im_path,
(input_size, input_size),
dataset['size'],
1,
rescale_masks=True,
with_file_names=True,
binary=True,
mean=np.array([[[0.36654497, 0.35386439, 0.30782658]]]),
std=np.array([[[0.19212837, 0.19031791, 0.18903286]]])
)
images, masks, file_names = next(generator)
m = model['method']((input_size, input_size), 1)
gpus = get_number_of_gpus()
if gpus > 1:
m = ModelMGPU(m, gpus)
m.compile(
optimizer=Adam(lr=1e-4),
loss=binary_soft_jaccard_loss,
metrics=['acc', binary_jaccard_distance_rounded])
weights_path = 'weights_train/weights.{}-{}-final-finetune.h5'.format(model['name'], dataset['name'])
m.load_weights(weights_path)
probs = m.predict(images, verbose=1)
probs = np.round(probs)
iou = batch_general_jaccard(masks, probs)
f1 = f1_score(masks, probs)
print('Mean IOU for {} on {}: {}'.format(model['name'], dataset['name'], np.mean(iou)))
print('F1 score for {} on {}: {}'.format(model['name'], dataset['name'], f1))
# wow such hack
from keras_utils.prediction import get_real_image, get_geo_frame, geo_reference_raster
for i, (prob, mask) in enumerate(zip(probs, masks)):
if i > 200:
break
raster = get_real_image(im_path, file_names[i], use_gdal=True)
R = raster.GetRasterBand(1).ReadAsArray()
G = raster.GetRasterBand(2).ReadAsArray()
B = raster.GetRasterBand(3).ReadAsArray()
img = np.zeros((512, 512, 3))
img[:, :, 0] = B
img[:, :, 1] = G
img[:, :, 2] = R
prob = np.round(prob)
prob = (prob[:, :, 0] * 255.).astype(np.uint8)
mask = (mask[:, :, 0] * 255.).astype(np.uint8)
pred_name = "pred-{}.tif".format(i)
out_path = '/data/finalpreds/{}/{}'.format(model['name'], dataset['name'])
pred_save_path = "{}/{}".format(out_path, pred_name)
cv2.imwrite(pred_save_path, prob)
cv2.imwrite("{}/image-{}.tif".format(out_path, i), img)
cv2.imwrite("{}/mask-{}.tif".format(out_path, i), mask)
try:
# Get coordinates for corresponding image
ulx, scalex, skewx, uly, skewy, scaley = get_geo_frame(raster)
# Geo reference newly created raster
geo_reference_raster(
pred_save_path,
[ulx, scalex, skewx, uly, skewy, scaley]
)
except ValueError as e:
print("Was not able to reference image at path: {}".format(pred_save_path))
K.clear_session()
def run():
parser = argparse.ArgumentParser()
parser.add_argument("--weights-path", type=str)
parser.add_argument("--epoch-number", type=int, default=5)
parser.add_argument("--test-images", type=str, default="")
parser.add_argument("--output-path", type=str, default="")
parser.add_argument("--input-size", type=int, default=713)
parser.add_argument("--batch-size", type=int, default=713)
parser.add_argument("--model-name", type=str, default="standard")
parser.add_argument("--save-imgs", type=bool, default=True)
args = parser.parse_args()
model_name = args.model_name
images_path = args.test_images
input_size = args.input_size
batch_size = args.batch_size
save_imgs = args.save_imgs
all_probs = {}
model_choices = {
'densenet': build_densenet,
'unet': build_unet
}
model_choice = model_choices[model_name]
generator, _ = create_generator(
images_path,
(input_size, input_size),
batch_size,
5,
rescale_masks=True,
with_file_names=True,
binary=False,
mean=np.array([[[0.36654497, 0.35386439, 0.30782658]]]),
std=np.array([[[0.19212837, 0.19031791, 0.18903286]]])
)
images, masks, file_names = next(generator)
for dataset in datasets:
model = model_choice((input_size, input_size), 1)
gpus = get_number_of_gpus()
print('Fund {} gpus'.format(gpus))
if gpus > 1:
model = ModelMGPU(model, gpus)
model.compile(
optimizer=Adam(lr=1e-4),
loss=binary_soft_jaccard_loss,
metrics=['acc', binary_jaccard_distance_rounded])
model.load_weights(args.weights_path.format(dataset))
probs = model.predict(images, verbose=1)
probs = np.round(probs)
# iou = batch_general_jaccard(masks, probs, binary=True)
# f1 = K.eval(f1_score(K.variable(masks), K.variable(probs)))
# print('Mean IOU for {}: {}'.format(dataset, np.mean(iou)))
# print('F1 score for {}: {}'.format(dataset, f1))
all_probs[dataset] = probs
final_prob = None
for i, key in enumerate(all_probs):
prob = all_probs[key]
prob[prob == 1] = scores[key]
if i == 0: # First iteration
final_prob = prob
continue
final_prob = np.maximum.reduce([final_prob, prob])
masks = np.argmax(masks, axis=2)
iou = batch_general_jaccard(masks, final_prob)
f1 = K.eval(f1_score(K.variable(masks), K.variable(final_prob)))
print('Mean IOU for {}: {}'.format('multiclass', np.mean(iou)))
print('F1 score for {}: {}'.format('multiclass', f1))
if not save_imgs:
return
# wow such hack
from keras_utils.prediction import get_real_image, get_geo_frame, geo_reference_raster
for i, prob in enumerate(final_prob):
mask = np.argmax(masks[i], axis=2)
raster = get_real_image(os.path.join(images_path, 'multiclass', 'test'), file_names[i], use_gdal=True)
R = raster.GetRasterBand(1).ReadAsArray()
G = raster.GetRasterBand(2).ReadAsArray()
B = raster.GetRasterBand(3).ReadAsArray()
img = np.zeros((512, 512, 3))
img[:, :, 0] = B
img[:, :, 1] = G
img[:, :, 2] = R
seg_pred = np.zeros((input_size, input_size, 3))
seg_mask = np.zeros((input_size, input_size, 3))
for c in range(5):
seg_pred[:, :, 0] += ((prob[:, :, 0] == c) * (class_color_map[c][2])).astype('uint8')
seg_pred[:, :, 1] += ((prob[:, :, 0] == c) * (class_color_map[c][1])).astype('uint8')
seg_pred[:, :, 2] += ((prob[:, :, 0] == c) * (class_color_map[c][0])).astype('uint8')
seg_mask[:, :, 0] += ((mask[:, :] == c) * (class_color_map[c][2])).astype('uint8')
seg_mask[:, :, 1] += ((mask[:, :] == c) * (class_color_map[c][1])).astype('uint8')
seg_mask[:, :, 2] += ((mask[:, :] == c) * (class_color_map[c][0])).astype('uint8')
pred_name = "pred-{}.tif".format(i)
pred_save_path = "{}/{}".format(args.output_path, pred_name)
cv2.imwrite(pred_save_path, seg_pred)
cv2.imwrite("{}/mask-{}.tif".format(args.output_path, i), seg_mask)
cv2.imwrite("{}/image-{}.tif".format(args.output_path, i), img)
try:
# Get coordinates for corresponding image
ulx, scalex, skewx, uly, skewy, scaley = get_geo_frame(raster)
# Geo reference newly created raster
geo_reference_raster(
pred_save_path,
[ulx, scalex, skewx, uly, skewy, scaley]
)
except ValueError as e:
print("Was not able to reference image at path: {}".format(pred_save_path))
def run():
np.random.seed(2)
tf.set_random_seed(2)
base_lr = 0.0002
max_lr = 0.002
data_dir = '/data/buildings/'
logs_dir = 'logs/{}'
weights_dir = 'weights_train'
weights_name = 'unet-experiment-{}'
input_size = (320, 320)
batch_size = 20
start_from = 6
binary = True
session_config()
train_generator, num_samples = create_generator(os.path.join(
data_dir, 'train'),
input_size,
batch_size,
nb_classes=1,
rescale_masks=False,
binary=binary,
augment=False)
val_generator, val_samples = create_generator(os.path.join(
data_dir, 'val'),
input_size,
batch_size,
nb_classes=1,
rescale_masks=False,
binary=binary,
augment=False)
for i, options in enumerate(experiments):
if i < start_from:
continue
print('Running experiment {} with options: {}'.format(str(i), options))
optimizer = get_optimizer(options['optimizer'], max_lr)
loss = get_loss(options['loss'])
dropout = options['dropout']
model = build_unet(input_size, nb_classes=1, dropout=dropout)
model.summary()
model.compile(optimizer=optimizer,
loss=loss,
metrics=['acc', binary_jaccard_distance_rounded])
steps_per_epoch = num_samples // batch_size
lr_opt = options['lr']
if lr_opt == 'annealing':
cyclic = None
else:
cyclic = lr_opt
print('Experiment {} using lr: {}'.format(str(i), cyclic))
model.fit_generator(
generator=train_generator,
validation_data=val_generator,
validation_steps=val_samples // batch_size,
steps_per_epoch=steps_per_epoch,
epochs=20,
verbose=True,
workers=8,
callbacks=callbacks(
logs_dir.format(str(i)),
filename=weights_name.format(str(i)),
weightsdir=weights_dir,
monitor_val='val_binary_jaccard_distance_rounded',
base_lr=base_lr,
max_lr=max_lr,
steps_per_epoch=steps_per_epoch,
cyclic=cyclic))
K.clear_session()
def train_densenet(data_dir, logdir, weights_dir, weights_name, input_size,
nb_classes, batch_size, config, initial_epoch,
pre_trained_weight, augment):
session_config()
model = build_densenet(input_size, nb_classes, config=config)
model.summary()
gpus = get_number_of_gpus()
print('Found {} gpus'.format(gpus))
if gpus > 1:
model = ModelMGPU(model, gpus)
is_binary = nb_classes == 1
if is_binary:
loss = binary_soft_jaccard_loss
else:
loss = soft_jaccard_loss
model.compile(optimizer=Adam(lr=1e-3),
loss=loss,
metrics=['acc', binary_jaccard_distance_rounded])
train_generator, num_samples = create_generator(
os.path.join(data_dir, 'train'),
input_size,
batch_size,
nb_classes,
rescale_masks=True,
binary=is_binary,
augment=False,
mean=np.array([[[0.42800662, 0.40565866, 0.3564895]]]),
std=np.array([[[0.19446792, 0.1984272, 0.19501258]]]))
val_generator, val_samples = create_generator(
os.path.join(data_dir, 'val'),
input_size,
batch_size,
nb_classes,
rescale_masks=True,
binary=is_binary,
augment=False,
mean=np.array([[[0.42800662, 0.40565866, 0.3564895]]]),
std=np.array([[[0.19446792, 0.1984272, 0.19501258]]]))
if pre_trained_weight:
print('Loading weights: {}'.format(pre_trained_weight))
model.load_weights(pre_trained_weight)
steps_per_epoch = num_samples // batch_size
if augment:
steps_per_epoch = steps_per_epoch * 4
base_lr = 0.00002
max_lr = 0.00055
cb = [
ValidationCallback(val_samples // batch_size, val_generator, is_binary)
] + callbacks(logdir,
filename=weights_name,
weightsdir=weights_dir,
monitor_val='mIOU',
base_lr=base_lr,
max_lr=max_lr,
steps_per_epoch=steps_per_epoch,
cyclic='triangular2')
model.fit_generator(generator=train_generator,
steps_per_epoch=steps_per_epoch,
epochs=100,
verbose=True,
workers=8,
callbacks=cb,
initial_epoch=initial_epoch)
def run(args):
n_classes = args.classes
model_name = args.model_name
images_path = args.test_images
input_size = args.input_size
batch_size = args.batch_size
no_save_imgs = args.no_save_imgs
model_choices = {
'unet': build_unet,
'densenet': build_densenet,
'pspnet': build_pspnet
}
model_choice = model_choices[model_name]
model = model_choice((input_size, input_size), n_classes)
gpus = get_number_of_gpus()
print('Fund {} gpus'.format(gpus))
if gpus > 1:
model = ModelMGPU(model, gpus)
model.compile(optimizer=Adam(lr=1e-4),
loss=binary_soft_jaccard_loss,
metrics=['acc', binary_jaccard_distance_rounded])
model.load_weights(args.weights_path)
binary = n_classes == 1
generator, samples = create_generator(
images_path, (input_size, input_size),
batch_size,
nb_classes=n_classes,
rescale_masks=False,
with_file_names=True,
binary=binary,
augment=False,
mean=np.array([[[0.36654497, 0.35386439, 0.30782658]]]),
std=np.array([[[0.19212837, 0.19031791, 0.18903286]]]))
steps = samples // batch_size
f1s = []
ious = []
for i in range(steps):
images, masks, file_names = next(generator)
probs = model.predict(images, verbose=1)
probs = np.argmax(probs, axis=3)
masks = np.argmax(masks, axis=3)
iou = batch_general_jaccard(masks, probs)
f1 = batch_classwise_f1_score(masks, probs)
ious.append(iou)
f1s.append(f1)
if no_save_imgs:
continue
for i, prob in enumerate(probs):
result = prob
mask_result = masks[i]
# img = get_real_image(images_path, file_names[i])
raster = get_real_image(images_path, file_names[i], use_gdal=True)
R = raster.GetRasterBand(1).ReadAsArray()
G = raster.GetRasterBand(2).ReadAsArray()
B = raster.GetRasterBand(3).ReadAsArray()
img = np.zeros((512, 512, 3))
img[:, :, 0] = B
img[:, :, 1] = G
img[:, :, 2] = R
seg_img = np.zeros((input_size, input_size, 3))
seg_mask = np.zeros((input_size, input_size, 3))
for c in range(n_classes):
seg_img[:, :, 0] += ((result[:, :] == c) *
(class_color_map[c][2])).astype('uint8')
seg_img[:, :, 1] += ((result[:, :] == c) *
(class_color_map[c][1])).astype('uint8')
seg_img[:, :, 2] += ((result[:, :] == c) *
(class_color_map[c][0])).astype('uint8')
seg_mask[:, :, 0] += ((mask_result[:, :] == c) *
(class_color_map[c][2])).astype('uint8')
seg_mask[:, :, 1] += ((mask_result[:, :] == c) *
(class_color_map[c][1])).astype('uint8')
seg_mask[:, :, 2] += ((mask_result[:, :] == c) *
(class_color_map[c][0])).astype('uint8')
pred_name = "pred-{}.tif".format(i)
pred_save_path = "{}/{}".format(args.output_path, pred_name)
cv2.imwrite(pred_save_path, seg_img)
cv2.imwrite("{}/mask-{}.tif".format(args.output_path, i), seg_mask)
cv2.imwrite("{}/image-{}.tif".format(args.output_path, i), img)
try:
# Get coordinates for corresponding image
ulx, scalex, skewx, uly, skewy, scaley = get_geo_frame(raster)
# Geo reference newly created raster
geo_reference_raster(pred_save_path,
[ulx, scalex, skewx, uly, skewy, scaley])
except ValueError as e:
print("Was not able to reference image at path: {}".format(
pred_save_path))
print('Mean IOU: {}'.format(np.mean(ious)))
print('F1 score: {}'.format(np.mean(f1s)))