forked from gaylejj/tensorflow-cyclegan
-
Notifications
You must be signed in to change notification settings - Fork 0
/
image.py
67 lines (59 loc) · 2.92 KB
/
image.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
import tensorflow as tf
class Images():
def __init__(self, tfrecords_file, image_size=256, batch_size=1, num_threads=2, shuffle=True, pipeline_tweaks=None, name=''):
self.tfrecords_file = tfrecords_file
self.image_size = image_size
self.batch_size = batch_size
self.num_threads = num_threads
self.shuffle = shuffle
self.mods = pipeline_tweaks
self.name = name
def feed(self):
with tf.name_scope(self.name):
filename_queue = tf.train.string_input_producer([self.tfrecords_file])
reader = tf.TFRecordReader()
_, serialized_example = reader.read(filename_queue)
features = tf.parse_single_example(
serialized_example,
features={
'image/file_name': tf.FixedLenFeature([], tf.string),
'image/encoded_image': tf.FixedLenFeature([], tf.string),
})
image_buffer = features['image/encoded_image']
image = tf.image.decode_jpeg(image_buffer, channels=3)
image = self.preprocess(image)
if self.shuffle:
images = tf.train.shuffle_batch(
[image], batch_size=self.batch_size, num_threads=self.num_threads,
capacity=100 + 3 * self.batch_size,
min_after_dequeue=100
)
else:
images = tf.train.batch(
[image], batch_size=self.batch_size, num_threads=self.num_threads,
capacity=100 + 3 * self.batch_size
)
tf.summary.image('_input', images)
return images
def preprocess(self, image):
if self.mods and self.mods['random_flip']:
image = tf.image.random_flip_left_right(image)
if self.mods and self.mods['random_saturation']:
image = tf.image.random_saturation(image, .95, 1.05)
if self.mods and self.mods['random_brightness']:
image = tf.image.random_brightness(image, .05)
if self.mods and self.mods['random_contrast']:
image = tf.image.random_contrast(image, .95, 1.05)
if self.mods and self.mods['crop_size'] > 0:
crop_size = self.mods['crop_size']
wiggle = 8
off_x, off_y = 25 - wiggle, 60 - wiggle
crop_size_plus = crop_size + 2 * wiggle
image = tf.image.resize_image_with_crop_or_pad(image, off_y + crop_size_plus, off_x + crop_size_plus)
image = tf.image.crop_to_bounding_box(image, off_y, off_x, crop_size_plus, crop_size_plus)
image = tf.random_crop(image, [crop_size, crop_size, 3])
image = tf.image.resize_images(image, size=(self.image_size, self.image_size))
image = tf.image.convert_image_dtype(image, dtype=tf.float32)
image = (image / 127.5) - 1.
image.set_shape([self.image_size, self.image_size, 3])
return image