Commit 2f3666ed authored by Duc Nguyen's avatar Duc Nguyen Committed by GitHub
Browse files

Change key of type 'tuple' to 'str'

dictionary having both 'tuple' and 'str' keys cannot be 'sorted'
parent 9b187704
...@@ -20,6 +20,8 @@ import tensorflow as tf ...@@ -20,6 +20,8 @@ import tensorflow as tf
from object_detection.core import prefetcher from object_detection.core import prefetcher
rt_shape_str = '_runtime_shapes'
class BatchQueue(object): class BatchQueue(object):
"""BatchQueue class. """BatchQueue class.
...@@ -81,7 +83,7 @@ class BatchQueue(object): ...@@ -81,7 +83,7 @@ class BatchQueue(object):
{key: tensor.get_shape() for key, tensor in tensor_dict.iteritems()}) {key: tensor.get_shape() for key, tensor in tensor_dict.iteritems()})
# Remember runtime shapes to unpad tensors after batching. # Remember runtime shapes to unpad tensors after batching.
runtime_shapes = collections.OrderedDict( runtime_shapes = collections.OrderedDict(
{(key + '_runtime_shapes'): tf.shape(tensor) {(key + rt_shape_str): tf.shape(tensor)
for key, tensor in tensor_dict.iteritems()}) for key, tensor in tensor_dict.iteritems()})
all_tensors = tensor_dict all_tensors = tensor_dict
all_tensors.update(runtime_shapes) all_tensors.update(runtime_shapes)
...@@ -112,8 +114,8 @@ class BatchQueue(object): ...@@ -112,8 +114,8 @@ class BatchQueue(object):
for key, batched_tensor in batched_tensors.iteritems(): for key, batched_tensor in batched_tensors.iteritems():
unbatched_tensor_list = tf.unstack(batched_tensor) unbatched_tensor_list = tf.unstack(batched_tensor)
for i, unbatched_tensor in enumerate(unbatched_tensor_list): for i, unbatched_tensor in enumerate(unbatched_tensor_list):
if '_runtime_shapes' in key: if rt_shape_str in key:
shapes[(key[:-15], i)] = unbatched_tensor shapes[(key[:-len(rt_shape_str)], i)] = unbatched_tensor
else: else:
tensors[(key, i)] = unbatched_tensor tensors[(key, i)] = unbatched_tensor
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment