Skip to content

GitLab

Commit bf65145d authored by Derek Chow's avatar Derek Chow
Browse files

Fix a bug in create_pet_tf_record.py related to mask creation.

parent 0ae575f2
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 2 deletions
research/object_detection/dataset_tools/create_pet_tf_record.py
View file @ bf65145d
...@@ -160,8 +160,6 @@ def dict_to_tf_example(data, ...@@ -160,8 +160,6 @@ def dict_to_tf_example(data,
if not faces_only: if not faces_only:
mask_remapped = mask_np != 2 mask_remapped = mask_np != 2
masks.append(mask_remapped) masks.append(mask_remapped)
mask_stack = np.stack(masks).astype(np.float32)
masks_flattened = np.reshape(mask_stack, [-1])
feature_dict = { feature_dict = {
'image/height': dataset_util.int64_feature(height), 'image/height': dataset_util.int64_feature(height),
...@@ -184,8 +182,11 @@ def dict_to_tf_example(data, ...@@ -184,8 +182,11 @@ def dict_to_tf_example(data,
'image/object/view': dataset_util.bytes_list_feature(poses), 'image/object/view': dataset_util.bytes_list_feature(poses),
} }
if not faces_only: if not faces_only:
mask_stack = np.stack(masks).astype(np.float32)
masks_flattened = np.reshape(mask_stack, [-1])
feature_dict['image/object/mask'] = ( feature_dict['image/object/mask'] = (
dataset_util.float_list_feature(masks_flattened.tolist())) dataset_util.float_list_feature(masks_flattened.tolist()))
example = tf.train.Example(features=tf.train.Features(feature=feature_dict)) example = tf.train.Example(features=tf.train.Features(feature=feature_dict))
return example return example
......
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