def predict_video(args, image_shape=None):
if args.video_file_in is None:
print("for video processing need --video_file_in")
return
if args.video_file_out is None:
print("for video processing need --video_file_out")
return
def process_frame(image):
if image_shape is not None:
image = scipy.misc.imresize(image, image_shape)
segmented_image, tf_time_ms, img_time_ms = predict_image(
sess, model, image, colors)
return segmented_image
tf.reset_default_graph()
with tf.Session(config=session_config(args)) as sess:
model = fcn16vgg16.FCN16_VGG16(define_graph=False)
model.load_model(
sess,
'trained_model' if args.model_dir is None else args.model_dir)
print('Running on video {}, output to: {}'.format(
args.video_file_in, args.video_file_out))
colors = get_colors()
input_clip = VideoFileClip(args.video_file_in)
annotated_clip = input_clip.fl_image(process_frame)
annotated_clip.write_videofile(args.video_file_out, audio=False)
def optimise_graph(args):
""" optimize frozen graph for inference """
if args.frozen_model_dir is None:
print("for optimise need --frozen_model_dir")
return
if args.optimised_model_dir is None:
print("for optimise need --optimised_model_dir")
return
print('calling c++ implementation of graph transform')
os.system('./optimise.sh {}'.format(args.frozen_model_dir))
# reading optimised graph
tf.reset_default_graph()
gd = tf.GraphDef()
output_graph_file = args.frozen_model_dir + "/optimised_graph.pb"
with tf.gfile.Open(output_graph_file, 'rb') as f:
gd.ParseFromString(f.read())
tf.import_graph_def(gd, name='')
print("{} ops in the optimised graph".format(len(gd.node)))
# save model in same format as usual
shutil.rmtree(args.optimised_model_dir, ignore_errors=True)
#if not os.path.exists(args.optimised_model_dir):
# os.makedirs(args.optimised_model_dir)
print('saving optimised model as saved_model to {}'.format(
args.optimised_model_dir))
model = fcn16vgg16.FCN16_VGG16(define_graph=False)
tf.reset_default_graph()
tf.import_graph_def(gd, name='')
with tf.Session() as sess:
model.save_model(sess, args.optimised_model_dir)
shutil.move(args.frozen_model_dir + '/optimised_graph.pb',
args.optimised_model_dir)
def predict_files(args, image_shape):
tf.reset_default_graph()
with tf.Session(config=session_config(args)) as sess:
model = fcn16vgg16.FCN16_VGG16(define_graph=False)
model.load_model(
sess,
'trained_model' if args.model_dir is None else args.model_dir)
# Make folder for current run
output_dir = os.path.join(args.runs_dir,
time.strftime("%Y%m%d_%H%M%S"))
if os.path.exists(output_dir):
shutil.rmtree(output_dir)
os.makedirs(output_dir)
print('Predicting on test images {} to: {}'.format(
args.images_paths, output_dir))
colors = get_colors()
images_pbar = tqdm(glob.glob(args.images_paths),
desc='Predicting (last tf call __ ms)',
unit='images')
tf_total_duration = 0.
img_total_duration = 0.
tf_count = 0.
img_count = 0.
for image_file in images_pbar:
image = scipy.misc.imresize(scipy.misc.imread(image_file),
image_shape)
segmented_image, tf_time_ms, img_time_ms = predict_image(
sess, model, image, colors)
if tf_count > 0:
tf_total_duration += tf_time_ms
tf_count += 1
tf_avg_ms = int(tf_total_duration /
(tf_count - 1 if tf_count > 1 else 1))
if img_count > 0:
img_total_duration += img_time_ms
img_count += 1
img_avg_ms = int(img_total_duration /
(img_count - 1 if img_count > 1 else 1))
images_pbar.set_description(
'Predicting (last tf call {} ms, avg tf {} ms, last img {} ms, avg {} ms)'
.format(tf_time_ms, tf_avg_ms, img_time_ms, img_avg_ms))
# tf timings:
# mac cpu inference is 670ms on trained but unoptimized graph. tf 1.3
# ubuntu cpu inference is 1360ms on pip tf-gpu 1.3.
# ubuntu cpu inference is 560ms on custom built tf-gpu 1.3 (cuda+xla).
# ubuntu gpu inference is 18ms on custom built tf-gpu 1.3 (cuda+xla). 580ms total per image. 1.7 fps
# quantize_weights increases inference to 50ms
# final performance on ubuntu/1080ti with ssd, including time to load/save is 3 fps
scipy.misc.imsave(
os.path.join(output_dir, os.path.basename(image_file)),
segmented_image)
def freeze_graph(args):
# based on https://blog.metaflow.fr/tensorflow-how-to-freeze-a-model-and-serve-it-with-a-python-api-d4f3596b3adc
if args.ckpt_dir is None:
print("for freezing need --ckpt_dir")
return
if args.frozen_model_dir is None:
print("for freezing need --frozen_model_dir")
return
checkpoint = tf.train.get_checkpoint_state(args.ckpt_dir)
input_checkpoint = checkpoint.model_checkpoint_path
print("freezing from {}".format(input_checkpoint))
saver = tf.train.import_meta_graph(input_checkpoint + '.meta',
clear_devices=True)
graph = tf.get_default_graph()
input_graph_def = graph.as_graph_def()
print("{} ops in the input graph".format(len(input_graph_def.node)))
output_node_names = "predictions/prediction_class"
# freeze graph
with tf.Session() as sess:
saver.restore(sess, input_checkpoint)
# use a built-in TF helper to export variables to constants
output_graph_def = tf_graph_util.convert_variables_to_constants(
sess, input_graph_def, output_node_names.split(","))
print("{} ops in the frozen graph".format(len(output_graph_def.node)))
if os.path.exists(args.frozen_model_dir):
shutil.rmtree(args.frozen_model_dir)
# save model in same format as usual
print('saving frozen model as saved_model to {}'.format(
args.frozen_model_dir))
model = fcn16vgg16.FCN16_VGG16(define_graph=False)
tf.reset_default_graph()
tf.import_graph_def(output_graph_def, name='')
with tf.Session() as sess:
model.save_model(sess, args.frozen_model_dir)
print('saving frozen model as graph.pb (for transforms) to {}'.format(
args.frozen_model_dir))
with tf.gfile.GFile(args.frozen_model_dir + '/graph.pb', "wb") as f:
f.write(output_graph_def.SerializeToString())
def train(args, image_shape):
config = session_config(args)
# extract pre-trained VGG weights
with tf.Session(config=config) as sess:
var_values = load_trained_vgg_vars(sess)
tf.reset_default_graph()
with tf.Session(config=config) as sess:
# define our FCN
num_classes = len(cityscape_labels.labels)
model = fcn16vgg16.FCN16_VGG16(num_classes)
# variables initialization
sess.run(tf.global_variables_initializer())
model.restore_variables(sess, var_values)
# Create batch generator
# TODO: Augment Images for better results
# https://datascience.stackexchange.com/questions/5224/how-to-prepare-augment-images-for-neural-network
train_batches_fn, num_samples = get_train_batch_generator_cityscapes(
args.images_paths, args.labels_paths, image_shape)
time_str = time.strftime("%Y%m%d_%H%M%S")
run_name = "/{}_ep{}_b{}_lr{:.6f}_kp{}".format(time_str, args.epochs,
args.batch_size,
args.learning_rate,
args.keep_prob)
start_time = time.time()
final_loss = model.train(sess, args.epochs, args.batch_size,
train_batches_fn, num_samples, args.keep_prob,
args.learning_rate, args.ckpt_dir,
args.summary_dir + run_name)
# Make folder for current run
output_dir = os.path.join(args.runs_dir, time_str)
if os.path.exists(output_dir):
shutil.rmtree(output_dir)
os.makedirs(output_dir)
# save training details to text file
with open(os.path.join(output_dir, "params.txt"), "w") as f:
f.write('keep_prob={}\n'.format(args.keep_prob))
f.write('images_paths={}\n'.format(args.images_paths))
f.write('num_samples={}\n'.format(num_samples))
f.write('batch_size={}\n'.format(args.batch_size))
f.write('epochs={}\n'.format(args.epochs))
f.write('gpu={}\n'.format(args.gpu))
f.write('gpu_mem={}\n'.format(args.gpu_mem))
f.write('learning_rate={}\n'.format(args.learning_rate))
f.write('final_loss={}\n'.format(final_loss))
duration = time.time() - start_time
f.write('total_time_hrs={}\n'.format(duration / 3600))
# save model
""" save trained model using SavedModelBuilder """
if args.model_dir is None:
model_dir = os.path.join(output_dir, 'model')
else:
model_dir = args.model_dir
print('saving trained model to {}'.format(model_dir))
model.save_model(sess, model_dir)