def writeImage(result, file_name):
file_name = utils.filepath_to_name(args.image)
cv2.imwrite("%s_pred.png" % (file_name),
cv2.cvtColor(np.uint8(out_vis_image), cv2.COLOR_RGB2BGR))
print("")
print("Finished!")
print("Wrote image " + "%s_pred.png" % (file_name))
def getCoordinatesFromFiles(images,
outputPath=False,
downscale_factor=False,
background='white',
saveMorphImage=True):
coordinates = dict()
for img in images:
if not os.path.isfile(img):
coordinates[utils.filepath_to_name(img, True)] = []
continue
image = cv2.cvtColor(cv2.imread(img), cv2.COLOR_RGB2GRAY)
if np.min(image) == np.max(image):
#Whole image is same color i.e. everything is background
coordinates[utils.filepath_to_name(img, True)] = []
continue
if background == 'white':
foregroundcolor = [np.min(image)]
elif background == 'black':
foregroundcolor = [np.max(image)]
else:
raise ValueError(
'Only black and white are allowed as values for background (it was {})'
.format(background))
coordinates[utils.filepath_to_name(
img, True)], morphImage = utils.getCoordsFromPrediction(
image, foregroundcolor, downscale_factor, close_dim=384)
#open_dim=8,,open2_dim=32
if saveMorphImage:
outMorph = np.ones(morphImage.shape, dtype=np.uint8) * np.iinfo(
np.uint8).max
outMorph[morphImage == 0] = 0
cv2.imwrite(img[0:-4] + '_morph.tif',
cv2.cvtColor(np.uint8(outMorph), cv2.COLOR_RGB2BGR))
if outputPath:
if not outputPath.endswith('.json'):
outputPath = outputPath + '.json'
with open(outputPath, 'w') as f:
json.dump(coordinates, f)
return coordinates, images
sess.run(tf.global_variables_initializer())
print('Loading model checkpoint weights')
saver = tf.train.Saver(max_to_keep=1000)
saver.restore(sess, args.checkpoint_path)
print("Testing image " + args.image)
loaded_image = utils.load_image(args.image)
resized_image = cv2.resize(loaded_image, (args.crop_width, args.crop_width))
input_image = np.expand_dims(np.float32(
resized_image[:args.crop_height, :args.crop_width]),
axis=0) / 255.0
st = time.time()
output_image = sess.run(network, feed_dict={net_input: input_image})
run_time = time.time() - st
output_image = np.array(output_image[0, :, :, :])
output_image = helpers.reverse_one_hot(output_image)
out_vis_image = helpers.colour_code_segmentation(output_image, label_values)
file_name = utils.filepath_to_name(args.image)
cv2.imwrite("%s_pred.png" % (file_name),
cv2.cvtColor(np.uint8(out_vis_image), cv2.COLOR_RGB2BGR))
print("")
print("Finished!")
print("Wrote image " + "%s_pred.png" % (file_name))
gt = helpers.reverse_one_hot(helpers.one_hot_it(gt, label_values))
# st = time.time()
output_image = sess.run(network,
feed_dict={net_input: input_image})
output_image = np.array(output_image[0, :, :, :])
output_image = helpers.reverse_one_hot(output_image)
out_vis_image = helpers.colour_code_segmentation(
output_image, label_values)
accuracy, class_accuracies, prec, rec, f1, iou = utils.evaluate_segmentation(
pred=output_image, label=gt, num_classes=num_classes)
file_name = utils.filepath_to_name(val_input_names[ind])
target.write("%s, %f, %f, %f, %f, %f" %
(file_name, accuracy, prec, rec, f1, iou))
for item in class_accuracies:
target.write(", %f" % (item))
target.write("\n")
scores_list.append(accuracy)
class_scores_list.append(class_accuracies)
precision_list.append(prec)
recall_list.append(rec)
f1_list.append(f1)
iou_list.append(iou)
gt = helpers.colour_code_segmentation(gt, label_values)
out_fore_clean = helpers.remove_small(out_fore_bin, min_size=300)
out_fore_clean = helpers.remove_edge_seg(out_fore_clean)
gt = helpers.remove_edge_seg(gt)
out_vis_image = helpers.colour_code_segmentation(output_image, label_values)
out_fore_vis_image = helpers.colour_code_segmentation(out_fore_clean, label_values)
accuracy, class_accuracies, prec, rec, f1, iou = utils.evaluate_segmentation(pred=output_image, label=gt, num_classes=num_classes)
accuracy_fore, class_accuracies_fore, prec_fore, rec_fore, f1_fore, iou_fore = utils.evaluate_segmentation(pred=out_fore_clean, label=gt, num_classes=num_classes)
file_name = utils.filepath_to_name(test_input_names[ind])
target.write("%s, %f, %f, %f, %f, %f"%(file_name, accuracy, prec, rec, f1, iou))
for item in class_accuracies:
target.write(", %f"%(item))
target.write("\n")
scores_list.append(accuracy)
class_scores_list.append(class_accuracies)
precision_list.append(prec)
recall_list.append(rec)
f1_list.append(f1)
iou_list.append(iou)
scores_list_fore.append(accuracy_fore)
class_scores_list_fore.append(class_accuracies_fore)
precision_list_fore.append(prec_fore)
def main(args, _log):
def data_augmentation(input_image, output_image):
# Data augmentation
input_image, output_image = utils.random_crop(input_image,
output_image,
args["crop_height"],
args["crop_width"])
if args["h_flip"] and random.randint(0, 1):
input_image = cv2.flip(input_image, 1)
output_image = cv2.flip(output_image, 1)
if args["v_flip"] and random.randint(0, 1):
input_image = cv2.flip(input_image, 0)
output_image = cv2.flip(output_image, 0)
if args["brightness"]:
factor = 1.0 + random.uniform(-1.0 * args["brightness"],
args["brightness"])
table = np.array([((i / 255.0) * factor) * 255
for i in np.arange(0, 256)]).astype(np.uint8)
input_image = cv2.LUT(input_image, table)
if args["rotation"]:
angle = random.uniform(-1 * args["rotation"], args["rotation"])
if args["rotation"]:
M = cv2.getRotationMatrix2D(
(input_image.shape[1] // 2, input_image.shape[0] // 2), angle,
1.0)
input_image = cv2.warpAffine(
input_image,
M, (input_image.shape[1], input_image.shape[0]),
flags=cv2.INTER_NEAREST)
output_image = cv2.warpAffine(
output_image,
M, (output_image.shape[1], output_image.shape[0]),
flags=cv2.INTER_NEAREST)
return input_image, output_image
print("Args:", args)
# Get the names of the classes so we can record the evaluation results
class_names_list, label_values = helpers.get_label_info(
os.path.join(args["dataset"], "class_dict.csv"))
class_names_string = ""
for class_name in class_names_list:
if not class_name == class_names_list[-1]:
class_names_string = class_names_string + class_name + ", "
else:
class_names_string = class_names_string + class_name
num_classes = len(label_values)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
# Compute your softmax cross entropy loss
net_input = tf.placeholder(tf.float32, shape=[None, None, None, 3])
net_output = tf.placeholder(tf.float32,
shape=[None, None, None, num_classes])
network, init_fn = model_builder.build_model(
model_name=args["model"],
frontend=args["frontend"],
net_input=net_input,
num_classes=num_classes,
crop_width=args["crop_width"],
crop_height=args["crop_height"],
is_training=True)
loss = tf.reduce_mean(
tf.nn.softmax_cross_entropy_with_logits(logits=network,
labels=net_output))
opt = tf.train.RMSPropOptimizer(
learning_rate=0.0001, decay=0.995).minimize(
loss, var_list=[var for var in tf.trainable_variables()])
saver = tf.train.Saver(max_to_keep=1000)
sess.run(tf.global_variables_initializer())
utils.count_params()
# If a pre-trained ResNet is required, load the weights.
# This must be done AFTER the variables are initialized with sess.run(tf.global_variables_initializer())
if init_fn is not None:
init_fn(sess)
# Load a previous checkpoint if desired
model_checkpoint_name = "checkpoints/latest_model_" + args[
"model"] + "_" + basename(normpath(args["dataset"])) + ".ckpt"
if args["continue_training"]:
_log.info('Loaded latest model checkpoint')
saver.restore(sess, model_checkpoint_name)
# Load the data
_log.info("Loading the data ...")
train_input_names, train_output_names, val_input_names, val_output_names, test_input_names, test_output_names = utils.prepare_data(
dataset_dir=args["dataset"])
_log.info("\n***** Begin training *****")
_log.debug("Dataset -->", args["dataset"])
_log.debug("Model -->", args["model"])
_log.debug("Crop Height -->", args["crop_height"])
_log.debug("Crop Width -->", args["crop_width"])
_log.debug("Num Epochs -->", args["num_epochs"])
_log.debug("Batch Size -->", args["batch_size"])
_log.debug("Num Classes -->", num_classes)
_log.debug("Data Augmentation:")
_log.debug("\tVertical Flip -->", args["v_flip"])
_log.debug("\tHorizontal Flip -->", args["h_flip"])
_log.debug("\tBrightness Alteration -->", args["brightness"])
_log.debug("\tRotation -->", args["rotation"])
avg_loss_per_epoch = []
avg_scores_per_epoch = []
avg_iou_per_epoch = []
# Which validation images do we want
val_indices = []
num_vals = min(args["num_val_images"], len(val_input_names))
# Set random seed to make sure models are validated on the same validation images.
# So you can compare the results of different models more intuitively.
random.seed(16)
val_indices = random.sample(range(0, len(val_input_names)), num_vals)
# Do the training here
for epoch in range(args["epoch_start_i"], args["num_epochs"]):
current_losses = []
cnt = 0
# Equivalent to shuffling
id_list = np.random.permutation(len(train_input_names))
num_iters = int(np.floor(len(id_list) / args["batch_size"]))
st = time.time()
epoch_st = time.time()
for i in range(num_iters):
# st=time.time()
input_image_batch = []
output_image_batch = []
# Collect a batch of images
for j in range(args["batch_size"]):
index = i * args["batch_size"] + j
id = id_list[index]
input_image = utils.load_image(train_input_names[id],
args["crop_width"],
args["crop_height"])
output_image = utils.load_image(train_output_names[id],
args["crop_width"],
args["crop_height"])
with tf.device('/cpu:0'):
input_image, output_image = data_augmentation(
input_image, output_image)
# Prep the data. Make sure the labels are in one-hot format
input_image = np.float32(input_image) / 255.0
output_image = np.float32(
helpers.one_hot_it(label=output_image,
label_values=label_values))
input_image_batch.append(
np.expand_dims(input_image, axis=0))
output_image_batch.append(
np.expand_dims(output_image, axis=0))
if args["batch_size"] == 1:
input_image_batch = input_image_batch[0]
output_image_batch = output_image_batch[0]
else:
input_image_batch = np.squeeze(
np.stack(input_image_batch, axis=1))
output_image_batch = np.squeeze(
np.stack(output_image_batch, axis=1))
# Do the training
_, current = sess.run([opt, loss],
feed_dict={
net_input: input_image_batch,
net_output: output_image_batch
})
current_losses.append(current)
cnt = cnt + args["batch_size"]
if cnt % 20 == 0:
string_print = "Epoch = %d Count = %d Iter:(%d of %d) Current_Loss = %.4f Time = %.2f" % (
epoch, cnt, i, num_iters, current, time.time() - st)
utils.LOG(string_print)
st = time.time()
mean_loss = np.mean(current_losses)
avg_loss_per_epoch.append(mean_loss)
# Create directories if needed
if not os.path.isdir("%s/%04d" % ("checkpoints", epoch)):
os.makedirs("%s/%04d" % ("checkpoints", epoch))
# Save latest checkpoint to same file name
_log.info("Saving latest checkpoint")
saver.save(sess, model_checkpoint_name)
if val_indices != 0 and epoch % args["checkpoint_step"] == 0:
_log.info("Saving checkpoint for this epoch")
saver.save(sess, "%s/%04d/model.ckpt" % ("checkpoints", epoch))
if epoch % args["validation_step"] == 0:
_log.info("Performing validation")
target_path = "%s/%04d/val_scores.csv" % ("checkpoints", epoch)
target = open(target_path, 'w')
target.write(
"val_name, avg_accuracy, precision, recall, f1 score, mean iou, %s\n"
% (class_names_string))
scores_list = []
class_scores_list = []
precision_list = []
recall_list = []
f1_list = []
iou_list = []
# Do the validation on a small set of validation images
for ind in val_indices:
input_image = np.expand_dims(np.float32(
utils.load_image(val_input_names[ind], args["crop_width"],
args["crop_height"])
[:args["crop_height"], :args["crop_width"]]),
axis=0) / 255.0
gt = utils.load_image(
val_output_names[ind], args["crop_width"],
args["crop_height"]
)[:args["crop_height"], :args["crop_width"]]
gt = helpers.reverse_one_hot(
helpers.one_hot_it(gt, label_values))
# st = time.time()
output_image = sess.run(network,
feed_dict={net_input: input_image})
output_image = np.array(output_image[0, :, :, :])
output_image = helpers.reverse_one_hot(output_image)
out_vis_image = helpers.colour_code_segmentation(
output_image, label_values)
accuracy, class_accuracies, prec, rec, f1, iou = utils.evaluate_segmentation(
pred=output_image, label=gt, num_classes=num_classes)
file_name = utils.filepath_to_name(val_input_names[ind])
target.write("%s, %f, %f, %f, %f, %f" %
(file_name, accuracy, prec, rec, f1, iou))
for item in class_accuracies:
target.write(", %f" % (item))
target.write("\n")
scores_list.append(accuracy)
class_scores_list.append(class_accuracies)
precision_list.append(prec)
recall_list.append(rec)
f1_list.append(f1)
iou_list.append(iou)
gt = helpers.colour_code_segmentation(gt, label_values)
file_name = os.path.basename(val_input_names[ind])
file_name = os.path.splitext(file_name)[0]
pred_img_path = "%s/%04d/%s_pred.png" % ("checkpoints", epoch,
file_name)
gt_img_path = "%s/%04d/%s_gt.png" % ("checkpoints", epoch,
file_name)
cv2.imwrite(
pred_img_path,
cv2.cvtColor(np.uint8(out_vis_image), cv2.COLOR_RGB2BGR))
cv2.imwrite(gt_img_path,
cv2.cvtColor(np.uint8(gt), cv2.COLOR_RGB2BGR))
# Send the first 16 images to sacred
if ind in val_indices[:16]:
ex.add_artifact(gt_img_path, "GtImage_%d" % ind)
ex.add_artifact(pred_img_path, "PredImage_%d" % ind)
target.close()
ex.add_artifact(target_path)
avg_score = np.mean(scores_list)
class_avg_scores = np.mean(class_scores_list, axis=0)
avg_scores_per_epoch.append(avg_score)
avg_precision = np.mean(precision_list)
avg_recall = np.mean(recall_list)
avg_f1 = np.mean(f1_list)
avg_iou = np.mean(iou_list)
avg_iou_per_epoch.append(avg_iou)
# Sacred info dict gets sent every heartbeat (10s)
ex.info["avg_score"] = avg_score
ex.info["class_avg_scores"] = class_avg_scores
ex.info["avg_precision"] = avg_precision
ex.info["avg_recall"] = avg_recall
ex.info["avg_f1"] = avg_f1
ex.info["avg_iou"] = avg_iou
_log.debug("\nAverage validation accuracy for epoch # %04d = %f" %
(epoch, avg_score))
_log.debug(
"Average per class validation accuracies for epoch # %04d:" %
(epoch))
for index, item in enumerate(class_avg_scores):
_log.debug("%s = %f" % (class_names_list[index], item))
_log.debug("Validation precision = ", avg_precision)
_log.debug("Validation recall = ", avg_recall)
_log.debug("Validation F1 score = ", avg_f1)
_log.debug("Validation IoU score = ", avg_iou)
epoch_time = time.time() - epoch_st
remain_time = epoch_time * (args["num_epochs"] - 1 - epoch)
m, s = divmod(remain_time, 60)
h, m = divmod(m, 60)
if s != 0:
train_time = "Remaining training time = %d hours %d minutes %d seconds\n" % (
h, m, s)
else:
train_time = "Remaining training time : Training completed.\n"
utils.LOG(train_time)
scores_list = []
fig1, ax1 = plt.subplots(figsize=(11, 8))
ax1.plot(range(epoch + 1), avg_scores_per_epoch)
ax1.set_title("Average validation accuracy vs epochs")
ax1.set_xlabel("Epoch")
ax1.set_ylabel("Avg. val. accuracy")
plt.savefig('accuracy_vs_epochs.png')
ex.add_artifact("accuracy_vs_epochs.png")
plt.clf()
fig2, ax2 = plt.subplots(figsize=(11, 8))
ax2.plot(range(epoch + 1), avg_loss_per_epoch)
ax2.set_title("Average loss vs epochs")
ax2.set_xlabel("Epoch")
ax2.set_ylabel("Current loss")
plt.savefig('loss_vs_epochs.png')
ex.add_artifact("loss_vs_epochs.png")
plt.clf()
fig3, ax3 = plt.subplots(figsize=(11, 8))
ax3.plot(range(epoch + 1), avg_iou_per_epoch)
ax3.set_title("Average IoU vs epochs")
ax3.set_xlabel("Epoch")
ax3.set_ylabel("Current IoU")
plt.savefig('iou_vs_epochs.png')
ex.add_artifact("iou_vs_epochs.png")
sess.run(tf.global_variables_initializer())
print('Loading model checkpoint weights')
saver=tf.train.Saver(max_to_keep=1000)
saver.restore(sess, args.checkpoint_path)
for ind in pre_indices:
print("Testing image " + pre_input_names[ind])
loaded_image = utils.load_image(pre_input_names[ind])
padding_image = cv2.copyMakeBorder(loaded_image, 14, 13, 14, 13, cv2.BORDER_REFLECT)
input_image = np.expand_dims(np.float32(padding_image),axis=0)/255.0
st = time.time()
output_image = sess.run(network,feed_dict={net_input:input_image})
run_time = time.time()-st
output_image = np.array(output_image[0,:,:,:])
output_image = helpers.reverse_one_hot(output_image)
out_vis_image = helpers.colour_code_segmentation(output_image, label_values)[14:-13,14:-13,:]
file_name = utils.filepath_to_name(pre_input_names[ind])
if not os.path.isdir("%s/predict/%s"%(args.dataset,args.model)):
os.makedirs("%s/predict/%s"%(args.dataset,args.model))
cv2.imwrite("%s/predict/%s/%s_pred.png"%(args.dataset,args.model,file_name),cv2.cvtColor(np.uint8(out_vis_image), cv2.COLOR_RGB2BGR))
print("Wrote image " + "%s/predict/%s/%s_pred.png"%(args.dataset,args.model,file_name))
print("")
print("Finished!")
def segment(file_list):
parser = argparse.ArgumentParser()
image = file_list[0]
checkpoint_path = "latest_model_DeepLabV3_plus_CamVid_people.ckpt"
crop_height = 512
crop_width = 512
model = "DeepLabV3_plus"
dataset = "CamVid_people"
class_names_list, label_values = helpers.get_label_info(
os.path.join(args.dataset, "class_dict.csv"))
num_classes = len(label_values)
#num_classes = 2
print("\n***** Begin prediction *****")
print("Dataset -->", args.dataset)
print("Model -->", args.model)
print("Crop Height -->", args.crop_height)
print("Crop Width -->", args.crop_width)
print("Num Classes -->", num_classes)
print("Image -->", args.image)
# Initializing network
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
net_input = tf.placeholder(tf.float32, shape=[None, None, None, 3])
net_output = tf.placeholder(tf.float32,
shape=[None, None, None, num_classes])
network, _ = model_builder.build_model(net_input=net_input,
num_classes=num_classes,
is_training=False)
sess.run(tf.global_variables_initializer())
print('Loading model checkpoint weights')
saver = tf.train.Saver(max_to_keep=1000)
saver.restore(sess, args.checkpoint_path)
print("Testing image " + args.image)
#loaded_image = utils.load_image(args.image)
images = [
"./CamVid_people/test/0001TP_007170.png",
"./CamVid_people/test/0001TP_007860.png",
"./CamVid_people/test/0001TP_008070.png"
]
for image in images:
loaded_image = utils.load_image(image)
resized_image = cv2.resize(loaded_image,
(args.crop_width, args.crop_height))
input_image = np.expand_dims(np.float32(
resized_image[:args.crop_height, :args.crop_width]),
axis=0) / 255.0
st = time.time()
output_image = sess.run(network, feed_dict={net_input: input_image})
run_time = time.time() - st
output_image = np.array(output_image[0, :, :, :])
output_image = helpers.reverse_one_hot(output_image)
out_vis_image = helpers.colour_code_segmentation(
output_image, label_values)
file_name = utils.filepath_to_name(args.image)
cv2.imwrite("%s_pred.png" % (file_name),
cv2.cvtColor(np.uint8(out_vis_image), cv2.COLOR_RGB2BGR))
print("")
print("Finished!")
print("time elapsed", run_time)
print("Wrote image " + "%s_pred.png" % (file_name))
(args.crop_width, args.crop_height))
input_image = np.expand_dims(np.float32(
resized_image[:args.crop_height, :args.crop_width]),
axis=0) / 255.0
st = time.time()
output_image = sess.run(network, feed_dict={net_input: input_image})
run_time = time.time() - st
output_image = np.array(output_image[0, :, :, :])
output_image = helpers.reverse_one_hot(output_image)
out_vis_image = helpers.colour_code_segmentation(output_image,
label_values)
overlayed_img = 0.35 * resized_image + 0.65 * out_vis_image
# overlayed_img = overlayed_img.astype(np.uint8)
file_name = utils.filepath_to_name(image)
cv2.imwrite("preds/%s_pred.png" % (file_name),
cv2.cvtColor(np.uint8(overlayed_img), cv2.COLOR_RGB2BGR))
cv2.imwrite(
"preds/%s_mask.png" % (file_name),
cv2.cvtColor((cv2.resize(np.uint8(out_vis_image),
(img_shape[0], img_shape[1]))),
cv2.COLOR_RGB2BGR))
print("")
print("Finished!")
print("Wrote image " + "%s_pred.png" % (file_name))
def main(args=None):
#If args is None, will parse command line arguments
#Otherwise args needs to be list with arguments, like
# ['--dataset', 'CamVid', '--model', 'FC-DenseNet103']
parser = argparse.ArgumentParser()
parser.add_argument('--num_epochs',
type=int,
default=300,
help='Number of epochs to train for')
parser.add_argument('--epoch_start_i',
type=int,
default=0,
help='Start counting epochs from this number')
parser.add_argument('--checkpoint_step',
type=int,
default=5,
help='How often to save checkpoints (epochs)')
parser.add_argument('--validation_step',
type=int,
default=1,
help='How often to perform validation (epochs)')
parser.add_argument('--continue_training',
type=utils.str2bool,
default=False,
help='Whether to continue training from a checkpoint')
parser.add_argument(
'--continue_from',
type=str,
default='',
help='From which checkpoint to continue. Only relevant with darea_call.'
)
parser.add_argument('--dataset',
type=str,
default="CamVid",
help='Dataset you are using.')
parser.add_argument('--dataset_path',
type=str,
default="",
help='Path to Dataset folder.')
parser.add_argument(
'--image_suffix',
type=str,
default='',
required=False,
help=
'Only files with this extension should be included. You should specify it if some non-image files will be in the same folder'
)
parser.add_argument('--crop_height',
type=int,
default=512,
help='Height of cropped input image to network')
parser.add_argument('--crop_width',
type=int,
default=512,
help='Width of cropped input image to network')
parser.add_argument(
'--biased_crop',
type=float,
default=0,
help=
'Probability of making a biased cropped. Biased crops always contain some foreground portion. Only works if one of the classes is named "Background".'
)
parser.add_argument(
'--downscale_factor',
type=float,
default=0,
required=False,
help=
'Shrink image by this factor. E.g. if image is 1024x1024 and downscale_factor is 0.5, downscaled image will be 512x512. This is applied before cropping.'
)
parser.add_argument('--batch_size',
type=int,
default=1,
help='Number of images in each batch')
parser.add_argument(
'--num_val_images',
type=int,
default=20,
help='The number of images to used for validations. If -1 -> use all')
parser.add_argument(
'--h_flip',
type=utils.str2bool,
default=False,
help=
'Whether to randomly flip the image horizontally for data augmentation'
)
parser.add_argument(
'--v_flip',
type=utils.str2bool,
default=False,
help=
'Whether to randomly flip the image vertically for data augmentation')
parser.add_argument(
'--brightness',
type=float,
default=None,
help=
'Whether to randomly change the image brightness for data augmentation. Specifies the max bightness change as a factor between 0.0 and 1.0. For example, 0.1 represents a max brightness change of 10%% (+-).'
)
parser.add_argument(
'--rotation',
type=float,
default=None,
help=
'DOES NOT WORK! Whether to randomly rotate the image for data augmentation. Specifies the max rotation angle in degrees.'
)
parser.add_argument('--rotation_perpendicular',
type=utils.str2bool,
default=False,
help='Randomly rotates by 0, 90, 180 or 270 degrees')
parser.add_argument(
'--model',
type=str,
default="FC-DenseNet103",
help=
'The model you are using. See model_builder.py for supported models')
parser.add_argument(
'--frontend',
type=str,
default="ResNet101",
help=
'The frontend you are using. See frontend_builder.py for supported models'
)
parser.add_argument(
'--save_best',
type=utils.str2bool,
default=False,
help='Saves model with smallest loss rather than last model')
parser.add_argument('--learn_rate',
type=float,
default=0.0001,
help='The learning rate')
parser.add_argument(
'--chkpt_prefix',
type=str,
default='',
help=
'Prefix in front of checkpoint (intended for distinguishing same models ran with different parameters)'
)
parser.add_argument(
'--darea_call',
type=utils.str2bool,
default=False,
required=False,
help='Set to true when you call it from Darea software')
parser.add_argument(
'--save_path',
type=str,
default='checkpoints',
required=False,
help='Name of saved model. Only used when darea_call is true.')
parser.add_argument('--makePlots',
type=utils.str2bool,
default=True,
required=False,
help='Whether plots should be made')
if args is None:
args = parser.parse_args()
else:
args = parser.parse_args(args)
if args.darea_call: os.chdir(os.path.dirname(os.path.realpath(__file__)))
# Get the names of the classes so we can record the evaluation results
class_names_list, label_values = helpers.get_label_info(
os.path.join(args.dataset_path, args.dataset, "class_dict.csv"))
class_names_string = ', '.join(class_names_list)
if 'Background' not in class_names_list:
args.biased_crop = 0
backgroundValue = None
else:
backgroundValue = label_values[class_names_list.index('Background')]
num_classes = len(label_values)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
# Compute your softmax cross entropy loss
net_input = tf.placeholder(tf.float32, shape=[
None, None, None, 3
]) #Try setting to 1 for grayscale, think theres no other changes needed.
net_output = tf.placeholder(tf.float32,
shape=[None, None, None, num_classes])
network, init_fn = model_builder.build_model(model_name=args.model,
frontend=args.frontend,
net_input=net_input,
num_classes=num_classes,
crop_width=args.crop_width,
crop_height=args.crop_height,
is_training=True)
loss = tf.reduce_mean(
tf.nn.softmax_cross_entropy_with_logits_v2(logits=network,
labels=net_output))
opt = tf.train.RMSPropOptimizer(
learning_rate=args.learn_rate, decay=0.995).minimize(
loss, var_list=[var for var in tf.trainable_variables()])
saver = tf.train.Saver(max_to_keep=500)
sess.run(tf.global_variables_initializer())
utils.count_params()
# If a pre-trained ResNet is required, load the weights.
# This must be done AFTER the variables are initialized with sess.run(tf.global_variables_initializer())
if init_fn is not None:
init_fn(sess)
# Load a previous checkpoint if desired
if args.darea_call:
if args.continue_training and args.continue_from:
model_checkpoint_name = '../../deepLearning/checkpoints/' + args.continue_from + '.ckpt'
if os.path.isfile(model_checkpoint_name + '.index'):
print('Loading latest model checkpoint...')
saver.restore(sess, model_checkpoint_name)
print('successfully loaded {}'.format(model_checkpoint_name))
else:
print('Specified checkpoint {} not found. Starting fresh one'.
format(os.path.abspath(model_checkpoint_name)))
if args.save_path:
model_checkpoint_name = os.path.join(args.save_path,
args.dataset + '.ckpt')
else:
if args.continue_training:
if args.continue_from:
model_checkpoint_name = args.continue_from
if not model_checkpoint_name.endswith('.ckpt'):
model_checkpoint_name += '.ckpt'
else:
model_checkpoint_name = os.path.join(
args.save_path, "latest_model_" + args.chkpt_prefix +
args.model + "_" + args.dataset + ".ckpt")
if os.path.isfile(model_checkpoint_name + '.index'):
print('Loading latest model checkpoint...')
saver.restore(sess, model_checkpoint_name)
print('successfully loaded {}'.format(model_checkpoint_name))
else:
print('{} not found. Starting a fresh training'.format(
args.continue_from))
model_checkpoint_name = os.path.join(
args.save_path, "latest_model_" + args.chkpt_prefix + args.model +
"_" + args.dataset + ".ckpt")
# Load the data
print("Loading the data ...")
train_input_names,train_output_names, val_input_names, val_output_names, test_input_names, test_output_names = \
utils.prepare_data(dataset_dir=os.path.join(args.dataset_path,args.dataset),image_suffix=args.image_suffix)
print("\n***** Begin training *****")
print("Dataset -->", args.dataset)
print("Model -->", args.model)
if args.downscale_factor:
print("Downscale Factor -->", args.downscale_factor)
print("Crop Height -->", args.crop_height)
print("Crop Width -->", args.crop_width)
print("Num Epochs -->", args.num_epochs)
print("Batch Size -->", args.batch_size)
print("Num Classes -->", num_classes)
print("Learn Rate -->", args.learn_rate)
print("Validation Step -->", args.validation_step)
print("Data Augmentation:")
print("\tVertical Flip -->", args.v_flip)
print("\tHorizontal Flip -->", args.h_flip)
print("\tBrightness Alteration -->", args.brightness)
print("\tRotation -->", args.rotation)
print("\tPerpendicular Rotation -->", args.rotation_perpendicular)
print("", flush=True)
avg_loss_per_epoch = []
avg_scores_per_epoch = []
avg_iou_per_epoch = []
if args.num_val_images == -1:
#Use all validation images if so wanted by users
val_indices = range(0, len(val_input_names))
else:
# Which validation images do we want
val_indices = []
num_vals = min(args.num_val_images, len(val_input_names))
# Set random seed to make sure models are validated on the same validation images.
# So you can compare the results of different models more intuitively.
random.seed(16)
val_indices = random.sample(range(0, len(val_input_names)), num_vals)
#Copy class file with same name as checkpoint
shutil.copyfile(
os.path.join(args.dataset_path, args.dataset, "class_dict.csv"),
os.path.splitext(model_checkpoint_name)[0] + '.classes')
# Do the training here
for epoch in range(args.epoch_start_i, args.num_epochs):
current_losses = []
cnt = 0
# Equivalent to shuffling
id_list = np.random.permutation(len(train_input_names))
num_iters = int(np.floor(len(id_list) / args.batch_size))
st = time.time()
epoch_st = time.time()
for i in range(num_iters):
# st=time.time()
input_image_batch = []
output_image_batch = []
# Collect a batch of images
for j in range(args.batch_size):
index = i * args.batch_size + j
id = id_list[index]
input_image = utils.load_image(train_input_names[id])
output_image = utils.load_image(train_output_names[id])
with tf.device('/cpu:0'):
input_image, output_image = data_augmentation(
input_image, output_image, args, backgroundValue)
if 0:
#Debugging:
try:
os.mkdir('imagesinTrain')
'kj'
except:
pass
try:
os.mkdir('imagesinTrain/%04d' % (epoch))
except:
pass
file_name = utils.filepath_to_name(
train_input_names[id])
file_name = os.path.splitext(file_name)[0]
cv2.imwrite(
"%s/%04d/%s_im.png" %
("imagesinTrain", epoch, file_name),
cv2.cvtColor(np.uint8(input_image),
cv2.COLOR_RGB2BGR))
cv2.imwrite(
"%s/%04d/%s_gt.png" %
("imagesinTrain", epoch, file_name),
cv2.cvtColor(np.uint8(output_image),
cv2.COLOR_RGB2BGR))
# Prep the data. Make sure the labels are in one-hot format
input_image = np.float32(input_image) / 255.0
output_image = np.float32(
helpers.one_hot_it(label=output_image,
label_values=label_values))
if 0:
#Debugging:
try:
os.mkdir('imagesinTrain')
except:
pass
try:
os.mkdir('imagesinTrain/%04d' % (epoch))
except:
pass
file_name = utils.filepath_to_name(
train_input_names[id])
file_name = os.path.splitext(file_name)[0]
print("%s/%04d/%s_im.png" %
("imagesinTrain", epoch, file_name))
cv2.imwrite(
"%s/%04d/%s_im.png" %
("imagesinTrain", epoch, file_name),
cv2.cvtColor(np.uint8(input_image),
cv2.COLOR_RGB2BGR))
#cv2.imwrite("%s/%04d/%s_gt.png"%("imagesinTrain",epoch, file_name),cv2.cvtColor(np.uint8(output_image), cv2.COLOR_RGB2BGR))
input_image_batch.append(
np.expand_dims(input_image, axis=0))
output_image_batch.append(
np.expand_dims(output_image, axis=0))
if args.batch_size == 1:
input_image_batch = input_image_batch[0]
output_image_batch = output_image_batch[0]
else:
input_image_batch = np.squeeze(
np.stack(input_image_batch, axis=1))
output_image_batch = np.squeeze(
np.stack(output_image_batch, axis=1))
# Do the training
_, current, output_image = sess.run([opt, loss, network],
feed_dict={
net_input:
input_image_batch,
net_output:
output_image_batch
})
current_losses.append(current)
cnt = cnt + args.batch_size
if cnt % 25 == 0:
string_print = "Epoch = %d Count = %d Current_Loss = %.4f Time = %.2f" % (
epoch, cnt, current, time.time() - st)
utils.LOG(string_print)
st = time.time()
if 0:
#For Debugging
output_image = np.array(output_image[0, :, :, :])
output_image = helpers.reverse_one_hot(output_image)
out_vis_image = helpers.colour_code_segmentation(
output_image, label_values)
cv2.imwrite(
"%s/%04d/%s_pred.png" %
("imagesinTrain", epoch, file_name),
cv2.cvtColor(np.uint8(out_vis_image), cv2.COLOR_RGB2BGR))
mean_loss = np.mean(current_losses)
avg_loss_per_epoch.append(mean_loss)
# Create directories if needed
if not os.path.isdir("%s/%04d" % ("checkpoints", epoch)):
os.makedirs("%s/%04d" % ("checkpoints", epoch))
# If latest checkpoint should be saved, save now to same file name,
if not args.save_best:
print("Saving latest checkpoint")
saver.save(sess, model_checkpoint_name)
if val_indices != 0 and epoch % args.checkpoint_step == 0:
print("Saving checkpoint for this epoch")
saver.save(sess, "%s/%04d/model.ckpt" % (args.save_path, epoch))
if epoch % args.validation_step == 0:
if epoch == 0:
best_avg_iou = 0
print("Performing validation")
target = open("%s/%04d/val_scores.csv" % ("checkpoints", epoch),
'w')
target.write(
"val_name, avg_accuracy, precision, recall, f1 score, mean iou, %s\n"
% (class_names_string))
scores_list = []
class_scores_list = []
precision_list = []
recall_list = []
f1_list = []
iou_list = []
# Do the validation on a small set of validation images
for ind in val_indices:
input_image = utils.load_image(val_input_names[ind])
gt = utils.load_image(val_output_names[ind])
if args.downscale_factor and args.downscale_factor != 1:
dim = (int(input_image.shape[0] * args.downscale_factor),
int(input_image.shape[1] * args.downscale_factor))
input_image = cv2.resize(input_image,
dim,
interpolation=cv2.INTER_CUBIC)
gt = cv2.resize(gt, dim, interpolation=cv2.INTER_NEAREST)
#input_image, gt = data_augmentation(input_image, gt, args)
gt = helpers.reverse_one_hot(
helpers.one_hot_it(gt, label_values))
crop_height = args.crop_height
crop_width = args.crop_width
if input_image.shape[0] > crop_height or input_image.shape[
1] > crop_width:
#rectangle in bottom right corner smaller than cropped im will not be used
nrCroppingsY = input_image.shape[0] // crop_height
nrCroppingsX = input_image.shape[1] // crop_width
output_image = np.zeros([
nrCroppingsY * crop_height, nrCroppingsX * crop_width
])
gt = gt[0:nrCroppingsY * crop_height,
0:nrCroppingsX * crop_width]
for yi in range(nrCroppingsY):
row = np.zeros(
[crop_height, nrCroppingsX * crop_width])
for xi in range(nrCroppingsX):
inputIm = input_image[yi * crop_height:(1 + yi) *
crop_height,
xi * crop_width:(1 + xi) *
crop_width, :]
inputIm = np.expand_dims(np.float32(inputIm) /
255.0,
axis=0)
out = sess.run(network,
feed_dict={net_input: inputIm})
out = np.array(out[0, :, :, :])
out = helpers.reverse_one_hot(out)
row[:, xi * crop_width:(1 + xi) * crop_width] = out
output_image[yi * crop_height:(1 + yi) *
crop_height, :] = row
# st = time.time()
else:
input_image = np.expand_dims(np.float32(input_image) /
255.0,
axis=0)
output_image = sess.run(network,
feed_dict={net_input: input_image})
output_image = np.array(output_image[0, :, :, :])
output_image = helpers.reverse_one_hot(output_image)
out_vis_image = helpers.colour_code_segmentation(
output_image, label_values)
accuracy, class_accuracies, prec, rec, f1, iou, class_iou = utils.evaluate_segmentation(
pred=output_image, label=gt, num_classes=num_classes)
file_name = utils.filepath_to_name(val_input_names[ind])
target.write("%s, %f, %f, %f, %f, %f" %
(file_name, accuracy, prec, rec, f1, iou))
for item in class_accuracies:
target.write(", %f" % (item))
target.write("\n")
scores_list.append(accuracy)
class_scores_list.append(class_accuracies)
precision_list.append(prec)
recall_list.append(rec)
f1_list.append(f1)
iou_list.append(iou)
gt = helpers.colour_code_segmentation(gt, label_values)
file_name = os.path.basename(val_input_names[ind])
file_name = os.path.splitext(file_name)[0]
cv2.imwrite(
"%s/%04d/%s_pred.png" % ("checkpoints", epoch, file_name),
cv2.cvtColor(np.uint8(out_vis_image), cv2.COLOR_RGB2BGR))
cv2.imwrite(
"%s/%04d/%s_gt.png" % ("checkpoints", epoch, file_name),
cv2.cvtColor(np.uint8(gt), cv2.COLOR_RGB2BGR))
target.close()
avg_score = np.mean(scores_list)
class_avg_scores = np.mean(class_scores_list, axis=0)
avg_scores_per_epoch.append(avg_score)
avg_precision = np.mean(precision_list)
avg_recall = np.mean(recall_list)
avg_f1 = np.mean(f1_list)
avg_iou = np.mean(iou_list)
avg_iou_per_epoch.append(avg_iou)
print("\nAverage validation accuracy for epoch # %04d = %f" %
(epoch, avg_score))
print("Average per class validation accuracies for epoch # %04d:" %
(epoch))
for index, item in enumerate(class_avg_scores):
print("%s = %f" % (class_names_list[index], item))
print("Validation precision = ", avg_precision)
print("Validation recall = ", avg_recall)
print("Validation F1 score = ", avg_f1)
print("Validation IoU score = ", avg_iou, flush=True)
if args.save_best and avg_iou > best_avg_iou:
#Save model if best model is wanted and it is the best or if it is first evaluation
print("Saving best checkpoint with iou {}".format(avg_iou))
saver.save(sess, model_checkpoint_name)
best_avg_iou = avg_iou
#Save an info file
with open(model_checkpoint_name[:-5] + '.info', 'w') as f:
f.write('Epoch\t{}\nValidation IoU score\t{}'.format(
epoch, avg_iou))
epoch_time = time.time() - epoch_st
remain_time = epoch_time * (args.num_epochs - 1 - epoch)
m, s = divmod(remain_time, 60)
h, m = divmod(m, 60)
if s != 0:
train_time = "Remaining training time = %d hours %d minutes %d seconds\n" % (
h, m, s)
else:
train_time = "Remaining training time : Training completed.\n"
utils.LOG(train_time)
scores_list = []
with open(model_checkpoint_name[:-5] + '.info', 'a+') as f:
#Save some info on filesizes
f.write('\nimageSize\t{}'.format([args.crop_height, args.crop_width]))
f.write('\nTraining completed\t{}'.format(
datetime.datetime.now().strftime("%d-%b-%Y %H:%M:%S")))
if not args.darea_call and args.makePlots:
fig1, ax1 = plt.subplots(figsize=(11, 8))
ax1.plot(range(epoch + 1), avg_scores_per_epoch)
ax1.set_title("Average validation accuracy vs epochs")
ax1.set_xlabel("Epoch")
ax1.set_ylabel("Avg. val. accuracy")
plt.savefig('accuracy_vs_epochs.png')
plt.clf()
fig2, ax2 = plt.subplots(figsize=(11, 8))
ax2.plot(range(epoch + 1), avg_loss_per_epoch)
ax2.set_title("Average loss vs epochs")
ax2.set_xlabel("Epoch")
ax2.set_ylabel("Current loss")
plt.savefig('loss_vs_epochs.png')
plt.clf()
fig3, ax3 = plt.subplots(figsize=(11, 8))
ax3.plot(range(epoch + 1), avg_iou_per_epoch)
ax3.set_title("Average IoU vs epochs")
ax3.set_xlabel("Epoch")
ax3.set_ylabel("Current IoU")
plt.savefig('iou_vs_epochs.png')
def main(args=None):
#If args is None, will parse command line arguments
#Otherwise args needs to be list with arguments, like
# ['--image', 'imageFolder', '--model', 'FC-DenseNet103']
parser = argparse.ArgumentParser()
parser.add_argument(
'--image',
type=str,
default=None,
required=True,
help='The image or folder with images you want to predict on. ')
parser.add_argument(
'--checkpoint_path',
type=str,
default=None,
required=False,
help=
'The path to the latest checkpoint weights for your model. Will guess based on model and dataset if not specified'
)
parser.add_argument('--crop_height',
type=int,
default=512,
help='Height of cropped input image to network')
parser.add_argument('--crop_width',
type=int,
default=512,
help='Width of cropped input image to network')
parser.add_argument(
'--downscale_factor',
type=float,
default=0,
required=False,
help=
'Shrink image by this factor. E.g. if image is 1024x1024 and downscale_factor is 0.5, downscaled image will be 512x512. This is applied before cropping.'
)
parser.add_argument('--model',
type=str,
default="FC-DenseNet103",
required=False,
help='The model you are using')
parser.add_argument('--dataset',
type=str,
default="CamVid",
required=False,
help='The dataset you are using')
parser.add_argument(
'--image_suffix',
type=str,
default='',
required=False,
help=
'Only files with this extension should be included. You should specify it if some non-image files will be in the same folder'
)
parser.add_argument(
'--outpath',
type=str,
default='./',
required=False,
help='Folder where predicted images should be saved to')
parser.add_argument(
'--file_suffix',
type=str,
default='_pred.png',
required=False,
help='Suffix appended to input image name for output image')
parser.add_argument(
'--darea_call',
type=utils.str2bool,
default=False,
required=False,
help='Set to true when you call it from Darea software')
parser.add_argument('--save_predictionImage',
type=utils.str2bool,
default=True,
required=False,
help='Whether predictions should be saved as images')
parser.add_argument(
'--save_coordinates',
type=utils.str2bool,
default=False,
required=False,
help='Whether coordinates of predicted structures should be saved')
if args is None:
args = parser.parse_args()
else:
args = parser.parse_args(args)
if args.darea_call:
os.chdir(os.path.dirname(os.path.realpath(__file__)))
class_names_list, label_values = helpers.get_label_info(
os.path.join(args.dataset, "class_dict.csv"))
num_classes = len(label_values)
print("\n***** Begin prediction *****")
print("Dataset -->", args.dataset)
print("Model -->", args.model)
print("Crop Height -->", args.crop_height)
print("Crop Width -->", args.crop_width)
print("Num Classes -->", num_classes)
print("Image -->", args.image)
if args.checkpoint_path is None:
args.checkpoint_path = "checkpoints/latest_model_" + args.model + "_" + args.dataset + ".ckpt"
# Initializing network
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
net_input = tf.placeholder(tf.float32, shape=[None, None, None, 3])
net_output = tf.placeholder(tf.float32,
shape=[None, None, None, num_classes])
network, _ = model_builder.build_model(args.model,
net_input=net_input,
num_classes=num_classes,
crop_width=args.crop_width,
crop_height=args.crop_height,
is_training=False)
sess.run(tf.global_variables_initializer())
print('Loading model checkpoint weights')
saver = tf.train.Saver(max_to_keep=1000)
saver.restore(sess, args.checkpoint_path)
if os.path.isdir(args.image):
folder = args.image
#Get list of all images in the folder, ignore hidden files, subfolders and files not ending with correct suffix
images = [
os.path.join(folder, file) for file in os.listdir(folder)
if not file.startswith('.') and not os.path.isdir(file)
and file.endswith(args.image_suffix)
]
else:
images = [args.image]
if not os.path.isdir(args.outpath):
os.mkdir(args.outpath)
if args.save_coordinates:
coordinates = dict()
print('Performing predictions...\n')
progress_string = ''
for index, image in enumerate(images):
if args.darea_call:
if index > 1:
print('\b' * (len(progress_string) +
2)) #Deletes previous output in matlab console
progress_string = 'Predicting on image {} / {}'.format(
index + 1, len(images))
print(progress_string)
else:
print('Predicting on image {} / {}\r'.format(
index + 1, len(images)))
input_image = utils.load_image(image)
if args.downscale_factor and args.downscale_factor != 1:
dim = (int(input_image.shape[0] * args.downscale_factor),
int(input_image.shape[1] * args.downscale_factor))
input_image = cv2.resize(input_image,
dim,
interpolation=cv2.INTER_CUBIC)
st = time.time()
crop_height = args.crop_height
crop_width = args.crop_width
if input_image.shape[0] > crop_height or input_image.shape[
1] > crop_width:
#rectangle in bottom right corner smaller than cropped im will not be used
nrCroppingsY = m.ceil(input_image.shape[0] / crop_height)
nrCroppingsX = m.ceil(input_image.shape[1] / crop_width)
output_image = np.zeros(
[nrCroppingsY * crop_height, nrCroppingsX * crop_width])
for yi in range(nrCroppingsY):
row = np.zeros([crop_height, nrCroppingsX * crop_width])
cropYstart = yi * crop_height
cropYstop = (1 + yi) * crop_height
if cropYstop >= input_image.shape[0]:
cropYstop = input_image.shape[0] - 1
cropYstart = cropYstop - crop_height
for xi in range(nrCroppingsX):
cropXstart = xi * crop_width
cropXstop = (1 + xi) * crop_width
if cropXstop >= input_image.shape[1]:
cropXstop = input_image.shape[1] - 1
cropXstart = cropXstop - crop_width
inputIm = input_image[cropYstart:cropYstop,
cropXstart:cropXstop, :]
inputIm = np.expand_dims(np.float32(inputIm) / 255.0,
axis=0)
#print(inputIm.shape)
out = sess.run(network, feed_dict={net_input: inputIm})
out = np.array(out[0, :, :, :])
out = helpers.reverse_one_hot(out)
if (1 + xi) * crop_width >= input_image.shape[1]:
row[:, xi * crop_width:] = out
else:
row[:, xi * crop_width:(1 + xi) * crop_width] = out
if (1 + yi) * crop_height >= input_image.shape[0]:
output_image[yi * crop_height:, :] = row
else:
output_image[yi * crop_height:(1 + yi) *
crop_height, :] = row
# st = time.time()
else:
input_image = np.expand_dims(np.float32(input_image) / 255.0,
axis=0)
output_image = sess.run(network,
feed_dict={net_input: input_image})
output_image = np.array(output_image[0, :, :, :])
output_image = helpers.reverse_one_hot(output_image)
run_time = time.time() - st
out_vis_image = helpers.colour_code_segmentation(
output_image, label_values)
if args.save_predictionImage:
file_name = "{}{}".format(utils.filepath_to_name(image),
args.file_suffix)
cv2.imwrite(
os.path.join(args.outpath, file_name),
cv2.cvtColor(np.uint8(out_vis_image), cv2.COLOR_RGB2BGR))
if args.save_coordinates:
coordinates[utils.filepath_to_name(
image)] = utils.getCoordsFromPrediction(
cv2.cvtColor(np.uint8(out_vis_image), cv2.COLOR_RGB2GRAY),
args.downscale_factor)
if not args.darea_call:
print("Wrote image {}".format(file_name))
print("")
if args.save_coordinates:
with open(os.path.join(args.outpath, 'centroids.json'), 'w') as f:
json.dump(coordinates, f)
if not args.darea_call:
print("Finished!")
def main(args=None):
#If args is None, will parse command line arguments
#Otherwise args needs to be list with arguments, like
# ['--image', 'imageFolder', '--model', 'FC-DenseNet103']
parser = argparse.ArgumentParser()
parser.add_argument('--image', type=str, default=None, required=True, help='The image or folder with images you want to predict on. ')
parser.add_argument('--classes', type=str, required=False, default=None, help='Path to class_dict.csv. If None, expects it to be same as .ckpt but with .classes extension')
parser.add_argument('--checkpoint_path', type=str, required=True, help='The path to the latest checkpoint weights for your model.')
parser.add_argument('--crop_height', type=int, default=512, help='Height of cropped input image to network')
parser.add_argument('--crop_width', type=int, default=512, help='Width of cropped input image to network')
parser.add_argument('--downscale_factor', type=float, default=0, required=False, help='Shrink image by this factor. E.g. if image is 1024x1024 and downscale_factor is 0.5, downscaled image will be 512x512. This is applied before cropping.')
parser.add_argument('--model', type=str, default="FC-DenseNet103", required=False, help='The model you are using')
parser.add_argument('--image_suffix', type=str, default='', required=False, help='Only files with this extension should be included. You should specify it if some non-image files will be in the same folder')
parser.add_argument('--outpath', type=str, default='./', required=False, help='Folder where predicted images should be saved to')
parser.add_argument('--file_suffix', type=str, default='_pred.png', required=False, help='Suffix appended to input image name for output image')
parser.add_argument('--save_predictionImage', type=utils.str2bool, default=True, required=False, help='Whether predictions should be saved as images')
parser.add_argument('--save_coordinates', type=utils.str2bool, default=True, required=False, help='Whether coordinates of predicted structures should be saved')
parser.add_argument('--coords_filename', type=str, default='coordinates',required=False, help='Filename of the coordinates file (without filextension).')
parser.add_argument('--coords_class', type=str, default=None, required=False, help='Specify class to generate coordinates from. (Default is all classes). Separate multiple classes by comma to pool them.')
parser.add_argument('--wait_time', type=float, default=1, required=False, help='Time (in s) to wait for new images to appear in folder')
if args is None:
args=parser.parse_args()
else:
args=parser.parse_args(args)
if args.classes is None:
args.classes=os.path.splitext(args.checkpoint_path)[0]+'.classes'
print(args.classes)
class_names_list, label_values = helpers.get_label_info(args.classes)
num_classes = len(label_values)
assert os.path.isdir(args.image), "Argument image needs to be a folder, got {}".format(args.image)
if args.save_coordinates:
if args.coords_class is None:
foregroundcolors=[c[0] for c in label_values[1:]]
#Labels are in RGB, but we only need grayscale value
else:
target_classes=args.coords_class.split(',')
foregroundcolors=[label_values[i][0] for i,c in enumerate(class_names_list) if c in target_classes]
if len(foregroundcolors)<1:
raise ValueError('None of the specified --coords_class found in class dict.\nSpecified classes were {}\nExisting classes are {}'.format(target_classes,class_names_list))
print("\n***** Begin prediction *****")
#print("Dataset -->", args.dataset)
print("Model -->", args.model)
print("Crop Height -->", args.crop_height)
print("Crop Width -->", args.crop_width)
print("Num Classes -->", num_classes)
print("Image -->", args.image)
if args.checkpoint_path is None:
args.checkpoint_path = "checkpoints/latest_model_" + args.model + "_" + args.dataset + ".ckpt"
# Initializing network
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess=tf.Session(config=config)
net_input = tf.placeholder(tf.float32,shape=[None,None,None,3])
net_output = tf.placeholder(tf.float32,shape=[None,None,None,num_classes])
network, _ = model_builder.build_model(args.model, net_input=net_input,
num_classes=num_classes,
crop_width=args.crop_width,
crop_height=args.crop_height,
is_training=False)
sess.run(tf.global_variables_initializer())
print('Loading model checkpoint weights')
saver=tf.train.Saver(max_to_keep=1000)
saver.restore(sess, args.checkpoint_path)
#time.sleep(args.wait_time*3)
folder=args.image
finished_images=[]
images=getUnfinishedImages(folder,finished_images,args.image_suffix)
if not os.path.isdir(args.outpath):
os.mkdir(args.outpath)
if args.save_coordinates:
coordinates=dict()
print('Performing predictions...')
while images:
time.sleep(args.wait_time)
for index,image in enumerate(images):
print("Predicting image {} (image {} out of {})".format(image,index+1, len(images)))
input_image = utils.load_image(image)
if args.downscale_factor and args.downscale_factor !=1:
dim=(int(input_image.shape[0]*args.downscale_factor), int(input_image.shape[1]*args.downscale_factor))
input_image=cv2.resize(input_image,dim,interpolation=cv2.INTER_CUBIC)
st = time.time()
crop_height=args.crop_height
crop_width=args.crop_width
if input_image.shape[0]>crop_height or input_image.shape[1]>crop_width:
#rectangle in bottom right corner smaller than cropped im will not be used
nrCroppingsY=m.ceil(input_image.shape[0]/crop_height)
nrCroppingsX=m.ceil(input_image.shape[1]/crop_width)
output_image=np.zeros([nrCroppingsY*crop_height,nrCroppingsX*crop_width])
for yi in range(nrCroppingsY):
row=np.zeros([crop_height,nrCroppingsX*crop_width])
cropYstart=yi*crop_height;
cropYstop=(1+yi)*crop_height
if cropYstop>=input_image.shape[0]:
cropYstop=input_image.shape[0]-1
cropYstart=cropYstop-crop_height
for xi in range(nrCroppingsX):
cropXstart=xi*crop_width;
cropXstop=(1+xi)*crop_width
if cropXstop>=input_image.shape[1]:
cropXstop=input_image.shape[1]-1
cropXstart=cropXstop-crop_width
inputIm=input_image[cropYstart:cropYstop,cropXstart:cropXstop,:]
inputIm=np.expand_dims(np.float32(inputIm)/255.0,axis=0)
#print(inputIm.shape)
out=sess.run(network,feed_dict={net_input:inputIm})
out = np.array(out[0,:,:,:])
out=helpers.reverse_one_hot(out)
if (1+xi)*crop_width>=input_image.shape[1]:
row[:,xi*crop_width:]=out
else:
row[:,xi*crop_width:(1+xi)*crop_width]=out
if (1+yi)*crop_height>=input_image.shape[0]:
output_image[yi*crop_height:,:]=row
else:
output_image[yi*crop_height:(1+yi)*crop_height,:]=row
# st = time.time()
else:
input_image=np.expand_dims(np.float32(input_image)/255.0,axis=0)
output_image = sess.run(network,feed_dict={net_input:input_image})
output_image = np.array(output_image[0,:,:,:])
output_image = helpers.reverse_one_hot(output_image)
run_time = time.time()-st
out_vis_image = helpers.colour_code_segmentation(output_image, label_values)
if args.save_predictionImage:
file_name = "{}{}".format(utils.filepath_to_name(image),args.file_suffix)
cv2.imwrite(os.path.join(args.outpath,file_name),cv2.cvtColor(np.uint8(out_vis_image), cv2.COLOR_RGB2BGR))
print("Wrote image {}".format(file_name))
if args.save_coordinates:
coordinates[utils.filepath_to_name(image)]=utils.getCoordsFromPrediction(cv2.cvtColor(np.uint8(out_vis_image), cv2.COLOR_RGB2GRAY),foregroundcolors,args.downscale_factor);
print("", flush=True)
#Finished predicting all images in list
finished_images+=images
#Now check if new images appeared in folder, if so, repeat, else wait, retry then finish
images=getUnfinishedImages(folder,finished_images,args.image_suffix)
if images:
continue
print("Waiting for new images", flush=True)
time.sleep(args.wait_time)
images=getUnfinishedImages(folder,finished_images,args.image_suffix)
if not images:
break
if args.save_coordinates:
with open(os.path.join(args.outpath,'{}.json'.format(args.coords_filename)), 'w') as f:
json.dump(coordinates,f)
print("Finished!")
run_times_list.append(time.time()-st)
#print('Predicted image, took {}'.format(time.time()-st))
#st = time.time()
output_image = np.array(output_image[0,:,:,:])
output_image = helpers.reverse_one_hot(output_image)
out_vis_image = helpers.colour_code_segmentation(output_image, label_values)
#print('Computed output images, took {}'.format(time.time()-st))
#st = time.time()
accuracy, class_accuracies, prec, rec, f1, iou, class_iou = utils.evaluate_segmentation(pred=output_image, label=gt, num_classes=num_classes)
#print('Computed metrics, took {}'.format(time.time()-st))
#st = time.time()
file_name = utils.filepath_to_name(input_name)
target.write("\n%s,%g,%g,%g,%g,%g"%(file_name, accuracy, prec, rec, f1, iou))
for item in class_accuracies:
target.write(",%g"%(item))
for item in class_iou:
target.write(",%g"%(item))
scores_list.append(accuracy)
class_scores_list.append(class_accuracies)
precision_list.append(prec)
recall_list.append(rec)
f1_list.append(f1)
iou_list.append(iou)
#print('Wrote metrics, took {}'.format(time.time()-st))
#st = time.time()
def predict_segment(images, save=False):
checkpoint_path = "./weights/weight3/DeepLabV3_plus_CamVid_people_weight3.ckpt"
model = "DeepLabV3_plus"
dataset = "CamVid_people"
crop_height = 512
crop_width = 512
class_names_list, label_values = helpers.get_label_info(os.path.join(dataset, "class_dict.csv"))
num_classes = len(label_values)
#num_classes = 2
print("\n***** Begin prediction *****")
print("Dataset -->", dataset)
print("Model -->", model)
print("Crop Height -->", crop_height)
print("Crop Width -->", crop_width)
print("Num Classes -->", num_classes)
print("Image -->", images)
# Initializing network
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess=tf.Session(config=config)
net_input = tf.placeholder(tf.float32,shape=[None,None,None,3])
net_output = tf.placeholder(tf.float32,shape=[None,None,None,num_classes])
network, _ = model_builder.build_model(model, net_input=net_input,
num_classes=num_classes,
crop_width=crop_width,
crop_height=crop_height,
is_training=False)
sess.run(tf.global_variables_initializer())
print('Loading model checkpoint weights')
saver=tf.train.Saver(max_to_keep=1000)
saver.restore(sess, checkpoint_path)
print("Testing image", images)
input_image = [cv2.resize(utils.load_image(image), (crop_width, crop_height))[:crop_height, :crop_width] for image in images]
input_image = np.array(input_image, dtype = np.float32)/255.0
st = time.time()
output_images = sess.run(network,feed_dict={net_input:input_image})
run_time = time.time()-st
output_images = helpers.reverse_one_hot(output_images)
out_vis_images = helpers.colour_binary_segmentation(output_images)
if save:
for i in range(len(images)):
out_vis_image = np.array(out_vis_images[i,:,:])
#output_image = helpers.reverse_one_hot(output_image)
#out_vis_image = helpers.colour_code_segmentation(output_image, label_values)
image = images[i]
file_name = utils.filepath_to_name(image)
cv2.imwrite("%s_pred.png"%(file_name),cv2.cvtColor(np.uint8(out_vis_image), cv2.COLOR_RGB2BGR))
print("Wrote image " + "%s_pred.png"%(file_name))
#output_images = helpers.reverse_one_hot(output_images)
#out_vis_images = helpers.colour_binary_segmentation(output_images)
print("")
print("segment Finished!")
print("time elapsed",run_time)
return out_vis_images
is_training=False)
sess.run(tf.global_variables_initializer())
print('Loading model checkpoint weights')
saver=tf.train.Saver(max_to_keep=1000)
saver.restore(sess, args.checkpoint_path)
for image in os.listdir(args.input_folder):
if not os.path.splitext(image)[-1] == '.png':
continue
image_path = os.path.join(args.input_folder, image)
loaded_image = utils.load_image(image_path)
resized_image =cv2.resize(loaded_image, (args.crop_width, args.crop_height))
input_image = np.expand_dims(np.float32(resized_image[:args.crop_height, :args.crop_width]),axis=0)/255.0
st = time.time()
output_image = sess.run(network,feed_dict={net_input:input_image})
run_time = time.time()-st
output_image = np.array(output_image[0,:,:,:])
output_image = helpers.reverse_one_hot(output_image)
out_vis_image = helpers.colour_code_segmentation(output_image, label_values)
image_path_no_ext = os.path.splitext(image_path)[0]
file_name = utils.filepath_to_name(image_path_no_ext)
cv2.imwrite(os.path.join(args.output_folder, "%s.png"%(file_name)),cv2.cvtColor(np.uint8(out_vis_image), cv2.COLOR_RGB2BGR))
def pred(img):
'''
parser = argparse.ArgumentParser()
parser.add_argument('--image', type=str, default=None, required=True, help='The image you want to predict on. ')
parser.add_argument('--checkpoint_path', type=str, default=None, required=True, help='The path to the latest checkpoint weights for your model.')
parser.add_argument('--crop_height', type=int, default=512, help='Height of cropped input image to network')
parser.add_argument('--crop_width', type=int, default=512, help='Width of cropped input image to network')
parser.add_argument('--model', type=str, default=None, required=True, help='The model you are using')
parser.add_argument('--dataset', type=str, default="CamVid", required=False, help='The dataset you are using')
'''
print("Testing image " + img)
loaded_image = utils.load_image(img)
resized_image = cv2.resize(loaded_image, (crop_width, crop_height))
input_image = np.expand_dims(
np.float32(resized_image[:crop_height, :crop_width]), axis=0) / 255.0
st = time.time()
output_image = sess.run(network, feed_dict={net_input: input_image})
run_time = time.time() - st
output_image = np.array(output_image[0, :, :, :])
output_image = helpers.reverse_one_hot(output_image)
out_vis_image = helpers.colour_code_segmentation(output_image,
label_values)
file_name = utils.filepath_to_name(img)
res = cv2.cvtColor(np.uint8(out_vis_image), cv2.COLOR_RGB2BGR)
cv2.imwrite("%s_pred.png" % (file_name), res)
final_out = cv2.addWeighted(resized_image, 1, res, 0.6, 0)
cv2.imwrite("%s_pred_overlayed.png" % (file_name), final_out)
print("")
print("Finished!")
print("Wrote image " + "%s_pred.png" % (file_name))
# USAGE
# python segment.py --model enet-cityscapes/enet-model.net --classes enet-cityscapes/enet-classes.txt --colors enet-cityscapes/enet-colors.txt --image images/example_01.png
# load the class label names
CLASSES = open(classes).read().strip().split("\n")
# if a colors file was supplied, load it from disk
if colors:
COLORS = open(colors).read().strip().split("\n")
COLORS = [np.array(c.split(",")).astype("int") for c in COLORS]
COLORS = np.array(COLORS, dtype="uint8")
#print(COLORS)
legend = np.zeros(((len(CLASSES) * 15) + 15, 300, 3), dtype="uint8") + 255
# loop over the class names + colors
for (i, (className, color)) in enumerate(zip(CLASSES, COLORS)):
# draw the class name + color on the legend
color = [int(c) for c in color]
cv2.putText(legend, ' ' + className,
(0, (i * 15) + 14), cv2.FONT_HERSHEY_SIMPLEX, 0.4,
tuple(color), 1, cv2.LINE_AA)
cv2.rectangle(legend, (100, (i * 15)), (300, (i * 15) + 15),
tuple(color), -1)
leg = cv2.cvtColor(np.uint8(legend), cv2.COLOR_RGB2BGR)
cv2.imwrite("legend.png", leg)
return [img, file_name + '_pred_overlayed.png', 'legend.png']