def __init__(self):
logger.info("Loading VITON_worker ...")
self.batch_size = 1
self.image_holder = \
tf.placeholder(tf.float32,
shape=[self.batch_size, 256, 192, 3])
self.prod_image_holder = tf.placeholder(
tf.float32, shape=[self.batch_size, 256, 192, 3])
self.body_segment_holder = tf.placeholder(
tf.float32, shape=[self.batch_size, 256, 192, 1])
self.prod_segment_holder = tf.placeholder(
tf.float32, shape=[self.batch_size, 256, 192, 1])
self.skin_segment_holder = tf.placeholder(
tf.float32, shape=[self.batch_size, 256, 192, 3])
self.pose_map_holder = \
tf.placeholder(tf.float32,
shape=[self.batch_size, 256, 192, 18])
self.viton_worker = create_model(self.prod_image_holder,
self.body_segment_holder,
self.skin_segment_holder,
self.pose_map_holder,
self.prod_segment_holder,
self.image_holder)
saver = tf.train.Saver()
self.sess = tf.Session()
logger.info("loading model from checkpoint")
checkpoint = tf.train.latest_checkpoint(FLAGS.checkpoint)
if checkpoint is None:
checkpoint = FLAGS.checkpoint
logger.info("Checkpoint: {}".format(checkpoint))
saver.restore(self.sess, checkpoint)
logger.info("Initialization done")
def main(unused_argv):
try:
os.mkdir(FLAGS.result_dir)
except:
pass
try:
os.mkdir(FLAGS.result_dir + "/images/")
except:
pass
try:
os.mkdir(FLAGS.result_dir + "/tps/")
except:
pass
# batch inference, can also be done one image per time.
batch_size = 1
image_holder = tf.placeholder(tf.float32, shape=[batch_size, 256, 192, 3])
prod_image_holder = tf.placeholder(tf.float32,
shape=[batch_size, 256, 192, 3])
body_segment_holder = tf.placeholder(tf.float32,
shape=[batch_size, 256, 192, 1])
prod_segment_holder = tf.placeholder(tf.float32,
shape=[batch_size, 256, 192, 1])
skin_segment_holder = tf.placeholder(tf.float32,
shape=[batch_size, 256, 192, 3])
pose_map_holder = tf.placeholder(tf.float32,
shape=[batch_size, 256, 192, 18])
model = create_model(prod_image_holder, body_segment_holder,
skin_segment_holder, pose_map_holder,
prod_segment_holder, image_holder)
images = np.zeros((batch_size, 256, 192, 3))
prod_images = np.zeros((batch_size, 256, 192, 3))
body_segments = np.zeros((batch_size, 256, 192, 1))
prod_segments = np.zeros((batch_size, 256, 192, 1))
skin_segments = np.zeros((batch_size, 256, 192, 3))
pose_raws = np.zeros((batch_size, 256, 192, 18))
saver = tf.train.Saver()
with tf.Session() as sess:
print("loading model from checkpoint")
checkpoint = tf.train.latest_checkpoint(FLAGS.checkpoint)
if checkpoint == None:
checkpoint = FLAGS.checkpoint
print(checkpoint)
saver.restore(sess, checkpoint)
# reading input data
test_info = open(FLAGS.test_label).read().splitlines()
for i in range(FLAGS.begin, FLAGS.end, batch_size):
# loading batch data
image_names = []
product_image_names = []
for j in range(i, i + batch_size):
info = test_info[j].split()
print(info)
image_name = info[0]
product_image_name = info[1]
image_names.append(image_name)
product_image_names.append(product_image_name)
(image, prod_image, pose_raw, body_segment, prod_segment,
skin_segment) = _process_image(image_name, product_image_name,
sess)
images[j - i] = image
prod_images[j - i] = prod_image
body_segments[j - i] = body_segment
prod_segments[j - i] = prod_segment
skin_segments[j - i] = skin_segment
pose_raws[j - i] = pose_raw
# inference
feed_dict = {
image_holder: images,
prod_image_holder: prod_images,
body_segment_holder: body_segments,
skin_segment_holder: skin_segments,
prod_segment_holder: prod_segments,
pose_map_holder: pose_raws,
}
[image_output, mask_output, loss,
step] = sess.run([
model.image_outputs, model.mask_outputs,
model.gen_loss_content_L1, model.global_step
],
feed_dict=feed_dict)
# write results
for j in range(batch_size):
imageio.imwrite(
FLAGS.result_dir + ("images/%08d_" % step) +
image_names[j] + "_" + product_image_names[j] + '.png',
(image_output[j] / 2.0 + 0.5))
imageio.imwrite(
FLAGS.result_dir + ("images/%08d_" % step) +
image_names[j] + "_" + product_image_names[j] +
'_mask.png', np.squeeze(mask_output[j]))
imageio.imwrite(FLAGS.result_dir + "images/" + image_names[j],
(images[j] / 2.0 + 0.5))
imageio.imwrite(
FLAGS.result_dir + "images/" + product_image_names[j],
(prod_images[j] / 2.0 + 0.5))
sio.savemat(
FLAGS.result_dir + "/tps/" + ("%08d_" % step) +
image_names[j] + "_" + product_image_names[j] +
"_mask.mat", {"mask": np.squeeze(mask_output[j])})
# write html
index_path = os.path.join(FLAGS.result_dir, "index.html")
if os.path.exists(index_path):
index = open(index_path, "a")
else:
index = open(index_path, "w")
index.write("<html><body><table><tr>")
index.write("<th>step</th>")
index.write("<th>name</th><th>input</th>"
"<th>output</th><th>target</th></tr>")
for j in range(batch_size):
index.write("<tr>")
index.write("<td>%d %d</td>" % (step, i + j))
index.write("<td>%s %s</td>" %
(image_names[j], product_image_names[j]))
index.write("<td><img src='images/%s'></td>" % image_names[j])
index.write("<td><img src='images/%s'></td>" %
product_image_names[j])
index.write("<td><img src='images/%08d_%s'></td>" %
(step, image_names[j] + "_" +
product_image_names[j] + '.png'))
index.write("<td><img src='images/%08d_%s'></td>" %
(step, image_names[j] + "_" +
product_image_names[j] + '_mask.png'))
index.write("</tr>")