Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
124 changes: 63 additions & 61 deletions build_data.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,114 +3,116 @@
import os

try:
from os import scandir
from os import scandir
except ImportError:
# Python 2 polyfill module
from scandir import scandir

# Python 2 polyfill module
from scandir import scandir

FLAGS = tf.flags.FLAGS
FLAGS = tf.compat.v1.flags.FLAGS

tf.flags.DEFINE_string('X_input_dir', 'data/apple2orange/trainA',
tf.compat.v1.flags.DEFINE_string('X_input_dir', 'data/apple2orange/trainA',
'X input directory, default: data/apple2orange/trainA')
tf.flags.DEFINE_string('Y_input_dir', 'data/apple2orange/trainB',
tf.compat.v1.flags.DEFINE_string('Y_input_dir', 'data/apple2orange/trainB',
'Y input directory, default: data/apple2orange/trainB')
tf.flags.DEFINE_string('X_output_file', 'data/tfrecords/apple.tfrecords',
tf.compat.v1.flags.DEFINE_string('X_output_file', 'data/tfrecords/apple.tfrecords',
'X output tfrecords file, default: data/tfrecords/apple.tfrecords')
tf.flags.DEFINE_string('Y_output_file', 'data/tfrecords/orange.tfrecords',
tf.compat.v1.flags.DEFINE_string('Y_output_file', 'data/tfrecords/orange.tfrecords',
'Y output tfrecords file, default: data/tfrecords/orange.tfrecords')


def data_reader(input_dir, shuffle=True):
"""Read images from input_dir then shuffle them
"""Read images from input_dir then shuffle them
Args:
input_dir: string, path of input dir, e.g., /path/to/dir
Returns:
file_paths: list of strings
"""
file_paths = []
file_paths = []

for img_file in scandir(input_dir):
if img_file.name.endswith('.jpg') and img_file.is_file():
file_paths.append(img_file.path)
for img_file in scandir(input_dir):
if img_file.name.endswith('.jpg') and img_file.is_file():
file_paths.append(img_file.path)

if shuffle:
# Shuffle the ordering of all image files in order to guarantee
# random ordering of the images with respect to label in the
# saved TFRecord files. Make the randomization repeatable.
shuffled_index = list(range(len(file_paths)))
random.seed(12345)
random.shuffle(shuffled_index)
if shuffle:
# Shuffle the ordering of all image files in order to guarantee
# random ordering of the images with respect to label in the
# saved TFRecord files. Make the randomization repeatable.
shuffled_index = list(range(len(file_paths)))
random.seed(12345)
random.shuffle(shuffled_index)

file_paths = [file_paths[i] for i in shuffled_index]
file_paths = [file_paths[i] for i in shuffled_index]

return file_paths
return file_paths


def _int64_feature(value):
"""Wrapper for inserting int64 features into Example proto."""
if not isinstance(value, list):
value = [value]
return tf.train.Feature(int64_list=tf.train.Int64List(value=value))
"""Wrapper for inserting int64 features into Example proto."""
if not isinstance(value, list):
value = [value]
return tf.train.Feature(int64_list=tf.train.Int64List(value=value))


def _bytes_feature(value):
"""Wrapper for inserting bytes features into Example proto."""
return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
"""Wrapper for inserting bytes features into Example proto."""
return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))


def _convert_to_example(file_path, image_buffer):
"""Build an Example proto for an example.
"""Build an Example proto for an example.
Args:
file_path: string, path to an image file, e.g., '/path/to/example.JPG'
image_buffer: string, JPEG encoding of RGB image
Returns:
Example proto
"""
file_name = file_path.split('/')[-1]
file_name = file_path.split('/')[-1]

example = tf.train.Example(features=tf.train.Features(feature={
'image/file_name': _bytes_feature(tf.compat.as_bytes(os.path.basename(file_name))),
'image/encoded_image': _bytes_feature((image_buffer))
example = tf.train.Example(features=tf.train.Features(feature={
'image/file_name': _bytes_feature(tf.compat.as_bytes(os.path.basename(file_name))),
'image/encoded_image': _bytes_feature((image_buffer))
}))
return example
return example


def data_writer(input_dir, output_file):
"""Write data to tfrecords
"""Write data to tfrecords
"""
file_paths = data_reader(input_dir)
file_paths = data_reader(input_dir)

# create tfrecords dir if not exists
output_dir = os.path.dirname(output_file)
try:
os.makedirs(output_dir)
except os.error as e:
pass

# create tfrecords dir if not exists
output_dir = os.path.dirname(output_file)
try:
os.makedirs(output_dir)
except os.error as e:
pass
images_num = len(file_paths)

images_num = len(file_paths)
# dump to tfrecords file
writer = tf.compat.v1.python_io.TFRecordWriter(output_file)

# dump to tfrecords file
writer = tf.python_io.TFRecordWriter(output_file)
for i in range(len(file_paths)):
file_path = file_paths[i]

for i in range(len(file_paths)):
file_path = file_paths[i]
with tf.compat.v1.gfile.FastGFile(file_path, 'rb') as f:
image_data = f.read()

with tf.gfile.FastGFile(file_path, 'rb') as f:
image_data = f.read()
example = _convert_to_example(file_path, image_data)
writer.write(example.SerializeToString())

example = _convert_to_example(file_path, image_data)
writer.write(example.SerializeToString())
if i % 500 == 0:
print("Processed {}/{}.".format(i, images_num))
print("Done.")
writer.close()

if i % 500 == 0:
print("Processed {}/{}.".format(i, images_num))
print("Done.")
writer.close()

def main(unused_argv):
print("Convert X data to tfrecords...")
data_writer(FLAGS.X_input_dir, FLAGS.X_output_file)
print("Convert Y data to tfrecords...")
data_writer(FLAGS.Y_input_dir, FLAGS.Y_output_file)
print("Convert X data to tfrecords...")
data_writer(FLAGS.X_input_dir, FLAGS.X_output_file)
print("Convert Y data to tfrecords...")
data_writer(FLAGS.Y_input_dir, FLAGS.Y_output_file)


if __name__ == '__main__':
tf.app.run()
tf.compat.v1.app.run()
51 changes: 26 additions & 25 deletions discriminator.py
Original file line number Diff line number Diff line change
@@ -1,39 +1,40 @@
import tensorflow as tf
import ops


class Discriminator:
def __init__(self, name, is_training, norm='instance', use_sigmoid=False):
self.name = name
self.is_training = is_training
self.norm = norm
self.reuse = False
self.use_sigmoid = use_sigmoid
def __init__(self, name, is_training, norm='instance', use_sigmoid=False):
self.name = name
self.is_training = is_training
self.norm = norm
self.reuse = False
self.use_sigmoid = use_sigmoid

def __call__(self, input):
"""
def __call__(self, input):
"""
Args:
input: batch_size x image_size x image_size x 3
Returns:
output: 4D tensor batch_size x out_size x out_size x 1 (default 1x5x5x1)
filled with 0.9 if real, 0.0 if fake
"""
with tf.variable_scope(self.name):
# convolution layers
C64 = ops.Ck(input, 64, reuse=self.reuse, norm=None,
is_training=self.is_training, name='C64') # (?, w/2, h/2, 64)
C128 = ops.Ck(C64, 128, reuse=self.reuse, norm=self.norm,
is_training=self.is_training, name='C128') # (?, w/4, h/4, 128)
C256 = ops.Ck(C128, 256, reuse=self.reuse, norm=self.norm,
is_training=self.is_training, name='C256') # (?, w/8, h/8, 256)
C512 = ops.Ck(C256, 512,reuse=self.reuse, norm=self.norm,
is_training=self.is_training, name='C512') # (?, w/16, h/16, 512)
with tf.compat.v1.variable_scope(self.name):
# convolution layers
C64 = ops.Ck(input, 64, reuse=self.reuse, norm=None,
is_training=self.is_training, name='C64') # (?, w/2, h/2, 64)
C128 = ops.Ck(C64, 128, reuse=self.reuse, norm=self.norm,
is_training=self.is_training, name='C128') # (?, w/4, h/4, 128)
C256 = ops.Ck(C128, 256, reuse=self.reuse, norm=self.norm,
is_training=self.is_training, name='C256') # (?, w/8, h/8, 256)
C512 = ops.Ck(C256, 512, reuse=self.reuse, norm=self.norm,
is_training=self.is_training, name='C512') # (?, w/16, h/16, 512)

# apply a convolution to produce a 1 dimensional output (1 channel?)
# use_sigmoid = False if use_lsgan = True
output = ops.last_conv(C512, reuse=self.reuse,
use_sigmoid=self.use_sigmoid, name='output') # (?, w/16, h/16, 1)
# apply a convolution to produce a 1 dimensional output (1 channel?)
# use_sigmoid = False if use_lsgan = True
output = ops.last_conv(C512, reuse=self.reuse,
use_sigmoid=self.use_sigmoid, name='output') # (?, w/16, h/16, 1)

self.reuse = True
self.variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope=self.name)
self.reuse = True
self.variables = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.TRAINABLE_VARIABLES, scope=self.name)

return output
return output
65 changes: 34 additions & 31 deletions export_graph.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,48 +13,51 @@
from model import CycleGAN
import utils

FLAGS = tf.flags.FLAGS
FLAGS = tf.compat.v1.flags.FLAGS

tf.flags.DEFINE_string('checkpoint_dir', '', 'checkpoints directory path')
tf.flags.DEFINE_string('XtoY_model', 'apple2orange.pb', 'XtoY model name, default: apple2orange.pb')
tf.flags.DEFINE_string('YtoX_model', 'orange2apple.pb', 'YtoX model name, default: orange2apple.pb')
tf.flags.DEFINE_integer('image_size', '256', 'image size, default: 256')
tf.flags.DEFINE_integer('ngf', 64,
tf.compat.v1.flags.DEFINE_string('checkpoint_dir', '', 'checkpoints directory path')
tf.compat.v1.flags.DEFINE_string('XtoY_model', 'apple2orange.pb', 'XtoY model name, default: apple2orange.pb')
tf.compat.v1.flags.DEFINE_string('YtoX_model', 'orange2apple.pb', 'YtoX model name, default: orange2apple.pb')
tf.compat.v1.flags.DEFINE_integer('image_size', '256', 'image size, default: 256')
tf.compat.v1.flags.DEFINE_integer('ngf', 64,
'number of gen filters in first conv layer, default: 64')
tf.flags.DEFINE_string('norm', 'instance',
tf.compat.v1.flags.DEFINE_string('norm', 'instance',
'[instance, batch] use instance norm or batch norm, default: instance')


def export_graph(model_name, XtoY=True):
graph = tf.Graph()
graph = tf.Graph()

with graph.as_default():
cycle_gan = CycleGAN(ngf=FLAGS.ngf, norm=FLAGS.norm, image_size=FLAGS.image_size)

with graph.as_default():
cycle_gan = CycleGAN(ngf=FLAGS.ngf, norm=FLAGS.norm, image_size=FLAGS.image_size)
input_image = tf.compat.v1.placeholder(tf.float32, shape=[FLAGS.image_size, FLAGS.image_size, 3], name='input_image')
cycle_gan.model()
if XtoY:
output_image = cycle_gan.G.sample(tf.expand_dims(input_image, 0))
else:
output_image = cycle_gan.F.sample(tf.expand_dims(input_image, 0))

input_image = tf.placeholder(tf.float32, shape=[FLAGS.image_size, FLAGS.image_size, 3], name='input_image')
cycle_gan.model()
if XtoY:
output_image = cycle_gan.G.sample(tf.expand_dims(input_image, 0))
else:
output_image = cycle_gan.F.sample(tf.expand_dims(input_image, 0))
output_image = tf.identity(output_image, name='output_image')
restore_saver = tf.compat.v1.train.Saver()
export_saver = tf.compat.v1.train.Saver()

output_image = tf.identity(output_image, name='output_image')
restore_saver = tf.train.Saver()
export_saver = tf.train.Saver()
with tf.compat.v1.Session(graph=graph) as sess:
sess.run(tf.compat.v1.global_variables_initializer())
latest_ckpt = tf.train.latest_checkpoint(FLAGS.checkpoint_dir)
restore_saver.restore(sess, latest_ckpt)
output_graph_def = tf.compat.v1.graph_util.convert_variables_to_constants(
sess, graph.as_graph_def(), [output_image.op.name])

with tf.Session(graph=graph) as sess:
sess.run(tf.global_variables_initializer())
latest_ckpt = tf.train.latest_checkpoint(FLAGS.checkpoint_dir)
restore_saver.restore(sess, latest_ckpt)
output_graph_def = tf.graph_util.convert_variables_to_constants(
sess, graph.as_graph_def(), [output_image.op.name])
tf.compat.v1.train.write_graph(output_graph_def, 'pretrained', model_name, as_text=False)

tf.train.write_graph(output_graph_def, 'pretrained', model_name, as_text=False)

def main(unused_argv):
print('Export XtoY model...')
export_graph(FLAGS.XtoY_model, XtoY=True)
print('Export YtoX model...')
export_graph(FLAGS.YtoX_model, XtoY=False)
print('Export XtoY model...')
export_graph(FLAGS.XtoY_model, XtoY=True)
print('Export YtoX model...')
export_graph(FLAGS.YtoX_model, XtoY=False)


if __name__ == '__main__':
tf.app.run()
tf.compat.v1.app.run()
Loading