def train_autoencoder(model_path,
train_data_path,
image_size,
report_rate=100,
learning_rate=1e-4,
num_epoch=50,
batch_size=10,
capacity=3000,
min_after_dequeue=800):
from data.common import TfReader
with tf.Graph().as_default():
# Read training data.
train_data = TfReader(data_path=train_data_path,
size=(image_size, image_size),
num_epochs=num_epoch)
images, classes = train_data.read(batch_size=batch_size,
capacity=capacity,
min_after_dequeue=min_after_dequeue)
y, z = build_autoencoder(images)
# Cost function measures pixel-wise difference
cost = tf.reduce_sum(tf.square(y - images))
optimize(cost=cost,
save_path=model_path,
report_rate=report_rate,
scope=None,
learning_rate=learning_rate)
def test_autoencoder(model_path,
test_data_path,
image_size,
report_rate=100,
batch_size=1,
capacity=3000,
min_after_dequeue=800):
from PIL import Image
import numpy as np
import os
from data.common import TfReader
with tf.Graph().as_default():
# Read test data.
test_data = TfReader(data_path=test_data_path,
size=(image_size, image_size))
images, classes = test_data.read(batch_size=batch_size,
capacity=capacity,
min_after_dequeue=min_after_dequeue)
y, z = build_autoencoder(images)
# Cost function measures pixel-wise difference
cost = tf.reduce_sum(tf.square(y - images))
# Run session.
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
saver = tf.train.Saver()
with open(os.path.join(model_path, "checkpoint")) as checkpoint_file:
checkpoint_name = checkpoint_file.readline().strip().split(
':', 1)[1].strip()[1:-1]
print "Using model: " + checkpoint_name
with tf.Session() as sess:
sess.run(init_op)
saver.restore(sess, os.path.join(model_path, checkpoint_name))
tf.train.start_queue_runners(sess)
overall_cost = 0.0
step_count = 0.0
try:
while True:
original, decoded, predictions, current_cost = sess.run(
[images, y, z, cost])
step_count += 1.0
overall_cost += current_cost
if step_count % report_rate == 0:
Image.fromarray(np.uint8(decoded[0]), "RGB") \
.save("decoded_step_{0}.jpg".format(step_count))
Image.fromarray(np.uint8(original[0]), "RGB") \
.save("original_step_{0}.jpg".format(step_count))
print "{0} steps passed with cost of {1}/{2}." \
.format(step_count, current_cost, overall_cost / step_count)
except tf.errors.OutOfRangeError:
print "All images used in {0} steps.".format(step_count)
finally:
print "Final cost: {0}.".format(overall_cost / step_count)
def extract_image(input_path,
resize,
limit,
output_path,
regression,
report_rate=50):
from data.common import TfReader
from PIL import Image
import numpy as np
import tensorflow as tf
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
reader = TfReader(data_path=input_path,
regression=regression,
size=(resize, resize))
img_op, label_op = reader.read(batch_size=1)
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
with tf.Session() as sess:
sess.run([init_op])
tf.train.start_queue_runners(sess)
step_count = 0
label_wrote_count = {}
try:
while True:
step_count += 1
img, label = sess.run([img_op, label_op])
img = img[0]
label = label[0]
output_dir = os.path.join(output_path, str(label))
# Count name wrote times.
if label not in label_wrote_count:
label_wrote_count[label] = 0
if not os.path.isdir(output_dir):
os.mkdir(output_dir)
# Skip extra images.
elif limit and label_wrote_count[label] >= limit:
continue
Image.fromarray(np.uint8(img), "RGB").save(
os.path.join(output_dir,
"step_{0}.jpg".format(step_count)))
label_wrote_count[label] += 1
if step_count % report_rate == 0:
print "{0} steps passed with label wrote: ".format(
step_count),
print label_wrote_count
except tf.errors.OutOfRangeError:
print "All images used in {0} steps.".format(step_count)
def train(model_path,
train_data_path,
class_count,
image_size,
image_channel=3,
report_rate=100,
build=build_cnn,
regression=False,
scope=None,
learning_rate=1e-4,
num_epoch=50,
batch_size=10,
capacity=3000,
min_after_dequeue=800):
from data.common import TfReader
with tf.Graph().as_default():
# Read training data.
train_data = TfReader(data_path=train_data_path,
regression=regression,
size=(image_size, image_size),
num_epochs=num_epoch)
images, classes = train_data.read(batch_size=batch_size,
capacity=capacity,
min_after_dequeue=min_after_dequeue)
y, y_pred_cls = build(input_tensor=images,
num_class=class_count,
image_size=image_size,
image_channel=image_channel)
if regression:
cost = tf.reduce_sum(tf.pow(tf.transpose(y) - classes,
2)) / (2 * batch_size)
else:
# Calculate cross-entropy for each image.
# This function calculates the softmax internally, so use the output layer directly.
# The input label is of type int in [0, num_class).
cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(
logits=y, labels=classes)
# Calculate average cost across all images.
cost = tf.reduce_mean(cross_entropy)
optimize(cost=cost,
save_path=model_path,
report_rate=report_rate,
scope=scope,
learning_rate=learning_rate)
def test(model_path,
test_data_path,
class_count,
image_size,
image_channel=3,
report_rate=100,
build=build_cnn,
regression=False,
batch_size=1,
capacity=3000,
min_after_dequeue=800):
from data.common import TfReader
import os
with tf.Graph().as_default():
# Read test data.
train_data = TfReader(data_path=test_data_path,
regression=regression,
size=(image_size, image_size))
images, classes = train_data.read(batch_size=batch_size,
capacity=capacity,
min_after_dequeue=min_after_dequeue)
y, y_pred_cls = build(input_tensor=images,
num_class=class_count,
image_size=image_size,
image_channel=image_channel)
if regression:
prediction_op = tf.squeeze(y)
correct_prediction = tf.abs(tf.transpose(y) - classes)
statistics = RegressionBias()
else:
prediction_op = y_pred_cls
correct_prediction = tf.equal(y_pred_cls, classes)
statistics = ConfusionMatrix(class_count=class_count)
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
# Run session.
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
saver = tf.train.Saver()
with open(os.path.join(model_path, "checkpoint")) as checkpoint_file:
checkpoint_name = checkpoint_file.readline().strip().split(
':', 1)[1].strip()[1:-1]
print "Using model: " + checkpoint_name
with tf.Session() as sess:
sess.run(init_op)
saver.restore(sess, os.path.join(model_path, checkpoint_name))
tf.train.start_queue_runners(sess)
overall_accuracy = 0.0
step_count = 0.0
try:
while True:
predictions, truth, current_accuracy = sess.run(
[prediction_op, classes, accuracy])
step_count += 1.0
overall_accuracy += current_accuracy
statistics.update(predictions=predictions, truth=truth)
if step_count % report_rate == 0:
print "{0} steps passed with result of {1}/{2}."\
.format(step_count, current_accuracy, overall_accuracy / step_count)
print "Sample predictions: {0}".format(predictions)
except tf.errors.OutOfRangeError:
print "All images used in {0} steps.".format(step_count)
finally:
print "Final result: {0}.".format(overall_accuracy /
step_count)
statistics.print_result()
return overall_accuracy / step_count