Spaces:
Running
Running
File size: 3,020 Bytes
0b8359d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
# Copyright 2016 The TensorFlow Authors All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Generate examples of two objects moving in different directions."""
import random
import sys
import numpy as np
from six.moves import xrange
import tensorflow as tf
tf.flags.DEFINE_string('out_file', '',
'Output file for the tfrecords.')
def _add_object(obj_type, image, image2, xpos, ypos):
"""Add a moving obj to two consecutive images."""
obj_size = random.randint(8, 10)
channel = random.randint(0, 2)
move = random.randint(6, 10)
obj = np.zeros([obj_size, obj_size, 3])
if obj_type == 'rectangle':
xpos2 = xpos + move
ypos2 = ypos
for i in xrange(obj_size):
obj[i, 0:i+1, channel] = [1.0 for _ in xrange(i+1)]
elif obj_type == 'square':
xpos2 = xpos
ypos2 = ypos + move
obj[:, :, channel] = 1.0
for x in xrange(obj_size):
for y in xrange(obj_size):
if obj[x, y, channel] == 1.0:
image[xpos+x, ypos+y, channel] = 1.0
image2[xpos2+x, ypos2+y, channel] = 1.0
def _images_to_example(image, image2):
"""Convert two consecutive images to SequenceExample."""
example = tf.SequenceExample()
feature_list = example.feature_lists.feature_list['moving_objs']
feature = feature_list.feature.add()
feature.float_list.value.extend(np.reshape(image, [-1]).tolist())
feature = feature_list.feature.add()
feature.float_list.value.extend(np.reshape(image2, [-1]).tolist())
return example
def generate_input():
"""Generate tfrecords."""
writer = tf.python_io.TFRecordWriter(tf.flags.FLAGS.out_file)
writer2 = tf.python_io.TFRecordWriter(tf.flags.FLAGS.out_file + '_test')
examples = []
for xpos in xrange(0, 40, 3):
for ypos in xrange(0, 40, 3):
for xpos2 in xrange(0, 40, 3):
for ypos2 in xrange(0, 40, 3):
image = np.zeros([64, 64, 3])
image2 = np.zeros([64, 64, 3])
_add_object('rectangle', image, image2, xpos, ypos)
_add_object('square', image, image2, xpos2, ypos2)
examples.append(_images_to_example(image, image2))
sys.stderr.write('Finish generating examples.\n')
random.shuffle(examples)
for count, ex in enumerate(examples):
if count % 10 == 0:
writer2.write(ex.SerializeToString())
else:
writer.write(ex.SerializeToString())
def main(_):
generate_input()
if __name__ == '__main__':
tf.app.run()
|