def test_hierarical_prediction(self, input_folder=None, save_path=None):
self.saver = tf.train.Saver()
_, restore_model_path = model_utils.pre_load_checkpoint(MODEL_DIR)
print restore_model_path
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
with tf.Session(config=config) as sess:
self.saver.restore(sess, restore_model_path)
total_time = 0
samples = glob(input_folder)
samples.sort()
for point_path in samples:
edge_path = None
print point_path, edge_path
start = time.time()
gm = GKNN(point_path, edge_path, patch_size=NUM_POINT, patch_num=30,add_noise=False,normalization=True)
##get the edge information
_,pred,pred_edge = self.pc_prediction(gm,sess,patch_num_ratio=3, edge_threshold=0.05)
end = time.time()
print "total time: ",end-start
#path = os.path.join(save_path, point_path.split('/')[-1][:-4] + "_input.xyz")
#data_provider.save_xyz(path, gm.data)
path = os.path.join(save_path, point_path.split('/')[-1][:-4] + "_output.xyz")
data_provider.save_xyz(path, pred)
path = os.path.join(save_path, point_path.split('/')[-1][:-4] + "_outputedge.ply")
data_provider.save_ply(path, pred_edge)
print total_time/len(samples)
def visualize(self, input_folder=None, save_path=None):
_, restore_model_path = model_utils.pre_load_checkpoint(MODEL_DIR)
logger.info(restore_model_path, bold=True)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
with tf.Session(config=config) as sess:
tf_util.optimistic_restore(sess, restore_model_path)
samples = glob(input_folder, recursive=True)
samples.sort()
if len(samples) > 100:
samples = samples[:100]
for point_path in samples:
start = time.time()
data = pc_util.load(point_path, count=NUM_SHAPE_POINT)
num_shape_point = data.shape[0]
data = data[:, 0:3]
is_2D = np.all(data[:, 2] == 0)
data, centroid, furthest_distance = pc_util.normalize_point_cloud(
data)
if FLAGS.drop_out < 1:
idx = farthest_point_sample(
int(num_shape_point * FLAGS.drop_out),
data[np.newaxis, ...]).eval()[0]
data = data[idx, 0:3]
if JITTER:
data = pc_util.jitter_perturbation_point_cloud(
data[np.newaxis, ...],
sigma=FLAGS.jitter_sigma,
clip=FLAGS.jitter_max,
is_2D=is_2D)
data = data[0, ...]
## get the edge information
logger.info(os.path.basename(point_path))
mid = data[(np.abs(data[:, 2]) < np.amax(data[:, 2]) * 0.2), :]
idx = farthest_point_sample(5, mid[np.newaxis, ...]).eval()[0]
# idx = np.random.choice(data.shape[0], 5, replace=False)
patches = pc_util.extract_knn_patch(mid[idx, :], data,
NUM_POINT)
end = time.time()
print("total time: ", end - start)
path = os.path.join(save_path,
point_path.split('/')[-1][:-4] + '.ply')
total_levels = int(np.log2(UP_RATIO) / np.log2(STEP_RATIO))
for p in range(patches.shape[0]):
patch = patches[p]
for i in range(1, total_levels + 1):
patch_result = self.patch_prediction(
patch, sess, STEP_RATIO**i)
pc_util.save_ply(
(patch_result * furthest_distance) + centroid,
path[:-4] + "_p_%d_%d.ply" % (p, i))
pc_util.save_ply((patch * furthest_distance) + centroid,
path[:-4] + "_p_%d_%d.ply" % (p, 0))
pc_util.save_ply((data * furthest_distance) + centroid,
path[:-4] + "_input.ply")
def init_network(self):
self.model_path = os.path.join(ROOT_DIR,
'log/PartialNetNew4/model.ckpt')
with tf.Graph().as_default():
with tf.device('/gpu:' + str(0)):
pointclouds_pl = tf.placeholder(tf.float32, shape=(1, None, 3))
pointclouds_pl_big = tf.placeholder(tf.float32,
shape=(1, None, 3))
pointclouds_gt_small = tf.placeholder(tf.float32,
shape=(1, None, 3))
is_training_pl = tf.placeholder(tf.bool, shape=())
pred_angle = PartialNet.get_model(pointclouds_pl,
is_training_pl,
bn_decay=0.5)
# with tf.variable_scope(tf.get_variable_scope()):
# with tf.device('/gpu:0'), tf.name_scope('gpu_0') as scope:
# # Evenly split input data to each GPU
# input_pc_batch = tf.slice(pointclouds_pl,
# [0, 0, 0], [1, -1, -1])
#
# pred_angle = PartialNet.get_model_new(input_pc_batch, is_training_pl,bn_decay=0.5)
# Get training operator
saver = tf.train.Saver()
# Create a session
config = tf.ConfigProto()
# config.gpu_options.visible_device_list = '1'
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
config.log_device_placement = False
# config.gpu_options.per_process_gpu_memory_fraction = 0.5
self.sess = tf.Session(config=config)
restore_epoch, checkpoint_path = model_utils.pre_load_checkpoint(
os.path.dirname(self.model_path))
# Restore variables from disk.
model_name = checkpoint_path.split('/')[-1]
checkpoint_path = os.path.join(os.path.dirname(self.model_path),
model_name)
print('model_name:', model_name, checkpoint_path)
saver.restore(self.sess, checkpoint_path)
print('Model restored.')
self.ops = {
'pointclouds_pl': pointclouds_pl,
'pointclouds_pl_big': pointclouds_pl_big,
'pointclouds_gt_small': pointclouds_gt_small,
'is_training_pl': is_training_pl,
'pred_angle': pred_angle
}
def prediction_whole_model(use_normal=False):
data_folder = FLAGS.dataset
phase = data_folder.split('/')[-2] + data_folder.split('/')[-1]
save_path = os.path.join(MODEL_DIR, 'result/' + phase)
if not os.path.exists(save_path):
os.makedirs(save_path)
samples = glob(data_folder + "/*.xyz")
samples.sort(reverse=True)
input = np.loadtxt(samples[0])
if use_normal:
pointclouds_ipt = tf.placeholder(tf.float32,
shape=(1, input.shape[0], 6))
else:
pointclouds_ipt = tf.placeholder(tf.float32,
shape=(1, input.shape[0], 3))
pred, _ = MODEL_GEN.get_gen_model(pointclouds_ipt,
is_training=False,
scope='generator',
reuse=False,
use_normal=use_normal,
use_bn=False,
use_ibn=False,
bn_decay=0.95,
up_ratio=UP_RATIO)
saver = tf.train.Saver()
_, restore_model_path = model_utils.pre_load_checkpoint(MODEL_DIR)
print(restore_model_path)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
with tf.Session(config=config) as sess:
saver.restore(sess, restore_model_path)
for i, item in enumerate(samples):
input = np.loadtxt(item)
input = np.expand_dims(input, axis=0)
if not use_normal:
input = input[:, :, 0:3]
print(item, input.shape)
pred_pl = sess.run(pred, feed_dict={pointclouds_ipt: input})
path = os.path.join(save_path, item.split('/')[-1])
if use_normal:
norm_pl = np.zeros_like(pred_pl)
data_provider.save_pl(
path, np.hstack((pred_pl[0, ...], norm_pl[0, ...])))
else:
data_provider.save_pl(path, pred_pl[0, ...])
path = path[:-4] + '_input.xyz'
data_provider.save_pl(path, input[0])
def train(self, assign_model_path=None):
# Create a session
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
# config.log_device_placement = False
with tf.Session(config=config) as self.sess:
self.train_writer = tf.summary.FileWriter(
os.path.join(MODEL_DIR, 'train'), self.sess.graph)
init = tf.global_variables_initializer()
self.sess.run(init)
# restore the model
saver = tf.train.Saver(max_to_keep=6)
restore_epoch, checkpoint_path = model_utils.pre_load_checkpoint(
MODEL_DIR)
global LOG_FOUT
if restore_epoch == 0:
LOG_FOUT = open(os.path.join(MODEL_DIR, 'log_train.txt'), 'w')
LOG_FOUT.write(str(socket.gethostname()) + '\n')
LOG_FOUT.write(str(FLAGS) + '\n')
else:
LOG_FOUT = open(os.path.join(MODEL_DIR, 'log_train.txt'), 'a')
saver.restore(self.sess, checkpoint_path)
###assign the generator with another model file
if assign_model_path is not None:
print "Load pre-train model from %s" % (assign_model_path)
assign_saver = tf.train.Saver(var_list=[
var for var in tf.trainable_variables()
if var.name.startswith("generator")
])
assign_saver.restore(self.sess, assign_model_path)
##read data
self.fetchworker = data_provider.Fetcher(BATCH_SIZE, NUM_POINT,
USE_DATA_NORM)
self.fetchworker.start()
for epoch in tqdm(range(restore_epoch, MAX_EPOCH + 1), ncols=55):
log_string('**** EPOCH %03d ****\t' % (epoch))
self.train_one_epoch()
if epoch % 20 == 0:
saver.save(self.sess,
os.path.join(MODEL_DIR, "model"),
global_step=epoch)
self.fetchworker.shutdown()
def test(self, show=False, use_normal=False):
data_folder = '../../PointSR_data/CAD/mesh_MC16k'
phase = data_folder.split('/')[-2] + data_folder.split('/')[-1]
save_path = os.path.join(MODEL_DIR, 'result/' + phase)
self.saver = tf.train.Saver()
_, restore_model_path = model_utils.pre_load_checkpoint(MODEL_DIR)
print restore_model_path
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
with tf.Session(config=config) as sess:
self.saver.restore(sess, restore_model_path)
samples = glob(data_folder + "/.xyz")
samples.sort()
total_time = 0
#input, dist, edge, data_radius, name = data_provider.load_patch_data(NUM_POINT, True, 30)
#edge = np.reshape(edge,[-1,NUM_EDGE,6])
for i, item in tqdm(enumerate(samples)):
input = np.loadtxt(item)
edge = np.loadtxt(
item.replace('mesh_MC16k',
'mesh_edge').replace('.xyz', '_edge.xyz'))
idx = np.all(edge[:, 0:3] == edge[:, 3:6], axis=-1)
edge = edge[idx == False]
l = len(edge)
idx = range(l) * (1300 / l) + list(
np.random.permutation(l)[:1300 % l])
edge = edge[idx]
# # coord = input[:, 0:3]
# # centroid = np.mean(coord, axis=0, keepdims=True)
# # coord = coord - centroid
# # furthest_distance = np.amax(np.sqrt(np.sum(abs(coord) ** 2, axis=-1)))
# # coord = coord / furthest_distance
# # input[:, 0:3] = coord
input = np.expand_dims(input, axis=0)
# input = data_provider.jitter_perturbation_point_cloud(input, sigma=0.01, clip=0.02)
start_time = time.time()
edge_pl = tf.placeholder(tf.float32, [1, edge.shape[0], 6])
dist_gt_pl = tf.sqrt(
tf.reduce_min(model_utils.distance_point2edge(
self.pred, edge_pl),
axis=-1))
pred, pred_dist, dist_gt = sess.run(
[self.pred, self.pred_dist, dist_gt_pl],
feed_dict={
self.pointclouds_input: input[:, :, 0:3],
self.pointclouds_radius: np.ones(BATCH_SIZE),
edge_pl: np.expand_dims(edge, axis=0)
})
total_time += time.time() - start_time
norm_pl = np.zeros_like(pred)
##--------------visualize predicted point cloud----------------------
if show:
f, axis = plt.subplots(3)
axis[0].imshow(
pc_util.point_cloud_three_views(input[:, 0:3],
diameter=5))
axis[1].imshow(
pc_util.point_cloud_three_views(pred[0, :, :],
diameter=5))
axis[2].imshow(
pc_util.point_cloud_three_views(gt[:, 0:3],
diameter=5))
plt.show()
path = os.path.join(save_path,
item.split('/')[-1][:-4] + ".ply")
# rgba =data_provider.convert_dist2rgba(pred_dist2,scale=10)
# data_provider.save_ply(path, np.hstack((pred[0, ...],rgba,pred_dist2.reshape(NUM_ADDPOINT,1))))
path = os.path.join(save_path,
item.split('/')[-1][:-4] + "_gt.ply")
rgba = data_provider.convert_dist2rgba(dist_gt[0], scale=5)
data_provider.save_ply(
path,
np.hstack(
(pred[0, ...], rgba, dist_gt.reshape(NUM_ADDPOINT,
1))))
path = path.replace(phase, phase + "_input")
path = path.replace('xyz', 'ply')
rgba = data_provider.convert_dist2rgba(pred_dist[0], scale=5)
data_provider.save_ply(
path,
np.hstack((input[0], rgba, pred_dist.reshape(NUM_POINT,
1))))
print total_time / len(samples)
def train(assign_model_path=None):
is_training = True
bn_decay = 0.95
step = tf.Variable(0, trainable=False)
learning_rate = BASE_LEARNING_RATE
tf.summary.scalar('bn_decay', bn_decay)
tf.summary.scalar('learning_rate', learning_rate)
# get placeholder
pointclouds_pl, pointclouds_gt, pointclouds_gt_normal, pointclouds_radius = MODEL_GEN.placeholder_inputs(
BATCH_SIZE, NUM_POINT, UP_RATIO)
#create the generator model
pred, _ = MODEL_GEN.get_gen_model(pointclouds_pl,
is_training,
scope='generator',
bradius=pointclouds_radius,
reuse=None,
use_normal=False,
use_bn=False,
use_ibn=False,
bn_decay=bn_decay,
up_ratio=UP_RATIO)
#get emd loss
gen_loss_emd, matchl_out = model_utils.get_emd_loss(
pred, pointclouds_gt, pointclouds_radius)
#get repulsion loss
if USE_REPULSION_LOSS:
gen_repulsion_loss = model_utils.get_repulsion_loss4(pred)
tf.summary.scalar('loss/gen_repulsion_loss', gen_repulsion_loss)
else:
gen_repulsion_loss = 0.0
#get total loss function
pre_gen_loss = 100 * gen_loss_emd + gen_repulsion_loss + tf.losses.get_regularization_loss(
)
# create pre-generator ops
gen_update_ops = [
op for op in tf.get_collection(tf.GraphKeys.UPDATE_OPS)
if op.name.startswith("generator")
]
gen_tvars = [
var for var in tf.trainable_variables()
if var.name.startswith("generator")
]
with tf.control_dependencies(gen_update_ops):
pre_gen_train = tf.train.AdamOptimizer(
learning_rate,
beta1=0.9).minimize(pre_gen_loss,
var_list=gen_tvars,
colocate_gradients_with_ops=True,
global_step=step)
# merge summary and add pointclouds summary
tf.summary.scalar('loss/gen_emd', gen_loss_emd)
tf.summary.scalar('loss/regularation', tf.losses.get_regularization_loss())
tf.summary.scalar('loss/pre_gen_total', pre_gen_loss)
pretrain_merged = tf.summary.merge_all()
pointclouds_image_input = tf.placeholder(tf.float32,
shape=[None, 500, 1500, 1])
pointclouds_input_summary = tf.summary.image('pointcloud_input',
pointclouds_image_input,
max_outputs=1)
pointclouds_image_pred = tf.placeholder(tf.float32,
shape=[None, 500, 1500, 1])
pointclouds_pred_summary = tf.summary.image('pointcloud_pred',
pointclouds_image_pred,
max_outputs=1)
pointclouds_image_gt = tf.placeholder(tf.float32,
shape=[None, 500, 1500, 1])
pointclouds_gt_summary = tf.summary.image('pointcloud_gt',
pointclouds_image_gt,
max_outputs=1)
image_merged = tf.summary.merge([
pointclouds_input_summary, pointclouds_pred_summary,
pointclouds_gt_summary
])
# Create a session
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
config.log_device_placement = False
with tf.Session(config=config) as sess:
train_writer = tf.summary.FileWriter(os.path.join(MODEL_DIR, 'train'),
sess.graph)
init = tf.global_variables_initializer()
sess.run(init)
ops = {
'pointclouds_pl': pointclouds_pl,
'pointclouds_gt': pointclouds_gt,
'pointclouds_gt_normal': pointclouds_gt_normal,
'pointclouds_radius': pointclouds_radius,
'pointclouds_image_input': pointclouds_image_input,
'pointclouds_image_pred': pointclouds_image_pred,
'pointclouds_image_gt': pointclouds_image_gt,
'pretrain_merged': pretrain_merged,
'image_merged': image_merged,
'gen_loss_emd': gen_loss_emd,
'pre_gen_train': pre_gen_train,
'pred': pred,
'step': step,
}
#restore the model
saver = tf.train.Saver(max_to_keep=6)
restore_epoch, checkpoint_path = model_utils.pre_load_checkpoint(
MODEL_DIR)
global LOG_FOUT
if restore_epoch == 0:
LOG_FOUT = open(os.path.join(MODEL_DIR, 'log_train.txt'), 'w')
LOG_FOUT.write(str(socket.gethostname()) + '\n')
LOG_FOUT.write(str(FLAGS) + '\n')
else:
LOG_FOUT = open(os.path.join(MODEL_DIR, 'log_train.txt'), 'a')
saver.restore(sess, checkpoint_path)
###assign the generator with another model file
if assign_model_path is not None:
print "Load pre-train model from %s" % (assign_model_path)
assign_saver = tf.train.Saver(var_list=[
var for var in tf.trainable_variables()
if var.name.startswith("generator")
])
assign_saver.restore(sess, assign_model_path)
##read data
input_data, gt_data, data_radius, _ = data_provider.load_patch_data(
skip_rate=1,
num_point=NUM_POINT,
norm=USE_DATA_NORM,
use_randominput=USE_RANDOM_INPUT)
fetchworker = data_provider.Fetcher(input_data, gt_data, data_radius,
BATCH_SIZE, NUM_POINT,
USE_RANDOM_INPUT, USE_DATA_NORM)
fetchworker.start()
for epoch in tqdm(range(restore_epoch, MAX_EPOCH + 1), ncols=55):
log_string('**** EPOCH %03d ****\t' % (epoch))
train_one_epoch(sess, ops, fetchworker, train_writer)
if epoch % 20 == 0:
saver.save(sess,
os.path.join(MODEL_DIR, "model"),
global_step=epoch)
fetchworker.shutdown()
def test_hierarical_prediction(self):
data_folder = '../../PointSR_data/virtualscan/chair_test1/*_noise_half.xyz'
# data_folder = '../../PointSR_data/rawscan/aaa.xyz'
# data_folder = '/home/lqyu/chair/tmp.xyz'
phase = data_folder.split('/')[-3] + "_" + data_folder.split('/')[-2]
save_path = os.path.join(
MODEL_DIR,
'result/' + 'halfnoise_' + phase + '_512_0.05_dynamic_96')
data_folder = '../../PointSR_data/realscan/ToyTurtle_clean_simply_simply.xyz'
save_path = os.path.join(
'../../PointSR_data/tmp/realscan_simply_simply_residual_2048_0.05')
self.saver = tf.train.Saver()
_, restore_model_path = model_utils.pre_load_checkpoint(MODEL_DIR)
print restore_model_path
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
with tf.Session(config=config) as sess:
self.saver.restore(sess, restore_model_path)
total_time = 0
samples = glob(data_folder)
samples.sort()
for point_path in samples:
if 'no_noise' in point_path:
continue
edge_path = point_path.replace('new_simu_noise',
'mesh_edge').replace(
'_noise_double.xyz',
'_edge.xyz')
edge_path = None
print point_path, edge_path
gm = GKNN_realscan(point_path,
edge_path,
patch_size=NUM_POINT,
patch_num=30,
add_noise=False,
normalization=False)
##get the edge information
_, pred, pred_edge = self.pc_prediction(gm,
sess,
patch_num_ratio=3,
edge_threshold=0.05)
## re-prediction with edge information
# input, pred,pred_edge = self.pc_prediction(gm,sess,patch_num_ratio=3, edge_threshold=0.05,edge=pred_edge[:,0:3])
path = os.path.join(
save_path,
point_path.split('/')[-1][:-4] + "_input.xyz")
data_provider.save_xyz(path, gm.data)
path = os.path.join(
save_path,
point_path.split('/')[-1][:-4] + "_output.xyz")
data_provider.save_xyz(path, pred)
path = os.path.join(
save_path,
point_path.split('/')[-1][:-4] + "_outputedge.ply")
data_provider.save_ply(path, pred_edge)
print total_time / len(samples)
def evaluate():
with tf.device('/gpu:' + str(GPU_INDEX)):
test_data_placeholder = tf.placeholder(tf.float32,
shape=(BATCH_SIZE, NUM_POINT,
data_dim))
is_training_placeholder = tf.placeholder(tf.bool)
per_point_val, _, spatial_weights = MODEL.get_model(
test_data_placeholder, is_training_placeholder)
saver = tf.train.Saver()
# Create a session
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
config.log_device_placement = True
sess = tf.Session(config=config)
# Restore variables from disk.
_, restore_model_path = model_utils.pre_load_checkpoint(LOG_DIR)
# restore_model_path = LOG_DIR + '/model.ckpt-200'
saver.restore(sess, restore_model_path)
print("Model restored: %s" % restore_model_path)
ops = {
'pointclouds_pl': test_data_placeholder,
'is_training_pl': is_training_placeholder,
'per_point_val': per_point_val
}
# 'spatial_weights': spatial_weights}
is_training = False
num_batches = test_data_size // BATCH_SIZE
test_data_size_inuse = int(num_batches * BATCH_SIZE)
per_point_vals = np.zeros(shape=(test_data_size_inuse, NUM_POINT, 1))
per_point_features = np.zeros(shape=(test_data_size_inuse, NUM_POINT, 128))
for batch_idx in tqdm(range(num_batches)):
start_idx = batch_idx * BATCH_SIZE
end_idx = (batch_idx + 1) * BATCH_SIZE
current_test_data = test_data[start_idx:end_idx, :, :]
feed_dict = {
ops['pointclouds_pl']: current_test_data,
ops['is_training_pl']: is_training,
}
per_point_val = sess.run([ops['per_point_val']], feed_dict=feed_dict)
per_point_features[start_idx:end_idx, :, :] = per_point_val[0]
per_point_vals[start_idx:end_idx, :, 0] = np.max(per_point_val[0],
axis=-1)
# per_point_vals[start_idx:end_idx, :, 0] = np.mean(per_point_val[0], axis=-1)
# per_point_vals[start_idx:end_idx, :, 0] = np.linalg.norm(per_point_val[0], ord=2, axis=-1)
## save detected point saliency probability
output_path = LOG_DIR + '/' + 'saliency_prob/'
if not os.path.exists(output_path):
os.makedirs(output_path)
for i in range(test_data_size):
detected_saliency_prob = np.concatenate(
[test_data[i, :, :], per_point_vals[i, :, :]], axis=-1)
model_name = name[i] + '.txt'
model_path = os.path.join(output_path, model_name)
np.savetxt(model_path, detected_saliency_prob, fmt='%.3f')
def prediction_whole_model(data_folder=None, show=False, use_normal=False):
data_folder = '../data/test_data/our_collected_data/MC_5k'
phase = data_folder.split('/')[-2] + data_folder.split('/')[-1]
save_path = os.path.join(MODEL_DIR, 'result/' + phase)
if not os.path.exists(save_path):
os.makedirs(save_path)
samples = glob(data_folder + "/*.xyz")
samples.sort(reverse=True)
input = np.loadtxt(samples[0])
if use_normal:
pointclouds_ipt = tf.placeholder(tf.float32,
shape=(1, input.shape[0], 6))
else:
pointclouds_ipt = tf.placeholder(tf.float32,
shape=(1, input.shape[0], 3))
pred, _ = MODEL_GEN.get_gen_model(pointclouds_ipt,
is_training=False,
scope='generator',
bradius=1.0,
reuse=None,
use_normal=use_normal,
use_bn=False,
use_ibn=False,
bn_decay=0.95,
up_ratio=UP_RATIO)
saver = tf.train.Saver()
_, restore_model_path = model_utils.pre_load_checkpoint(MODEL_DIR)
print restore_model_path
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
with tf.Session(config=config) as sess:
saver.restore(sess, restore_model_path)
samples = glob(data_folder + "/*.xyz")
samples.sort()
total_time = 0
for i, item in enumerate(samples):
input = np.loadtxt(item)
gt = input
# input = data_provider.jitter_perturbation_point_cloud(np.expand_dims(input,axis=0),sigma=0.003,clip=0.006)
input = np.expand_dims(input, axis=0)
if not use_normal:
input = input[:, :, 0:3]
gt = gt[:, 0:3]
print item, input.shape
start_time = time.time()
pred_pl = sess.run(pred, feed_dict={pointclouds_ipt: input})
total_time += time.time() - start_time
norm_pl = np.zeros_like(pred_pl)
##--------------visualize predicted point cloud----------------------
path = os.path.join(save_path, item.split('/')[-1])
if show:
f, axis = plt.subplots(3)
axis[0].imshow(
pc_util.point_cloud_three_views(input[0, :, 0:3],
diameter=5))
axis[1].imshow(
pc_util.point_cloud_three_views(pred_pl[0, :, :],
diameter=5))
axis[2].imshow(
pc_util.point_cloud_three_views(gt[:, 0:3], diameter=5))
plt.show()
data_provider.save_pl(
path, np.hstack((pred_pl[0, ...], norm_pl[0, ...])))
path = path[:-4] + '_input.xyz'
data_provider.save_pl(path, input[0])
print total_time / 20
def train(assign_model_path=None, bn_decay=0.95):
step = tf.Variable(0, trainable=False)
learning_rate = BASE_LEARNING_RATE
# get placeholder
pointclouds_pl, pointclouds_gt, pointclouds_gt_normal, pointclouds_radius = MODEL_GEN.placeholder_inputs(
BATCH_SIZE, NUM_POINT, UP_RATIO)
# create discriminator
if FLAGS.gan:
d_real = MODEL_GEN.get_discriminator(pointclouds_gt,
True,
'd_1',
reuse=False,
use_bn=False,
use_ibn=False,
use_normal=False,
bn_decay=bn_decay)
# create the generator model
pred, _ = MODEL_GEN.get_gen_model(pointclouds_pl,
True,
scope='generator',
reuse=False,
use_normal=False,
use_bn=False,
use_ibn=False,
bn_decay=bn_decay,
up_ratio=UP_RATIO)
if FLAGS.gan:
d_fake = MODEL_GEN.get_discriminator(pred,
True,
'd_1',
reuse=True,
use_bn=False,
use_ibn=False,
use_normal=False,
bn_decay=bn_decay)
# get cd loss
gen_loss_cd, _ = model_utils.get_cd_loss(pred, pointclouds_gt,
pointclouds_radius, 1.0)
# get gan loss
if FLAGS.gan:
d_loss_real = tf.reduce_mean((d_real - 1)**2)
d_loss_fake = tf.reduce_mean(d_fake**2)
d_loss = 0.5 * (d_loss_real + d_loss_fake)
# get loss for generator
g_loss = tf.reduce_mean((d_fake - 1)**2)
# get total loss function
pre_gen_loss = gen_loss_cd
if FLAGS.gan:
pre_gen_loss = g_loss + FLAGS.lambd * pre_gen_loss
""" Training """
# divide trainable variables into a group for D and a group for G
t_vars = tf.trainable_variables()
if FLAGS.gan:
d_vars = [var for var in t_vars if 'd_' in var.name]
g_vars = [var for var in t_vars if 'generator' in var.name]
# optimizers
with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS)):
if FLAGS.gan:
d_optim = tf.train.AdamOptimizer(
learning_rate,
beta1=0.5).minimize(d_loss,
var_list=d_vars,
colocate_gradients_with_ops=True)
if assign_model_path:
learning_rate = learning_rate / 10
pre_gen_train = tf.train.AdamOptimizer(
learning_rate,
beta1=0.9).minimize(pre_gen_loss,
var_list=g_vars,
colocate_gradients_with_ops=True,
global_step=step)
# weight clipping
if FLAGS.gan:
clip_D = [p.assign(tf.clip_by_value(p, -0.01, 0.01)) for p in d_vars]
# merge summary and add pointclouds summary
tf.summary.scalar('bn_decay', bn_decay)
tf.summary.scalar('learning_rate', learning_rate)
tf.summary.scalar('loss/gen_cd', gen_loss_cd)
tf.summary.scalar('loss/regularation', tf.losses.get_regularization_loss())
tf.summary.scalar('loss/pre_gen_total', pre_gen_loss)
if FLAGS.gan:
tf.summary.scalar('loss/d_loss_real', d_loss_real)
tf.summary.scalar('loss/d_loss_fake', d_loss_fake)
tf.summary.scalar('loss/d_loss', d_loss)
tf.summary.scalar('loss/g_loss', g_loss)
pretrain_merged = tf.summary.merge_all()
pointclouds_image_input = tf.placeholder(tf.float32,
shape=[None, 500, 1500, 1])
pointclouds_input_summary = tf.summary.image('pointcloud_input',
pointclouds_image_input,
max_outputs=1)
pointclouds_image_pred = tf.placeholder(tf.float32,
shape=[None, 500, 1500, 1])
pointclouds_pred_summary = tf.summary.image('pointcloud_pred',
pointclouds_image_pred,
max_outputs=1)
pointclouds_image_gt = tf.placeholder(tf.float32,
shape=[None, 500, 1500, 1])
pointclouds_gt_summary = tf.summary.image('pointcloud_gt',
pointclouds_image_gt,
max_outputs=1)
image_merged = tf.summary.merge([
pointclouds_input_summary, pointclouds_pred_summary,
pointclouds_gt_summary
])
# Create a session
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
with tf.Session(config=config) as sess:
train_writer = tf.summary.FileWriter(os.path.join(MODEL_DIR, 'train'),
sess.graph)
init = tf.global_variables_initializer()
sess.run(init)
ops = {
'pointclouds_pl': pointclouds_pl,
'pointclouds_gt': pointclouds_gt,
'pointclouds_gt_normal': pointclouds_gt_normal,
'pointclouds_radius': pointclouds_radius,
'pointclouds_image_input': pointclouds_image_input,
'pointclouds_image_pred': pointclouds_image_pred,
'pointclouds_image_gt': pointclouds_image_gt,
'pretrain_merged': pretrain_merged,
'image_merged': image_merged,
'gen_loss_cd': gen_loss_cd,
'pre_gen_train': pre_gen_train,
'd_optim': d_optim if FLAGS.gan else None,
'pred': pred,
'step': step,
'clip': clip_D if FLAGS.gan else None,
}
# restore the model
saver = tf.train.Saver(max_to_keep=6)
restore_epoch, checkpoint_path = model_utils.pre_load_checkpoint(
MODEL_DIR)
if restore_epoch == 0:
LOG_FOUT = open(os.path.join(MODEL_DIR, 'log_train.txt'), 'w')
LOG_FOUT.write(str(socket.gethostname()) + '\n')
LOG_FOUT.write(str(FLAGS) + '\n')
else:
LOG_FOUT = open(os.path.join(MODEL_DIR, 'log_train.txt'), 'a')
saver.restore(sess, checkpoint_path)
###assign the generator with another model file
if assign_model_path is not None:
print("Load pre-train model from %s" % assign_model_path)
assign_saver = tf.train.Saver(var_list=[
var for var in tf.trainable_variables()
if var.name.startswith("generator")
])
assign_saver.restore(sess, assign_model_path)
##read data
input_data, gt_data, data_radius, _ = data_provider.load_patch_data(
FLAGS.dataset,
skip_rate=1,
norm=USE_DATA_NORM,
use_randominput=USE_RANDOM_INPUT)
fetchworker = data_provider.Fetcher(input_data, gt_data, data_radius,
BATCH_SIZE, NUM_POINT,
USE_RANDOM_INPUT, USE_DATA_NORM)
fetchworker.start()
for epoch in tqdm(range(restore_epoch, MAX_EPOCH + 1), ncols=55):
log_string(LOG_FOUT, '**** EPOCH %03d ****\t' % epoch)
train_one_epoch(sess, ops, fetchworker, train_writer, LOG_FOUT,
FLAGS.gan)
if epoch % 20 == 0:
saver.save(sess,
os.path.join(MODEL_DIR, "model"),
global_step=epoch)
fetchworker.shutdown()
LOG_FOUT.close()
def train():
with tf.Graph().as_default():
with tf.device('/gpu:' + str(GPU_INDEX)):
pointclouds_pl, pointclouds_angle, pointclouds_gt = MODEL.placeholder_inputs(
BATCH_SIZE, NUM_POINT)
is_training_pl = tf.placeholder(tf.bool, shape=())
# Note the global_step=batch parameter to minimize.
# That tells the optimizer to helpfully increment the 'batch' parameter
# for you every time it trains.
batch = tf.get_variable('batch', [],
initializer=tf.constant_initializer(0),
trainable=False)
bn_decay = get_bn_decay(batch)
tf.summary.scalar('bn_decay', bn_decay)
# Get model and loss
pred_angle = MODEL.get_model(pointclouds_pl,
is_training_pl,
bn_decay=bn_decay)
# loss_util.get_transform_loss2(pred_angle, pointclouds_angle)
cd_dists, knn_dists = loss_util.get_partialNet_loss(
pred_angle, pointclouds_angle, pointclouds_pl, pointclouds_gt)
losses = tf.get_collection('losses')
total_loss = tf.add_n(losses, name='total_loss')
tf.summary.scalar('total_loss', total_loss)
for l in losses + [total_loss]:
tf.summary.scalar(l.op.name, l)
print("--- Get training operator")
# Get training operator
learning_rate = get_learning_rate(batch)
tf.summary.scalar('learning_rate', learning_rate)
if OPTIMIZER == 'momentum':
optimizer = tf.train.MomentumOptimizer(learning_rate,
momentum=MOMENTUM)
elif OPTIMIZER == 'adam':
optimizer = tf.train.AdamOptimizer(learning_rate)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
# train_op = optimizer.minimize(loss)
train_op = optimizer.minimize(total_loss, global_step=batch)
# train_op = optimizer.minimize(total_loss, global_step=batch)
# Add ops to save and restore all the variables.
saver = tf.train.Saver(max_to_keep=1)
# Create a session
config = tf.ConfigProto()
# config.gpu_options.visible_device_list = '1'
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
config.log_device_placement = False
# config.gpu_options.per_process_gpu_memory_fraction = 0.5
sess = tf.Session(config=config)
# Add summary writers
merged = tf.summary.merge_all()
train_writer = tf.summary.FileWriter(os.path.join(LOG_DIR, 'train'),
sess.graph)
test_writer = tf.summary.FileWriter(os.path.join(LOG_DIR, 'test'),
sess.graph)
# Init variables
init = tf.global_variables_initializer()
sess.run(init)
ops = {
'pointclouds_pl': pointclouds_pl,
'pointclouds_gt': pointclouds_gt,
'pointclouds_angle': pointclouds_angle,
'is_training_pl': is_training_pl,
'pred_angle': pred_angle,
'loss': total_loss,
'train_op': train_op,
'merged': merged,
'step': batch,
'cd_dists': cd_dists,
'knn_dists': knn_dists
}
restore_epoch = 0
restore_epoch, checkpoint_path = model_utils.pre_load_checkpoint(
LOG_DIR)
global LOG_FOUT
if restore_epoch == 0:
LOG_FOUT = open(os.path.join(LOG_DIR, 'log_train.txt'), 'w')
LOG_FOUT.write(str(socket.gethostname()) + '\n')
LOG_FOUT.write(str(FLAGS) + '\n')
else:
LOG_FOUT = open(os.path.join(LOG_DIR, 'log_train.txt'), 'a')
saver.restore(sess, checkpoint_path)
# restore_epoch = 0
for epoch in range(restore_epoch, MAX_EPOCH):
log_string('**** EPOCH %03d ****' % (epoch))
sys.stdout.flush()
train_one_epoch(sess, ops, train_writer)
# eval_one_epoch(sess, ops, test_writer)
# Save the variables to disk.
if epoch % 10 == 0:
eval_one_epoch(sess, ops, test_writer)
save_path = saver.save(sess,
os.path.join(LOG_DIR, "model.ckpt"))
log_string("Model saved in file: %s" % save_path)
def train(self, assign_model_path=None):
# Create a session
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
if TRAIN_H5 is not None:
if NUM_SHAPE_POINT:
input_data, gt_data, data_radius, _ = data_provider.load_patch_data(
h5_filename=TRAIN_H5,
up_ratio=UP_RATIO,
step_ratio=STEP_RATIO,
num_point=
NUM_SHAPE_POINT, # roughly 1/4 of the complete model
norm=USE_DATA_NORM,
use_randominput=USE_RANDOM_INPUT)
else:
input_data, gt_data, data_radius, _ = data_provider.load_patch_data(
h5_filename=TRAIN_H5,
up_ratio=UP_RATIO,
step_ratio=STEP_RATIO,
num_point=NUM_SHAPE_POINT,
patch_size=NUM_POINT,
norm=USE_DATA_NORM,
use_randominput=USE_RANDOM_INPUT)
self.fetchworker = data_provider.Fetcher(
input_data,
gt_data,
data_radius,
BATCH_SIZE,
num_in_point=NUM_POINT,
step_ratio=STEP_RATIO,
up_ratio=UP_RATIO,
jitter=JITTER,
jitter_sigma=FLAGS.jitter_sigma,
jitter_max=FLAGS.jitter_max,
drop_out=FLAGS.drop_out)
elif TRAIN_RECORD is not None:
self.fetchworker = data_provider.Fetcher(
TRAIN_RECORD,
batch_size=BATCH_SIZE,
step_ratio=STEP_RATIO,
up_ratio=UP_RATIO,
num_in_point=NUM_POINT,
num_shape_point=NUM_SHAPE_POINT,
jitter=JITTER,
drop_out=FLAGS.drop_out,
jitter_max=FLAGS.jitter_max,
jitter_sigma=FLAGS.jitter_sigma)
else:
raise (ValueError)
with tf.Session(
config=config) as self.sess, self.train_writer.as_default():
tf.global_variables_initializer().run()
tf.contrib.summary.initialize(graph=tf.get_default_graph())
### assign the generator with another model file
# restore the model
self.saver = tf.train.Saver(max_to_keep=None)
if assign_model_path is not None:
logger.info("Load pre-train model from %s" %
(assign_model_path),
bold=True)
# assign_saver = tf.train.Saver(
# var_list=[var for var in tf.trainable_variables() if var.name.startswith("generator")])
# assign_saver.restore(self.sess, assign_model_path)
# self.saver.restore(self.sess, assign_model_path)
tf_util.optimistic_restore(self.sess, assign_model_path)
self.restore_epoch = RESTORE_EPOCH or 0
logger.info(("Resume training from %d epoch model from %s" %
(self.restore_epoch, assign_model_path)),
bold=True)
if RESTORE_EPOCH is not None:
tf.assign(self.step, RESTORE_EPOCH *
self.fetchworker.num_batches).eval()
else:
self.restore_epoch, checkpoint_path = model_utils.pre_load_checkpoint(
MODEL_DIR)
if checkpoint_path is not None:
# self.saver.restore(self.sess, checkpoint_path)
tf_util.optimistic_restore(self.sess, checkpoint_path)
try:
self.restore_epoch = RESTORE_EPOCH or int(
self.step.eval() / self.fetchworker.num_batches)
except Exception:
self.restore_epoch = RESTORE_EPOCH or 0
logger.info(
("Resume training from %d epoch model from %s" %
(self.restore_epoch, checkpoint_path)),
bold=True)
global LOG_FOUT
if self.restore_epoch == 0:
LOG_FOUT = open(os.path.join(MODEL_DIR, 'log_train.txt'), 'w')
LOG_FOUT.write(str(socket.gethostname()) + '\n')
LOG_FOUT.write(str(FLAGS) + '\n')
else:
LOG_FOUT = open(os.path.join(MODEL_DIR, 'log_train.txt'), 'a')
self.total_steps = self.fetchworker.num_batches * MAX_EPOCH
self.stage_steps = STAGE_STEPS
self.last_max_ratio = self.get_next_ratio(self.step.eval())
self.last_is_combined = self.blend.eval(
feed_dict={self.model_up_ratio: self.last_max_ratio}) >= 1.0
self.fetchworker.initialize(self.sess, self.last_max_ratio,
self.last_is_combined)
for epoch in tqdm(list(range(self.restore_epoch, MAX_EPOCH + 1))):
log_string('**** EPOCH %03d ****\t' % (epoch))
self.train_one_epoch(epoch)
if epoch % 10 == 0:
self.saver.save(self.sess,
os.path.join(MODEL_DIR, "model"))
# self.eval_per_epoch(epoch, TEST_DATA)
self.saver.save(self.sess, os.path.join(MODEL_DIR, "final"))
def test_hierarical_prediction(self):
data_folder = '../../PointSR_data/CAD_imperfect/simu_noise'
phase = data_folder.split('/')[-2] + "_" + data_folder.split('/')[-1]
save_path = os.path.join(MODEL_DIR, 'result/' + phase)
self.saver = tf.train.Saver()
_, restore_model_path = model_utils.pre_load_checkpoint(MODEL_DIR)
print restore_model_path
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
with tf.Session(config=config) as sess:
self.saver.restore(sess, restore_model_path)
total_time = 0
samples = glob(data_folder + "/*.xyz")
samples.sort()
for point_path in samples:
print point_path
edge_path = point_path.replace('simu_noise', 'mesh_edge')
edge_path = edge_path.replace('_noise.xyz', '_edge.xyz')
gm = GKNN(point_path,
edge_path,
patch_size=NUM_POINT,
patch_num=30,
add_noise=False)
l = gm.edge.shape[0]
idxx = range(l) * (NUM_EDGE / l) + list(
np.random.permutation(l)[:NUM_EDGE % l])
edge = gm.edge[idxx]
## get patch seed from farthestsampling
points = tf.convert_to_tensor(np.expand_dims(gm.data, axis=0),
dtype=tf.float32)
start = time.time()
seed = farthest_point_sample(gm.data.shape[0] / 2,
points).eval()[0]
print "aaaaa", time.time() - start
seed1_num = int(gm.data.shape[0] / NUM_POINT * 30)
seed_list1 = seed[:seed1_num]
seed_left = seed[seed1_num:]
# seed2_num = int(gm.data.shape[0] / NUM_POINT * 1)
# seed_list2 = gm.get_seed_fromdensity(seed2_num)
# seed_list = np.concatenate([seed_list1, seed_list2])
seed_list = np.unique(seed_list1)
inputs = []
up_point_list = []
up_edge_list = []
up_edgedist_list = []
input_edge_list = []
input_edgedist_list = []
fail = 0
for seed in tqdm(seed_list):
try:
patch_size = NUM_POINT * np.random.randint(1, 5)
point = gm.bfs_knn(seed, patch_size)
idx = np.random.permutation(patch_size)[:NUM_POINT]
idx.sort()
point = point[idx]
except:
fail = fail + 1
continue
#get the idx
idx1 = np.reshape(np.arange(NUM_POINT), [1, NUM_POINT])
idx0 = np.zeros((1, NUM_POINT))
idx = np.stack((idx0, idx1), axis=-1)
up_point, up_edgepoint, up_edgedist, input_edge, input_edgedist = self.patch_prediction(
point, edge, idx, sess)
inputs.append(point)
up_point_list.append(up_point)
up_edge_list.append(up_edgepoint)
up_edgedist_list.append(up_edgedist)
input_edge_list.append(input_edge)
input_edgedist_list.append(input_edgedist)
print "total %d fails" % fail
input = np.concatenate(inputs, axis=0)
path = os.path.join(
save_path,
point_path.split('/')[-1][:-4] + "_input.xyz")
data_provider.save_xyz(path, gm.data)
input_edge = np.concatenate(input_edge_list, axis=0)
input_edgedist = np.concatenate(input_edgedist_list, axis=0)
rgba = data_provider.convert_dist2rgba(input_edgedist,
scale=10)
path = os.path.join(
save_path,
point_path.split('/')[-1][:-4] + "_inputedge.ply")
data_provider.save_ply(
path,
np.hstack((input_edge, rgba, input_edgedist.reshape(-1,
1))))
pred = np.concatenate(up_point_list, axis=0)
path = os.path.join(
save_path,
point_path.split('/')[-1][:-4] + "_output.xyz")
data_provider.save_xyz(path, pred)
pred_edge = np.concatenate(up_edge_list, axis=0)
t1 = time.time()
print "total %d edgepoint" % pred_edge.shape[0]
edge_dist = np.zeros(pred_edge.shape[0])
for sid in range(0, pred_edge.shape[0], 20000):
eid = np.minimum(pred_edge.shape[0], sid + 20000)
tf_point = tf.placeholder(tf.float32, [1, eid - sid, 3])
tf_edge = tf.placeholder(tf.float32,
[1, gm.edge.shape[0], 6])
pred_edge_dist_tf = model_utils.distance_point2edge(
tf_point, tf_edge)
pred_edge_dist_tf = tf.sqrt(
tf.reduce_min(pred_edge_dist_tf, axis=-1))
edge_dist[sid:eid] = sess.run(pred_edge_dist_tf,
feed_dict={
tf_point:
np.expand_dims(
pred_edge[sid:eid],
axis=0),
tf_edge:
np.expand_dims(gm.edge,
axis=0)
})
t3 = time.time()
print "tf time %f" % (t3 - t1)
rgba = data_provider.convert_dist2rgba(edge_dist, scale=10)
path = os.path.join(
save_path,
point_path.split('/')[-1][:-4] + "_outputedgeerror.ply")
data_provider.save_ply(
path, np.hstack((pred_edge, rgba, edge_dist.reshape(-1,
1))))
pred_edgedist = np.concatenate(up_edgedist_list, axis=0)
rgba = data_provider.convert_dist2rgba(pred_edgedist, scale=10)
path = os.path.join(
save_path,
point_path.split('/')[-1][:-4] + "_outputedge.ply")
data_provider.save_ply(
path,
np.hstack((pred_edge, rgba, pred_edgedist.reshape(-1, 1))))
print total_time / len(samples)
def train():
with tf.Graph().as_default():
with tf.device('/cpu:0'):
pointclouds_pl, pointclouds_angle, pointclouds_gt,pointclouds_gt_big = MODEL.placeholder_inputs(BATCH_SIZE, NUM_POINT,CD_Ratio,K=4)
pointclouds_gt_small = tf.placeholder(tf.float32, shape=(BATCH_SIZE, 1024, 3))
is_training_pl = tf.placeholder(tf.bool, shape=())
# Note the global_step=batch parameter to minimize.
# That tells the optimizer to helpfully increment the 'batch' parameter
# for you every time it trains.
batch = tf.get_variable('batch', [],
initializer=tf.constant_initializer(0), trainable=False)
bn_decay = get_bn_decay(batch)
tf.summary.scalar('bn_decay', bn_decay)
print("--- Get training operator")
# Get training operator
learning_rate = get_learning_rate(batch)
tf.summary.scalar('learning_rate', learning_rate)
if OPTIMIZER == 'momentum':
optimizer = tf.train.MomentumOptimizer(learning_rate, momentum=MOMENTUM)
elif OPTIMIZER == 'adam':
optimizer = tf.train.AdamOptimizer(learning_rate)
#MODEL.get_model(pointclouds_pl, is_training_pl, bn_decay=bn_decay)
tower_grads = []
pred_gpu = []
cd_dists_gpu = []
knn_dists_gpu = []
total_loss_gpu = []
for i in range(NUM_GPUS):
with tf.variable_scope(tf.get_variable_scope(), reuse=bool(i != 0)):
with tf.device('/gpu:%d' % (i)), tf.name_scope('gpu_%d' % (i)) as scope:
# Evenly split input data to each GPU
input_pc_batch = tf.slice(pointclouds_pl,
[i * DEVICE_BATCH_SIZE, 0, 0], [DEVICE_BATCH_SIZE, -1, -1])
input_angle_batch = tf.slice(pointclouds_angle,
[i * DEVICE_BATCH_SIZE, 0], [DEVICE_BATCH_SIZE, -1])
input_gt_batch = tf.slice(pointclouds_gt,
[i * DEVICE_BATCH_SIZE, 0, 0], [DEVICE_BATCH_SIZE, -1, -1])
input_gt_big_batch = tf.slice(pointclouds_gt_big,
[i * DEVICE_BATCH_SIZE, 0, 0], [DEVICE_BATCH_SIZE, -1, -1])
input_gt_small_batch = tf.slice(pointclouds_gt_small,
[i * DEVICE_BATCH_SIZE, 0, 0], [DEVICE_BATCH_SIZE, -1, -1])
pred_result = MODEL.get_model_new(input_pc_batch, is_training=is_training_pl, bn_decay=bn_decay)
cd_dists_g, knn_dists_g = loss_util.get_partialNet_loss(pred_result, input_angle_batch, input_pc_batch,
input_gt_batch, input_gt_big_batch,input_gt_small_batch)
losses = tf.get_collection('losses', scope)
total_loss = tf.add_n(losses, name='total_loss')
for l in losses + [total_loss]:
tf.summary.scalar(l.op.name, l)
update_op = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
update_ops = tf.group(*update_op)
with tf.control_dependencies([update_ops]):
pc_l2_loss = tf.losses.get_regularization_loss()
pc_l2_loss = 0.00001 * pc_l2_loss
total_loss = total_loss+pc_l2_loss
grads = optimizer.compute_gradients(total_loss)
tower_grads.append(grads)
pred_gpu.append(pred_result)
total_loss_gpu.append(total_loss)
cd_dists_gpu.append(cd_dists_g)
knn_dists_gpu.append(knn_dists_g)
# Merge pred and losses from multiple GPUs
pred_angle = tf.concat(pred_gpu, 0)
cd_dists = tf.concat(cd_dists_gpu,0)
knn_dists = tf.concat(knn_dists_gpu, 0)
total_loss = tf.reduce_mean(total_loss_gpu)
# Get training operator
grads = average_gradients(tower_grads)
apply_gradient_op = optimizer.apply_gradients(grads, global_step=batch)
variable_averages = tf.train.ExponentialMovingAverage(bn_decay, batch)
variables_averages_op = variable_averages.apply(tf.trainable_variables())
train_op = tf.group(apply_gradient_op, variables_averages_op)
# Add ops to save and restore all the variables.
saver = tf.train.Saver(max_to_keep=2)
# Create a session
config = tf.ConfigProto()
# config.gpu_options.visible_device_list = '1'
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
config.log_device_placement = False
# config.gpu_options.per_process_gpu_memory_fraction = 0.5
sess = tf.Session(config=config)
# Add summary writers
merged = tf.summary.merge_all()
train_writer = tf.summary.FileWriter(os.path.join(LOG_DIR, 'train'), sess.graph)
test_writer = tf.summary.FileWriter(os.path.join(LOG_DIR, 'test'), sess.graph)
# Init variables
init = tf.global_variables_initializer()
sess.run(init)
ops = {'pointclouds_pl' : pointclouds_pl,
'pointclouds_gt' : pointclouds_gt,
'pointclouds_gt_big': pointclouds_gt_big,
'pointclouds_gt_small': pointclouds_gt_small,
'pointclouds_angle': pointclouds_angle,
'is_training_pl' : is_training_pl,
'pred_angle' : pred_angle,
'loss' : total_loss,
'train_op' : train_op,
'merged' : merged,
'step' : batch,
'cd_dists' : cd_dists,
'knn_dists' : knn_dists}
restore_epoch, checkpoint_path = model_utils.pre_load_checkpoint(LOG_DIR)
global LOG_FOUT
if restore_epoch == 0:
LOG_FOUT = open(os.path.join(LOG_DIR, 'log_train.txt'), 'w')
LOG_FOUT.write(str(socket.gethostname()) + '\n')
LOG_FOUT.write(str(FLAGS) + '\n')
else:
LOG_FOUT = open(os.path.join(LOG_DIR, 'log_train.txt'), 'a')
saver.restore(sess, checkpoint_path)
# restore_epoch = 0
for epoch in tqdm(range(restore_epoch,MAX_EPOCH)):
log_string('**** EPOCH %03d ****' % (epoch))
sys.stdout.flush()
train_one_epoch(sess, ops, train_writer)
# eval_one_epoch(sess, ops, test_writer)
# Save the variables to disk.
if epoch % 10 == 0:
#eval_one_epoch(sess, ops, test_writer)
save_path = saver.save(sess, os.path.join(LOG_DIR, "model.ckpt"),global_step=epoch)
log_string("Model saved in file: %s" % save_path)
if epoch % 50 == 0:
eval_one_epoch(sess, ops, test_writer)
def train(palette):
with tf.Graph().as_default():
with tf.device('/cpu:0'):
###################################################
### placeholders # triplet input
data_train_placeholder = tf.placeholder(tf.float32,
shape=(BATCH_SIZE,
NUM_POINT,
data_dim))
data_positive_train_placeholder = tf.placeholder(tf.float32,
shape=(BATCH_SIZE,
NUM_POINT,
data_dim))
data_negative_train_placeholder = tf.placeholder(tf.float32,
shape=(BATCH_SIZE,
NUM_POINT,
data_dim))
is_training_placeholder = tf.placeholder(tf.bool, shape=())
cluster_mean_placeholder = tf.placeholder(tf.float32,
shape=(cluster_num, 128))
indices_placeholder = tf.placeholder(tf.int64, shape=(BATCH_SIZE))
batch = tf.get_variable('batch', [],
initializer=tf.constant_initializer(0),
trainable=False)
bn_decay = get_bn_decay(batch)
tf.summary.scalar('bn_decay', bn_decay)
###################################################
### get training operator
learning_rate = get_learning_rate(batch)
tf.summary.scalar('learning_rate', learning_rate)
if OPTIMIZER == 'momentum':
optimizer = tf.train.MomentumOptimizer(learning_rate,
momentum=MOMENTUM)
elif OPTIMIZER == 'adam':
optimizer = tf.train.AdamOptimizer(learning_rate)
###################################################
tower_grads = []
feature_anchor_gpu = []
net_anchor_gpu = []
net_positive_gpu = []
net_negative_gpu = []
cluster_loss_gpu = []
contrastive_loss_gpu = []
total_loss_gpu = []
l2_loss_gpu = []
all_loss_gpu = []
### get model and loss
for i in range(NUM_GPUS):
with tf.variable_scope(tf.get_variable_scope(),
reuse=bool(i != 0)):
with tf.device('/gpu:%d' % (i)), tf.name_scope(
'gpu_%d' % (i)) as scope:
pc_patch = tf.slice(data_train_placeholder,
[i * DEVICE_BATCH_SIZE, 0, 0],
[DEVICE_BATCH_SIZE, -1, -1])
pc_patch_positive = tf.slice(
data_positive_train_placeholder,
[i * DEVICE_BATCH_SIZE, 0, 0],
[DEVICE_BATCH_SIZE, -1, -1])
pc_patch_negative = tf.slice(
data_negative_train_placeholder,
[i * DEVICE_BATCH_SIZE, 0, 0],
[DEVICE_BATCH_SIZE, -1, -1])
pc_features, pc_net, _ = MODEL.get_model(
pc_patch, is_training_placeholder, bn_decay)
tf.get_variable_scope().reuse_variables()
pc_features_positive, pc_net_positive, _ = MODEL.get_model(
pc_patch_positive, is_training_placeholder,
bn_decay)
pc_features_negative, pc_net_negative, _ = MODEL.get_model(
pc_patch_negative, is_training_placeholder,
bn_decay)
pc_indices = tf.slice(indices_placeholder,
[i * DEVICE_BATCH_SIZE],
[DEVICE_BATCH_SIZE])
pc_cluster_loss, _, _, pc_contrastive_loss, pc_total_loss = MODEL.get_loss(
pc_net, pc_net_positive, pc_net_negative,
cluster_mean_placeholder, pc_indices)
update_op = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
update_ops = tf.group(*update_op)
with tf.control_dependencies([update_ops]):
pc_l2_loss = tf.losses.get_regularization_loss()
pc_l2_loss = 0.00001 * pc_l2_loss
pc_all_loss = pc_total_loss + pc_l2_loss
grads = optimizer.compute_gradients(pc_all_loss)
tower_grads.append(grads)
cluster_loss_gpu.append(pc_cluster_loss)
contrastive_loss_gpu.append(pc_contrastive_loss)
total_loss_gpu.append(pc_total_loss)
l2_loss_gpu.append(pc_l2_loss)
all_loss_gpu.append(pc_all_loss)
feature_anchor_gpu.append(pc_features)
net_anchor_gpu.append(pc_net)
net_positive_gpu.append(pc_net_positive)
net_negative_gpu.append(pc_net_negative)
cluster_loss = tf.reduce_mean(cluster_loss_gpu)
contrastive_loss = tf.reduce_mean(contrastive_loss_gpu)
total_loss = tf.reduce_mean(total_loss_gpu)
l2_loss = tf.reduce_mean(l2_loss_gpu)
all_loss = tf.reduce_mean(all_loss_gpu)
tf.summary.scalar('cluster_loss', cluster_loss)
tf.summary.scalar('contrastive_loss', contrastive_loss)
tf.summary.scalar('l2_loss', l2_loss)
tf.summary.scalar('total loss', all_loss)
feature_anchor = tf.concat(feature_anchor_gpu, 0)
net_anchor = tf.concat(net_anchor_gpu, 0)
net_positive = tf.concat(net_positive_gpu, 0)
net_negative = tf.concat(net_negative_gpu, 0)
grads = average_gradients(tower_grads)
apply_gradient_op = optimizer.apply_gradients(grads,
global_step=batch)
variable_averages = tf.train.ExponentialMovingAverage(
bn_decay, batch)
variables_averages_op = variable_averages.apply(
tf.trainable_variables())
train_op = tf.group(apply_gradient_op, variables_averages_op)
saver = tf.train.Saver(max_to_keep=11)
# Create a session
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
config.log_device_placement = False
sess = tf.Session(config=config)
# Add summary writers
merged = tf.summary.merge_all()
train_writer = tf.summary.FileWriter(os.path.join(LOG_DIR, 'train'),
sess.graph)
# Init variables
init = tf.global_variables_initializer()
sess.run(init, {is_training_placeholder: True})
restore_epoch, checkpoint_path = model_utils.pre_load_checkpoint(
LOG_DIR)
global LOG_FOUT
if restore_epoch == 0:
LOG_FOUT = open(os.path.join(LOG_DIR, 'log_train.txt'), 'w')
LOG_FOUT.write(str(socket.gethostname()) + '\n')
LOG_FOUT.write(str(FLAGS) + '\n')
else:
LOG_FOUT = open(os.path.join(LOG_DIR, 'log_train.txt'), 'a')
saver.restore(sess, checkpoint_path)
ops = {
'pointclouds_pl': data_train_placeholder,
'pointclouds_pl_positive': data_positive_train_placeholder,
'pointclouds_pl_negative': data_negative_train_placeholder,
'is_training_pl': is_training_placeholder,
'buffer': cluster_mean_placeholder,
'indics': indices_placeholder,
'feature_anchor': feature_anchor,
'net_anchor': net_anchor,
'net_positive': net_positive,
'net_negative': net_negative,
'cluster_loss': cluster_loss,
'contrastive_loss': contrastive_loss,
'total_loss': total_loss,
'l2_loss': l2_loss,
'all_loss': all_loss,
'train_op': train_op,
'merged': merged,
'step': batch,
}
feature_buffer = np.random.rand(file_size, 128)
feature_buffer = trafo.unit_vector(feature_buffer, axis=1) # normalize
for epoch in tqdm(range(restore_epoch, MAX_EPOCH)):
# for epoch in tqdm(range(MAX_EPOCH)):
print('******* EPOCH %03d *******' % (epoch))
sys.stdout.flush()
## clustering
# cluster the memory buffer
cluster_pred = SpectralClustering(
n_clusters=cluster_num, gamma=1).fit_predict(feature_buffer)
# cluster_pred = KMeans(n_clusters=cluster_num).fit_predict(feature_buffer)
# cluster_pred = AgglomerativeClustering(n_clusters=cluster_num).fit_predict(feature_buffer)
cluster_mean = np.zeros(shape=(cluster_num, 128))
for cluster_idx in range(cluster_num):
indices = np.where(cluster_pred == cluster_idx)[0]
cluster_avg = np.mean(feature_buffer[indices, :], axis=0)
cluster_mean[cluster_idx, :] = cluster_avg
cluster_mean = trafo.unit_vector(cluster_mean, axis=1)
train_one_epoch(sess, cluster_mean, cluster_pred, data, ops,
train_writer)
## get the feature buffer after each epoch
list = range(pts_num)
seed = random.sample(list, NUM_POINT)
data_2048 = data[:, seed, :]
num_batches = file_size // BATCH_SIZE
feature_buffer = np.zeros(shape=(file_size, 128))
for batch_idx in range(num_batches + 1):
if batch_idx != num_batches:
start_idx = batch_idx * BATCH_SIZE
end_idx = (batch_idx + 1) * BATCH_SIZE
else:
start_idx = file_size - BATCH_SIZE
end_idx = file_size
data_input = data_2048[start_idx:end_idx, :, :]
feed_dict = {
ops['pointclouds_pl']: data_input,
ops['is_training_pl']: False,
}
features = sess.run([ops['net_anchor']], feed_dict=feed_dict)
feature_buffer[start_idx:end_idx, :] = features[0]
if epoch % 50 == 0:
# start_time = time.time()
# tsne_visualization_color(feature_buffer, cluster_pred, palette, epoch)
# end_time = time.time()
# print 'tsne running time is %f' % (end_time-start_time)
# start_time = time.time()
# pca_visualization_color(feature_buffer, cluster_pred, palette, epoch)
# end_time = time.time()
# print 'pca running time is %f' % (end_time-start_time)
# if epoch % 10 == 0:
save_path = saver.save(sess,
os.path.join(LOG_DIR, "model.ckpt"),
global_step=epoch)
print("Model saved in file: %s" % save_path)
def test_hierarical_prediction(self, input_folder=None, save_path=None):
_, restore_model_path = model_utils.pre_load_checkpoint(MODEL_DIR)
logger.info(restore_model_path, bold=True)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
with tf.Session(config=config) as sess:
tf_util.optimistic_restore(sess, restore_model_path)
total_time = 0
samples = glob(input_folder, recursive=True)
samples.sort()
# if len(samples)>100:
# samples = samples[:100]
for point_path in samples:
start = time.time()
data = pc_util.load(point_path, count=NUM_SHAPE_POINT)
num_shape_point = data.shape[0]
data = data[:, 0:3]
is_2D = np.all(data[:, 2] == 0)
data, centroid, furthest_distance = pc_util.normalize_point_cloud(
data)
if FLAGS.drop_out < 1:
idx = farthest_point_sample(
int(num_shape_point * FLAGS.drop_out),
data[np.newaxis, ...]).eval()[0]
data = data[idx, 0:3]
if JITTER:
data = pc_util.jitter_perturbation_point_cloud(
data[np.newaxis, ...],
sigma=FLAGS.jitter_sigma,
clip=FLAGS.jitter_max,
is_2D=is_2D)
data = data[0, ...]
## get the edge information
logger.info(os.path.basename(point_path))
input_list, pred_list = self.pc_prediction(
data, sess, patch_num_ratio=PATCH_NUM_RATIO)
end = time.time()
print("total time: ", end - start)
pred_pc = np.concatenate(pred_list, axis=0)
pred_pc = (pred_pc * furthest_distance) + centroid
data = (data * furthest_distance) + centroid
folder = os.path.basename(os.path.dirname(point_path))
path = os.path.join(save_path, folder,
point_path.split('/')[-1][:-4] + '.ply')
# pc_util.save_ply(pred_pc, path[:-4]+'_overlapped.ply')
pc_util.save_ply(data, path[:-4] + '_input.ply')
idx = farthest_point_sample(
int(num_shape_point * FLAGS.drop_out) * UP_RATIO,
pred_pc[np.newaxis, ...]).eval()[0]
pred_pc = pred_pc[idx, 0:3]
# pred_pc, _, _ = pc_util.normalize_point_cloud(pred_pc)
# pred_pc = (pred_pc * furthest_distance) + centroid
pc_util.save_ply(pred_pc, path[:-4] + '.ply')
# if len(input_list) > 1:
# counter = 0
# for in_p, pred_p in zip(input_list, pred_list):
# pc_util.save_ply(in_p*furthest_distance+centroid, path[:-4]+"_input_patch_%d.ply" % counter)
# pc_util.save_ply(pred_p*furthest_distance+centroid, path[:-4]+"_pred_patch_%d.ply" % counter)
# counter += 1
print(total_time / len(samples))
