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 = fcn8vgg16.FCN8_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 predict_files(args, image_shape):
tf.reset_default_graph()
with tf.Session(config=session_config(args)) as sess:
model = fcn8vgg16.FCN8_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 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 = fcn8vgg16.FCN8_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 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 = fcn8vgg16.FCN8_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)
def predict_files(args, image_shape):
tf.reset_default_graph()
with tf.Session(config=session_config(args)) as sess:
model = fcn8vgg16.FCN8_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))
segmented_image = cv2.resize(segmented_image, (2048, 1024),
interpolation=cv2.INTER_NEAREST)
scipy.misc.toimage(segmented_image, cmin=0, cmax=255,
mode='I').save(
os.path.join(output_dir,
os.path.basename(image_file)))
def predict_video(args, image_shape=None, force_reshape=True):
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
# Initializing tracker object to track movable vehicles
tracker_labels = [19, 52, 54, 55, 56, 57, 59, 60, 61, 62]
trackers = []
for tracker_label in tracker_labels:
trackers.append(object_tracking.Tracker(tracker_label))
# The actual frame processing is dealt with in this function
def process_frame(image):
segmented_image, tf_time_ms, img_time_ms = predict_image(
sess, model, image, colors, trackers)
return segmented_image
def process_frame_with_reshape(image):
if image_shape is not None:
# Apply intrisic camera calibration (undistort) --> not needed as I've already undistorted the test videos beforehand
# image = camera_calibration.undistort_image(image, mtx, dist) # adds about 14% of overhead
image = scipy.misc.imresize(image, image_shape)
segmented_image, tf_time_ms, img_time_ms = predict_image(
sess, model, image, colors, trackers)
return segmented_image
tf.reset_default_graph()
with tf.Session(config=session_config(args)) as sess:
model = fcn8vgg16.FCN8_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)
if force_reshape:
if args.video_start_second is None or args.video_end_second is None:
annotated_clip = input_clip.fl_image(
process_frame_with_reshape)
else:
annotated_clip = input_clip.fl_image(
process_frame_with_reshape).subclip(
args.video_start_second, args.video_end_second)
else:
if args.video_start_second is None or args.video_end_second is None:
annotated_clip = input_clip.fl_image(process_frame)
else:
annotated_clip = input_clip.fl_image(process_frame).subclip(
args.video_start_second, args.video_end_second)
annotated_clip.write_videofile(args.video_file_out, audio=False)
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 = fcn8vgg16.FCN8_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):
# DOCUMENTATION
# https://www.tensorflow.org/api_docs/python/tf/train/batch
# https://www.tensorflow.org/programmers_guide/datasets
# http://ischlag.github.io/2016/06/19/tensorflow-input-pipeline-example/
# https://kratzert.github.io/2017/06/15/example-of-tensorflows-new-input-pipeline.html
config = session_config(args)
# extract pre-trained VGG weights
with tf.Session(config=config) as sess:
var_values = load_vgg(sess)
tf.reset_default_graph()
# This portion of the code is enabled only if USE_TF_BATCHING is enabled
if USE_TF_BATCHING:
# Load paths and images in memory (they're only string, there's space)
images_path_pattern = args.images_paths
labels_path_pattern = args.labels_paths
image_paths = glob.glob(images_path_pattern)
if dataset == "cityscapes":
label_paths_tmp = {
re.sub('_gtFine_labelTrainIds', '_leftImg8bit',
os.path.basename(path)): path
for path in glob.glob(labels_path_pattern)
}
elif dataset == "mapillary":
label_paths_tmp = {
re.sub('_instance', '_image', os.path.basename(path)): path
for path in glob.glob(labels_path_pattern)
}
num_classes = len(dataset_labels.labels)
num_samples = len(image_paths)
assert len(image_paths) == len(label_paths_tmp)
print("num_classes={}".format(num_classes))
print("num_samples={}".format(num_samples))
label_paths = []
for image_file in image_paths:
label_paths.append(label_paths_tmp[os.path.basename(image_file)])
# Now we have a list of image_paths and one of label_paths
# Convert string into tensors
all_images = ops.convert_to_tensor(image_paths, dtype=dtypes.string)
all_labels = ops.convert_to_tensor(label_paths, dtype=dtypes.string)
# Create input queues
train_input_queue = tf.train.slice_input_producer(
[all_images, all_labels], shuffle=True)
# Process path and string tensor into an image and a label
train_image = tf.image.decode_png(tf.read_file(train_input_queue[0]),
channels=3)
train_label = tf.image.decode_png(tf.read_file(train_input_queue[1]),
channels=1)
# Define tensor shape
train_image.set_shape([256, 512, 3])
train_label.set_shape([256, 512, 1])
# One hot on the label
train_label_one_hot = tf.one_hot(tf.squeeze(train_label), num_classes)
# Collect batches of images before processing
train_image_batch, train_label_batch = tf.train.batch(
[train_image, train_label_one_hot],
batch_size=args.batch_size,
allow_smaller_final_batch=True,
num_threads=1)
def dummy_batch_fn(sess):
tmp_image_batch, tmp_label_batch = sess.run(
[train_image_batch, train_label_batch])
"""
print(debug)
for i in range(tmp_label_batch.shape[0]):
print("R={:6d} S={:6d} T={:6d}".format(np.count_nonzero(tmp_label_batch[i,:,:,13]),
np.count_nonzero(tmp_label_batch[i,:,:,27]),
np.count_nonzero(tmp_label_batch[i,:,:,30])))
"""
return np.array(tmp_image_batch), np.array(tmp_label_batch)
# TF session
with tf.Session(config=config) as sess:
# define our FCN
num_classes = len(dataset_labels.labels)
model = fcn8vgg16.FCN8_VGG16(num_classes)
# variables initialization
sess.run(tf.global_variables_initializer())
model.restore_variables(sess, var_values)
if USE_TF_BATCHING:
# initialize the queue threads to start to shovel data
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(coord=coord)
else:
train_batches_fn, num_samples = get_train_batch_generator(
args.images_paths, args.labels_paths, image_shape)
# Details for saving
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()
if USE_TF_BATCHING:
final_loss = model.train2(sess, args.epochs, args.batch_size,
dummy_batch_fn, num_samples,
args.keep_prob, args.learning_rate,
args.ckpt_dir,
args.summary_dir + run_name)
else:
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
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)
# stop our queue threads and properly close the session
coord.request_stop()
coord.join(threads)
sess.close()
