default="tfrecords/pascalvoc2012.tfrecords")
parser.add_argument('--train_dir',
help="where to log training",
default="train_log")
parser.add_argument('--batch_size',
help="batch size",
type=int,
default=10)
parser.add_argument('--num_epochs',
help="number of epochs.",
type=int,
default=50)
parser.add_argument('--lr', help="learning rate", type=float, default=1e-6)
args = parser.parse_args()
trn_images_batch, trn_segmentations_batch = input_pipeline(
args.train_record, args.batch_size, args.num_epochs)
deconvnet = DeconvNet(trn_images_batch,
trn_segmentations_batch,
use_cpu=False)
logits = deconvnet.logits
cross_entropy=tf.nn.sparse_softmax_cross_entropy_with_logits(logits = logits, \
labels = tf.cast(tf.reshape(deconvnet.y, [-1]), tf.int64), name='x_entropy')
loss_mean = tf.reduce_mean(cross_entropy, name='x_entropy_mean')
train_step = tf.train.AdamOptimizer(args.lr).minimize(loss_mean)
summary_op = tf.summary.merge_all() # v0.12
values = tf.reshape(x1, [-1])
return tf.scatter_nd(indices, values, tf.to_int64(top_shape))
if __name__ == '__main__':
# Using argparse over tf.FLAGS as I find they behave better in ipython
parser = argparse.ArgumentParser()
parser.add_argument('--train_record', help="training tfrecord file", default="tfrecords/pascalvoc2012.tfrecords")
parser.add_argument('--train_dir', help="where to log training", default="train_log")
parser.add_argument('--batch_size', help="batch size", type=int, default=10)
parser.add_argument('--num_epochs', help="number of epochs.", type=int, default=50)
parser.add_argument('--lr',help="learning rate",type=float, default=1e-6)
args = parser.parse_args()
trn_images_batch, trn_segmentations_batch = input_pipeline(
args.train_record,
args.batch_size,
args.num_epochs)
deconvnet = DeconvNet(trn_images_batch, trn_segmentations_batch, use_cpu=False)
logits=deconvnet.logits
cross_entropy=tf.nn.sparse_softmax_cross_entropy_with_logits(logits, \
tf.cast(tf.reshape(deconvnet.y, [-1]), tf.int64), name='x_entropy')
loss_mean=tf.reduce_mean(cross_entropy, name='x_entropy_mean')
train_step=tf.train.AdamOptimizer(args.lr).minimize(loss_mean)
summary_op = tf.summary.merge_all() # v0.12
from matplotlib import pyplot as plt
from utils import rgb2yuv,yuv2rgb,input_pipeline,concat_images
from net import color_net
if __name__ == "__main__":
filenames = glob.glob("demo/*")
with open("model/vgg16-20160129.tfmodel", mode='rb') as f:
fileContent = f.read()
graph_def = tf.GraphDef()
graph_def.ParseFromString(fileContent)
batch_size = 4
num_epochs = 1e+9
colorimage = input_pipeline(filenames, batch_size, num_epochs=num_epochs)
grayscale = tf.image.rgb_to_grayscale(colorimage)
grayscale_rgb = tf.image.grayscale_to_rgb(grayscale)
grayscale_yuv = rgb2yuv(grayscale_rgb)
grayscale = tf.concat([grayscale, grayscale, grayscale],3)
tf.import_graph_def(graph_def, input_map={"images": grayscale})
graph = tf.get_default_graph()
# Saver.
with tf.Session() as sess:
saver = tf.train.import_meta_graph('checkpoints/color_net_model200500.ckpt.meta')
saver.restore(sess, "checkpoints/color_net_model200500.ckpt")
def just_train():
logdir = os.path.join(board_logs_path,
datetime.datetime.now().strftime("%Y%m%d-%H%M%S"))
low_resolution_shape = (64, 64, 3)
high_resolution_shape = (256, 256, 3)
dataloader = iter(
input_pipeline(training_images_path, 1, high_resolution_shape[:2],
low_resolution_shape[:2]))
# Build and compile VGG19 network to extract features
epochs = 30001
batch_size = 1
vgg = build_vgg()
gan, generator, discriminator = build_gan()
writer = tf.summary.create_file_writer(logdir)
for epoch in range(epochs):
"""
Train the discriminator network
"""
# Sample a batch of images
# high_resolution_images, low_resolution_images = next(highres_dataloader) , next(lowres_dataloader)
high_resolution_images, low_resolution_images = next(dataloader)
# Generate high-resolution images from low-resolution images
generated_high_resolution_images = generator.predict(
low_resolution_images)
# Generate batch of real and fake labels
# real_labels = np.ones((batch_size, 16, 16, 1))
# fake_labels = np.zeros((batch_size, 16, 16, 1))
real_labels = np.ones((batch_size, 1))
fake_labels = np.zeros((batch_size, 1))
# Train the discriminator network on real and fake images
d_loss_real = discriminator.train_on_batch(high_resolution_images,
real_labels)
d_loss_fake = discriminator.train_on_batch(
generated_high_resolution_images, fake_labels)
# Calculate total discriminator loss
d_loss = 0.5 * np.add(d_loss_real, d_loss_fake)
# print("d_loss:", d_loss)
"""
Train the generator network
"""
# Sample a batch of images
high_resolution_images, low_resolution_images = next(dataloader)
# Extract feature maps for real high-resolution images
image_features = vgg.predict(high_resolution_images)
# Train the generator network
g_loss = gan.train_on_batch(
[low_resolution_images, high_resolution_images],
[real_labels, image_features])
# print("g_loss:", g_loss)
print("Epoch {} : g_loss: {} , d_loss: {}".format(
epoch, g_loss[0], d_loss[0]))
# Sample and save images after every 100 epochs
with writer.as_default():
tf.summary.scalar('g_loss', g_loss[0], step=epoch)
tf.summary.scalar('d_loss', d_loss[0], step=epoch)
writer.flush()
print("Epoch {}: Gloss : {} , Dloss {}".format(
epoch + 1, g_loss, d_loss))
if (epoch) % 100 == 0:
high_resolution_images, low_resolution_images = next(dataloader)
generated_images = generator.predict_on_batch(
low_resolution_images)
for index, img in enumerate(generated_images):
save_images(res_im_path + "img_{}_{}".format(epoch, index),
low_resolution_images[index],
generated_images[index],
high_resolution_images[index])
if (epoch) % 1000 == 0:
# Save models
generator.save_weights(models_path +
"generator_{}.h5".format(epoch))
discriminator.save_weights(models_path +
"discriminator_{}.h5".format(epoch))