import argparse
import subprocess
import tensorflow as tf
import numpy as np
from datetime import datetime
import json
import os
import sys
import datetime
import time
import collections
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.append(BASE_DIR)
import P2PNET
import ioUtil
# DEFAULT SETTINGS
parser = argparse.ArgumentParser()
parser.add_argument('--train_hdf5', default='training data file name(*.hdf5)' )
parser.add_argument('--test_hdf5', default='test data file name(*.hdf5)' )
parser.add_argument('--domain_A', default='skeleton', help='name of domain A')
parser.add_argument('--domain_B', default='surface', help='name of domain B')
parser.add_argument('--mode', type=str, default='train', help='train or test')
parser.add_argument('--gpu', type=int, default=0, help='which GPU to use [default: 0]')
parser.add_argument('--batch_size', type=int, default=4, help='Batch Size during training [default: 4]')
parser.add_argument('--epoch', type=int, default=200, help='number of epoches to run [default: 200]')
parser.add_argument('--decayEpoch', type=int, default=50, help='steps(how many epoches) for decaying learning rate')
parser.add_argument("--densityWeight", type=float, default=1.0, help="density weight [default: 1.0]")
parser.add_argument("--regularWeight", type=float, default=0.1, help="regularization weight [default: 0.1]")
parser.add_argument("--nnk", type=int, default=8, help="density: number of nearest neighbours [default: 8]")
parser.add_argument("--range_max", type=float, default=1.0, help="max length of point displacement[default: 1.0]")
parser.add_argument("--radiusScal", type=float, default=1.0, help="a constant for scaling radii in pointnet++ [default: 1.0]")
parser.add_argument("--noiseLength", type=int, default=32, help="length of point-wise noise vector [default: 32]")
parser.add_argument('--checkpoint', default=None, help='epoch_##.ckpt')
### None None None
parser.add_argument('--point_num', type=int, default=None, help='do not set the argument')
parser.add_argument('--example_num', type=int, default=None, help='do not set the argument')
parser.add_argument('--output_dir', type=str, default=None, help='do not set the argument')
FLAGS = parser.parse_args()
Train_examples = ioUtil.load_examples(FLAGS.train_hdf5, FLAGS.domain_A, FLAGS.domain_B, 'names')
Test_examples = ioUtil.load_examples(FLAGS.test_hdf5, FLAGS.domain_A, FLAGS.domain_B, 'names')
FLAGS.point_num = Train_examples.pointsets_A.shape[1]
POINT_NUM = FLAGS.point_num
Example_NUM = Train_examples.pointsets_A.shape[0]
FLAGS.example_num = Example_NUM
TRAINING_EPOCHES = FLAGS.epoch
batch_size = FLAGS.batch_size
if Train_examples.pointsets_B.shape[1] != POINT_NUM \
or Test_examples.pointsets_A.shape[1] != POINT_NUM \
or Test_examples.pointsets_B.shape[1] != POINT_NUM :
print( 'point number inconsistent in the data set.')
exit()
########## create output folders
datapath, basefname = os.path.split( FLAGS.train_hdf5 )
output_dir = 'output_' + basefname[0:basefname.index('_')] + '_' + FLAGS.domain_A + '-' + FLAGS.domain_B ## + '_noise' + str(FLAGS.noiseLength) + '_dw' + str(FLAGS.densityWeight)+ '_rw' + str(FLAGS.regularWeight)
if not os.path.exists(output_dir):
os.mkdir(output_dir)
MODEL_STORAGE_PATH = os.path.join(output_dir, 'trained_models')
if not os.path.exists(MODEL_STORAGE_PATH):
os.mkdir(MODEL_STORAGE_PATH)
SUMMARIES_FOLDER = os.path.join(output_dir, 'summaries')
if not os.path.exists(SUMMARIES_FOLDER):
os.mkdir(SUMMARIES_FOLDER)
########## Save test input
ioUtil.output_point_cloud_ply( Test_examples.pointsets_A, Test_examples.names, output_dir, 'gt_'+FLAGS.domain_A)
ioUtil.output_point_cloud_ply( Test_examples.pointsets_B, Test_examples.names, output_dir, 'gt_'+FLAGS.domain_B)
# print arguments
for k, v in FLAGS._get_kwargs():
print(k + ' = ' + str(v) )
def train():
with tf.Graph().as_default():
with tf.device('/gpu:' + str(FLAGS.gpu)):
model = P2PNET.create_model(FLAGS)
########## Init and Configuration ##########
saver = tf.train.Saver( max_to_keep=5 )
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
sess = tf.Session(config=config)
init = tf.global_variables_initializer()
sess.run(init)
# Restore variables from disk.
Start_epoch_number = 1
if FLAGS.checkpoint is not None:
print('load checkpoint: ' + FLAGS.checkpoint )
saver.restore(sess, FLAGS.checkpoint )
fname = os.path.basename( FLAGS.checkpoint )
Start_epoch_number = int( fname[6:-5] ) + 1
print( 'Start_epoch_number = ' + str(Start_epoch_number) )
train_writer = tf.summary.FileWriter(SUMMARIES_FOLDER + '/train', sess.graph)
test_writer = tf.summary.FileWriter(SUMMARIES_FOLDER + '/test')
fcmd = open(os.path.join(output_dir, 'arguments.txt'), 'w')
fcmd.write(str(FLAGS))
fcmd.close()
########## Training one epoch ##########
def train_one_epoch(epoch_num):
now = datetime.datetime.now()
print(now.strftime("%Y-%m-%d %H:%M:%S"))
start_time = time.time()
is_training = True
Train_examples_shuffled = ioUtil.shuffle_examples(Train_examples)
pointsets_A = Train_examples_shuffled.pointsets_A
pointsets_B = Train_examples_shuffled.pointsets_B
names = Train_examples_shuffled.names
num_data = pointsets_A.shape[0]
num_batch = num_data // batch_size
total_data_loss_A = 0.0
total_shape_loss_A = 0.0
total_density_loss_A = 0.0
total_data_loss_B = 0.0
total_shape_loss_B = 0.0
total_density_loss_B = 0.0
total_reg_loss = 0.0
for j in range(num_batch):
begidx = j * batch_size
endidx = (j + 1) * batch_size
feed_dict = {
model.pointSet_A_ph: pointsets_A[begidx: endidx, ...],
model.pointSet_B_ph: pointsets_B[begidx: endidx, ...],
model.is_training_ph: is_training,
}
fetches = {
"train": model.total_train,
"shapeLoss_A": model.shapeLoss_A,
"densityLoss_A": model.densityLoss_A,
"shapeLoss_B": model.shapeLoss_B,
"densityLoss_B": model.densityLoss_B,
"data_loss_A": model.data_loss_A,
"data_loss_B": model.data_loss_B,
"regul_loss": model.regul_loss,
"learning_rate": model.learning_rate,
"global_step": model.global_step,
}
results = sess.run(fetches, feed_dict=feed_dict)
total_data_loss_A += results["data_loss_A"]
total_shape_loss_A += results["shapeLoss_A"]
total_density_loss_A += results["densityLoss_A"]
total_data_loss_B += results["data_loss_B"]
total_shape_loss_B += results["shapeLoss_B"]
total_density_loss_B += results["densityLoss_B"]
total_reg_loss += results["regul_loss"]
if j % 50 == 0:
print(' ' + str(j) + '/' + str(num_batch) + ': ' )
print(' data_loss_A = {:.4f},'.format(results["data_loss_A"] ) + \
' shape = {:.4f},'.format(results["shapeLoss_A"] ) + \
' density = {:.4f}'.format(results["densityLoss_A"] ) )
print(' data_loss_B = {:.4f},'.format(results["data_loss_B"] ) + \
' shape = {:.4f},'.format(results["shapeLoss_B"] ) + \
' density = {:.4f}'.format(results["densityLoss_B"] ) )
print(' regul_loss = {:.4f}\n'.format(results["regul_loss"] ) )
print(' learning_rate = {:.6f}'.format(results["learning_rate"] ) )
print(' global_step = {0}'.format(results["global_step"] ) )
total_data_loss_A /= num_batch
total_shape_loss_A /= num_batch
total_density_loss_A /= num_batch
total_data_loss_B /= num_batch
total_shape_loss_B /= num_batch
total_density_loss_B /= num_batch
total_reg_loss /= num_batch
# evaluate summaries
training_sum = sess.run( model.training_sum_ops, \
feed_dict={model.train_dataloss_A_ph: total_data_loss_A, \
model.train_dataloss_B_ph: total_data_loss_B, \
model.train_regul_ph: total_reg_loss,\
})
train_writer.add_summary(training_sum, epoch_num)
print( '\tData_loss_A = %.4f,' % total_data_loss_A + \
' shape = %.4f,' % total_shape_loss_A + \
' density = %.4f' % total_density_loss_A )
print( '\tData_loss_B = %.4f,' % total_data_loss_B + \
' shape = %.4f,' % total_shape_loss_B + \
' density = %.4f' % total_density_loss_B )
print( '\tReg_loss: %.4f\n' % total_reg_loss)
elapsed_time = time.time() - start_time
print( '\tply/sec:' + str( round(num_data/elapsed_time) ) )
print( '\tduration of this epoch:' + str(round(elapsed_time/60) ) + ' min' )
print( '\testimated finishing time:' + str(round(elapsed_time/60.0 * (TRAINING_EPOCHES-epoch_num-1)) ) + ' min' )
################## end of train function #################### end of train function ##########
def eval_one_epoch(epoch_num, mustSavePly=False):
is_training = False
pointsets_A = Test_examples.pointsets_A
pointsets_B = Test_examples.pointsets_B
names = Test_examples.names
num_data = pointsets_A.shape[0]
num_batch = num_data // batch_size
total_data_loss_A = 0.0
total_shape_loss_A = 0.0
total_density_loss_A = 0.0
total_data_loss_B = 0.0
total_shape_loss_B = 0.0
total_density_loss_B = 0.0
total_reg_loss = 0.0
for j in range(num_batch):
begidx = j * batch_size
endidx = (j + 1) * batch_size
feed_dict = {
model.pointSet_A_ph: pointsets_A[begidx: endidx, ...],
model.pointSet_B_ph: pointsets_B[begidx: endidx, ...],
model.is_training_ph: is_training,
}
fetches = {
"shapeLoss_A": model.shapeLoss_A,
"densityLoss_A": model.densityLoss_A,
"shapeLoss_B": model.shapeLoss_B,
"densityLoss_B": model.densityLoss_B,
"data_loss_A": model.data_loss_A,
"data_loss_B": model.data_loss_B,
"regul_loss": model.regul_loss,
"Predicted_A": model.Predicted_A,
"Predicted_B": model.Predicted_B,
}
results = sess.run(fetches, feed_dict=feed_dict)
total_data_loss_A += results["data_loss_A"]
total_shape_loss_A += results["shapeLoss_A"]
total_density_loss_A += results["densityLoss_A"]
total_data_loss_B += results["data_loss_B"]
total_shape_loss_B += results["shapeLoss_B"]
total_density_loss_B += results["densityLoss_B"]
total_reg_loss += results["regul_loss"]
# write test results
if epoch_num % 20 == 0 or mustSavePly:
# save predicted point sets with 1 single feeding pass
nametosave = names[begidx: endidx, ...]
Predicted_A_xyz = np.squeeze(np.array(results["Predicted_A"]))
Predicted_B_xyz = np.squeeze(np.array(results["Predicted_B"]))
ioUtil.output_point_cloud_ply(Predicted_A_xyz, nametosave, output_dir,
'Ep' + str(epoch_num) + '_predicted_' + FLAGS.domain_A + 'X1')
ioUtil.output_point_cloud_ply(Predicted_B_xyz, nametosave, output_dir,
'Ep' + str(epoch_num) + '_predicted_' + FLAGS.domain_B + 'X1')
# save predicted point sets with 4 feeding passes
for i in range(3):
results = sess.run(fetches, feed_dict=feed_dict)
Predicted_A_xyz__ = np.squeeze(np.array(results["Predicted_A"]))
Predicted_B_xyz__ = np.squeeze(np.array(results["Predicted_B"]))
Predicted_A_xyz = np.concatenate((Predicted_A_xyz, Predicted_A_xyz__), axis=1)
Predicted_B_xyz = np.concatenate((Predicted_B_xyz, Predicted_B_xyz__), axis=1)
ioUtil.output_point_cloud_ply(Predicted_A_xyz, nametosave, output_dir,
'Ep' + str(epoch_num) + '_predicted_' + FLAGS.domain_A + 'X4')
ioUtil.output_point_cloud_ply(Predicted_B_xyz, nametosave, output_dir,
'Ep' + str(epoch_num) + '_predicted_' + FLAGS.domain_B + 'X4')
# save predicted point sets with 8 feeding passes
for i in range(4):
results = sess.run(fetches, feed_dict=feed_dict)
Predicted_A_xyz__ = np.squeeze(np.array(results["Predicted_A"]))
Predicted_B_xyz__ = np.squeeze(np.array(results["Predicted_B"]))
Predicted_A_xyz = np.concatenate((Predicted_A_xyz, Predicted_A_xyz__), axis=1)
Predicted_B_xyz = np.concatenate((Predicted_B_xyz, Predicted_B_xyz__), axis=1)
ioUtil.output_point_cloud_ply( Predicted_A_xyz, nametosave, output_dir,
'Ep' + str(epoch_num) + '_predicted_' + FLAGS.domain_A + 'X8')
ioUtil.output_point_cloud_ply( Predicted_B_xyz, nametosave, output_dir,
'Ep' + str(epoch_num) + '_predicted_' + FLAGS.domain_B + 'X8')
total_data_loss_A /= num_batch
total_shape_loss_A /= num_batch
total_density_loss_A /= num_batch
total_data_loss_B /= num_batch
total_shape_loss_B /= num_batch
total_density_loss_B /= num_batch
total_reg_loss /= num_batch
# evaluate summaries
testing_sum = sess.run( model.testing_sum_ops, \
feed_dict={model.test_dataloss_A_ph: total_data_loss_A, \
model.test_dataloss_B_ph: total_data_loss_B, \
model.test_regul_ph: total_reg_loss,\
})
test_writer.add_summary(testing_sum, epoch_num)
print('\tData_loss_A = %.4f,' % total_data_loss_A + \
' shape = %.4f,' % total_shape_loss_A + \
' density = %.4f' % total_density_loss_A)
print('\tData_loss_B = %.4f,' % total_data_loss_B + \
' shape = %.4f,' % total_shape_loss_B + \
' density = %.4f' % total_density_loss_B)
print('\tReg_loss: %.4f\n' % total_reg_loss)
################## end of test function #################### end of test function ##########
if not os.path.exists(MODEL_STORAGE_PATH):
os.mkdir(MODEL_STORAGE_PATH)
if FLAGS.mode=='train':
for epoch in range(Start_epoch_number, TRAINING_EPOCHES+1):
print( '\n>>> Training for the epoch %d/%d ...' % (epoch, TRAINING_EPOCHES))
train_one_epoch(epoch)
if epoch % 20 == 0:
cp_filename = saver.save(sess, os.path.join(MODEL_STORAGE_PATH, 'epoch_' + str(epoch) + '.ckpt'))
print( 'Successfully store the checkpoint model into ' + cp_filename)
print('\n<<< Testing on the test dataset...')
eval_one_epoch(epoch, mustSavePly=True)
else:
print( '\n<<< Testing on the test dataset ...')
eval_one_epoch(Start_epoch_number, mustSavePly=True)
if __name__ == '__main__':
train()
