def __init__(self, args, is_training):
df_train, self.num_train = lmdb_dataflow(args.lmdb_train, args.batch_size,
args.num_input_points, args.num_gt_points, is_training=True)
batch_train = get_queued_data(df_train.get_data(), [tf.string, tf.float32, tf.float32],
[[args.batch_size],
[args.batch_size, args.num_input_points, 3],
[args.batch_size, args.num_gt_points, 3]])
df_valid, self.num_valid = lmdb_dataflow(args.lmdb_valid, args.batch_size,
args.num_input_points, args.num_gt_points, is_training=False)
batch_valid = get_queued_data(df_valid.get_data(), [tf.string, tf.float32, tf.float32],
[[args.batch_size],
[args.batch_size, args.num_input_points, 3],
[args.batch_size, args.num_gt_points, 3]])
self.batch_data = tf.cond(is_training, lambda: batch_train, lambda: batch_valid)
def __init__(self, args, is_training):
df_test, self.num_test = lmdb_dataflow(args.lmdb_test,
args.batch_size,
args.num_input_points,
args.num_gt_points,
is_training=False)
batch_test = get_queued_data(
df_test.get_data(), [tf.string, tf.float32, tf.float32],
[[args.batch_size], [args.batch_size, args.num_input_points, 3],
[args.batch_size, args.num_gt_points, 3]])
self.batch_data = batch_test
avg_acc, acc_update = tf.metrics.accuracy(labels_pl, prediction)
avg_iou, iou_update = tf.metrics.mean_iou(labels_pl, prediction, n_parts)
tf.summary.scalar('train/learning rate', learning_rate, collections=['train'])
tf.summary.scalar('train/gradient norm', global_norm, collections=['train'])
tf.summary.scalar('train/bn decay', bn_decay, collections=['train'])
tf.summary.scalar('valid/loss', avg_loss, collections=['valid'])
tf.summary.scalar('valid/accuracy', avg_acc, collections=['valid'])
tf.summary.scalar('valid/mean iou', avg_iou, collections=['valid'])
update_ops = [loss_update, acc_update, iou_update]
train_summary = tf.summary.merge_all('train')
valid_summary = tf.summary.merge_all('valid')
lmdb_train = os.path.join(args.data_dir, 'train.lmdb')
lmdb_valid = os.path.join(args.data_dir, 'valid.lmdb')
df_train, num_train = lmdb_dataflow(lmdb_train, args.batch_size, args.num_points,
shuffle=True, render=True, task='seg')
df_valid, num_valid = lmdb_dataflow(lmdb_valid, args.batch_size, args.num_points,
shuffle=False, task='seg')
train_gen = df_train.get_data()
valid_gen = df_valid.get_data()
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
sess = tf.Session(config=config)
sess.run(tf.global_variables_initializer())
writer = create_log_dir(args, sess)
saver = tf.train.Saver()
step = sess.run(global_step)
def train(args):
is_training_pl = tf.placeholder(tf.bool, shape=(), name='is_training')
global_step = tf.Variable(0, trainable=False, name='global_step')
alpha = tf.train.piecewise_constant(global_step, [10000, 20000, 50000],
[0.01, 0.1, 0.5, 1.0], 'alpha_op')
inputs_pl = tf.placeholder(tf.float32, (1, None, 3), 'inputs')
npts_pl = tf.placeholder(tf.int32, (args.batch_size, ), 'num_points')
gt_pl = tf.placeholder(tf.float32,
(args.batch_size, args.num_gt_points, 3),
'ground_truths')
model_module = importlib.import_module('.%s' % args.model_type, 'models')
model = model_module.Model(inputs_pl, npts_pl, gt_pl, alpha)
add_train_summary('alpha', alpha)
if args.lr_decay:
learning_rate = tf.train.exponential_decay(args.base_lr,
global_step,
args.lr_decay_steps,
args.lr_decay_rate,
staircase=True,
name='lr')
learning_rate = tf.maximum(learning_rate, args.lr_clip)
add_train_summary('learning_rate', learning_rate)
else:
learning_rate = tf.constant(args.base_lr, name='lr')
train_summary = tf.summary.merge_all('train_summary')
valid_summary = tf.summary.merge_all('valid_summary')
trainer = tf.train.AdamOptimizer(learning_rate)
train_op = trainer.minimize(model.loss, global_step)
df_train, num_train = lmdb_dataflow(args.lmdb_train,
args.batch_size,
args.num_input_points,
args.num_gt_points,
is_training=True)
train_gen = df_train.get_data()
df_valid, num_valid = lmdb_dataflow(args.lmdb_valid,
args.batch_size,
args.num_input_points,
args.num_gt_points,
is_training=False)
valid_gen = df_valid.get_data()
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
sess = tf.Session(config=config)
saver = tf.train.Saver()
if args.restore:
saver.restore(sess, tf.train.latest_checkpoint(args.log_dir))
writer = tf.summary.FileWriter(args.log_dir)
else:
sess.run(tf.global_variables_initializer())
if os.path.exists(args.log_dir):
delete_key = input(
colored('%s exists. Delete? [y (or enter)/N]' % args.log_dir,
'white', 'on_red'))
if delete_key == 'y' or delete_key == "":
os.system('rm -rf %s/*' % args.log_dir)
os.makedirs(os.path.join(args.log_dir, 'plots'))
else:
os.makedirs(os.path.join(args.log_dir, 'plots'))
with open(os.path.join(args.log_dir, 'args.txt'), 'w') as log:
for arg in sorted(vars(args)):
log.write(arg + ': ' + str(getattr(args, arg)) +
'\n') # log of arguments
os.system('cp models/%s.py %s' %
(args.model_type, args.log_dir)) # bkp of model def
os.system('cp train.py %s' % args.log_dir) # bkp of train procedure
writer = tf.summary.FileWriter(args.log_dir, sess.graph)
total_time = 0
train_start = time.time()
init_step = sess.run(global_step)
for step in range(init_step + 1, args.max_step + 1):
epoch = step * args.batch_size // num_train + 1
ids, inputs, npts, gt = next(train_gen)
start = time.time()
feed_dict = {
inputs_pl: inputs,
npts_pl: npts,
gt_pl: gt,
is_training_pl: True
}
_, loss, summary = sess.run([train_op, model.loss, train_summary],
feed_dict=feed_dict)
total_time += time.time() - start
writer.add_summary(summary, step)
if step % args.steps_per_print == 0:
print('epoch %d step %d loss %.8f - time per batch %.4f' %
(epoch, step, loss, total_time / args.steps_per_print))
total_time = 0
if step % args.steps_per_eval == 0:
print(colored('Testing...', 'grey', 'on_green'))
num_eval_steps = num_valid // args.batch_size
total_loss = 0
total_time = 0
sess.run(tf.local_variables_initializer())
for i in range(num_eval_steps):
start = time.time()
ids, inputs, npts, gt = next(valid_gen)
feed_dict = {
inputs_pl: inputs,
npts_pl: npts,
gt_pl: gt,
is_training_pl: False
}
loss, _ = sess.run([model.loss, model.update],
feed_dict=feed_dict)
total_loss += loss
total_time += time.time() - start
summary = sess.run(valid_summary,
feed_dict={is_training_pl: False})
writer.add_summary(summary, step)
print(
colored(
'epoch %d step %d loss %.8f - time per batch %.4f' %
(epoch, step, total_loss / num_eval_steps,
total_time / num_eval_steps), 'grey', 'on_green'))
total_time = 0
if step % args.steps_per_visu == 0:
all_pcds = sess.run(model.visualize_ops, feed_dict=feed_dict)
for i in range(0, args.batch_size, args.visu_freq):
plot_path = os.path.join(
args.log_dir, 'plots',
'epoch_%d_step_%d_%s.png' % (epoch, step, ids[i]))
pcds = [x[i] for x in all_pcds]
plot_pcd_three_views(plot_path, pcds,
model.visualize_titles)
if step % args.steps_per_save == 0:
saver.save(sess, os.path.join(args.log_dir, 'model'), step)
print(
colored('Model saved at %s' % args.log_dir, 'white',
'on_blue'))
print('Total time', datetime.timedelta(seconds=time.time() - train_start))
sess.close()
def get_data(args):
# Speicfy batch size just equals to 1
#NOTE: specifically put is_training to false, because if I put
# is_training=True, it will result to double ids in an epoch for some reason idk
df_train, num_train = lmdb_dataflow(args.lmdb_train,
1,
args.num_input_points,
args.num_gt_points,
is_training=False)
train_gen = df_train.get_data()
df_valid, num_valid = lmdb_dataflow(args.lmdb_valid,
1,
args.num_input_points,
args.num_gt_points,
is_training=False)
valid_gen = df_valid.get_data()
print('==================================================================')
print('We have {} train files and {} val files'.format(
num_train, num_valid))
print('==================================================================')
###########################################################################
# SAVE TRAINING
id_train_list = []
for step in tqdm(range(num_train)):
ids, inputs, npts, gt = next(train_gen)
# Sanity check
assert inputs.shape[1] == npts[0], 'number of points do not match'
# Save input
with open(
os.path.join(SAVE_INPUT_PATH_TRAIN, '{:08d}.npy'.format(step)),
'wb') as f:
np.save(f, inputs.reshape(-1, 3))
# Save gt
with open(os.path.join(SAVE_GT_PATH_TRAIN, '{:08d}.npy'.format(step)),
'wb') as f:
np.save(f, gt.reshape(-1, 3))
#assert ids[0] not in id_train_list, 'Double id at step {}, id {}'.format(step, ids[0])
if ids[0] in id_train_list:
print('step {} is the same with step {}'.format(
step, id_train_list.index(ids[0])))
id_train_list.append(ids[0])
print('==================================================================')
print('Finish saving {} train files'.format(num_train))
print('==================================================================')
###########################################################################
# SAVE VALIDATION
id_val_list = []
for step in tqdm(range(num_valid)):
ids, inputs, npts, gt = next(valid_gen)
# Sanity check
assert inputs.shape[1] == npts[0], 'number of points do not match'
print(inputs.shape, gt.shape)
# Save input
with open(os.path.join(SAVE_INPUT_PATH_VAL, '{:08d}.npy'.format(step)),
'wb') as f:
np.save(f, inputs.reshape(-1, 3))
# Save gt
with open(os.path.join(SAVE_GT_PATH_VAL, '{:08d}.npy'.format(step)),
'wb') as f:
np.save(f, gt.reshape(-1, 3))
assert ids[0] not in id_val_list, 'Double id at step {}, id {}'.format(
step, ids[0])
id_val_list.append(ids[0])
print('==================================================================')
print('Finish saving {} val files'.format(num_valid))
print('==================================================================')
def get_transformed_test_data(
args,
xmax=40.,
zmax=40.,
):
"""Make and save a test data for PCN:
apply random y rotation & random xz translation
Save the point cloud, and also the pose"""
# Speicfy batch size just equals to 1
df_valid, num_valid = lmdb_dataflow(args.lmdb_valid,
1,
args.num_input_points,
args.num_gt_points,
is_training=False)
valid_gen = df_valid.get_data()
print('==================================================================')
print('We have {} val files'.format(num_valid))
print('==================================================================')
###########################################################################
# SAVE VALIDATION
id_val_list = []
for step in tqdm(range(num_valid)):
ids, inputs, npts, gt = next(valid_gen)
# inputs, gt ---- (1,N_input,3), (1,N_gt,3)
# Sanity check
assert inputs.shape[1] == npts[0], 'number of points do not match'
# (1,?,3) -> (?,3)
inputs = inputs.reshape(-1, 3)
gt = gt.reshape(-1, 3)
print(inputs.shape, gt.shape)
# Random y-rotation
ang = np.random.uniform(low=0., high=2. * np.pi)
R = np.eye(3)
# First col (x)
R[0, 0] = np.cos(ang)
R[2, 0] = -np.sin(ang)
# Third col (z)
R[0, 2] = np.sin(ang)
R[2, 2] = np.cos(ang)
# Note: before multiplied by rotation need to be (3,N), then
# move back to (N,3)
gt_transformed = (R @ gt.transpose()).transpose()
inputs_transformed = (R @ inputs.transpose()).transpose()
# Random x-ztranslation
xtrans = np.random.uniform(low=-xmax, high=xmax)
ztrans = np.random.uniform(low=-zmax, high=zmax)
gt_transformed[:, 0] = gt_transformed[:, 0] + xtrans
gt_transformed[:, 2] = gt_transformed[:, 2] + ztrans
inputs_transformed[:, 0] = inputs_transformed[:, 0] + xtrans
inputs_transformed[:, 2] = inputs_transformed[:, 2] + ztrans
# The pose
pose = np.array([xtrans, 0, ztrans, ang])
# Save input
with open(
os.path.join(SAVE_INPUT_PATH_VAL_TRANSFORMED,
'{:08d}.npy'.format(step)), 'wb') as f:
np.save(f, inputs_transformed)
# Save original gt
with open(
os.path.join(SAVE_GT_PATH_VAL_ORIGINAL,
'{:08d}.npy'.format(step)), 'wb') as f:
np.save(f, gt)
# Save gt
with open(
os.path.join(SAVE_GT_PATH_VAL_TRANSFORMED,
'{:08d}.npy'.format(step)), 'wb') as f:
np.save(f, gt_transformed)
# Save pose
with open(os.path.join(SAVE_POSE_PATH, '{:08d}.npy'.format(step)),
'wb') as f:
np.save(f, pose)
assert ids[0] not in id_val_list, 'Double id at step {}, id {}'.format(
step, ids[0])
id_val_list.append(ids[0])
print('==================================================================')
print('Finish saving {} val files'.format(num_valid))
print('==================================================================')
f.close()
print("=== ModelNet40 %s %s Done ===\n" % (t, res))
elif args.dataset == 'shapenet8':
print("\n\n=== ShapeNet8 ===\n")
for t in ['train', 'valid']:
sum_dict = json.loads(
open(os.path.join(SHAPENET8_PATH, 'keys.json')).read())
for key in sum_dict.keys():
sum_dict[key] = np.zeros(
3) # num of objects, num of points, average
# the data stored in the lmdb files is with varying number of points
df, num = lmdb_dataflow(lmdb_path=os.path.join(
SHAPENET8_PATH, '%s.lmdb' % t),
batch_size=1,
input_size=1000000,
output_size=1,
is_training=False)
data_gen = df.get_data()
for _ in tqdm(range(num)):
ids, _, npts, _ = next(data_gen)
model_name = ids[0][:8]
sum_dict[model_name][1] += npts[0]
sum_dict[model_name][0] += 1
sum_dict[model_name][
2] = sum_dict[model_name][1] / sum_dict[model_name][0]
f = open("./dump_sum_points/shapenet8_%s.json" % t, "w+")
for key in sum_dict.keys():
parser.add_argument("--hdf5_path",
type=str,
default=r'../data/modelnet40_pcn/hdf5_partial_1024')
parser.add_argument("--partial",
action='store_true',
help='store partial scan or not')
parser.add_argument('--num_per_obj', type=int, default=1024)
parser.add_argument('--num_scan', type=int, default=10)
args = parser.parse_args()
lmdb_file = os.path.join(args.lmdb_path, args.f_name + '.lmdb')
os.system('mkdir -p %s' % args.hdf5_path)
df_train, num_train = lmdb_dataflow(lmdb_path=lmdb_file,
batch_size=1,
input_size=args.num_per_obj,
output_size=args.num_per_obj,
is_training=False)
if args.partial:
print('Now we generate point cloud from partial observed objects.')
file_per_h5 = 2048 * 4 # of objects within each hdf5 file
data_gen = df_train.get_data()
idx = 0
data_np = np.zeros((file_per_h5, args.num_per_obj, 3))
label_np = np.zeros((file_per_h5, ), dtype=np.int32)
ids_np = np.chararray((file_per_h5, ), itemsize=32)
# convert label string to integers
def get_transformed_test_data(
args,
xmax=40.,
zmax=40.,
):
"""Make and save a test data for PCN:
apply random y rotation & random xz translation
Save the point cloud, and also the pose"""
# Speicfy batch size just equals to 1
df_valid, num_valid = lmdb_dataflow(args.lmdb_valid,
1,
args.num_input_points,
args.num_gt_points,
is_training=False)
valid_gen = df_valid.get_data()
posename = 'pose1'
posename = 'pose_trans'
posename = 'pose_rot'
print('==================================================================')
print('We have {} val files'.format(num_valid))
print('==================================================================')
###########################################################################
# SAVE VALIDATION
id_val_list = []
for step in tqdm(range(num_valid)):
ids, inputs, npts, gt = next(valid_gen)
# inputs, gt ---- (1,N_input,3), (1,N_gt,3)
# Sanity check
assert inputs.shape[1] == npts[0], 'number of points do not match'
# (1,?,3) -> (?,3)
inputs = inputs.reshape(-1, 3)
gt = gt.reshape(-1, 3)
print(inputs.shape, gt.shape)
# TODO
save_pose_path = '/home/josephinemonica/Documents/gpu_link/joint_pose_and_shape_estimation/data/data-PCN/my_data'
save_pose_path = os.path.join(save_pose_path, posename)
with open(os.path.join(save_pose_path, '{:08d}.npy'.format(step)),
'rb') as f:
pose_random = np.load(f)
# Random y-rotation
ang = pose_random[3]
R = np.eye(3)
# First col (x)
R[0, 0] = np.cos(ang)
R[2, 0] = -np.sin(ang)
# Third col (z)
R[0, 2] = np.sin(ang)
R[2, 2] = np.cos(ang)
# Note: before multiplied by rotation need to be (3,N), then
# move back to (N,3)
gt_transformed = (R @ gt.transpose()).transpose()
inputs_transformed = (R @ inputs.transpose()).transpose()
# Random x-ztranslation
xtrans = pose_random[0]
ztrans = pose_random[2]
gt_transformed[:, 0] = gt_transformed[:, 0] + xtrans
gt_transformed[:, 2] = gt_transformed[:, 2] + ztrans
inputs_transformed[:, 0] = inputs_transformed[:, 0] + xtrans
inputs_transformed[:, 2] = inputs_transformed[:, 2] + ztrans
# Save input
with open(
os.path.join(SAVE_INPUT_PATH_VAL_TRANSFORMED,
'{:08d}.npy'.format(step)), 'wb') as f:
np.save(f, inputs_transformed)
# Save original gt
with open(
os.path.join(SAVE_GT_PATH_VAL_ORIGINAL,
'{:08d}.npy'.format(step)), 'wb') as f:
np.save(f, gt)
# Save gt
with open(
os.path.join(SAVE_GT_PATH_VAL_TRANSFORMED,
'{:08d}.npy'.format(step)), 'wb') as f:
np.save(f, gt_transformed)
assert ids[0] not in id_val_list, 'Double id at step {}, id {}'.format(
step, ids[0])
id_val_list.append(ids[0])
print('==================================================================')
print('Finish saving {} val files'.format(num_valid))
print('==================================================================')
