def ops_init():
with tf.device('/gpu:' + str(gpu_to_use)):
pointclouds_ph, input_label_ph = placeholder_inputs()
is_training_ph = tf.placeholder(tf.bool, shape=())
# simple model
pred, seg_pred, end_points = model.get_model(pointclouds_ph, input_label_ph, \
cat_num=NUM_OBJ_CATS, part_num=NUM_PART_CATS, is_training=is_training_ph, \
batch_size=batch_size, num_point=point_num, weight_decay=0.0, bn_decay=None)
ops = {}
ops['pred'] = pred
ops['seg_pred'] = seg_pred
return ops
def predict():
is_training = False
with tf.device('/gpu:'+str(gpu_to_use)):
pointclouds_ph, input_label_ph = placeholder_inputs()
is_training_ph = tf.placeholder(tf.bool, shape=())
# simple model
pred, seg_pred, end_points = model.get_model(pointclouds_ph, input_label_ph, \
cat_num=NUM_OBJ_CATS, part_num=NUM_PART_CATS, is_training=is_training_ph, \
batch_size=batch_size, num_point=point_num, weight_decay=0.0, bn_decay=None)
# Add ops to save and restore all the variables.
saver = tf.train.Saver()
# Later, launch the model, use the saver to restore variables from disk, and
# do some work with the model.
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
with tf.Session(config=config) as sess:
if not os.path.exists(output_dir):
os.mkdir(output_dir)
flog = open(os.path.join(output_dir, 'log.txt'), 'w')
# Restore variables from disk.
printout(flog, 'Loading model %s' % pretrained_model_path)
saver.restore(sess, pretrained_model_path)
printout(flog, 'Model restored.')
# Note: the evaluation for the model with BN has to have some statistics
# Using some test datas as the statistics
batch_data = np.zeros([batch_size, point_num, 3]).astype(np.float32)
total_acc = 0.0
total_seen = 0
total_acc_iou = 0.0
total_per_cat_acc = np.zeros((NUM_OBJ_CATS)).astype(np.float32)
total_per_cat_iou = np.zeros((NUM_OBJ_CATS)).astype(np.float32)
total_per_cat_seen = np.zeros((NUM_OBJ_CATS)).astype(np.int32)
ffiles = open(test_file_list, 'r')
files = [line.rstrip() for line in ffiles.readlines()]
ffiles.close()
len_pts_files = len(files)
data, labels, segs = provider.load_h5_data_label_seg(
os.path.join(hdf5_data_dir, files[0]))
for i in range(len(data)):
if i % 100 == 0:
printout(flog, '%d/%d ...' % (i, len(data)))
cur_gt_label = labels[i]
cur_label_one_hot = np.zeros((1, NUM_OBJ_CATS), dtype=np.float32)
cur_label_one_hot[0, cur_gt_label] = 1
# pts_file_to_load = os.path.join(ply_data_dir, pts_files[shape_idx])
# seg_file_to_load = os.path.join(ply_data_dir, seg_files[shape_idx])
pts = data[i]
seg = segs[i]
ori_point_num = len(seg)
batch_data[0, ...] = pc_augment_to_point_num(pc_normalize(pts), point_num)
label_pred_val, seg_pred_res = sess.run([pred, seg_pred], feed_dict={
pointclouds_ph: batch_data,
input_label_ph: cur_label_one_hot,
is_training_ph: is_training,
})
label_pred_val = np.argmax(label_pred_val[0, :])
seg_pred_res = seg_pred_res[0, ...]
iou_oids = object2setofoid[objcats[cur_gt_label]]
non_cat_labels = list(set(np.arange(NUM_PART_CATS)).difference(set(iou_oids)))
mini = np.min(seg_pred_res)
seg_pred_res[:, non_cat_labels] = mini - 1000
seg_pred_val = np.argmax(seg_pred_res, axis=1)[:ori_point_num]
seg_acc = np.mean(seg_pred_val == seg)
total_acc += seg_acc
total_seen += 1
total_per_cat_seen[cur_gt_label] += 1
total_per_cat_acc[cur_gt_label] += seg_acc
mask = np.int32(seg_pred_val == seg)
total_iou = 0.0
iou_log = ''
for oid in iou_oids:
n_pred = np.sum(seg_pred_val == oid)
n_gt = np.sum(seg == oid)
n_intersect = np.sum(np.int32(seg == oid) * mask)
n_union = n_pred + n_gt - n_intersect
iou_log += '_' + str(n_pred)+'_'+str(n_gt)+'_'+str(n_intersect)+'_'+str(n_union)+'_'
if n_union == 0:
total_iou += 1
iou_log += '_1\n'
else:
total_iou += n_intersect * 1.0 / n_union
iou_log += '_'+str(n_intersect * 1.0 / n_union)+'\n'
avg_iou = total_iou / len(iou_oids)
total_acc_iou += avg_iou
total_per_cat_iou[cur_gt_label] += avg_iou
if output_verbose:
output_color_point_cloud(pts, seg, os.path.join(output_dir, str(i)+'_gt.obj'))
output_color_point_cloud(pts, seg_pred_val, os.path.join(output_dir, str(i)+'_pred.obj'))
output_color_point_cloud_red_blue(pts, np.int32(seg == seg_pred_val),
os.path.join(output_dir, str(i)+'_diff.obj'))
with open(os.path.join(output_dir, str(i)+'.log'), 'w') as fout:
fout.write('Total Point: %d\n\n' % ori_point_num)
fout.write('Ground Truth: %s\n' % objnames[cur_gt_label])
fout.write('Predict: %s\n\n' % objnames[label_pred_val])
fout.write('Accuracy: %f\n' % seg_acc)
fout.write('IoU: %f\n\n' % avg_iou)
fout.write('IoU details: %s\n' % iou_log)
printout(flog, 'Accuracy: %f' % (total_acc / total_seen))
printout(flog, 'IoU: %f' % (total_acc_iou / total_seen))
for cat_idx in range(NUM_OBJ_CATS):
printout(flog, '\t ' + objcats[cat_idx] + ' Total Number: ' + str(total_per_cat_seen[cat_idx]))
if total_per_cat_seen[cat_idx] > 0:
printout(flog, '\t ' + objcats[cat_idx] + ' Accuracy: ' + \
str(total_per_cat_acc[cat_idx] / total_per_cat_seen[cat_idx]))
printout(flog, '\t ' + objcats[cat_idx] + ' IoU: '+ \
str(total_per_cat_iou[cat_idx] / total_per_cat_seen[cat_idx]))
def train():
with tf.Graph().as_default():
with tf.device('/gpu:'+str(FLAGS.gpu)):
pointclouds_ph, input_label_ph, labels_ph, seg_ph = placeholder_inputs()
is_training_ph = tf.placeholder(tf.bool, shape=())
batch = tf.Variable(0, trainable=False)
learning_rate = tf.train.exponential_decay(
BASE_LEARNING_RATE, # base learning rate
batch * batch_size, # global_var indicating the number of steps
DECAY_STEP, # step size
DECAY_RATE, # decay rate
staircase=True # Stair-case or continuous decreasing
)
learning_rate = tf.maximum(learning_rate, LEARNING_RATE_CLIP)
bn_momentum = tf.train.exponential_decay(
BN_INIT_DECAY,
batch*batch_size,
BN_DECAY_DECAY_STEP,
BN_DECAY_DECAY_RATE,
staircase=True)
bn_decay = tf.minimum(BN_DECAY_CLIP, 1 - bn_momentum)
lr_op = tf.summary.scalar('learning_rate', learning_rate)
batch_op = tf.summary.scalar('batch_number', batch)
bn_decay_op = tf.summary.scalar('bn_decay', bn_decay)
labels_pred, seg_pred, end_points = model.get_model(pointclouds_ph, input_label_ph, \
is_training=is_training_ph, bn_decay=bn_decay, cat_num=NUM_CATEGORIES, \
part_num=NUM_PART_CATS, batch_size=batch_size, num_point=point_num, weight_decay=FLAGS.wd)
# model.py defines both classification net and segmentation net, which share the common global feature extractor network.
# In model.get_loss, we define the total loss to be weighted sum of the classification and segmentation losses.
# Here, we only train for segmentation network. Thus, we set weight to be 1.0.
loss, label_loss, per_instance_label_loss, seg_loss, per_instance_seg_loss, per_instance_seg_pred_res \
= model.get_loss(labels_pred, seg_pred, labels_ph, seg_ph, 1.0, end_points)
total_training_loss_ph = tf.placeholder(tf.float32, shape=())
total_testing_loss_ph = tf.placeholder(tf.float32, shape=())
label_training_loss_ph = tf.placeholder(tf.float32, shape=())
label_testing_loss_ph = tf.placeholder(tf.float32, shape=())
seg_training_loss_ph = tf.placeholder(tf.float32, shape=())
seg_testing_loss_ph = tf.placeholder(tf.float32, shape=())
label_training_acc_ph = tf.placeholder(tf.float32, shape=())
label_testing_acc_ph = tf.placeholder(tf.float32, shape=())
label_testing_acc_avg_cat_ph = tf.placeholder(tf.float32, shape=())
seg_training_acc_ph = tf.placeholder(tf.float32, shape=())
seg_testing_acc_ph = tf.placeholder(tf.float32, shape=())
seg_testing_acc_avg_cat_ph = tf.placeholder(tf.float32, shape=())
total_train_loss_sum_op = tf.summary.scalar('total_training_loss', total_training_loss_ph)
total_test_loss_sum_op = tf.summary.scalar('total_testing_loss', total_testing_loss_ph)
label_train_loss_sum_op = tf.summary.scalar('label_training_loss', label_training_loss_ph)
label_test_loss_sum_op = tf.summary.scalar('label_testing_loss', label_testing_loss_ph)
seg_train_loss_sum_op = tf.summary.scalar('seg_training_loss', seg_training_loss_ph)
seg_test_loss_sum_op = tf.summary.scalar('seg_testing_loss', seg_testing_loss_ph)
label_train_acc_sum_op = tf.summary.scalar('label_training_acc', label_training_acc_ph)
label_test_acc_sum_op = tf.summary.scalar('label_testing_acc', label_testing_acc_ph)
label_test_acc_avg_cat_op = tf.summary.scalar('label_testing_acc_avg_cat', label_testing_acc_avg_cat_ph)
seg_train_acc_sum_op = tf.summary.scalar('seg_training_acc', seg_training_acc_ph)
seg_test_acc_sum_op = tf.summary.scalar('seg_testing_acc', seg_testing_acc_ph)
seg_test_acc_avg_cat_op = tf.summary.scalar('seg_testing_acc_avg_cat', seg_testing_acc_avg_cat_ph)
train_variables = tf.trainable_variables()
trainer = tf.train.AdamOptimizer(learning_rate)
train_op = trainer.minimize(loss, var_list=train_variables, global_step=batch)
saver = tf.train.Saver()
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)
train_writer = tf.summary.FileWriter(SUMMARIES_FOLDER + '/train', sess.graph)
test_writer = tf.summary.FileWriter(SUMMARIES_FOLDER + '/test')
train_file_list = provider.getDataFiles(TRAINING_FILE_LIST)
num_train_file = len(train_file_list)
test_file_list = provider.getDataFiles(TESTING_FILE_LIST)
num_test_file = len(test_file_list)
fcmd = open(os.path.join(LOG_STORAGE_PATH, 'cmd.txt'), 'w')
fcmd.write(str(FLAGS))
fcmd.close()
# write logs to the disk
flog = open(os.path.join(LOG_STORAGE_PATH, 'log.txt'), 'w')
def train_one_epoch(train_file_idx, epoch_num):
is_training = True
for i in range(num_train_file):
cur_train_filename = os.path.join(hdf5_data_dir, train_file_list[train_file_idx[i]])
printout(flog, 'Loading train file ' + cur_train_filename)
cur_data, cur_labels, cur_seg = provider.loadDataFile_with_seg(cur_train_filename)
cur_data, cur_labels, order = provider.shuffle_data(cur_data, np.squeeze(cur_labels))
cur_seg = cur_seg[order, ...]
cur_labels_one_hot = convert_label_to_one_hot(cur_labels)
num_data = len(cur_labels)
num_batch = num_data // batch_size
total_loss = 0.0
total_label_loss = 0.0
total_seg_loss = 0.0
total_label_acc = 0.0
total_seg_acc = 0.0
for j in range(num_batch):
begidx = j * batch_size
endidx = (j + 1) * batch_size
feed_dict = {
pointclouds_ph: cur_data[begidx: endidx, ...],
labels_ph: cur_labels[begidx: endidx, ...],
input_label_ph: cur_labels_one_hot[begidx: endidx, ...],
seg_ph: cur_seg[begidx: endidx, ...],
is_training_ph: is_training,
}
_, loss_val, label_loss_val, seg_loss_val, per_instance_label_loss_val, \
per_instance_seg_loss_val, label_pred_val, seg_pred_val, pred_seg_res \
= sess.run([train_op, loss, label_loss, seg_loss, per_instance_label_loss, \
per_instance_seg_loss, labels_pred, seg_pred, per_instance_seg_pred_res], \
feed_dict=feed_dict)
per_instance_part_acc = np.mean(pred_seg_res == cur_seg[begidx: endidx, ...], axis=1)
average_part_acc = np.mean(per_instance_part_acc)
total_loss += loss_val
total_label_loss += label_loss_val
total_seg_loss += seg_loss_val
per_instance_label_pred = np.argmax(label_pred_val, axis=1)
total_label_acc += np.mean(np.float32(per_instance_label_pred == cur_labels[begidx: endidx, ...]))
total_seg_acc += average_part_acc
total_loss = total_loss * 1.0 / num_batch
total_label_loss = total_label_loss * 1.0 / num_batch
total_seg_loss = total_seg_loss * 1.0 / num_batch
total_label_acc = total_label_acc * 1.0 / num_batch
total_seg_acc = total_seg_acc * 1.0 / num_batch
lr_sum, bn_decay_sum, batch_sum, train_loss_sum, train_label_acc_sum, \
train_label_loss_sum, train_seg_loss_sum, train_seg_acc_sum = sess.run(\
[lr_op, bn_decay_op, batch_op, total_train_loss_sum_op, label_train_acc_sum_op, \
label_train_loss_sum_op, seg_train_loss_sum_op, seg_train_acc_sum_op], \
feed_dict={total_training_loss_ph: total_loss, label_training_loss_ph: total_label_loss, \
seg_training_loss_ph: total_seg_loss, label_training_acc_ph: total_label_acc, \
seg_training_acc_ph: total_seg_acc})
train_writer.add_summary(train_loss_sum, i + epoch_num * num_train_file)
train_writer.add_summary(train_label_loss_sum, i + epoch_num * num_train_file)
train_writer.add_summary(train_seg_loss_sum, i + epoch_num * num_train_file)
train_writer.add_summary(lr_sum, i + epoch_num * num_train_file)
train_writer.add_summary(bn_decay_sum, i + epoch_num * num_train_file)
train_writer.add_summary(train_label_acc_sum, i + epoch_num * num_train_file)
train_writer.add_summary(train_seg_acc_sum, i + epoch_num * num_train_file)
train_writer.add_summary(batch_sum, i + epoch_num * num_train_file)
printout(flog, '\tTraining Total Mean_loss: %f' % total_loss)
printout(flog, '\t\tTraining Label Mean_loss: %f' % total_label_loss)
printout(flog, '\t\tTraining Label Accuracy: %f' % total_label_acc)
printout(flog, '\t\tTraining Seg Mean_loss: %f' % total_seg_loss)
printout(flog, '\t\tTraining Seg Accuracy: %f' % total_seg_acc)
def eval_one_epoch(epoch_num):
is_training = False
total_loss = 0.0
total_label_loss = 0.0
total_seg_loss = 0.0
total_label_acc = 0.0
total_seg_acc = 0.0
total_seen = 0
total_label_acc_per_cat = np.zeros((NUM_CATEGORIES)).astype(np.float32)
total_seg_acc_per_cat = np.zeros((NUM_CATEGORIES)).astype(np.float32)
total_seen_per_cat = np.zeros((NUM_CATEGORIES)).astype(np.int32)
for i in range(num_test_file):
cur_test_filename = os.path.join(hdf5_data_dir, test_file_list[i])
printout(flog, 'Loading test file ' + cur_test_filename)
cur_data, cur_labels, cur_seg = provider.loadDataFile_with_seg(cur_test_filename)
cur_labels = np.squeeze(cur_labels)
cur_labels_one_hot = convert_label_to_one_hot(cur_labels)
num_data = len(cur_labels)
num_batch = num_data // batch_size
for j in range(num_batch):
begidx = j * batch_size
endidx = (j + 1) * batch_size
feed_dict = {
pointclouds_ph: cur_data[begidx: endidx, ...],
labels_ph: cur_labels[begidx: endidx, ...],
input_label_ph: cur_labels_one_hot[begidx: endidx, ...],
seg_ph: cur_seg[begidx: endidx, ...],
is_training_ph: is_training,
}
loss_val, label_loss_val, seg_loss_val, per_instance_label_loss_val, \
per_instance_seg_loss_val, label_pred_val, seg_pred_val, pred_seg_res \
= sess.run([loss, label_loss, seg_loss, per_instance_label_loss, \
per_instance_seg_loss, labels_pred, seg_pred, per_instance_seg_pred_res], \
feed_dict=feed_dict)
per_instance_part_acc = np.mean(pred_seg_res == cur_seg[begidx: endidx, ...], axis=1)
average_part_acc = np.mean(per_instance_part_acc)
total_seen += 1
total_loss += loss_val
total_label_loss += label_loss_val
total_seg_loss += seg_loss_val
per_instance_label_pred = np.argmax(label_pred_val, axis=1)
total_label_acc += np.mean(np.float32(per_instance_label_pred == cur_labels[begidx: endidx, ...]))
total_seg_acc += average_part_acc
for shape_idx in range(begidx, endidx):
total_seen_per_cat[cur_labels[shape_idx]] += 1
total_label_acc_per_cat[cur_labels[shape_idx]] += np.int32(per_instance_label_pred[shape_idx-begidx] == cur_labels[shape_idx])
total_seg_acc_per_cat[cur_labels[shape_idx]] += per_instance_part_acc[shape_idx - begidx]
total_loss = total_loss * 1.0 / total_seen
total_label_loss = total_label_loss * 1.0 / total_seen
total_seg_loss = total_seg_loss * 1.0 / total_seen
total_label_acc = total_label_acc * 1.0 / total_seen
total_seg_acc = total_seg_acc * 1.0 / total_seen
test_loss_sum, test_label_acc_sum, test_label_loss_sum, test_seg_loss_sum, test_seg_acc_sum = sess.run(\
[total_test_loss_sum_op, label_test_acc_sum_op, label_test_loss_sum_op, seg_test_loss_sum_op, seg_test_acc_sum_op], \
feed_dict={total_testing_loss_ph: total_loss, label_testing_loss_ph: total_label_loss, \
seg_testing_loss_ph: total_seg_loss, label_testing_acc_ph: total_label_acc, seg_testing_acc_ph: total_seg_acc})
test_writer.add_summary(test_loss_sum, (epoch_num+1) * num_train_file-1)
test_writer.add_summary(test_label_loss_sum, (epoch_num+1) * num_train_file-1)
test_writer.add_summary(test_seg_loss_sum, (epoch_num+1) * num_train_file-1)
test_writer.add_summary(test_label_acc_sum, (epoch_num+1) * num_train_file-1)
test_writer.add_summary(test_seg_acc_sum, (epoch_num+1) * num_train_file-1)
printout(flog, '\tTesting Total Mean_loss: %f' % total_loss)
printout(flog, '\t\tTesting Label Mean_loss: %f' % total_label_loss)
printout(flog, '\t\tTesting Label Accuracy: %f' % total_label_acc)
printout(flog, '\t\tTesting Seg Mean_loss: %f' % total_seg_loss)
printout(flog, '\t\tTesting Seg Accuracy: %f' % total_seg_acc)
for cat_idx in range(NUM_CATEGORIES):
if total_seen_per_cat[cat_idx] > 0:
printout(flog, '\n\t\tCategory %s Object Number: %d' % (all_obj_cats[cat_idx][0], total_seen_per_cat[cat_idx]))
printout(flog, '\t\tCategory %s Label Accuracy: %f' % (all_obj_cats[cat_idx][0], total_label_acc_per_cat[cat_idx]/total_seen_per_cat[cat_idx]))
printout(flog, '\t\tCategory %s Seg Accuracy: %f' % (all_obj_cats[cat_idx][0], total_seg_acc_per_cat[cat_idx]/total_seen_per_cat[cat_idx]))
if not os.path.exists(MODEL_STORAGE_PATH):
os.mkdir(MODEL_STORAGE_PATH)
for epoch in range(TRAINING_EPOCHES):
printout(flog, '\n<<< Testing on the test dataset ...')
eval_one_epoch(epoch)
printout(flog, '\n>>> Training for the epoch %d/%d ...' % (epoch, TRAINING_EPOCHES))
train_file_idx = np.arange(0, len(train_file_list))
np.random.shuffle(train_file_idx)
train_one_epoch(train_file_idx, epoch)
if (epoch+1) % 10 == 0:
cp_filename = saver.save(sess, os.path.join(MODEL_STORAGE_PATH, 'epoch_' + str(epoch+1)+'.ckpt'))
printout(flog, 'Successfully store the checkpoint model into ' + cp_filename)
flog.flush()
flog.close()
def predict():
is_training = False
with tf.device('/gpu:'+str(gpu_to_use)):
pointclouds_ph, input_label_ph = placeholder_inputs()
is_training_ph = tf.placeholder(tf.bool, shape=())
# simple model
pred, seg_pred, end_points = model.get_model(pointclouds_ph, input_label_ph, \
cat_num=NUM_OBJ_CATS, part_num=NUM_PART_CATS, is_training=is_training_ph, \
batch_size=batch_size, num_point=point_num, weight_decay=0.0, bn_decay=None)
# Add ops to save and restore all the variables.
saver = tf.train.Saver()
# Later, launch the model, use the saver to restore variables from disk, and
# do some work with the model.
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
with tf.Session(config=config) as sess:
if not os.path.exists(output_dir):
os.mkdir(output_dir)
flog = open(os.path.join(output_dir, 'log.txt'), 'w')
# Restore variables from disk.
printout(flog, 'Loading model %s' % pretrained_model_path)
saver.restore(sess, pretrained_model_path)
printout(flog, 'Model restored.')
# Note: the evaluation for the model with BN has to have some statistics
# Using some test datas as the statistics
batch_data = np.zeros([batch_size, point_num, 3]).astype(np.float32)
total_acc = 0.0
total_seen = 0
total_acc_iou = 0.0
total_per_cat_acc = np.zeros((NUM_OBJ_CATS)).astype(np.float32)
total_per_cat_iou = np.zeros((NUM_OBJ_CATS)).astype(np.float32)
total_per_cat_seen = np.zeros((NUM_OBJ_CATS)).astype(np.int32)
ffiles = open(test_file_list, 'r')
lines = [line.rstrip() for line in ffiles.readlines()]
pts_files = [line.split()[0] for line in lines]
seg_files = [line.split()[1] for line in lines]
labels = [line.split()[2] for line in lines]
ffiles.close()
len_pts_files = len(pts_files)
for shape_idx in range(len_pts_files):
if shape_idx % 100 == 0:
printout(flog, '%d/%d ...' % (shape_idx, len_pts_files))
cur_gt_label = on2oid[labels[shape_idx]]
cur_label_one_hot = np.zeros((1, NUM_OBJ_CATS), dtype=np.float32)
cur_label_one_hot[0, cur_gt_label] = 1
pts_file_to_load = os.path.join(ply_data_dir, pts_files[shape_idx])
seg_file_to_load = os.path.join(ply_data_dir, seg_files[shape_idx])
pts, seg = load_pts_seg_files(pts_file_to_load, seg_file_to_load, objcats[cur_gt_label])
ori_point_num = len(seg)
batch_data[0, ...] = pc_augment_to_point_num(pc_normalize(pts), point_num)
label_pred_val, seg_pred_res = sess.run([pred, seg_pred], feed_dict={
pointclouds_ph: batch_data,
input_label_ph: cur_label_one_hot,
is_training_ph: is_training,
})
label_pred_val = np.argmax(label_pred_val[0, :])
seg_pred_res = seg_pred_res[0, ...]
iou_oids = object2setofoid[objcats[cur_gt_label]]
non_cat_labels = list(set(np.arange(NUM_PART_CATS)).difference(set(iou_oids)))
mini = np.min(seg_pred_res)
seg_pred_res[:, non_cat_labels] = mini - 1000
seg_pred_val = np.argmax(seg_pred_res, axis=1)[:ori_point_num]
seg_acc = np.mean(seg_pred_val == seg)
total_acc += seg_acc
total_seen += 1
total_per_cat_seen[cur_gt_label] += 1
total_per_cat_acc[cur_gt_label] += seg_acc
mask = np.int32(seg_pred_val == seg)
total_iou = 0.0
iou_log = ''
for oid in iou_oids:
n_pred = np.sum(seg_pred_val == oid)
n_gt = np.sum(seg == oid)
n_intersect = np.sum(np.int32(seg == oid) * mask)
n_union = n_pred + n_gt - n_intersect
iou_log += '_' + str(n_pred)+'_'+str(n_gt)+'_'+str(n_intersect)+'_'+str(n_union)+'_'
if n_union == 0:
total_iou += 1
iou_log += '_1\n'
else:
total_iou += n_intersect * 1.0 / n_union
iou_log += '_'+str(n_intersect * 1.0 / n_union)+'\n'
avg_iou = total_iou / len(iou_oids)
total_acc_iou += avg_iou
total_per_cat_iou[cur_gt_label] += avg_iou
if output_verbose:
output_color_point_cloud(pts, seg, os.path.join(output_dir, str(shape_idx)+'_gt.obj'))
output_color_point_cloud(pts, seg_pred_val, os.path.join(output_dir, str(shape_idx)+'_pred.obj'))
output_color_point_cloud_red_blue(pts, np.int32(seg == seg_pred_val),
os.path.join(output_dir, str(shape_idx)+'_diff.obj'))
with open(os.path.join(output_dir, str(shape_idx)+'.log'), 'w') as fout:
fout.write('Total Point: %d\n\n' % ori_point_num)
fout.write('Ground Truth: %s\n' % objnames[cur_gt_label])
fout.write('Predict: %s\n\n' % objnames[label_pred_val])
fout.write('Accuracy: %f\n' % seg_acc)
fout.write('IoU: %f\n\n' % avg_iou)
fout.write('IoU details: %s\n' % iou_log)
printout(flog, 'Accuracy: %f' % (total_acc / total_seen))
printout(flog, 'IoU: %f' % (total_acc_iou / total_seen))
for cat_idx in range(NUM_OBJ_CATS):
printout(flog, '\t ' + objcats[cat_idx] + ' Total Number: ' + str(total_per_cat_seen[cat_idx]))
if total_per_cat_seen[cat_idx] > 0:
printout(flog, '\t ' + objcats[cat_idx] + ' Accuracy: ' + \
str(total_per_cat_acc[cat_idx] / total_per_cat_seen[cat_idx]))
printout(flog, '\t ' + objcats[cat_idx] + ' IoU: '+ \
str(total_per_cat_iou[cat_idx] / total_per_cat_seen[cat_idx]))
def train():
with tf.Graph().as_default():
with tf.device('/gpu:' + str(FLAGS.gpu)):
pointclouds_ph, input_label_ph, labels_ph, seg_ph, pairwise_distances_ph = placeholder_inputs(
)
is_training_ph = tf.placeholder(tf.bool, shape=())
batch = tf.Variable(0, trainable=False)
learning_rate = tf.train.exponential_decay(
BASE_LEARNING_RATE, # base learning rate
batch *
batch_size, # global_var indicating the number of steps
DECAY_STEP, # step size
DECAY_RATE, # decay rate
staircase=True # Stair-case or continuous decreasing
)
learning_rate = tf.maximum(learning_rate, LEARNING_RATE_CLIP)
bn_momentum = tf.train.exponential_decay(BN_INIT_DECAY,
batch * batch_size,
BN_DECAY_DECAY_STEP,
BN_DECAY_DECAY_RATE,
staircase=True)
bn_decay = tf.minimum(BN_DECAY_CLIP, 1 - bn_momentum)
lr_op = tf.summary.scalar('learning_rate', learning_rate)
batch_op = tf.summary.scalar('batch_number', batch)
bn_decay_op = tf.summary.scalar('bn_decay', bn_decay)
labels_pred, seg_pred, end_points, feature_output = model.get_model(pointclouds_ph, input_label_ph, \
is_training=is_training_ph, bn_decay=bn_decay, cat_num=NUM_CATEGORIES, \
part_num=NUM_PART_CATS, batch_size=batch_size, num_point=point_num, weight_decay=FLAGS.wd)
#labels_pred, seg_pred, end_points = model.get_last_layers(labels_pred,seg_pred,end_points,part_num=NUM_PART_CATS, batch_size=batch_size, num_point=point_num,weight_decay=FLAGS.wd)
# model.py defines both classification net and segmentation net, which share the common global feature extractor network.
# In model.get_loss, we define the total loss to be weighted sum of the classification and segmentation losses.
# Here, we only train for segmentation network. Thus, we set weight to be 1.0.
loss, label_loss, per_instance_label_loss, seg_loss, per_instance_seg_loss, per_instance_seg_pred_res, debug_pred, debug_label \
= model.get_loss_finetune(labels_pred, seg_pred, labels_ph, seg_ph, 1.0, end_points, pairwise_distances_ph)
total_training_loss_ph = tf.placeholder(tf.float32, shape=())
total_testing_loss_ph = tf.placeholder(tf.float32, shape=())
label_training_loss_ph = tf.placeholder(tf.float32, shape=())
label_testing_loss_ph = tf.placeholder(tf.float32, shape=())
seg_training_loss_ph = tf.placeholder(tf.float32, shape=())
seg_testing_loss_ph = tf.placeholder(tf.float32, shape=())
label_training_acc_ph = tf.placeholder(tf.float32, shape=())
label_testing_acc_ph = tf.placeholder(tf.float32, shape=())
label_testing_acc_avg_cat_ph = tf.placeholder(tf.float32, shape=())
seg_training_acc_ph = tf.placeholder(tf.float32, shape=())
seg_testing_acc_ph = tf.placeholder(tf.float32, shape=())
seg_testing_acc_avg_cat_ph = tf.placeholder(tf.float32, shape=())
total_train_loss_sum_op = tf.summary.scalar(
'total_training_loss', total_training_loss_ph)
total_test_loss_sum_op = tf.summary.scalar('total_testing_loss',
total_testing_loss_ph)
label_train_loss_sum_op = tf.summary.scalar(
'label_training_loss', label_training_loss_ph)
label_test_loss_sum_op = tf.summary.scalar('label_testing_loss',
label_testing_loss_ph)
seg_train_loss_sum_op = tf.summary.scalar('seg_training_loss',
seg_training_loss_ph)
seg_test_loss_sum_op = tf.summary.scalar('seg_testing_loss',
seg_testing_loss_ph)
label_train_acc_sum_op = tf.summary.scalar('label_training_acc',
label_training_acc_ph)
label_test_acc_sum_op = tf.summary.scalar('label_testing_acc',
label_testing_acc_ph)
label_test_acc_avg_cat_op = tf.summary.scalar(
'label_testing_acc_avg_cat', label_testing_acc_avg_cat_ph)
seg_train_acc_sum_op = tf.summary.scalar('seg_training_acc',
seg_training_acc_ph)
seg_test_acc_sum_op = tf.summary.scalar('seg_testing_acc',
seg_testing_acc_ph)
seg_test_acc_avg_cat_op = tf.summary.scalar(
'seg_testing_acc_avg_cat', seg_testing_acc_avg_cat_ph)
train_variables = tf.trainable_variables()
trainer = tf.train.AdamOptimizer(learning_rate)
train_op = trainer.minimize(loss,
var_list=train_variables,
global_step=batch)
variables = tf.contrib.framework.get_variables_to_restore()
print('Number of output variables for last layer {0}'.format(
variables[121].shape[3]))
if NUM_PART_CATS == variables[121].shape[
3] and initialize_last_layer == False:
print("Not setting the last layers weights to random weights")
variables_to_restore = variables
else:
print("Loading all the weights except last layer weights")
#variables_to_restore = [variables[0]]
indices_to_take = range(121)
indices_to_take.append(range(123, 289))
variables_to_restore = variables[0:121]
next_set_variables_to_restore = variables[123:289]
variables_to_restore.extend(next_set_variables_to_restore)
saver = tf.train.Saver(variables_to_restore)
# saver = tf.train.Saver()
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
sess = tf.Session(config=config)
saver.restore(sess, pretrained_model_path)
init = tf.global_variables_initializer()
sess.run(init)
train_writer = tf.summary.FileWriter(SUMMARIES_FOLDER + '/train',
sess.graph)
test_writer = tf.summary.FileWriter(SUMMARIES_FOLDER + '/test')
#train_file_list = provider.getDataFiles(TRAINING_FILE_LIST)
train_file_list = [TRAINING_FILE_LIST]
num_train_file = len(train_file_list)
#test_file_list = provider.getDataFiles(TESTING_FILE_LIST)
#num_test_file = len(test_file_list)
fcmd = open(os.path.join(LOG_STORAGE_PATH, 'cmd.txt'), 'w')
fcmd.write(str(FLAGS))
fcmd.close()
# write logs to the disk
flog = open(os.path.join(LOG_STORAGE_PATH, 'log.txt'), 'w')
def train_one_epoch(train_file_idx, epoch_num):
is_training = True
for i in range(num_train_file):
#cur_train_filename = os.path.join(hdf5_data_dir, train_file_list[train_file_idx[i]])
cur_train_filename = os.path.join(point_cloud_file_path_,
training_data_file_name)
printout(flog, 'Loading train file ' + cur_train_filename)
cur_data, cur_labels, cur_seg = provider.loadDataFile_with_seg(
cur_train_filename)
cur_data, cur_labels, order = provider.shuffle_data(
cur_data, np.squeeze(cur_labels, axis=0))
cur_seg = cur_seg[order, ...]
cur_labels_one_hot = convert_label_to_one_hot(cur_labels)
num_data = len(cur_labels)
num_batch = num_data // batch_size
total_loss = 0.0
total_label_loss = 0.0
total_seg_loss = 0.0
total_label_acc = 0.0
total_seg_acc = 0.0
for j in range(num_batch):
begidx = j * batch_size
endidx = (j + 1) * batch_size
calc_pairwise_distance = calculatePairWiseDistanceEnergies(
cur_data[begidx:endidx, ...])
feed_dict = {
pointclouds_ph: cur_data[begidx:endidx, ...],
labels_ph: cur_labels[begidx:endidx, ...],
input_label_ph: cur_labels_one_hot[begidx:endidx, ...],
seg_ph: cur_seg[begidx:endidx, ...],
is_training_ph: is_training,
pairwise_distances_ph: calc_pairwise_distance
}
_, loss_val, label_loss_val, seg_loss_val, per_instance_label_loss_val, \
per_instance_seg_loss_val, label_pred_val, seg_pred_val, pred_seg_res, feature_output_val\
= sess.run([train_op, loss, label_loss, seg_loss, per_instance_label_loss, \
per_instance_seg_loss, labels_pred, seg_pred, per_instance_seg_pred_res, feature_output], \
feed_dict=feed_dict)
if epoch_num == TRAINING_EPOCHES - 1:
text_file = open(point_cloud_file_path, "r")
input = text_file.read().split('\n')
l = []
for i in range(len(input)):
x = list(input[i].split())
l.append(x)
l = l[:-1]
points = np.array(l, dtype=np.float)
point_data = points[:, 0:3]
seg_pred_output = np.argmax(np.squeeze(seg_pred_val),
axis=-1)
colors = [
label_color_dict[x + 1] for x in seg_pred_output
]
output = np.column_stack(
(np.squeeze(point_data), colors))
np.savetxt(labelled_cloud_file_name, output)
# There is a bug here and I fully don't understand why it is crashing here
per_instance_part_acc = np.mean(
pred_seg_res == cur_seg[begidx:endidx, ...])
average_part_acc = np.mean(per_instance_part_acc)
total_loss += loss_val
total_label_loss += label_loss_val
total_seg_loss += seg_loss_val
per_instance_label_pred = np.argmax(label_pred_val, axis=1)
total_label_acc += np.mean(
np.float32(per_instance_label_pred == cur_labels[
begidx:endidx, ...]))
total_seg_acc += average_part_acc
total_loss = total_loss * 1.0 / num_batch
total_label_loss = total_label_loss * 1.0 / num_batch
total_seg_loss = total_seg_loss * 1.0 / num_batch
total_label_acc = total_label_acc * 1.0 / num_batch
total_seg_acc = total_seg_acc * 1.0 / num_batch
lr_sum, bn_decay_sum, batch_sum, train_loss_sum, train_label_acc_sum, \
train_label_loss_sum, train_seg_loss_sum, train_seg_acc_sum = sess.run(\
[lr_op, bn_decay_op, batch_op, total_train_loss_sum_op, label_train_acc_sum_op, \
label_train_loss_sum_op, seg_train_loss_sum_op, seg_train_acc_sum_op], \
feed_dict={total_training_loss_ph: total_loss, label_training_loss_ph: total_label_loss, \
seg_training_loss_ph: total_seg_loss, label_training_acc_ph: total_label_acc, \
seg_training_acc_ph: total_seg_acc})
train_writer.add_summary(train_loss_sum,
i + epoch_num * num_train_file)
train_writer.add_summary(train_label_loss_sum,
i + epoch_num * num_train_file)
train_writer.add_summary(train_seg_loss_sum,
i + epoch_num * num_train_file)
train_writer.add_summary(lr_sum,
i + epoch_num * num_train_file)
train_writer.add_summary(bn_decay_sum,
i + epoch_num * num_train_file)
train_writer.add_summary(train_label_acc_sum,
i + epoch_num * num_train_file)
train_writer.add_summary(train_seg_acc_sum,
i + epoch_num * num_train_file)
train_writer.add_summary(batch_sum,
i + epoch_num * num_train_file)
printout(flog, '\tTraining Total Mean_loss: %f' % total_loss)
printout(flog,
'\t\tTraining Label Mean_loss: %f' % total_label_loss)
printout(flog,
'\t\tTraining Label Accuracy: %f' % total_label_acc)
printout(flog,
'\t\tTraining Seg Mean_loss: %f' % total_seg_loss)
printout(flog, '\t\tTraining Seg Accuracy: %f' % total_seg_acc)
def eval_one_epoch(epoch_num):
is_training = False
total_loss = 0.0
total_label_loss = 0.0
total_seg_loss = 0.0
total_label_acc = 0.0
total_seg_acc = 0.0
total_seen = 0
total_label_acc_per_cat = np.zeros(
(NUM_CATEGORIES)).astype(np.float32)
total_seg_acc_per_cat = np.zeros(
(NUM_CATEGORIES)).astype(np.float32)
total_seen_per_cat = np.zeros((NUM_CATEGORIES)).astype(np.int32)
for i in range(num_test_file):
cur_test_filename = os.path.join(hdf5_data_dir,
test_file_list[i])
printout(flog, 'Loading test file ' + cur_test_filename)
cur_data, cur_labels, cur_seg = provider.loadDataFile_with_seg(
cur_test_filename)
cur_labels = np.squeeze(cur_labels)
cur_labels_one_hot = convert_label_to_one_hot(cur_labels)
num_data = len(cur_labels)
num_batch = num_data // batch_size
for j in range(num_batch):
begidx = j * batch_size
endidx = (j + 1) * batch_size
feed_dict = {
pointclouds_ph: cur_data[begidx:endidx, ...],
labels_ph: cur_labels[begidx:endidx, ...],
input_label_ph: cur_labels_one_hot[begidx:endidx, ...],
seg_ph: cur_seg[begidx:endidx, ...],
is_training_ph: is_training,
}
loss_val, label_loss_val, seg_loss_val, per_instance_label_loss_val, \
per_instance_seg_loss_val, label_pred_val, seg_pred_val, pred_seg_res \
= sess.run([loss, label_loss, seg_loss, per_instance_label_loss, \
per_instance_seg_loss, labels_pred, seg_pred, per_instance_seg_pred_res], \
feed_dict=feed_dict)
per_instance_part_acc = np.mean(
pred_seg_res == cur_seg[begidx:endidx, ...], axis=1)
average_part_acc = np.mean(per_instance_part_acc)
total_seen += 1
total_loss += loss_val
total_label_loss += label_loss_val
total_seg_loss += seg_loss_val
per_instance_label_pred = np.argmax(label_pred_val, axis=1)
total_label_acc += np.mean(
np.float32(per_instance_label_pred == cur_labels[
begidx:endidx, ...]))
total_seg_acc += average_part_acc
for shape_idx in range(begidx, endidx):
total_seen_per_cat[cur_labels[shape_idx]] += 1
total_label_acc_per_cat[
cur_labels[shape_idx]] += np.int32(
per_instance_label_pred[shape_idx - begidx] ==
cur_labels[shape_idx])
total_seg_acc_per_cat[
cur_labels[shape_idx]] += per_instance_part_acc[
shape_idx - begidx]
total_loss = total_loss * 1.0 / total_seen
total_label_loss = total_label_loss * 1.0 / total_seen
total_seg_loss = total_seg_loss * 1.0 / total_seen
total_label_acc = total_label_acc * 1.0 / total_seen
total_seg_acc = total_seg_acc * 1.0 / total_seen
test_loss_sum, test_label_acc_sum, test_label_loss_sum, test_seg_loss_sum, test_seg_acc_sum = sess.run(\
[total_test_loss_sum_op, label_test_acc_sum_op, label_test_loss_sum_op, seg_test_loss_sum_op, seg_test_acc_sum_op], \
feed_dict={total_testing_loss_ph: total_loss, label_testing_loss_ph: total_label_loss, \
seg_testing_loss_ph: total_seg_loss, label_testing_acc_ph: total_label_acc, seg_testing_acc_ph: total_seg_acc})
test_writer.add_summary(test_loss_sum,
(epoch_num + 1) * num_train_file - 1)
test_writer.add_summary(test_label_loss_sum,
(epoch_num + 1) * num_train_file - 1)
test_writer.add_summary(test_seg_loss_sum,
(epoch_num + 1) * num_train_file - 1)
test_writer.add_summary(test_label_acc_sum,
(epoch_num + 1) * num_train_file - 1)
test_writer.add_summary(test_seg_acc_sum,
(epoch_num + 1) * num_train_file - 1)
printout(flog, '\tTesting Total Mean_loss: %f' % total_loss)
printout(flog,
'\t\tTesting Label Mean_loss: %f' % total_label_loss)
printout(flog, '\t\tTesting Label Accuracy: %f' % total_label_acc)
printout(flog, '\t\tTesting Seg Mean_loss: %f' % total_seg_loss)
printout(flog, '\t\tTesting Seg Accuracy: %f' % total_seg_acc)
for cat_idx in range(NUM_CATEGORIES):
if total_seen_per_cat[cat_idx] > 0:
printout(
flog, '\n\t\tCategory %s Object Number: %d' %
(all_obj_cats[cat_idx][0],
total_seen_per_cat[cat_idx]))
printout(
flog, '\t\tCategory %s Label Accuracy: %f' %
(all_obj_cats[cat_idx][0],
total_label_acc_per_cat[cat_idx] /
total_seen_per_cat[cat_idx]))
printout(
flog, '\t\tCategory %s Seg Accuracy: %f' %
(all_obj_cats[cat_idx][0],
total_seg_acc_per_cat[cat_idx] /
total_seen_per_cat[cat_idx]))
if not os.path.exists(MODEL_STORAGE_PATH):
os.mkdir(MODEL_STORAGE_PATH)
for epoch in range(TRAINING_EPOCHES):
# Commenting out the validation / eval part because of one shot learning
'''
printout(flog, '\n<<< Testing on the test dataset ...')
eval_one_epoch(epoch)
'''
printout(
flog, '\n>>> Training for the epoch %d/%d ...' %
(epoch, TRAINING_EPOCHES))
train_file_idx = np.arange(0, len(train_file_list))
np.random.shuffle(train_file_idx)
train_one_epoch(train_file_idx, epoch)
'''
if (epoch+1) % 10 == 0:
cp_filename = saver.save(sess, os.path.join(MODEL_STORAGE_PATH, 'base_' + FLAGS.model_prefix + '.ckpt'))
printout(flog, 'Successfully store the checkpoint model into ' + cp_filename)
'''
flog.flush()
if FLAGS.update_model:
cp_filename = saver.save(
sess,
os.path.join(MODEL_STORAGE_PATH,
'base_' + FLAGS.model_prefix + '.ckpt'))
printout(
flog,
'Successfully store the checkpoint model into ' + cp_filename)
flog.close()