-
Notifications
You must be signed in to change notification settings - Fork 33
/
decoder.py
executable file
·127 lines (101 loc) · 4.67 KB
/
decoder.py
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
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
#!/usr/bin/python
#
# 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.
# ==============================================================================
r"""Neural Network Image Compression Decoder.
Decompress an image from the numpy's npz format generated by the encoder.
Example usage:
python decoder.py --input_codes=output_codes.pkl --iteration=15 \
--output_directory=/tmp/compression_output/ --model=residual_gru.pb
"""
import io
import os
import numpy as np
import tensorflow as tf
tf.flags.DEFINE_string('input_codes', None, 'Location of binary code file.')
tf.flags.DEFINE_integer('iteration', -1, 'The max quality level of '
'the images to output. Use -1 to infer from loaded '
' codes.')
tf.flags.DEFINE_string('output_directory', None, 'Directory to save decoded '
'images.')
tf.flags.DEFINE_string('model', None, 'Location of compression model.')
FLAGS = tf.flags.FLAGS
def get_input_tensor_names():
name_list = ['GruBinarizer/SignBinarizer/Sign:0']
for i in range(1, 16):
name_list.append('GruBinarizer/SignBinarizer/Sign_{}:0'.format(i))
return name_list
def get_output_tensor_names():
return ['loop_{0:02d}/add:0'.format(i) for i in range(0, 16)]
def main(_):
if (FLAGS.input_codes is None or FLAGS.output_directory is None or
FLAGS.model is None):
print('\nUsage: python decoder.py --input_codes=output_codes.pkl '
'--iteration=15 --output_directory=/tmp/compression_output/ '
'--model=residual_gru.pb\n\n')
return
if FLAGS.iteration < -1 or FLAGS.iteration > 15:
print('\n--iteration must be between 0 and 15 inclusive, or -1 to infer '
'from file.\n')
return
iteration = FLAGS.iteration
if not tf.gfile.Exists(FLAGS.output_directory):
tf.gfile.MkDir(FLAGS.output_directory)
if not tf.gfile.Exists(FLAGS.input_codes):
print('\nInput codes not found.\n')
return
contents = ''
with tf.gfile.FastGFile(FLAGS.input_codes, 'r') as code_file:
contents = code_file.read()
loaded_codes = np.load(io.BytesIO(contents))
assert ['codes', 'shape'] not in loaded_codes.files
loaded_shape = loaded_codes['shape']
loaded_array = loaded_codes['codes']
# Unpack and recover code shapes.
unpacked_codes = np.reshape(np.unpackbits(loaded_array)
[:np.prod(loaded_shape)],
loaded_shape)
numpy_int_codes = np.split(unpacked_codes, len(unpacked_codes))
if iteration == -1:
iteration = len(unpacked_codes) - 1
# Convert back to float and recover scale.
numpy_codes = [np.squeeze(x.astype(np.float32), 0) * 2 - 1 for x in
numpy_int_codes]
with tf.Graph().as_default() as graph:
# Load the inference model for decoding.
with tf.gfile.FastGFile(FLAGS.model, 'rb') as model_file:
graph_def = tf.GraphDef()
graph_def.ParseFromString(model_file.read())
_ = tf.import_graph_def(graph_def, name='')
# For encoding the tensors into PNGs.
input_image = tf.placeholder(tf.uint8)
encoded_image = tf.image.encode_png(input_image)
input_tensors = [graph.get_tensor_by_name(name) for name in
get_input_tensor_names()][0:iteration 1]
outputs = [graph.get_tensor_by_name(name) for name in
get_output_tensor_names()][0:iteration 1]
feed_dict = {key: value for (key, value) in zip(input_tensors,
numpy_codes)}
with tf.Session(graph=graph) as sess:
results = sess.run(outputs, feed_dict=feed_dict)
for index, result in enumerate(results):
img = np.uint8(np.clip(result 0.5, 0, 255))
img = img.squeeze()
png_img = sess.run(encoded_image, feed_dict={input_image: img})
with tf.gfile.FastGFile(os.path.join(FLAGS.output_directory,
'image_{0:02d}.png'.format(index)),
'w') as output_image:
output_image.write(png_img)
if __name__ == '__main__':
tf.app.run()