Tensorflow: working tf.while_loop does not work as part of Dataset API input pipeline - tensorflow

Question

My problem is an image keypoint recognition task on images of snails. I have found that although there are many prewritten image augmentation functions for classification tasks (such as Keras' ImageDataGenerator), there are none that I can find suitable for this problem, which requires changes to the output keypoints to match the random transformations of the image. Hence I am writing my own to be mapped onto the dataset as it is read from TFRecord.
The logic I am using involves a while loop which continues to generate random transformations (rotation + shift + zoom etc.) and apply them the real keypoints until it finds a set of transformations where the keypoints fit into the image. This is to avoid transformations that leave part of the snail outside the image. It would then apply those same transformations to the image and return them.
My problem is that, while I have successfully got this augmentation function to work on a single test set of keypoints, when I use the same function as part of my input pipeline, it does not work, throwing the following error: 'Merge can not have more than one valid input' (full trace included at end). I have not been able to find an explanation anywhere.
# Defining cond argument to while loop.'ph' are placeholders to match numbers of arguments for tf.while_loop
def not_fit_in_image(landmarks, ph2, ph3, ph4, ph5, ph6):
# tf logical operators to find if landmarks fit in image
return landmarks_not_fit_in_image
def augmentation_function(image, original_landmarks):
def body(ph1, ph2, ph3, ph4, ph5, ph6):
shift = tf.random_uniform([1, 2], -shift_max, shift_max, tf.float32)
landmarks = original_landmarks + shift
# More random transformations generated and applied
return landmarks, rotation, shift, zoom, y_over_x_proportion_change, shear
# placeholders to match number of arguments
ph_a = tf.constant(0, dtype=tf.float32)
landmarks, rotation, shift, zoom, y_over_x_proportion_change, shear = tf.while_loop(not_fit_in_image, body, [original_landmarks, ph_a, ph_b, ph_a, ph_a, ph_a])
# In future, would now apply these same transformations to image.
return image, landmarks
# Setting up input data pipeline using Dataset API
train = tf.data.TFRecordDataset(train_data_tfrecords).map(parse_function)
train = train.map(augmentation_function) # Using the above augmentation function
train = train.repeat().shuffle(buffer_size).batch(batch_size)
# ... Set up handle, iterator, init ops ... all works ...
with tf.Session() as sess:
train_handle = sess.run(train_iterator.string_handle())
sess.run(train_init_op)
train_images, train_landmarks = sess.run(next_batch, feed_dict={handle: train_handle})
The following error occurs:
2017-11-10 13:08:14.449612: W C:\tf_jenkins\home\workspace\rel-win\M\windows-gpu\PY\35\tensorflow\core\framework\op_kernel.cc:1192] Internal: Merge can not have more than one valid input.
[[Node: while/Merge_5 = Merge[N=2, T=DT_FLOAT](while/Enter_5, while/NextIteration_5)]]
Traceback (most recent call last):
File "C:\Users\hanne\Anaconda3\envs\tensorflow-gpu\lib\site-packages\tensorflow\python\client\session.py", line 1323, in _do_call
return fn(*args)
File "C:\Users\hanne\Anaconda3\envs\tensorflow-gpu\lib\site-packages\tensorflow\python\client\session.py", line 1302, in _run_fn
status, run_metadata)
File "C:\Users\hanne\Anaconda3\envs\tensorflow-gpu\lib\site-packages\tensorflow\python\framework\errors_impl.py", line 473, in __exit__
c_api.TF_GetCode(self.status.status))
tensorflow.python.framework.errors_impl.InternalError: Merge can not have more than one valid input.
[[Node: while/Merge_5 = Merge[N=2, T=DT_FLOAT](while/Enter_5, while/NextIteration_5)]]
[[Node: IteratorGetNext = IteratorGetNext[output_shapes=[[?,384,384], [?,15,2]], output_types=[DT_FLOAT, DT_FLOAT], _device="/job:localhost/replica:0/task:0/device:CPU:0"](IteratorFromStringHandle)]]
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
File "C:/Users/hanne/Documents/Tensorflow Projects/Snails/random_rotations_working_while_loop_experiments.py", line 143, in <module>
train_images, train_landmarks = sess.run(next_batch, feed_dict={handle: train_handle})
File "C:\Users\hanne\Anaconda3\envs\tensorflow-gpu\lib\site-packages\tensorflow\python\client\session.py", line 889, in run
run_metadata_ptr)
File "C:\Users\hanne\Anaconda3\envs\tensorflow-gpu\lib\site-packages\tensorflow\python\client\session.py", line 1120, in _run
feed_dict_tensor, options, run_metadata)
File "C:\Users\hanne\Anaconda3\envs\tensorflow-gpu\lib\site-packages\tensorflow\python\client\session.py", line 1317, in _do_run
options, run_metadata)
File "C:\Users\hanne\Anaconda3\envs\tensorflow-gpu\lib\site-packages\tensorflow\python\client\session.py", line 1336, in _do_call
raise type(e)(node_def, op, message)
tensorflow.python.framework.errors_impl.InternalError: Merge can not have more than one valid input.
[[Node: while/Merge_5 = Merge[N=2, T=DT_FLOAT](while/Enter_5, while/NextIteration_5)]]
[[Node: IteratorGetNext = IteratorGetNext[output_shapes=[[?,384,384], [?,15,2]], output_types=[DT_FLOAT, DT_FLOAT], _device="/job:localhost/replica:0/task:0/device:CPU:0"](IteratorFromStringHandle)]]
This is my first time asking a question on stack overflow, so any comments about how to write better questions are also very welcome! I have tried to strip down the code above as much as I can for brevity and it is hence minimal but NOT complete or verifiable - let me know if I should include more code.
EDIT
I was able to figure out what was wrong! tf.while_loop acts like a python while loop, checking the condition before each run of 'body', which includes THE VERY FIRST RUN. The argument 'loop_vars' takes the variables for this first check. I had entered placeholder values of the wrong format to 'loop_vars', which caused the error above. A good way around this, which worked for me, is to enter the result of a first run of 'body' to the loop_vars variable, as this is assured of being of the right form.