def train(): with tf.Graph().as_default(), tf.device('/cpu:0'): 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) trainer = tf.train.AdamOptimizer(learning_rate) # store tensors for different gpus tower_grads = [] pointclouds_phs = [] input_label_phs = [] seg_phs =[] is_training_phs =[] # 变量 with tf.variable_scope(tf.get_variable_scope()): for i in range(FLAGS.num_gpu): with tf.device('/gpu:%d' % i): with tf.name_scope('%s_%d' % (TOWER_NAME, i)) as scope: pointclouds_phs.append(tf.placeholder(tf.float32, shape=(batch_size, point_num, 3))) # for points 4*2048*3 input_label_phs.append(tf.placeholder(tf.float32, shape=(batch_size, NUM_CATEGORIES))) # for one-hot category label # 4*16 seg_phs.append(tf.placeholder(tf.int32, shape=(batch_size, point_num))) # for part labels # labels 4*2048 is_training_phs.append(tf.placeholder(tf.bool, shape=())) # 预测值 # seg_pred : B*N*part_num=4*2048*50 seg_pred = model.get_model(pointclouds_phs[-1], input_label_phs[-1], \ is_training=is_training_phs[-1], 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) # cat_num=16,part_num=50 # loss:这个batch的总的loss # per_instance_loss:每个点云的loss,shape=4 # per_instance_seg_pred_res : 每个点云的分割结果,由seg_pred取max得到的,shape=4*2048 loss, per_instance_seg_loss, per_instance_seg_pred_res \ = model.get_loss(seg_pred, seg_phs[-1]) # placeholder total_training_loss_ph = tf.placeholder(tf.float32, shape=()) total_testing_loss_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=()) # scalar 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) 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) tf.get_variable_scope().reuse_variables() grads = trainer.compute_gradients(loss) tower_grads.append(grads) grads = average_gradients(tower_grads) # 计算平均梯度 train_op = trainer.apply_gradients(grads, global_step=batch) saver = tf.train.Saver(tf.global_variables(), sharded=True, max_to_keep=20) # 保存变量 config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True sess = tf.Session(config=config) init = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) sess.run(init) train_writer = tf.summary.FileWriter(SUMMARIES_FOLDER + '/train', sess.graph) # 保存graph test_writer = tf.summary.FileWriter(SUMMARIES_FOLDER + '/test') train_file_list = provider.getDataFiles(TRAINING_FILE_LIST) # [train0.h5,train1.h5,...train5.h5] num_train_file = len(train_file_list) # 6 test_file_list = provider.getDataFiles(TESTING_FILE_LIST) # [ply_data_val0.h5] num_test_file = len(test_file_list) # 1 # 保存终端输入 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 # 遍历每一个train文件 for i in range(num_train_file): cur_train_filename = os.path.join(hdf5_data_dir, train_file_list[train_file_idx[i]]) # 获取当前的train的点云文件 printout(flog, 'Loading train file ' + cur_train_filename) cur_data, cur_labels, cur_seg = provider.load_h5_data_label_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 // (FLAGS.num_gpu * batch_size) # For all working gpus num——batch代表这个train文件分几个batch total_loss = 0.0 total_seg_acc = 0.0 # 对每一个batch for j in range(num_batch): begidx_0 = j * batch_size # 第一个gpu endidx_0 = (j + 1) * batch_size begidx_1 = (j + 1) * batch_size # 第二个gpu endidx_1 = (j + 2) * batch_size feed_dict = { # For the first gpu pointclouds_phs[0]: cur_data[begidx_0: endidx_0, ...], # 4*2048*3 input_label_phs[0]: cur_labels_one_hot[begidx_0: endidx_0, ...], # 4*16 seg_phs[0]: cur_seg[begidx_0: endidx_0, ...], # 4*2048 ,每一个数都在0-49之间 is_training_phs[0]: is_training, # # For the second gpu # pointclouds_phs[1]: cur_data[begidx_1: endidx_1, ...], # input_label_phs[1]: cur_labels_one_hot[begidx_1: endidx_1, ...], # seg_phs[1]: cur_seg[begidx_1: endidx_1, ...], # is_training_phs[1]: is_training, } # train_op is for both gpus, and the others are for gpu_1 # 每一个batch的平均损失、每一个点云的损失、分割预测值 _, loss_val, per_instance_seg_loss_val, seg_pred_val, pred_seg_res \ = sess.run([train_op, loss, per_instance_seg_loss, seg_pred, per_instance_seg_pred_res], \ feed_dict=feed_dict) # per_instance_part_acc = np.mean(pred_seg_res == cur_seg[begidx_1: endidx_1, ...], axis=1) per_instance_part_acc = np.mean(pred_seg_res == cur_seg[begidx_0: endidx_0, ...], axis=1) # 每一个点云的分割精度 average_part_acc = np.mean(per_instance_part_acc) # 当前batch的平均精度 total_loss += loss_val total_seg_acc += average_part_acc # 至此,一个train文件遍历完成 total_loss = total_loss * 1.0 / num_batch # 每一个train文件都得到一个loss和seg_acc total_seg_acc = total_seg_acc * 1.0 / num_batch # 绘制图 lr_sum, bn_decay_sum, batch_sum, train_loss_sum, train_seg_acc_sum = sess.run(\ [lr_op, bn_decay_op, batch_op, total_train_loss_sum_op, seg_train_acc_sum_op], \ feed_dict={total_training_loss_ph: total_loss, seg_training_acc_ph: total_seg_acc}) train_writer.add_summary(train_loss_sum, i + epoch_num * num_train_file) # epoch_num是一个不断变化的值 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_seg_acc_sum, i + epoch_num * num_train_file) train_writer.add_summary(batch_sum, i + epoch_num * num_train_file) printout(flog, '\tTanin_file: {},Training Total Mean_loss: {}'.format(i,total_loss)) # 每一个train文件的loss和acc printout(flog, '\t\tTanin_file: {} ,Training Seg Accuracy: {}'.format(i,total_seg_acc)) def eval_one_epoch(epoch_num): is_training = False total_loss = 0.0 total_seg_acc = 0.0 total_seen = 0 total_seg_acc_per_cat = np.zeros((NUM_CATEGORIES)).astype(np.float32) # [0. 0. ....0.] 16个0 total_seen_per_cat = np.zeros((NUM_CATEGORIES)).astype(np.int32) # [0 0 0 0 ...0] for i in range(num_test_file): # i= 0 num_test_file=1 cur_test_filename = os.path.join(hdf5_data_dir, test_file_list[i]) # ply_data_val0.h5 printout(flog, 'Loading test file ' + cur_test_filename) # ply_data_val0.h5 # 读取数据 # data = f['data'][:] # (1870, 2048, 3) # label = f['label'][:] # (1870,1) 1870个点云的类别0-15 # seg = f['pid'][:] # (1870, 2048) ,表示每一个点属于的类别0-49,一共50类点 cur_data, cur_labels, cur_seg = provider.load_h5_data_label_seg(cur_test_filename) cur_labels = np.squeeze(cur_labels) # shape:(1870,) cur_labels_one_hot = convert_label_to_one_hot(cur_labels) # 1870*16 num_data = len(cur_labels) # 1870个点云 num_batch = num_data // batch_size # 467个batch # Run on gpu_1, since the tensors used for evaluation are defined on gpu_1 for j in range(num_batch): # 0-466 begidx = j * batch_size endidx = (j + 1) * batch_size # feed_dict = { # pointclouds_phs[1]: cur_data[begidx: endidx, ...], # input_label_phs[1]: cur_labels_one_hot[begidx: endidx, ...], # seg_phs[1]: cur_seg[begidx: endidx, ...], # is_training_phs[1]: is_training} feed_dict = { pointclouds_phs[0]: cur_data[begidx: endidx, ...], # 4*2048*3 input_label_phs[0]: cur_labels_one_hot[begidx: endidx, ...], # 4*16 seg_phs[0]: cur_seg[begidx: endidx, ...], # 4*2048 is_training_phs[0]: is_training} # pred_seg_res:每个点的分割结果 shape:4*2048 # seg_pred_val:分割结果,概率编码 shape:4*2048*50 # per_instance_seg_loss_val: 每一个点云的损失 shape:(4,) # loss_val:当前batch的平均损失 shape:() loss_val, per_instance_seg_loss_val, seg_pred_val, pred_seg_res \ = sess.run([loss, per_instance_seg_loss, 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) # 4个点云的每一个的分割精度 shape:(4,) average_part_acc = np.mean(per_instance_part_acc) # 4个点云的平均分割精度 total_seen += 1 # 以batch为单位 total_loss += loss_val total_seg_acc += average_part_acc # average_part_acc:当前batch的分割精度 for shape_idx in range(begidx, endidx): # 对当前batch中的每一个点云 total_seen_per_cat[cur_labels[shape_idx]] += 1 # 统计看过的每一类点云的数量 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_seg_acc = total_seg_acc * 1.0 / total_seen # 总的平均分割精度,以batch为单位 # 绘制图 test_loss_sum, test_seg_acc_sum = sess.run(\ [total_test_loss_sum_op, seg_test_acc_sum_op], \ feed_dict={total_testing_loss_ph: total_loss, \ 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_seg_acc_sum, (epoch_num+1) * num_train_file-1) printout(flog, '\t\tTesting Total Mean_loss: %f' % total_loss) printout(flog, '\t\tTesting Seg Accuracy: %f' % total_seg_acc) for cat_idx in range(NUM_CATEGORIES): # 0-15 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 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 % 5 == 0: cp_filename = saver.save(sess, os.path.join(MODEL_STORAGE_PATH, 'epoch_' + str(epoch)+'.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=()) seg_pred = 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) saver = tf.train.Saver() 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'), 'a+') printout(flog, 'Loading model %s' % pretrained_model_path) saver.restore(sess, pretrained_model_path) printout(flog, 'Model restored.') 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]] # 0/1/.../15 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) seg_pred_res = sess.run(seg_pred, feed_dict={ pointclouds_ph: batch_data, input_label_ph: cur_label_one_hot, is_training_ph: is_training }) 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('Accuracy: %f\n' % seg_acc) fout.write('IoU: %f\n\n' % avg_iou) fout.write('IoU details: %s\n' % iou_log) printout(flog, pretrained_model_path) 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(), tf.device('/cpu:0'): 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) trainer = tf.train.AdamOptimizer(learning_rate) # store tensors for different gpus tower_grads = [] pointclouds_phs = [] input_label_phs = [] seg_phs = [] is_training_phs = [] with tf.variable_scope(tf.get_variable_scope()): for i in range(FLAGS.num_gpu): with tf.device('/gpu:%d' % i): with tf.name_scope('%s_%d' % (TOWER_NAME, i)) as scope: pointclouds_phs.append( tf.placeholder(tf.float32, shape=(batch_size, point_num, 3))) # for points input_label_phs.append( tf.placeholder( tf.float32, shape=(batch_size, NUM_CATEGORIES ))) # for one-hot category label seg_phs.append( tf.placeholder( tf.int32, shape=(batch_size, point_num))) # for part labels is_training_phs.append( tf.placeholder(tf.bool, shape=())) seg_pred = model.get_model(pointclouds_phs[-1], input_label_phs[-1], \ is_training=is_training_phs[-1], 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) loss, per_instance_seg_loss, per_instance_seg_pred_res \ = model.get_loss(seg_pred, seg_phs[-1]) total_training_loss_ph = tf.placeholder(tf.float32, shape=()) total_testing_loss_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) 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) tf.get_variable_scope().reuse_variables() grads = trainer.compute_gradients(loss) tower_grads.append(grads) grads = average_gradients(tower_grads) train_op = trainer.apply_gradients(grads, global_step=batch) saver = tf.train.Saver(tf.global_variables(), sharded=True, max_to_keep=20) config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True sess = tf.Session(config=config) init = tf.group(tf.global_variables_initializer(), tf.local_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.load_h5_data_label_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 // (FLAGS.num_gpu * batch_size ) # For all working gpus total_loss = 0.0 total_seg_acc = 0.0 for j in range(num_batch): begidx_0 = j * batch_size endidx_0 = (j + 1) * batch_size begidx_1 = (j + 1) * batch_size endidx_1 = (j + 2) * batch_size feed_dict = { # For the first gpu pointclouds_phs[0]: cur_data[begidx_0:endidx_0, ...], input_label_phs[0]: cur_labels_one_hot[begidx_0:endidx_0, ...], seg_phs[0]: cur_seg[begidx_0:endidx_0, ...], is_training_phs[0]: is_training, # For the second gpu pointclouds_phs[1]: cur_data[begidx_1:endidx_1, ...], input_label_phs[1]: cur_labels_one_hot[begidx_1:endidx_1, ...], seg_phs[1]: cur_seg[begidx_1:endidx_1, ...], is_training_phs[1]: is_training, } # train_op is for both gpus, and the others are for gpu_1 _, loss_val, per_instance_seg_loss_val, seg_pred_val, pred_seg_res \ = sess.run([train_op, loss, per_instance_seg_loss, seg_pred, per_instance_seg_pred_res], \ feed_dict=feed_dict) per_instance_part_acc = np.mean( pred_seg_res == cur_seg[begidx_1:endidx_1, ...], axis=1) average_part_acc = np.mean(per_instance_part_acc) total_loss += loss_val total_seg_acc += average_part_acc total_loss = total_loss * 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_seg_acc_sum = sess.run(\ [lr_op, bn_decay_op, batch_op, total_train_loss_sum_op, seg_train_acc_sum_op], \ feed_dict={total_training_loss_ph: total_loss, seg_training_acc_ph: total_seg_acc}) train_writer.add_summary(train_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_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 Seg Accuracy: %f' % total_seg_acc) def eval_one_epoch(epoch_num): is_training = False total_loss = 0.0 total_seg_acc = 0.0 total_seen = 0 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.load_h5_data_label_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 # Run on gpu_1, since the tensors used for evaluation are defined on gpu_1 for j in range(num_batch): begidx = j * batch_size endidx = (j + 1) * batch_size feed_dict = { pointclouds_phs[1]: cur_data[begidx:endidx, ...], input_label_phs[1]: cur_labels_one_hot[begidx:endidx, ...], seg_phs[1]: cur_seg[begidx:endidx, ...], is_training_phs[1]: is_training } loss_val, per_instance_seg_loss_val, seg_pred_val, pred_seg_res \ = sess.run([loss, per_instance_seg_loss, 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_seg_acc += average_part_acc for shape_idx in range(begidx, endidx): total_seen_per_cat[cur_labels[shape_idx]] += 1 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_seg_acc = total_seg_acc * 1.0 / total_seen test_loss_sum, test_seg_acc_sum = sess.run(\ [total_test_loss_sum_op, seg_test_acc_sum_op], \ feed_dict={total_testing_loss_ph: total_loss, \ 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_seg_acc_sum, (epoch_num + 1) * num_train_file - 1) printout(flog, '\tTesting Total Mean_loss: %f' % total_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 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 % 5 == 0: cp_filename = saver.save( sess, os.path.join(MODEL_STORAGE_PATH, 'epoch_' + str(epoch) + '.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.compat.v1.placeholder(tf.bool, shape=()) seg_pred = 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) saver = tf.train.Saver() 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'), 'a') printout(flog, 'Loading model %s' % pretrained_model_path) saver.restore(sess, pretrained_model_path) printout(flog, 'Model restored.') 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() ######################################################################################################## # load their data # 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]] # 0/1/.../15 # 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]) # batch_data[0, ...] = pc_augment_to_point_num(pc_normalize(pts), point_num) ######################################################################################################### # Load sapien data cur_label_one_hot = np.zeros( (1, NUM_OBJ_CATS), dtype=np.float32 ) # dummy info for pure part segmentation without cat info sapien_h5_files = os.listdir(sapien_h5_dir) for h5_file_path in sapien_h5_files: pts, seg = load_h5_data(os.path.join(sapien_h5_dir, h5_file_path), point_num) # only take one point cloud from the first camera angle at the first time instance pts = pts[0, :] seg = seg[0, :] batch_data[0, ...] = pc_normalize(pts) ######################################################################################################### ori_point_num = len(seg) seg_pred_res = sess.run(seg_pred, feed_dict={ pointclouds_ph: batch_data, input_label_ph: cur_label_one_hot, is_training_ph: is_training }) 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(h5_file_path) + '_gt.obj')) output_color_point_cloud( pts, seg_pred_val, os.path.join(output_dir, str(h5_file_path) + '_pred.obj')) output_color_point_cloud_red_blue( pts, np.int32(seg == seg_pred_val), os.path.join(output_dir, str(h5_file_path) + '_diff.obj')) with open(os.path.join(output_dir, str(h5_file_path) + '.log'), 'w') as fout: printout(flog, str(h5_file_path) + ':\n') fout.write('Total Point: %d\n\n' % ori_point_num) printout(flog, 'Total Point: %d\n\n' % ori_point_num) #fout.write('Ground Truth: %s\n' % objnames[cur_gt_label]) fout.write('Accuracy: %f\n' % seg_acc) printout(flog, '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))