Spaces:
Running
Running
File size: 3,224 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 95 96 97 98 99 100 |
# Copyright 2018 Google LLC
#
# 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
#
# https://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.
# =============================================================================
"""An example script to generate a tfrecord file from a folder containing the
renderings.
Example usage:
python gen_tfrecords.py --input=FOLDER --output=output.tfrecord
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import numpy as np
import os
from scipy import misc
import tensorflow as tf
FLAGS = tf.app.flags.FLAGS
tf.app.flags.DEFINE_string("input", "", "Input folder containing images")
tf.app.flags.DEFINE_string("output", "", "Output tfrecord.")
def get_matrix(lines):
return np.array([[float(y) for y in x.strip().split(" ")] for x in lines])
def read_model_view_matrices(filename):
with open(filename, "r") as f:
lines = f.readlines()
return get_matrix(lines[:4]), get_matrix(lines[4:])
def bytes_feature(values):
return tf.train.Feature(bytes_list=tf.train.BytesList(value=[values]))
def generate():
with tf.python_io.TFRecordWriter(FLAGS.output) as tfrecord_writer:
with tf.Graph().as_default():
im0 = tf.placeholder(dtype=tf.uint8)
im1 = tf.placeholder(dtype=tf.uint8)
encoded0 = tf.image.encode_png(im0)
encoded1 = tf.image.encode_png(im1)
with tf.Session() as sess:
count = 0
indir = FLAGS.input + "/"
while tf.gfile.Exists(indir + "%06d.txt" % count):
print("saving %06d" % count)
image0 = misc.imread(indir + "%06d.png" % (count * 2))
image1 = misc.imread(indir + "%06d.png" % (count * 2 + 1))
mat0, mat1 = read_model_view_matrices(indir + "%06d.txt" % count)
mati0 = np.linalg.inv(mat0).flatten()
mati1 = np.linalg.inv(mat1).flatten()
mat0 = mat0.flatten()
mat1 = mat1.flatten()
st0, st1 = sess.run([encoded0, encoded1],
feed_dict={im0: image0, im1: image1})
example = tf.train.Example(features=tf.train.Features(feature={
'img0': bytes_feature(st0),
'img1': bytes_feature(st1),
'mv0': tf.train.Feature(
float_list=tf.train.FloatList(value=mat0)),
'mvi0': tf.train.Feature(
float_list=tf.train.FloatList(value=mati0)),
'mv1': tf.train.Feature(
float_list=tf.train.FloatList(value=mat1)),
'mvi1': tf.train.Feature(
float_list=tf.train.FloatList(value=mati1)),
}))
tfrecord_writer.write(example.SerializeToString())
count += 1
def main(argv):
del argv
generate()
if __name__ == "__main__":
tf.app.run()
|