Skip to content

GitLab

Commit 7907ba50 authored by Fan Yang's avatar Fan Yang Committed by A. Unique TensorFlower
Browse files

Internal change 

PiperOrigin-RevId: 372036675
parent 64f6527c
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 1 deletion
official/vision/beta/data/tfrecord_lib.py
View file @ 7907ba50
...@@ -63,12 +63,14 @@ def convert_to_feature(value, value_type=None): ...@@ -63,12 +63,14 @@ def convert_to_feature(value, value_type=None):
return tf.train.Feature(int64_list=tf.train.Int64List(value=[value])) return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))
elif value_type == 'int64_list': elif value_type == 'int64_list':
value = np.asarray(value).astype(np.int64).reshape(-1)
return tf.train.Feature(int64_list=tf.train.Int64List(value=value)) return tf.train.Feature(int64_list=tf.train.Int64List(value=value))
elif value_type == 'float': elif value_type == 'float':
return tf.train.Feature(float_list=tf.train.FloatList(value=[value])) return tf.train.Feature(float_list=tf.train.FloatList(value=[value]))
elif value_type == 'float_list': elif value_type == 'float_list':
value = np.asarray(value).astype(np.float32).reshape(-1)
return tf.train.Feature(float_list=tf.train.FloatList(value=value)) return tf.train.Feature(float_list=tf.train.FloatList(value=value))
elif value_type == 'bytes': elif value_type == 'bytes':
...@@ -172,4 +174,3 @@ def check_and_make_dir(directory): ...@@ -172,4 +174,3 @@ def check_and_make_dir(directory):
"""Creates the directory if it doesn't exist.""" """Creates the directory if it doesn't exist."""
if not tf.io.gfile.isdir(directory): if not tf.io.gfile.isdir(directory):
tf.io.gfile.makedirs(directory) tf.io.gfile.makedirs(directory)
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