forked from vahidk/EffectiveTensorflow
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathutils.py
More file actions
70 lines (57 loc) · 1.63 KB
/
Copy pathutils.py
File metadata and controls
70 lines (57 loc) · 1.63 KB
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
"""Auxiliary functions."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import io
import numpy as np
import PIL
import multiprocessing as mp
import tensorflow as tf
def parallel_record_writer(iterator, create_example, path, num_threads=4):
"""Create a RecordIO file from data for efficient reading."""
def _queue(inputs):
for item in iterator:
inputs.put(item)
for _ in range(num_threads):
inputs.put(None)
def _map_fn(inputs, outputs):
while True:
item = inputs.get()
if item is None:
break
example = create_example(item)
outputs.put(example)
outputs.put(None)
# Read the inputs.
inputs = mp.Queue()
mp.Process(target=_queue, args=(inputs,)).start()
# Convert to tf.Example
outputs = mp.Queue()
for _ in range(num_threads):
mp.Process(target=_map_fn, args=(inputs, outputs)).start()
# Write the output to file.
writer = tf.python_io.TFRecordWriter(path)
counter = 0
while True:
example = outputs.get()
if example is None:
counter += 1
if counter == num_threads:
break
else:
continue
writer.write(example.SerializeToString())
writer.close()
def encode_image(data, format='png'):
"""Encodes a numpy array to string."""
im = PIL.Image.fromarray(data)
buf = io.BytesIO()
data = im.save(buf, format=format)
buf.seek(0)
return buf.getvalue()
def decode_image(data):
"""Decode the given image to a numpy array."""
buf = io.BytesIO(data)
im = PIL.Image.open(buf)
data = np.array(im.getdata()).reshape([im.height, im.width, -1])
return data