├── DeepMLS_Generation.py ├── LICENSE ├── Octree ├── ocnn-tf │ └── libs │ │ ├── ReadMe.txt │ │ └── __init__.py └── ocnn.py ├── Pretrained ├── Config_d6_1p_pretrained.json ├── Config_d7_1p_pretrained.json ├── d6_1p_pretrained.zip ├── d7_1p_pretrained.zip └── download_models.py ├── README.md ├── environment.yml ├── examples ├── Config_g2_bs32_1p_d6.json ├── d0fa70e45dee680fa45b742ddc5add59.ply ├── d3380ee3db68aefb3f214ef9c53ac06.ply └── fcfc935c2ff7c66c858699aaad4acee4.ply ├── media └── teaser.png ├── mls_marching_cubes.py ├── network_architecture.py ├── points3d-tf ├── points3d │ ├── __init__.py │ ├── build.sh │ ├── get_neighbor_points_spatial_grid_local_self.cc │ ├── get_neighbor_points_spatial_grid_local_self.cu.cc │ ├── get_neighbor_points_spatial_grid_radius_weighting_voting.cc │ └── get_neighbor_points_spatial_grid_radius_weighting_voting.cu.cc └── pointsdataset │ └── points_dataset_octreed7.py ├── scripts ├── Readme.md └── generate_tf_records.py ├── utils.py └── vdb_tsdf ├── ReadMe.MD └── main.cpp /DeepMLS_Generation.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import numpy as np 4 | import tensorflow as tf 5 | import argparse 6 | from utils import * 7 | import tqdm 8 | 9 | sys.path.append("points3d-tf") 10 | from pointsdataset.points_dataset_octreed7 import * 11 | from network_architecture import * 12 | from points3d import get_neighbor_spatial_grid_local_self 13 | from points3d import get_neighbor_spatial_grid_radius_v_voting 14 | 15 | parser = argparse.ArgumentParser(description='Point-based Shape Generation') 16 | parser.add_argument('config', type=str, metavar='N', help='config json file') 17 | parser.add_argument('--test', action='store_true', help='forward the test data(inference)') 18 | args = parser.parse_args() 19 | config = config_reader(args.config) 20 | 21 | DBL_EPSILON = 1E-22 22 | 23 | class param: 24 | def __init__(self, config): 25 | self.train_data = config['train_data'] 26 | self.train_batch_size = config['train_batch_size'] 27 | self.learning_rate = config['learning_rate'] 28 | if('learning_rate_lower_bound' in config): 29 | self.learning_rate_lower_bound = config['learning_rate_lower_bound'] 30 | else: 31 | self.learning_rate_lower_bound = 1e-4 32 | self.lr_decay_epochs = config['lr_decay_epochs'] 33 | 34 | #choose different sdf samples for training(generate on depth-6 grids or depth-7 grids) 35 | self.sdf_data_sources = config['sdf_data_sources'] 36 | assert(self.sdf_data_sources == 6 or self.sdf_data_sources == 7) 37 | if(self.sdf_data_sources == 6): 38 | self.max_training_epochs = config['max_training_epochs_d6'] 39 | elif(self.sdf_data_sources == 7): 40 | self.max_training_epochs = config['max_training_epochs_d7'] 41 | 42 | self.exp_folder = config['exp_name'] 43 | self.ckpt = config['ckpt'] 44 | 45 | self.test = args.test 46 | self.gpu = config['gpu'] 47 | self.num_of_gpus = len(self.gpu.split(",")) 48 | self.num_of_input_points = config['num_of_input_points'] 49 | self.num_neighbors_to_search = config['num_neighbors_to_search'] 50 | 51 | #loss weighting 52 | self.sdf_loss_weight = config['sdf_loss_weight'] 53 | self.sdf_grad_loss_weight = config['sdf_grad_loss_weight'] 54 | self.geo_reg = config['geo_reg'] 55 | self.repulsion_weight = config['repulsion_weight'] 56 | self.normal_norm_reg_weight = config['normal_norm_reg_weight'] 57 | self.patch_radius_smoothness = config['patch_radius_smoothness'] 58 | self.octree_split_loss_weighting = config['octree_split_loss_weighting'] 59 | self.weight_decay = config['weight_decay'] 60 | 61 | #for octree data 62 | self.input_normal_signals = config['input_normal_signals'] 63 | if(self.input_normal_signals): 64 | self.channel = 4 65 | else: 66 | self.channel = 3 67 | self.depth = config['octree_depth'] 68 | if("decoder_octree_depth" in config): 69 | self.decoder_octree_depth = config['decoder_octree_depth'] 70 | else: 71 | self.decoder_octree_depth = 6 72 | print("decoder_octree_depth setting to {}".format(self.decoder_octree_depth)) 73 | 74 | #predict how many mls points in each non-empty octree leaf node 75 | self.points_per_node = config['points_per_node'] 76 | self.constant_radius = config['constant_radius'] 77 | assert(self.sdf_data_sources == 6 or self.sdf_data_sources == 7) 78 | if(self.sdf_data_sources == 7): 79 | self.sdf_loss_weight *= 4 80 | 81 | self.sdf_samples_each_iter = config['sdf_samples_each_iter'] 82 | print("using {} sdf samples for evaluation in each iteration".format(self.sdf_samples_each_iter)) 83 | 84 | self.node_receptive_field = config['node_receptive_field'] 85 | print("node receptive field times: {}".format(self.node_receptive_field)) 86 | 87 | self.radius_range = config['radius_range'] 88 | print("radius range = {}".format(self.radius_range)) 89 | 90 | if("noise_stddev" in config): 91 | self.noise_stddev = config['noise_stddev'] 92 | else: 93 | #our model is normaled to [-1, 1]^3 bounding box with 5% padding 94 | #to achieve same noise level with conv-onet 95 | self.noise_stddev = 0.0095 96 | 97 | FLAGS = param(config) 98 | 99 | os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' 100 | os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu 101 | KNN_PREFERRED = FLAGS.num_neighbors_to_search 102 | print("KNN={}".format(KNN_PREFERRED)) 103 | assert(KNN_PREFERRED >= 2) 104 | 105 | train_batch_size_gpu = FLAGS.train_batch_size // FLAGS.num_of_gpus 106 | assert(train_batch_size_gpu * FLAGS.num_of_gpus == FLAGS.train_batch_size) 107 | print("utilize {} gpus with total training batch size={}".format(FLAGS.num_of_gpus, FLAGS.train_batch_size)) 108 | 109 | lr_placeholder = tf.placeholder(tf.float32) 110 | print('=====') 111 | sys.stdout.flush() 112 | 113 | #normal unit norm regularization to avoid degenerated normal (should be very gentle) 114 | def normal_unit_norm_regularization(normals, points_segment, points_num): 115 | #input shape [batch_size*n,3] 116 | per_vertex_loss = tf.math.square(tf.reduce_sum(tf.math.square(normals), axis=1) - 1) 117 | loss = tf.segment_sum(per_vertex_loss, points_segment) / tf.cast(points_num, normals.dtype) 118 | return tf.reduce_mean(loss) 119 | 120 | def self_regularization_loss(predict_position_matrix, predict_normal_normalized, points_segment, points_num, per_point_squared_ball_radius): 121 | #first get neighbor info and weights 122 | predict_points = tf.concat([tf.reshape(predict_position_matrix, [-1, 3]), tf.reshape(predict_normal_normalized, [-1, 3])], axis=1) 123 | 124 | p2p_neighbor = tf.reshape(get_neighbor_spatial_grid_local_self(predict_points, tf.cumsum(points_num), knn=KNN_PREFERRED), [-1,KNN_PREFERRED]) 125 | invalid_index_mask = tf.cast(tf.less(p2p_neighbor, 0), dtype=p2p_neighbor.dtype) 126 | #make -1 to 0 127 | p2p_neighbor += invalid_index_mask 128 | p2p_neighbor = tf.expand_dims(p2p_neighbor, axis=-1) 129 | p2p_neighbor = tf.stop_gradient(p2p_neighbor) 130 | 131 | p2p_patch_pos = tf.gather_nd(predict_position_matrix, p2p_neighbor) 132 | p2p_patch_normal = tf.gather_nd(predict_normal_normalized, p2p_neighbor) 133 | per_point_radius = tf.math.sqrt(per_point_squared_ball_radius + 1e-19) 134 | p2p_patch_radius = tf.gather_nd(tf.reshape(per_point_radius, [-1]), p2p_neighbor) 135 | 136 | p2p_patch_pos_diff = tf.tile(tf.reshape(predict_position_matrix, [-1,1,3]), multiples=[1,KNN_PREFERRED,1]) - p2p_patch_pos 137 | p2p_patch_pos_diff_norm_squared = rowwise_l2_norm_squared(p2p_patch_pos_diff) 138 | p2p_patch_pos_diff_norm = tf.math.sqrt(p2p_patch_pos_diff_norm_squared + 1e-10) 139 | 140 | #using dot product [batch_size, n] 141 | p2p_patch_normal_dot = tf.reduce_sum(tf.tile(tf.reshape(predict_normal_normalized, [-1,1,3]),multiples=[1,KNN_PREFERRED,1])*p2p_patch_normal, axis=-1) 142 | p2p_patch_distance = rowwise_l2_norm_squared(p2p_patch_pos_diff) 143 | 144 | if(FLAGS.constant_radius): 145 | squared_ball_radius = per_point_squared_ball_radius[0] 146 | else: 147 | squared_ball_radius = tf.gather_nd(per_point_squared_ball_radius, p2p_neighbor) 148 | 149 | valid_neighbor_mask = tf.cast(1 - invalid_index_mask, dtype=p2p_patch_distance.dtype) 150 | p2p_patch_distance += squared_ball_radius * (1 - p2p_patch_normal_dot) 151 | 152 | #mls weight is stored in matrix with shape=[batch_size*n, KNN] 153 | p2p_mls_weight = tf.math.exp(-p2p_patch_distance / squared_ball_radius) 154 | p2p_mls_weight = p2p_mls_weight*valid_neighbor_mask 155 | #mls weights should also be normalized 156 | #p2p_mls_weight = p2p_mls_weight / tf.tile(tf.reshape(tf.reduce_sum(p2p_mls_weight,axis=-1),[-1,1]),multiples=[1,KNN_PREFERRED]) 157 | 158 | stop_regularization_weight_gradient = True 159 | if(stop_regularization_weight_gradient): 160 | p2p_mls_weight = tf.stop_gradient(p2p_mls_weight) 161 | p2p_mls_weight_sum = tf.reduce_sum(p2p_mls_weight, axis=-1) 162 | p2p_mls_weight_sum_dim3 = tf.tile(tf.expand_dims(p2p_mls_weight_sum, axis=-1), multiples=[1,3]) 163 | 164 | Laplacian_radius = tf.math.square(p2p_mls_weight_sum*tf.reshape(per_point_radius, [-1]) - tf.reduce_sum(p2p_mls_weight*p2p_patch_radius, axis=-1)) / tf.math.square(p2p_mls_weight_sum + 1e-6) 165 | 166 | #local geometry regularization(planar shape) 167 | #get neighbor point distance to tangent plane in shape[batch_size*n_point, KNN_PREFERRED] 168 | p2p_patch_pos_diff_normal = tf.reduce_sum(p2p_patch_pos_diff*tf.tile(tf.reshape(predict_normal_normalized, [-1, 1, 3]), multiples=[1, KNN_PREFERRED, 1]), axis=-1) 169 | p2p_patch_pos_diff_tangent = p2p_patch_pos_diff - tf.tile(tf.expand_dims(p2p_patch_pos_diff_normal, axis=-1), multiples=[1,1,3])*tf.tile(tf.reshape(predict_normal_normalized, [-1, 1, 3]), multiples=[1, KNN_PREFERRED, 1]) 170 | p2p_patch_pos_diff_tangent_norm_squared = rowwise_l2_norm_squared(p2p_patch_pos_diff_tangent) 171 | p2p_patch_pos_diff_tangent_norm = tf.math.sqrt(p2p_patch_pos_diff_tangent_norm_squared + 1e-10) 172 | 173 | tangent_distance = tf.math.square(p2p_patch_pos_diff_normal) 174 | if(stop_regularization_weight_gradient): 175 | p2p_mls_weight_plane = tf.stop_gradient(p2p_mls_weight*tf.math.exp(-tangent_distance/squared_ball_radius)) 176 | else: 177 | p2p_mls_weight_plane = p2p_mls_weight*tf.math.exp(-tangent_distance/squared_ball_radius) 178 | 179 | tangent_distance *= p2p_mls_weight_plane 180 | 181 | #-d repulsion 182 | if(True): 183 | #repulsion in tangent direction 184 | repulsion_vec = - p2p_patch_pos_diff_tangent_norm / (octree_mls_points_squared_radius()**0.5) 185 | else: 186 | repulsion_vec = - p2p_patch_pos_diff_norm / (octree_mls_points_squared_radius()**0.5) 187 | 188 | #then we get a tensor with shape [\sum(all_mls_points)] 189 | repulsion_force = tf.reduce_mean(p2p_mls_weight * repulsion_vec, axis=-1) 190 | repulsion_force = tf.segment_sum(repulsion_force, points_segment) / tf.cast(points_num, repulsion_force.dtype) 191 | 192 | radius_smoothness_loss = tf.segment_sum(Laplacian_radius, points_segment) / tf.cast(points_num, Laplacian_radius.dtype) 193 | local_plane_regularization = tf.segment_sum(tf.reduce_sum(tangent_distance, axis=-1), points_segment) / tf.cast(points_num, repulsion_force.dtype) 194 | 195 | repulsion_force = tf.reduce_mean(repulsion_force) 196 | local_plane_regularization = tf.reduce_mean(local_plane_regularization) 197 | radius_smoothness_loss = tf.reduce_mean(radius_smoothness_loss) 198 | 199 | return repulsion_force, local_plane_regularization, radius_smoothness_loss 200 | 201 | def eval_sdf_from_mls(src_position, target_position, target_normal_normalized, per_point_squared_ball_radius, s2t_neighbor): 202 | invalid_index_mask = tf.cast(tf.less(s2t_neighbor, 0), dtype=s2t_neighbor.dtype) 203 | 204 | #make -1 to 0(index -1 indicates invalid neighbor index, which means we cannot find up to K neighbors) 205 | s2t_neighbor += invalid_index_mask 206 | s2t_neighbor = tf.expand_dims(s2t_neighbor, axis=-1) 207 | 208 | #get per vertex patch vertices position & normal [batch_size*n, KNN, 3]+[batch_size*n, KNN, 3] 209 | s2t_patch_pos = tf.gather_nd(target_position, s2t_neighbor) 210 | s2t_patch_normal = tf.gather_nd(target_normal_normalized, s2t_neighbor) 211 | #compute mls weights 212 | s2t_patch_pos_diff = tf.tile(tf.reshape(src_position, [-1,1,3]),multiples=[1,KNN_PREFERRED,1]) - s2t_patch_pos 213 | s2t_patch_distance = rowwise_l2_norm_squared(s2t_patch_pos_diff) 214 | 215 | valid_neighbor_mask = tf.cast(1 - invalid_index_mask, dtype=s2t_patch_distance.dtype) 216 | 217 | #avoid divide by zero error 218 | if(FLAGS.constant_radius): 219 | s2t_mls_weight = -s2t_patch_distance / per_point_squared_ball_radius[0] 220 | else: 221 | s2t_mls_weight = -s2t_patch_distance / tf.gather_nd(tf.reshape(per_point_squared_ball_radius, [-1]), s2t_neighbor) 222 | 223 | s2t_mls_weight -= tf.stop_gradient(tf.tile(tf.expand_dims(s2t_mls_weight[:,0], axis=-1), multiples=[1, KNN_PREFERRED])) 224 | 225 | #mls weight is stored in matrix with shape=[batch_size*n, KNN] 226 | s2t_mls_weight = s2t_mls_weight*valid_neighbor_mask 227 | s2t_mls_weight = tf.math.exp(s2t_mls_weight) 228 | s2t_mls_weight = s2t_mls_weight*valid_neighbor_mask 229 | 230 | #mls weights should also be normalized 231 | s2t_mls_weight = s2t_mls_weight / tf.tile(tf.reshape(tf.reduce_sum(s2t_mls_weight,axis=-1),[-1,1]),multiples=[1,KNN_PREFERRED]) 232 | 233 | s2t_sdf = tf.reduce_sum(s2t_mls_weight*tf.reduce_sum(s2t_patch_pos_diff*s2t_patch_normal, axis=-1), axis=-1) 234 | #get weighted average patch points 235 | s2t_patch_pos_mean = tf.reduce_sum(tf.tile(tf.expand_dims(s2t_mls_weight, axis=-1),multiples=[1,1,3])*s2t_patch_pos, axis=1) 236 | #get weighted average normal: gradient 237 | s2t_patch_normal_mean = tf.reduce_sum(tf.tile(tf.expand_dims(s2t_mls_weight, axis=-1),multiples=[1,1,3])*s2t_patch_normal, axis=1) 238 | s2t_sdf_grad = tf.math.l2_normalize(s2t_patch_normal_mean, axis=1) 239 | 240 | return s2t_sdf_grad, s2t_sdf 241 | 242 | def MLS_sdf_Loss_Pack(evaluate_position_matrix, predict_position_matrix, predict_normal_normalized, evaluate_points_num, ps2p_neighbor_index, \ 243 | per_point_squared_ball_radius, sdf_precompute): 244 | 245 | evaluate_position_matrix = tf.reshape(evaluate_position_matrix, [-1, 3]) 246 | sdf_precompute = tf.reshape(sdf_precompute, [-1, 4]) 247 | 248 | #select valid grid points out of padded data 249 | valid_data_mask = tf.cast(tf.greater(evaluate_position_matrix[:,0], -5), dtype=evaluate_position_matrix.dtype) 250 | 251 | predict_sdf_grad, predict_sdf = \ 252 | eval_sdf_from_mls(evaluate_position_matrix, predict_position_matrix, predict_normal_normalized, per_point_squared_ball_radius, ps2p_neighbor_index) 253 | 254 | valid_data_mask = tf.reshape(valid_data_mask, [-1, evaluate_points_num]) 255 | valid_data_mask = tf.stop_gradient(valid_data_mask) 256 | #select the normal for grid points 257 | predict_sdf_grad = tf.reshape(predict_sdf_grad, [-1, evaluate_points_num, 3]) 258 | predict_sdf = tf.reshape(predict_sdf, [-1, evaluate_points_num]) 259 | 260 | #using precomputed sdf value and gradients 261 | gt_sdf = tf.reshape(sdf_precompute[:,0], [-1, evaluate_points_num]) 262 | gt_sdf_grad = tf.reshape(sdf_precompute[:,1:], [-1, evaluate_points_num, 3]) 263 | 264 | gt_sdf = tf.stop_gradient(gt_sdf) 265 | gt_sdf_grad = tf.stop_gradient(tf.nn.l2_normalize(gt_sdf_grad, axis=-1)) 266 | 267 | #L2 loss of difference between predict sdf and gt sdf (as well as gradients) 268 | grid_sdf_squared_diff = valid_data_mask * tf.math.square(gt_sdf - predict_sdf) 269 | 270 | if(False): 271 | #l2 loss 272 | grid_sdf_grad_squared_diff = valid_data_mask*rowwise_l2_norm_squared(gt_sdf_grad - predict_sdf_grad) 273 | else: 274 | #using dot product 275 | grid_sdf_grad_squared_diff = valid_data_mask*(1 - tf.reduce_sum(gt_sdf_grad*predict_sdf_grad, axis=-1)) 276 | 277 | valid_count = tf.reduce_sum(tf.reshape(valid_data_mask, [-1, evaluate_points_num]), axis=-1) 278 | 279 | predict_sdf_diff_loss = tf.reduce_sum(grid_sdf_squared_diff, axis=-1) / valid_count 280 | predict_sdf_grad_loss = tf.reduce_sum(grid_sdf_grad_squared_diff, axis=-1) / valid_count 281 | 282 | predict_sdf_diff_loss = tf.reduce_mean(predict_sdf_diff_loss) 283 | predict_sdf_grad_loss = tf.reduce_mean(predict_sdf_grad_loss) 284 | 285 | return predict_sdf_diff_loss, predict_sdf_grad_loss 286 | 287 | def network_loss(predict_points, points_segment, points_num, per_point_squared_ball_radius, sampled_position_matrix, sdf_precompute, name="network_loss"): 288 | #prepare ingredients for cooking 289 | predict_position_matrix = predict_points[:,:3] 290 | predict_normal_matrix = predict_points[:,3:] 291 | predict_normal_normalized = tf.math.l2_normalize(predict_normal_matrix, axis=1) 292 | 293 | loss_dict = {} 294 | train_loss = 0 295 | 296 | sampled_points = sampled_position_matrix 297 | evaluated_points_num = tf.cast(tf.shape(sampled_points)[1], tf.int32) 298 | sampled_points = tf.reshape(tf.transpose(sampled_points, perm=[0,2,1]), [-1, evaluated_points_num*3]) 299 | 300 | #compute neighbor mls points for sdf samples 301 | ps2p_neighbor_index = tf.reshape(get_neighbor_spatial_grid_radius_v_voting(sampled_points, predict_points, tf.cumsum(points_num),\ 302 | per_point_squared_ball_radius, knn=KNN_PREFERRED),[-1, KNN_PREFERRED]) 303 | 304 | if(FLAGS.sdf_loss_weight > DBL_EPSILON or FLAGS.sdf_grad_loss_weight > DBL_EPSILON): 305 | mls_sdf_loss, mls_sdf_grad_loss = MLS_sdf_Loss_Pack(tf.reshape(sampled_position_matrix, [-1,3]), predict_position_matrix, predict_normal_normalized, evaluated_points_num, ps2p_neighbor_index, \ 306 | per_point_squared_ball_radius, sdf_precompute=sdf_precompute) 307 | mls_sdf_loss *= FLAGS.sdf_loss_weight 308 | mls_sdf_grad_loss *= FLAGS.sdf_grad_loss_weight 309 | 310 | if(FLAGS.sdf_loss_weight > DBL_EPSILON): 311 | loss_dict[name + "_mls_sdf_loss"] = mls_sdf_loss 312 | 313 | if(FLAGS.sdf_grad_loss_weight > DBL_EPSILON): 314 | loss_dict[name + "_mls_sdf_grad_loss"] = mls_sdf_grad_loss 315 | 316 | train_loss += mls_sdf_loss + mls_sdf_grad_loss 317 | 318 | wLop_repulsion_loss, local_plane_regularization, radius_smoothness_loss = \ 319 | self_regularization_loss(predict_position_matrix, predict_normal_normalized, points_segment, points_num, per_point_squared_ball_radius) 320 | 321 | if(FLAGS.repulsion_weight > DBL_EPSILON): 322 | LOP_regularization_loss = FLAGS.repulsion_weight * wLop_repulsion_loss 323 | train_loss += LOP_regularization_loss 324 | loss_dict[name + '_repulsion'] = LOP_regularization_loss 325 | 326 | if(FLAGS.normal_norm_reg_weight > DBL_EPSILON): 327 | normal_unit_norm_loss = FLAGS.normal_norm_reg_weight*normal_unit_norm_regularization(predict_normal_matrix, points_segment, points_num) 328 | loss_dict[name +'_normal_unit_norm_reg'] = normal_unit_norm_loss 329 | train_loss += normal_unit_norm_loss 330 | 331 | if(FLAGS.geo_reg > DBL_EPSILON): 332 | local_plane_regularization *= FLAGS.geo_reg 333 | train_loss += local_plane_regularization 334 | loss_dict[name + "_local_plane_regularization"] = local_plane_regularization 335 | 336 | if(not FLAGS.constant_radius and FLAGS.patch_radius_smoothness > DBL_EPSILON): 337 | radius_smoothness_loss *= FLAGS.patch_radius_smoothness 338 | train_loss += radius_smoothness_loss 339 | loss_dict["radius_smoothness"] = radius_smoothness_loss 340 | 341 | if(FLAGS.weight_decay > DBL_EPSILON): 342 | network_weights_decay_loss = l2_regularizer("ocnn", FLAGS.weight_decay) 343 | loss_dict['network_weight_decay_loss'] = network_weights_decay_loss 344 | train_loss += network_weights_decay_loss 345 | 346 | loss_dict[name + "_total_loss"] = train_loss 347 | return train_loss, loss_dict 348 | 349 | def train_data_loader(): 350 | data_list = points_dataset_AE_multiple_GPU(FLAGS.train_data, FLAGS.train_batch_size, points_num=FLAGS.num_of_input_points, depth=FLAGS.depth, \ 351 | gpu_num=FLAGS.num_of_gpus, sample_grid_points_number=FLAGS.sdf_samples_each_iter, data_sources=FLAGS.sdf_data_sources, noise_stddev=FLAGS.noise_stddev) 352 | 353 | assert(len(data_list) == FLAGS.num_of_gpus*2+5) 354 | train_record_num = data_list[0] 355 | octree_all_gpu = data_list[1:(FLAGS.num_of_gpus+1)] 356 | gt_octree_all_gpu = data_list[(FLAGS.num_of_gpus+1):(FLAGS.num_of_gpus*2+1)] 357 | points, filenames, sampled_points, sampled_sdf = data_list[(FLAGS.num_of_gpus*2+1):] 358 | 359 | points_all_gpu = [] 360 | sampled_points_all_gpu = [] 361 | sampled_sdf_all_gpu = [] 362 | for i in range(FLAGS.num_of_gpus): 363 | points_all_gpu.append(points[i*train_batch_size_gpu:(i+1)*train_batch_size_gpu]) 364 | sampled_points_all_gpu.append(sampled_points[i*train_batch_size_gpu:(i+1)*train_batch_size_gpu]) 365 | sampled_sdf_all_gpu.append(sampled_sdf[i*train_batch_size_gpu:(i+1)*train_batch_size_gpu]) 366 | return train_record_num, octree_all_gpu, gt_octree_all_gpu, points_all_gpu, sampled_points_all_gpu, sampled_sdf_all_gpu, filenames 367 | 368 | def train_network(): 369 | train_record_num, octree_all_gpu, gt_octree_all_gpu, points_all_gpu, sampled_points_all_gpu, sampled_sdf_all_gpu, filenames = train_data_loader() 370 | loss_all_gpu = [] 371 | gradients_all_gpu=[] 372 | solver = tf.train.AdamOptimizer(learning_rate=lr_placeholder) 373 | all_loss_dict = [] 374 | 375 | for g_id in range(FLAGS.num_of_gpus): 376 | with tf.device('/gpu:{}'.format(g_id)): 377 | with tf.name_scope('GPU_{}'.format(g_id)) as scope: 378 | if(g_id == 0): 379 | g_loss, predict_points, sampled_position_matrix, loss_dict = build_graph_gpu_octree_local(train_batch_size_gpu, octree_all_gpu[g_id], None if points_all_gpu is None else points_all_gpu[g_id], 380 | sampled_points_all_gpu[g_id], sampled_sdf_all_gpu[g_id], reuse=None, gt_octree= gt_octree_all_gpu[g_id]) 381 | else: 382 | g_loss, _, _, loss_dict = build_graph_gpu_octree_local(train_batch_size_gpu, octree_all_gpu[g_id], None if points_all_gpu is None else points_all_gpu[g_id], 383 | sampled_points_all_gpu[g_id], sampled_sdf_all_gpu[g_id], reuse=True, gt_octree= gt_octree_all_gpu[g_id]) 384 | loss_all_gpu.append(g_loss) 385 | gradients_all_gpu.append(solver.compute_gradients(g_loss)) 386 | all_loss_dict.append(loss_dict) 387 | #use last tower statistics to update the moving mean/variance 388 | if(g_id == 0): 389 | batchnorm_update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS, scope=scope) 390 | 391 | with tf.device('/cpu:0'): 392 | train_summary = tf_summary_from_dict(all_loss_dict, True) 393 | 394 | with tf.device('/gpu:0'): 395 | total_loss = sum(loss_all_gpu) / FLAGS.num_of_gpus 396 | with tf.control_dependencies(batchnorm_update_ops): 397 | if(FLAGS.num_of_gpus == 1): 398 | apply_gradient_op = solver.apply_gradients(gradients_all_gpu[0]) 399 | else: 400 | grad_avg = average_gradient(gradients_all_gpu) 401 | apply_gradient_op = solver.apply_gradients(grad_avg) 402 | 403 | points_all_gpu[0] = tf.transpose(tf.reshape(points_all_gpu[0],[-1, 6, FLAGS.num_of_input_points]), perm=[0,2,1]) 404 | 405 | return train_record_num, train_summary, total_loss, apply_gradient_op, predict_points, sampled_position_matrix, points_all_gpu[0], filenames 406 | 407 | def octree_mls_points_squared_radius(): 408 | return 0.25**(FLAGS.decoder_octree_depth - 1) / FLAGS.points_per_node 409 | 410 | def build_graph_gpu_octree_local(input_batch_size, octree, points, grid_position, grid_sdf, reuse, is_training=True, gt_octree=None): 411 | assert(gt_octree is not None) 412 | 413 | split_loss, split_acc, mls_points, points_segment, octree_node_xyz_valid = \ 414 | octree_network_unet(octree, FLAGS.depth, FLAGS.decoder_octree_depth, FLAGS.channel, FLAGS.points_per_node, training=is_training, reuse=reuse, node_receptive_field=FLAGS.node_receptive_field, predict_radius=not FLAGS.constant_radius, radius_range=FLAGS.radius_range, gt_octree=gt_octree) 415 | 416 | points_num = tf.cast(tf.segment_sum(tf.ones_like(points_segment, dtype=tf.float64), points_segment), tf.int32) 417 | constant_squared_radius = octree_mls_points_squared_radius() 418 | 419 | if(not FLAGS.constant_radius): 420 | per_point_squared_radius = constant_squared_radius* tf.math.square(mls_points[:,6]) 421 | mls_points = mls_points[:,:6] 422 | else: 423 | per_point_squared_radius = constant_squared_radius* tf.ones(tf.shape(mls_points)[0]) 424 | 425 | sampled_position_matrix = tf.reshape(grid_position, [-1, FLAGS.sdf_samples_each_iter, 3]) 426 | grid_sdf = tf.reshape(grid_sdf, [-1, FLAGS.sdf_samples_each_iter, 4]) 427 | 428 | sampled_position_matrix = tf.stop_gradient(sampled_position_matrix) 429 | grid_sdf = tf.stop_gradient(grid_sdf) 430 | 431 | final_loss, loss_dict = network_loss(mls_points, points_segment, points_num, per_point_squared_radius, sampled_position_matrix, grid_sdf, name="final_geometry") 432 | 433 | loss_dict['split_accuracy'] = tf.add_n(split_acc) / len(split_acc) 434 | loss_dict['octree_split_loss'] = FLAGS.octree_split_loss_weighting*tf.add_n(split_loss) 435 | train_loss = final_loss + FLAGS.octree_split_loss_weighting*tf.add_n(split_loss) 436 | 437 | predict_points_list = [mls_points, points_num, per_point_squared_radius] 438 | loss_dict['final_geometry_total_loss'] = train_loss 439 | 440 | return train_loss, predict_points_list, sampled_position_matrix, loss_dict 441 | 442 | def build_graph_gpu_octree_local_inference(octree, reuse=None, is_training=False, test_batch_size=1): 443 | assert(not is_training) 444 | print("inference using prediction octree") 445 | mls_points, points_segment, octree_node_xyz_valid = octree_network_unet_completion_decode_shape(octree, FLAGS.depth, FLAGS.decoder_octree_depth, FLAGS.channel, FLAGS.points_per_node, test_batch_size, training=is_training, reuse=reuse, node_receptive_field=FLAGS.node_receptive_field, predict_radius=not FLAGS.constant_radius, radius_range=FLAGS.radius_range) 446 | 447 | points_num = tf.cast(tf.segment_sum(tf.ones_like(points_segment, dtype=tf.float64), points_segment), tf.int32) 448 | constant_squared_radius = octree_mls_points_squared_radius() 449 | 450 | if(not FLAGS.constant_radius): 451 | per_point_squared_radius = constant_squared_radius* tf.math.square(mls_points[:,6]) 452 | mls_points = mls_points[:,:6] 453 | else: 454 | per_point_squared_radius = constant_squared_radius* tf.ones(tf.shape(mls_points)[0]) 455 | 456 | predict_points_list = [mls_points, points_num, per_point_squared_radius] 457 | 458 | return predict_points_list 459 | 460 | def write_visualization_results(predict_points, sampled_position_matrix, input_points, dir, iter, filenames=None): 461 | final_prediction = predict_points[0] 462 | assert(len(predict_points) == 3) 463 | per_point_radius = np.reshape(predict_points[2], [-1]) 464 | per_point_radius = np.split(per_point_radius, np.cumsum(predict_points[1])) 465 | per_point_radius.pop() 466 | final_prediction_ = np.split(np.reshape(final_prediction, [-1]), 6*np.cumsum(predict_points[1])) 467 | final_prediction = [np.reshape(item, [-1,6]) for item in final_prediction_] 468 | final_prediction.pop() 469 | batch_size = len(final_prediction) 470 | assert(len(per_point_radius) == batch_size) 471 | 472 | input_points = np.reshape(input_points, [-1, FLAGS.num_of_input_points, 6]) 473 | assert(batch_size == input_points.shape[0]) 474 | 475 | sampled_position_matrix = np.reshape(sampled_position_matrix, [batch_size, -1, 3]) 476 | 477 | if(filenames is not None): 478 | filename_text_file = open(os.path.join(dir, "filenames_{:06d}.txt".format(iter)), "w") 479 | for train_batch_idx in range(batch_size): 480 | words_arr = filenames[train_batch_idx].decode("utf8")#.split("\\") 481 | #filename_text_file.write("{}\t{}\n".format(words_arr[1], words_arr[3][:-4])) 482 | filename_text_file.write("{}\n".format(words_arr)) 483 | filename_text_file.close() 484 | 485 | for i in range(batch_size): 486 | np.savetxt(os.path.join(dir, 'input_{:06d}_{:04d}.xyz'.format(iter, i)), input_points[i], fmt='%0.4f') 487 | np.savetxt(os.path.join(dir, 'predict_pc_{:06d}_{:04d}.xyz'.format(iter, i)), final_prediction[i], fmt='%0.4f') 488 | if(iter == 0): 489 | cur_sample_position_matrix = sampled_position_matrix[i] 490 | cur_sample_position_matrix = cur_sample_position_matrix[np.where(cur_sample_position_matrix[:,0] > -50)[0]] 491 | np.savetxt(os.path.join(dir, 'sampled_{:06d}_{:04d}.xyz'.format(iter, i)), cur_sample_position_matrix) 492 | 493 | if(per_point_radius is not None): 494 | np.savetxt(os.path.join(dir, 'predict_pc_{:06d}_{:04d}_radius.txt'.format(iter, i)), per_point_radius[i]) 495 | 496 | def inference_from_inputs(): 497 | #placeholder for input pointcloud, should in shape [-1, 3] 498 | input_points = tf.placeholder(tf.float32) 499 | #convert input_points to octree 500 | input_points_str = points_new(input_points, [], input_points, []) 501 | input_octree = octree_batch([points2octree(input_points_str, depth=FLAGS.depth, full_depth=2, node_dis=False, split_label=True)]) 502 | mls_points = build_graph_gpu_octree_local_inference(input_octree) 503 | 504 | #finish build the graph and next we load checkpoint 505 | gvars = tf.global_variables() 506 | print("{} global variables".format(len(gvars))) 507 | print("{} model parameters total".format(get_num_params())) 508 | 509 | ckpt = tf.train.latest_checkpoint(FLAGS.ckpt) 510 | start_iters = 0 if not ckpt else int(ckpt[ckpt.find('iter') + 4:-5]) + 1 511 | if(ckpt): 512 | tf_restore_saver = tf.train.Saver(var_list=gvars, max_to_keep=100) 513 | else: 514 | print("Cannot load checkpoint from {}".format(FLAGS.ckpt)) 515 | raise NotImplementedError 516 | 517 | config = tf.ConfigProto() 518 | config.gpu_options.allow_growth = True 519 | with tf.Session(config=config) as sess: 520 | # initialize 521 | init = tf.compat.v1.global_variables_initializer() 522 | sess.run(init) 523 | 524 | print("===========================") 525 | print("restore training from iteration %d: %s" %(start_iters, ckpt)) 526 | print("===========================") 527 | tf_restore_saver.restore(sess, ckpt) 528 | 529 | if(True): 530 | #forward single ShapeNet object 531 | import plyfile 532 | #inference 533 | input_file = "examples/d0fa70e45dee680fa45b742ddc5add59.ply" 534 | assert(os.path.exists(input_file)) 535 | plydata = plyfile.PlyData.read(input_file) 536 | noisy_points = np.stack([plydata['vertex']['x'], plydata['vertex']['y'], plydata['vertex']['z']], axis=1) 537 | assert(noisy_points.shape[1] == 3) 538 | #first normalize noisy points to [-0.95, 0.95] 539 | scale = 0.95 / np.abs(noisy_points).max() 540 | noisy_points *= scale 541 | assert(noisy_points.min() >= -1 and noisy_points.max() <= 1) 542 | mls_points_prediction = sess.run(mls_points, feed_dict={input_points:noisy_points}) 543 | assert(len(mls_points_prediction) == 3) 544 | 545 | mls_points_geometry = mls_points_prediction[0] 546 | mls_points_radius = mls_points_prediction[2] 547 | 548 | mls_points_position = mls_points_geometry[:,:3] 549 | mls_points_normal = mls_points_geometry[:,3:] / np.linalg.norm(mls_points_geometry[:,3:], axis=1, keepdims=True) 550 | mls_points_geometry = np.concatenate([mls_points_position, mls_points_normal], axis=1) 551 | 552 | shape_scale = np.ones(1)*scale 553 | np.savetxt(input_file+".xyz", mls_points_geometry, fmt="%0.6f") 554 | np.savetxt(input_file+"_radius.txt", mls_points_radius, fmt="%0.6f") 555 | np.savetxt(input_file+"_scale.txt", shape_scale) 556 | return 557 | 558 | #Shape Completion of ShapeNet:on ShapeNet 13 classes 559 | import plyfile 560 | if not os.path.exists(FLAGS.exp_folder): 561 | os.mkdir(FLAGS.exp_folder) 562 | test_output_dir = os.path.join(FLAGS.exp_folder, 'test') 563 | if not os.path.exists(test_output_dir): os.makedirs(test_output_dir, exist_ok=True) 564 | 565 | #where input pointcloud lies 566 | input_dir = "your_directory/convolutional_occupancy_networks-master/out/pointcloud/shapenet_3plane/generation_pretrained/input" 567 | categories = os.listdir(input_dir) 568 | 569 | for cat in categories: 570 | cat_input_dir = os.path.join(input_dir, cat) 571 | filelist = os.listdir(cat_input_dir) 572 | cat_output_dir = os.path.join(test_output_dir, cat) 573 | if(not os.path.exists(cat_output_dir)): os.mkdir(cat_output_dir) 574 | for point_file in filelist: 575 | if(not os.path.isdir(point_file) and point_file.endswith(".ply")): 576 | sample_id = point_file.replace(".ply", "") 577 | print("{}:{}".format(cat, point_file)) 578 | #do inference 579 | plydata = plyfile.PlyData.read(os.path.join(cat_input_dir, point_file)) 580 | noisy_points = np.stack([plydata['vertex']['x'], plydata['vertex']['y'], plydata['vertex']['z']], axis=1) 581 | assert(noisy_points.shape[1] == 3) 582 | #first normalize noisy points to [-0.95, 0.95] 583 | scale = 0.95 / np.abs(noisy_points).max() 584 | noisy_points *= scale 585 | assert(noisy_points.min() >= -1 and noisy_points.max() <= 1) 586 | mls_points_prediction = sess.run(mls_points, feed_dict={input_points:noisy_points}) 587 | assert(len(mls_points_prediction) == 3) 588 | 589 | mls_points_geometry = mls_points_prediction[0] 590 | mls_points_radius = mls_points_prediction[2] 591 | 592 | mls_points_position = mls_points_geometry[:,:3] 593 | mls_points_normal = mls_points_geometry[:,3:] / np.linalg.norm(mls_points_geometry[:,3:], axis=1, keepdims=True) 594 | mls_points_geometry = np.concatenate([mls_points_position, mls_points_normal], axis=1) 595 | 596 | shape_scale = np.ones(1)*scale 597 | np.savetxt(os.path.join(cat_output_dir, sample_id+".xyz"), mls_points_geometry, fmt="%0.6f") 598 | np.savetxt(os.path.join(cat_output_dir, sample_id+"_radius.txt"), mls_points_radius, fmt="%0.6f") 599 | np.savetxt(os.path.join(cat_output_dir, sample_id+"_scale.txt"), shape_scale) 600 | 601 | def training_pipeline(): 602 | if not os.path.exists(FLAGS.exp_folder): 603 | os.mkdir(FLAGS.exp_folder) 604 | train_record_num, train_summary, total_loss, train_op, predict_points, sampled_position_matrix, input_points, train_filenames = train_network() 605 | 606 | # checkpoint 607 | ckpt = tf.train.latest_checkpoint(FLAGS.ckpt) 608 | start_iters = 0 if not ckpt else int(ckpt[ckpt.find('iter') + 4:-5]) + 1 609 | 610 | # saver 611 | gvars = tf.global_variables() 612 | print("{} global variables".format(len(gvars))) 613 | print("{} model parameters total".format(get_num_params())) 614 | 615 | save_vars = gvars 616 | save_vars_wo_adam = [var for var in gvars if 'Adam' not in var.name] 617 | tf_saver_model = tf.train.Saver(var_list = save_vars_wo_adam, max_to_keep=100) 618 | restore_vars = save_vars_wo_adam 619 | 620 | print("restore {} vars".format(len(restore_vars))) 621 | print("save {} vars".format(len(save_vars))) 622 | 623 | tf_saver = tf.train.Saver(var_list=save_vars, max_to_keep=100) 624 | if ckpt: 625 | tf_restore_saver = tf.train.Saver(var_list=restore_vars, max_to_keep=100) 626 | 627 | #train_record_num = sum(1 for _ in tf.python_io.tf_record_iterator(FLAGS.train_data)) 628 | print("train_record_num: {}".format(train_record_num)) 629 | iter_100_epoch = int(100.0 * train_record_num / FLAGS.train_batch_size) 630 | 631 | config = tf.ConfigProto() 632 | config.gpu_options.allow_growth = True 633 | with tf.Session(config=config) as sess: 634 | # initialize 635 | init = tf.compat.v1.global_variables_initializer() 636 | sess.run(init) 637 | 638 | if ckpt: 639 | print("===========================") 640 | print("restore training from iteration %d: %s" %(start_iters, ckpt)) 641 | print("===========================") 642 | tf_restore_saver.restore(sess, ckpt) 643 | 644 | #folders for training outputs 645 | obj_dir = os.path.join(FLAGS.exp_folder, 'obj') 646 | if not os.path.exists(obj_dir): os.makedirs(obj_dir) 647 | 648 | #folders for checkpoint 649 | if not os.path.exists(os.path.join(FLAGS.exp_folder, "model")): os.makedirs(os.path.join(FLAGS.exp_folder, "model"), exist_ok=True) 650 | 651 | # tf summary 652 | summary_writer = tf.summary.FileWriter(FLAGS.exp_folder, sess.graph) 653 | lr_decay_step = int(FLAGS.lr_decay_epochs*iter_100_epoch/100.0) 654 | print("lr decay 20% per {} iterations({} epochs)".format(lr_decay_step, FLAGS.lr_decay_epochs)) 655 | cur_lr = max(FLAGS.learning_rate_lower_bound, FLAGS.learning_rate*0.8**int(start_iters / lr_decay_step)) 656 | print("initial lr setting to {}".format(cur_lr)) 657 | 658 | max_iter = int(1.0 * FLAGS.max_training_epochs * train_record_num / FLAGS.train_batch_size) + 1 659 | print("max training iterations: {}".format(max_iter)) 660 | last_loss = 0 661 | # start training 662 | for i in tqdm.tqdm(range(start_iters, max_iter)): 663 | #cur_epochs = 1.0 * i * FLAGS.train_batch_size / train_record_num 664 | if(i % lr_decay_step == 0): 665 | cur_lr = max(FLAGS.learning_rate_lower_bound, FLAGS.learning_rate*0.8**int(i / lr_decay_step)) 666 | print("lr setting to {}".format(cur_lr)) 667 | 668 | if(i % 5000 == 0): 669 | #write out visualization results 670 | _, train_summary_fetch, train_loss_eval, predict_points_eval, sampled_position_matrix_eval, input_points_fetch, train_filenames_fetch = \ 671 | sess.run([train_op, train_summary, total_loss, predict_points, sampled_position_matrix, input_points, train_filenames], feed_dict={lr_placeholder: cur_lr}) 672 | write_visualization_results(predict_points_eval, sampled_position_matrix_eval, input_points_fetch, obj_dir, i, filenames=train_filenames_fetch) 673 | else: 674 | _, train_summary_fetch, train_loss_eval = sess.run([train_op, train_summary, total_loss], feed_dict={lr_placeholder: cur_lr}) 675 | 676 | #write summary to tfevents 677 | summary_writer.add_summary(train_summary_fetch, i) 678 | #save checkpoint 679 | if (i % 5000 == 0): 680 | tf_saver.save(sess, os.path.join(FLAGS.exp_folder, 'model/iter{:06d}.ckpt'.format(i))) 681 | #print logs to stdout 682 | if(i % 500 == 0): 683 | print("iteration {}: train loss = {}".format(i, train_loss_eval)) 684 | sys.stdout.flush() 685 | 686 | print("Program terminated.") 687 | 688 | def main(): 689 | if(FLAGS.test): 690 | return inference_from_inputs() 691 | return training_pipeline() 692 | 693 | if __name__ == '__main__': 694 | print('Program running...') 695 | main() -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 Andy97 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /Octree/ocnn-tf/libs/ReadMe.txt: -------------------------------------------------------------------------------- 1 | Please copy compiled "libocnn.so" here. -------------------------------------------------------------------------------- /Octree/ocnn-tf/libs/__init__.py: -------------------------------------------------------------------------------- 1 | import os 2 | import tensorflow as tf 3 | from tensorflow.python.framework import ops 4 | 5 | if 'OCTREE_KEY' in os.environ and os.environ['OCTREE_KEY'] == '64': 6 | print('INFO from ocnn: The octree key is 64 bits') 7 | octree_key64 = True 8 | tf_uintk = tf.uint64 9 | tf_uints = tf.uint16 10 | tf_intk = tf.int64 11 | else: 12 | print('INFO from ocnn: The octree key is 32 bits, ' 13 | 'the octree depth should be smaller than 8.') 14 | octree_key64 = False 15 | tf_uintk = tf.uint32 16 | tf_uints = tf.uint8 17 | tf_intk = tf.int32 18 | 19 | _current_path = os.path.dirname(os.path.realpath(__file__)) 20 | _tf_ocnn_module = tf.load_op_library(os.path.join(_current_path, 'libocnn.so')) 21 | 22 | bounding_sphere = _tf_ocnn_module.bounding_sphere 23 | points_property = _tf_ocnn_module.points_property 24 | transform_points = _tf_ocnn_module.transform_points 25 | normalize_points = _tf_ocnn_module.normalize_points 26 | points_new = _tf_ocnn_module.points_new 27 | points_set_property = _tf_ocnn_module.points_set_property 28 | octree_drop = _tf_ocnn_module.octree_drop 29 | octree_scan = _tf_ocnn_module.octree_scan 30 | octree_cast = _tf_ocnn_module.octree_cast 31 | octree_batch = _tf_ocnn_module.octree_batch 32 | points2octree = _tf_ocnn_module.points_to_octree 33 | octree_property = _tf_ocnn_module.octree_property 34 | octree_pad = _tf_ocnn_module.octree_pad 35 | octree_depad = _tf_ocnn_module.octree_depad 36 | octree2col = _tf_ocnn_module.octree_to_col 37 | col2octree = _tf_ocnn_module.col_to_octree 38 | octree_grow = _tf_ocnn_module.octree_grow 39 | octree_new = _tf_ocnn_module.octree_new 40 | octree_update = _tf_ocnn_module.octree_update 41 | octree_align = _tf_ocnn_module.octree_align 42 | octree_mask = _tf_ocnn_module.octree_mask 43 | octree_samples = _tf_ocnn_module.octree_samples 44 | octree_search = _tf_ocnn_module.octree_search 45 | octree_key2xyz = _tf_ocnn_module.octree_key_to_xyz 46 | octree_xyz2key = _tf_ocnn_module.octree_xyz_to_key 47 | octree_decode_key = _tf_ocnn_module.octree_decode_key 48 | octree_encode_key = _tf_ocnn_module.octree_encode_key 49 | octree_search_key = _tf_ocnn_module.octree_search_key 50 | octree_set_property = _tf_ocnn_module.octree_set_property 51 | octree_gather = _tf_ocnn_module.octree_gather 52 | octree_gatherbk = _tf_ocnn_module.octree_gatherbk 53 | _octree_max_pool = _tf_ocnn_module.octree_max_pool 54 | _octree_mask_pool = _tf_ocnn_module.octree_mask_pool 55 | _octree_max_unpool = _tf_ocnn_module.octree_max_unpool 56 | _octree_conv = _tf_ocnn_module.octree_conv 57 | _octree_deconv = _tf_ocnn_module.octree_deconv 58 | _octree_conv_grad = _tf_ocnn_module.octree_conv_grad 59 | _octree_deconv_grad = _tf_ocnn_module.octree_deconv_grad 60 | _octree_align_grad = _tf_ocnn_module.octree_align_grad 61 | _octree_bilinear = _tf_ocnn_module.octree_bilinear 62 | 63 | 64 | ops.NotDifferentiable('BoundingSphere') 65 | ops.NotDifferentiable('OctreeSetProperty') 66 | ops.NotDifferentiable('OctreeBatch') 67 | ops.NotDifferentiable('TransformPoints') 68 | ops.NotDifferentiable('NormalizePoints') 69 | ops.NotDifferentiable('PointsNew') 70 | ops.NotDifferentiable('PointsSetProperty') 71 | ops.NotDifferentiable('PointsToOctree') 72 | ops.NotDifferentiable('OctreeProperty') 73 | ops.NotDifferentiable('OctreeNew') 74 | ops.NotDifferentiable('OctreeUpdate') 75 | ops.NotDifferentiable('OctreeGrow') 76 | ops.NotDifferentiable('OctreeSamples') 77 | ops.NotDifferentiable('OctreeBilinear') 78 | ops.NotDifferentiable('OctreeKeyToXyz') 79 | ops.NotDifferentiable('OctreeXyzToKey') 80 | ops.NotDifferentiable('OctreeDecodeKey') 81 | ops.NotDifferentiable('OctreeEncodeKey') 82 | ops.NotDifferentiable('OctreeSearchKey') 83 | ops.NotDifferentiable('OctreeSearch') 84 | ops.NotDifferentiable('PointsProperty') 85 | ops.NotDifferentiable('OctreeScan') 86 | ops.NotDifferentiable('OctreeCast') 87 | ops.NotDifferentiable('OctreeDrop') 88 | 89 | 90 | @ops.RegisterGradient('OctreePad') 91 | def _OctreePadGrad(op, grad): 92 | grad_out = octree_depad(grad, op.inputs[1], op.get_attr('depth')) 93 | return [grad_out, None] 94 | 95 | 96 | @ops.RegisterGradient('OctreeDepad') 97 | def _OctreeDepadGrad(op, grad): 98 | grad_out = octree_pad(grad, op.inputs[1], op.get_attr('depth')) 99 | return [grad_out, None] 100 | 101 | 102 | @ops.RegisterGradient('OctreeToCol') 103 | def _OctreeToColGrad(op, grad): 104 | grad_out = col2octree(grad, op.inputs[1], op.get_attr('depth'), 105 | op.get_attr('kernel_size'), op.get_attr('stride')) 106 | return [grad_out, None] 107 | 108 | 109 | @ops.RegisterGradient('ColToOctree') 110 | def _ColToOctreeGrad(op, grad): 111 | grad_out = octree2col(grad, op.inputs[1], op.get_attr('depth'), 112 | op.get_attr('kernel_size'), op.get_attr('stride')) 113 | return [grad_out, None] 114 | 115 | 116 | @ops.RegisterGradient('OctreeMaxPool') 117 | def _OctreeMaxPoolGrad(op, *grad): 118 | grad_out = _octree_max_unpool(grad[0], op.outputs[1], op.inputs[1], 119 | op.get_attr('depth')) 120 | return [grad_out, None] 121 | 122 | 123 | @ops.RegisterGradient('OctreeMaxUnpool') 124 | def _OctreeMaxUnpoolGrad(op, grad): 125 | grad_out = _octree_mask_pool(grad, op.inputs[1], op.inputs[2], 126 | op.get_attr('depth')) 127 | return [grad_out, None, None] 128 | 129 | 130 | @ops.RegisterGradient('OctreeMaskPool') 131 | def _OctreeMaskPoolGrad(op, grad): 132 | grad_out = _octree_max_unpool(grad, op.inputs[1], op.inputs[2], 133 | op.get_attr('depth')) 134 | return [grad_out, None, None] 135 | 136 | 137 | @ops.RegisterGradient('OctreeConv') 138 | def _OctreeConvGrad(op, grad): 139 | grad_out = _octree_conv_grad(op.inputs[0], op.inputs[1], op.inputs[2], grad, 140 | op.get_attr('depth'), op.get_attr('num_output'), 141 | op.get_attr('kernel_size'), op.get_attr('stride')) 142 | return grad_out + (None, ) 143 | 144 | 145 | @ops.RegisterGradient('OctreeDeconv') 146 | def _OctreeDeconvGrad(op, grad): 147 | grad_out = _octree_deconv_grad(op.inputs[0], op.inputs[1], op.inputs[2], grad, 148 | op.get_attr('depth'), op.get_attr('num_output'), 149 | op.get_attr('kernel_size'), op.get_attr('stride')) 150 | return grad_out + (None, ) 151 | 152 | 153 | @ops.RegisterGradient('OctreeAlign') 154 | def _OctreeAlignGrad(op, *grad): 155 | grad_out = _octree_align_grad(grad[0], op.outputs[1]) 156 | return [grad_out, None, None] 157 | 158 | 159 | @ops.RegisterGradient('OctreeMask') 160 | def _OctreeMaskGrad(op, grad): 161 | grad_out = octree_mask(grad, op.inputs[1], op.get_attr('mask')) 162 | return [grad_out, None] 163 | 164 | 165 | @ops.RegisterGradient('OctreeGather') 166 | def _OctreeGatherGrad(op, grad): 167 | shape = tf.shape(op.inputs[0]) 168 | grad_out = octree_gatherbk(grad, op.inputs[1], shape) 169 | return [grad_out, None] 170 | 171 | 172 | def octree_max_pool(data, octree, depth): 173 | with tf.variable_scope('octree_max_pool'): 174 | data, mask = _octree_max_pool(data, octree, depth) # the bottom data depth 175 | data = octree_pad(data, octree, depth-1) # !!! depth-1 176 | return data, mask 177 | 178 | 179 | def octree_max_unpool(data, mask, octree, depth): 180 | with tf.variable_scope('octree_max_unpool'): 181 | data = octree_depad(data, octree, depth) # !!! depth 182 | data = _octree_max_unpool(data, mask, octree, depth) # the bottom data depth 183 | return data 184 | 185 | 186 | def octree_avg_pool(data, octree, depth): 187 | with tf.variable_scope('octree_avg_pool'): 188 | data = tf.reshape(data, [1, int(data.shape[1]), -1, 8]) 189 | data = tf.reduce_mean(data, axis=3, keepdims=True) 190 | data = octree_pad(data, octree, depth-1) # !!! depth-1 191 | return data 192 | 193 | 194 | # todo: merge octree_conv_fast and octree_conv_memory to reduce code redundancy 195 | def octree_conv_fast(data, octree, depth, channel, kernel_size=[3], stride=1): 196 | assert(type(kernel_size) is list and len(kernel_size) < 4) 197 | for i in range(len(kernel_size), 3): 198 | kernel_size.append(kernel_size[-1]) 199 | 200 | with tf.variable_scope('octree_conv'): 201 | dim = int(data.shape[1]) * kernel_size[0] * kernel_size[1] * kernel_size[2] 202 | kernel = tf.get_variable('weights', shape=[channel, dim], dtype=tf.float32, 203 | initializer=tf.contrib.layers.xavier_initializer()) 204 | col = octree2col(data, octree, depth, kernel_size, stride) 205 | col = tf.reshape(col, [dim, -1]) 206 | conv = tf.matmul(kernel, col) 207 | conv = tf.expand_dims(tf.expand_dims(conv, 0), -1) # [C, H] -> [1, C, H, 1] 208 | if stride == 2: 209 | conv = octree_pad(conv, octree, depth-1, 0) 210 | return conv 211 | 212 | 213 | def octree_conv_memory(data, octree, depth, channel, kernel_size=[3], stride=1): 214 | assert(type(kernel_size) is list and len(kernel_size) < 4) 215 | for i in range(len(kernel_size), 3): 216 | kernel_size.append(kernel_size[-1]) 217 | 218 | with tf.variable_scope('octree_conv'): 219 | dim = int(data.shape[1]) * kernel_size[0] * kernel_size[1] * kernel_size[2] 220 | kernel = tf.get_variable('weights', shape=[channel, dim], dtype=tf.float32, 221 | initializer=tf.contrib.layers.xavier_initializer()) 222 | conv = _octree_conv(data, kernel, octree, depth, channel, kernel_size, stride) 223 | if stride == 2: 224 | conv = octree_pad(conv, octree, depth-1) 225 | return conv 226 | 227 | 228 | def octree_deconv_fast(data, octree, depth, channel, kernel_size=[3], stride=1): 229 | assert(type(kernel_size) is list and len(kernel_size) < 4) 230 | for i in range(len(kernel_size), 3): 231 | kernel_size.append(kernel_size[-1]) 232 | 233 | with tf.variable_scope('octree_deconv'): 234 | kernel_sdim = kernel_size[0] * kernel_size[1] * kernel_size[2] 235 | dim = channel * kernel_sdim 236 | kernel = tf.get_variable('weights', shape=[int(data.shape[1]), dim], dtype=tf.float32, 237 | initializer=tf.contrib.layers.xavier_initializer()) 238 | if stride == 2: 239 | data = octree_depad(data, octree, depth) 240 | depth = depth + 1 241 | data = tf.squeeze(data, [0, 3]) 242 | deconv = tf.matmul(kernel, data, transpose_a=True, transpose_b=False) 243 | deconv = tf.reshape(deconv, [channel, kernel_sdim, -1]) 244 | col = col2octree(deconv, octree, depth, kernel_size, stride) 245 | return col 246 | 247 | 248 | def octree_deconv_memory(data, octree, depth, channel, kernel_size=[3], stride=1): 249 | assert(type(kernel_size) is list and len(kernel_size) < 4) 250 | for i in range(len(kernel_size), 3): 251 | kernel_size.append(kernel_size[-1]) 252 | 253 | with tf.variable_scope('octree_deconv'): 254 | kernel_sdim = kernel_size[0] * kernel_size[1] * kernel_size[2] 255 | dim = channel * kernel_sdim 256 | kernel = tf.get_variable('weights', shape=[int(data.shape[1]), dim], dtype=tf.float32, 257 | initializer=tf.contrib.layers.xavier_initializer()) 258 | if stride == 2: 259 | data = octree_depad(data, octree, depth) 260 | deconv = _octree_deconv(data, kernel, octree, depth, channel, kernel_size, stride) 261 | return deconv 262 | 263 | 264 | def octree_full_voxel(data, depth): 265 | height = 2 ** (3 * depth) 266 | channel = int(data.shape[1]) 267 | with tf.variable_scope('octree_full_voxel'): 268 | data = tf.reshape(data, [channel, -1, height]) # (1, C, H, 1) -> (C, batch_size, H1) 269 | data = tf.transpose(data, perm=[1, 0, 2]) 270 | return data 271 | 272 | 273 | def octree_tile(data, octree, depth): 274 | with tf.variable_scope('octree_tile'): 275 | data = octree_depad(data, octree, depth) # (1, C, H, 1) 276 | data = tf.tile(data, [1, 1, 1, 8]) # (1, C, H, 8) 277 | channel = int(data.shape[1]) 278 | output = tf.reshape(data, [1, channel, -1, 1]) 279 | return output 280 | 281 | 282 | def octree_global_pool(data, octree, depth): 283 | with tf.variable_scope('octree_global_pool'): 284 | segment_ids = octree_property(octree, property_name='index', dtype=tf.int32, 285 | depth=depth, channel=1) 286 | segment_ids = tf.reshape(segment_ids, [-1]) 287 | data = tf.squeeze(data, axis=[0, 3]) # (1, C, H, 1) -> (C, H) 288 | data = tf.transpose(data) # (C, H) -> (H, C) 289 | output = tf.math.segment_mean(data, segment_ids) # (H, C) -> (batch_size, C) 290 | return output 291 | 292 | 293 | def octree_bilinear_legacy(data, octree, depth, target_depth): 294 | with tf.variable_scope('octree_bilinear'): 295 | mask = tf.constant( 296 | [[0, 0, 0], [0, 0, 1], [0, 1, 0], [1, 0, 0], 297 | [0, 1, 1], [1, 0, 1], [1, 1, 0], [1, 1, 1]], dtype=tf.float32) 298 | index, fracs = _octree_bilinear(octree, depth, target_depth) 299 | feat = tf.transpose(tf.squeeze(data, [0, 3])) # (1, C, H, 1) -> (H, C) 300 | output = tf.zeros([tf.shape(index)[0], tf.shape(feat)[1]], dtype=tf.float32) 301 | norm = tf.zeros([tf.shape(index)[0], 1], dtype=tf.float32) 302 | for i in range(8): 303 | idxi = index[:, i] 304 | weight = tf.abs(tf.reduce_prod(mask[i, :] - fracs, axis=1, keepdims=True)) 305 | output += weight * tf.gather(feat, idxi) 306 | norm += weight * tf.expand_dims(tf.cast(idxi > -1, dtype=tf.float32), -1) 307 | output = tf.div(output, norm) 308 | output = tf.expand_dims(tf.expand_dims(tf.transpose(output), 0), -1) 309 | return output 310 | 311 | 312 | # pts: (N, 4), i.e. N x (x, y, z, id) 313 | # data: (1, C, H, 1) 314 | def octree_bilinear_v1(pts, data, octree, depth): 315 | with tf.variable_scope('octree_bilinear'): 316 | mask = tf.constant( 317 | [[0, 0, 0], [0, 0, 1], [0, 1, 0], [0, 1, 1], 318 | [1, 0, 0], [1, 0, 1], [1, 1, 0], [1, 1, 1]], dtype=tf.float32) 319 | 320 | xyzf, ids = tf.split(pts, [3, 1], 1) 321 | xyzf = xyzf - 0.5 # since the value is defined on the center of each voxel 322 | xyzi = tf.floor(xyzf) # the integer part 323 | frac = xyzf - xyzi # the fraction part 324 | 325 | feat = tf.transpose(tf.squeeze(data, [0, 3])) # (1, C, H, 1) -> (H, C) 326 | output = tf.zeros([tf.shape(xyzi)[0], tf.shape(feat)[1]], dtype=tf.float32) 327 | norm = tf.zeros([tf.shape(xyzi)[0], 1], dtype=tf.float32) 328 | 329 | for i in range(8): 330 | maski = mask[i, :] 331 | maskc = 1.0 - maski 332 | xyzm = xyzi + maski 333 | xyzm = tf.cast(tf.concat([xyzm, ids], axis=1), dtype=tf_uints) 334 | idxi = octree_search_key(octree_encode_key(xyzm), octree, depth, is_xyz=True) 335 | 336 | weight = tf.abs(tf.reduce_prod(maskc - frac, axis=1, keepdims=True)) 337 | output += weight * tf.gather(feat, idxi) 338 | norm += weight * tf.expand_dims(tf.cast(idxi > -1, dtype=tf.float32), -1) 339 | output = tf.div(output, norm) 340 | 341 | output = tf.expand_dims(tf.expand_dims(tf.transpose(output), 0), -1) 342 | frac = tf.expand_dims(tf.expand_dims(tf.transpose(frac), 0), -1) 343 | 344 | return output, frac 345 | 346 | # pts: (N, 4), i.e. N x (x, y, z, id) 347 | # data: (1, C, H, 1) 348 | def octree_bilinear_v2(pts, data, octree, depth): 349 | with tf.variable_scope('octree_bilinear'): 350 | mask = tf.constant( 351 | [[0, 0, 0], [0, 0, 1], [0, 1, 0], [0, 1, 1], 352 | [1, 0, 0], [1, 0, 1], [1, 1, 0], [1, 1, 1]], dtype=tf.float32) 353 | 354 | xyzf, ids = tf.split(pts, [3, 1], 1) 355 | xyzf = xyzf - 0.5 # since the value is defined on the center of each voxel 356 | xyzi = tf.floor(xyzf) # the integer part 357 | frac = xyzf - xyzi # the fraction part 358 | 359 | output = tf.zeros([1, tf.shape(data)[1], tf.shape(xyzi)[0], 1], dtype=tf.float32) 360 | norm = tf.zeros([tf.shape(xyzi)[0], 1], dtype=tf.float32) 361 | 362 | for i in range(8): 363 | maski = mask[i, :] 364 | maskc = 1.0 - maski 365 | xyzm = xyzi + maski 366 | xyzm = tf.cast(tf.concat([xyzm, ids], axis=1), dtype=tf_uints) 367 | # !!! Note some elements of idxi may be -1 368 | idxi = octree_search_key(octree_encode_key(xyzm), octree, depth, is_xyz=True) 369 | 370 | weight = tf.abs(tf.reduce_prod(maskc - frac, axis=1, keepdims=True)) 371 | # output += weight * tf.gather(data, idxi, axis=2) 372 | output += weight * octree_gather(data, idxi) 373 | norm += weight * tf.expand_dims(tf.cast(idxi > -1, dtype=tf.float32), -1) 374 | output = tf.div(output, norm) 375 | return output 376 | 377 | 378 | # pts: (N, 4), i.e. N x (x, y, z, id). 379 | # data: (1, C, H, 1) 380 | # !!! Note: the pts should be scaled into [0, 2^depth] 381 | def octree_bilinear_v3(pts, data, octree, depth): 382 | with tf.variable_scope('octree_linear'): 383 | mask = tf.constant( 384 | [[0, 0, 0], [0, 0, 1], [0, 1, 0], [0, 1, 1], 385 | [1, 0, 0], [1, 0, 1], [1, 1, 0], [1, 1, 1]], dtype=tf.float32) 386 | if octree_key64: 387 | masku = tf.constant([0, 4294967296, 65536, 4295032832, 388 | 1, 4294967297, 65537, 4295032833], dtype=tf.int64) 389 | else: 390 | masku = tf.constant([0, 65536, 256, 65792, 391 | 1, 65537, 257, 65793], dtype=tf.int32) 392 | 393 | maskc = 1 - mask 394 | 395 | xyzf, ids = tf.split(pts, [3, 1], 1) 396 | xyzf = xyzf - 0.5 # since the value is defined on the center of each voxel 397 | xyzi = tf.floor(xyzf) # the integer part (N, 3) 398 | frac = xyzf - xyzi # the fraction part (N, 3) 399 | 400 | key = tf.cast(tf.concat([xyzi, ids], axis=1), dtype=tf_uints) 401 | key = tf.cast(octree_encode_key(key), dtype=tf_intk) 402 | # Cast the key to `int32` since the `add` below does not support `uint64` 403 | # The size effect is that the batch_size must be smaller than 128 404 | key = tf.expand_dims(key, 1) + masku # (N, 8), 405 | key = tf.cast(tf.reshape(key, [-1]), dtype=tf_uintk) 406 | 407 | idx = octree_search_key(key, octree, depth) # (N*8,) 408 | flgs = idx > -1 # filtering flags 409 | idx = tf.boolean_mask(idx, flgs) 410 | 411 | npt = tf.shape(xyzi)[0] 412 | ids = tf.reshape(tf.range(npt), [-1, 1]) 413 | ids = tf.reshape(tf.tile(ids, [1, 8]), [-1]) # (N*8,) 414 | ids = tf.boolean_mask(ids, flgs) 415 | 416 | frac = maskc - tf.expand_dims(frac, axis=1) 417 | weight = tf.abs(tf.reshape(tf.reduce_prod(frac, axis=2), [-1])) 418 | weight = tf.boolean_mask(weight, flgs) 419 | 420 | indices = tf.concat([tf.expand_dims(ids, 1), tf.expand_dims(idx, 1)], 1) 421 | indices = tf.cast(indices, tf.int64) 422 | data = tf.squeeze(data, [0, 3]) # (C, H) 423 | h = tf.shape(data)[1] 424 | mat = tf.SparseTensor(indices=indices, values=weight, dense_shape=[npt, h]) 425 | 426 | # channel, max_channel = int(data.shape[0]), 512 427 | # if channel > max_channel: 428 | # num = channel // max_channel 429 | # remain = channel % max_channel 430 | # splits = [max_channel] * num 431 | # if remain != 0: 432 | # splits.append(remain) 433 | # num += 1 434 | # output_split = [None] * num 435 | # data_split = tf.split(data, splits, axis=0) 436 | # for i in range(num): 437 | # with tf.name_scope('mat_%d' % i): 438 | # output_split[i] = tf.sparse.sparse_dense_matmul( 439 | # mat, data_split[i], adjoint_a=False, adjoint_b=True) 440 | # output = tf.concat(output_split, axis=1) 441 | # else: 442 | # output = tf.sparse.sparse_dense_matmul(mat, data, adjoint_a=False, adjoint_b=True) 443 | 444 | output = tf.sparse.sparse_dense_matmul(mat, data, adjoint_a=False, adjoint_b=True) 445 | norm = tf.sparse.sparse_dense_matmul(mat, tf.ones([h, 1])) 446 | output = tf.div(output, norm + 1.0e-10) # avoid dividing by zeros 447 | output = tf.expand_dims(tf.expand_dims(tf.transpose(output), 0), -1) 448 | return output 449 | 450 | 451 | def octree_bilinear(data, octree, depth, target_depth, mask=None): 452 | with tf.name_scope('Octree_bilinear'): 453 | xyz = octree_property(octree, property_name='xyz', depth=target_depth, 454 | channel=1, dtype=tf_uintk) 455 | xyz = tf.reshape(xyz, [-1]) 456 | if mask is not None: 457 | xyz = tf.boolean_mask(xyz, mask) 458 | xyz = tf.cast(octree_decode_key(xyz), dtype=tf.float32) 459 | 460 | # Attention: displacement 0.5, scale 461 | scale = 2.0**(depth-target_depth) 462 | xyz += tf.constant([0.5, 0.5, 0.5, 0.0], dtype=tf.float32) 463 | xyz *= tf.constant([scale, scale, scale, 1.0], dtype=tf.float32) 464 | 465 | output = octree_bilinear_v3(xyz, data, octree, depth) 466 | return output 467 | 468 | 469 | # pts: (N, 4), i.e. N x (x, y, z, id) 470 | # data: (1, C, H, 1) 471 | def octree_nearest_interp(pts, data, octree, depth): 472 | with tf.variable_scope('octree_nearest_interp'): 473 | # The value is defined on the center of each voxel, 474 | # so we can get the closest grid point by simply casting the value to tf_uints 475 | pts = tf.cast(pts, dtype=tf_uints) 476 | key = tf.reshape(octree_encode_key(pts), [-1]) 477 | 478 | idx = octree_search_key(key, octree, depth) 479 | # !!! Note that some of idx may be -1 or over-bound 480 | # Use tf.gather may be problematic with some version of tensorflow 481 | # according to my experiments. So I implemented octree_gather to 482 | # replace the original tf.gather. If you encounter errors, please 483 | # use the octree_gather 484 | # output = tf.gather(data, idx, axis=2) 485 | output = octree_gather(data, idx) 486 | return output 487 | 488 | 489 | 490 | def octree_signal(octree, depth, channel): 491 | with tf.name_scope('octree_signal'): 492 | signal = octree_property(octree, property_name='feature', dtype=tf.float32, 493 | depth=depth, channel=channel) 494 | signal = tf.reshape(signal, [1, channel, -1, 1]) 495 | return signal 496 | 497 | 498 | def octree_xyz(octree, depth, decode=True): 499 | with tf.name_scope('octree_xyz'): 500 | xyz = octree_property(octree, property_name='xyz', dtype=tf_uintk, 501 | depth=depth, channel=1) 502 | xyz = tf.cast(xyz, tf.int64) 503 | xyz = tf.reshape(xyz, [-1]) # uint32, N 504 | xyz = tf.cast(xyz, tf_uintk) 505 | if decode: 506 | xyz = octree_decode_key(xyz) # uint8, Nx4 507 | return xyz 508 | 509 | 510 | def octree_child(octree, depth): 511 | with tf.name_scope('octree_child'): 512 | child = octree_property(octree, property_name='child', dtype=tf.int32, 513 | depth=depth, channel=1) 514 | child = tf.reshape(child, [-1]) 515 | return child 516 | 517 | 518 | def octree_split(octree, depth): 519 | with tf.name_scope('octree_split'): 520 | split = octree_property(octree, property_name='split', dtype=tf.float32, 521 | depth=depth, channel=1) 522 | split = tf.reshape(split, [-1]) 523 | return split -------------------------------------------------------------------------------- /Octree/ocnn.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import tensorflow as tf 3 | sys.path.append("..") 4 | from libs import * 5 | 6 | 7 | def get_variables_with_name(name=None, without=None, train_only=True, verbose=False): 8 | if name is None: 9 | raise Exception("please input a name") 10 | 11 | t_vars = tf.trainable_variables() if train_only else tf.all_variables() 12 | d_vars = [var for var in t_vars if name in var.name] 13 | 14 | if without is not None: 15 | d_vars = [var for var in d_vars if without not in var.name] 16 | 17 | if verbose: 18 | print(" [*] geting variables with %s" % name) 19 | for idx, v in enumerate(d_vars): 20 | print(" got {:3}: {:15} {}".format(idx, v.name, str(v.get_shape()))) 21 | 22 | return d_vars 23 | 24 | 25 | def dense(inputs, nout, use_bias=False): 26 | inputs = tf.layers.flatten(inputs) 27 | fc = tf.layers.dense(inputs, nout, use_bias=use_bias, 28 | kernel_initializer=tf.contrib.layers.xavier_initializer()) 29 | return fc 30 | 31 | 32 | def batch_norm(inputs, training, axis=1): 33 | return tf.layers.batch_normalization(inputs, axis=axis, training=training) 34 | 35 | 36 | def fc_bn_relu(inputs, nout, training): 37 | fc = dense(inputs, nout, use_bias=False) 38 | bn = batch_norm(fc, training) 39 | return tf.nn.relu(bn) 40 | 41 | 42 | def conv2d(inputs, nout, kernel_size, stride, padding='SAME', data_format='channels_first'): 43 | return tf.layers.conv2d(inputs, nout, kernel_size=kernel_size, strides=stride, 44 | padding=padding, data_format=data_format, use_bias=False, 45 | kernel_initializer=tf.contrib.layers.xavier_initializer()) 46 | 47 | 48 | def octree_conv1x1(inputs, nout, use_bias=False): 49 | outputs = tf.layers.conv2d(inputs, nout, kernel_size=1, strides=1, 50 | data_format='channels_first', use_bias=use_bias, 51 | kernel_initializer=tf.contrib.layers.xavier_initializer()) 52 | return outputs 53 | 54 | 55 | def octree_conv1x1(inputs, nout, use_bias=False): 56 | with tf.variable_scope('conv2d_1x1'): 57 | inputs = tf.squeeze(inputs, axis=[0, 3]) # (1, C, H, 1) -> (C, H) 58 | weights = tf.get_variable('weights', shape=[nout, int(inputs.shape[0])], 59 | dtype=tf.float32, initializer=tf.contrib.layers.xavier_initializer()) 60 | outputs = tf.matmul(weights, inputs) # (C, H) -> (nout, H) 61 | if use_bias: 62 | bias = tf.get_variable('bias', shape=[nout, 1], dtype=tf.float32, 63 | initializer=tf.contrib.layers.xavier_initializer()) 64 | outputs = bias + outputs 65 | outputs = tf.expand_dims(tf.expand_dims(outputs, axis=0), axis=-1) 66 | return outputs 67 | 68 | 69 | def octree_conv1x1_bn(inputs, nout, training): 70 | conv = octree_conv1x1(inputs, nout, use_bias=False) 71 | return batch_norm(conv, training) 72 | 73 | 74 | def octree_conv1x1_bn_relu(inputs, nout, training): 75 | conv = octree_conv1x1_bn(inputs, nout, training) 76 | return tf.nn.relu(conv) 77 | 78 | 79 | def conv2d_bn(inputs, nout, kernel_size, stride, training): 80 | conv = conv2d(inputs, nout, kernel_size, stride) 81 | return batch_norm(conv, training) 82 | 83 | 84 | def conv2d_bn_relu(inputs, nout, kernel_size, stride, training): 85 | conv = conv2d_bn(inputs, nout, kernel_size, stride, training) 86 | return tf.nn.relu(conv) 87 | 88 | 89 | def upsample(data, channel, training): 90 | deconv = tf.layers.conv2d_transpose(data, channel, kernel_size=[8, 1], 91 | strides=[8, 1], data_format='channels_first', use_bias=False, 92 | kernel_initializer=tf.contrib.layers.xavier_initializer()) 93 | bn = tf.layers.batch_normalization(deconv, axis=1, training=training) 94 | return tf.nn.relu(bn) 95 | 96 | 97 | def downsample(data, channel, training): 98 | deconv = tf.layers.conv2d(data, channel, kernel_size=[8, 1], 99 | strides=[8, 1], data_format='channels_first', use_bias=False, 100 | kernel_initializer=tf.contrib.layers.xavier_initializer()) 101 | bn = tf.layers.batch_normalization(deconv, axis=1, training=training) 102 | return tf.nn.relu(bn) 103 | 104 | 105 | def avg_pool2d(inputs, data_format='NCHW'): 106 | return tf.nn.avg_pool2d(inputs, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], 107 | padding='SAME', data_format=data_format) 108 | 109 | 110 | def global_pool(inputs, data_format='channels_first'): 111 | axis = [2, 3] if data_format == 'channels_first' else [1, 2] 112 | return tf.reduce_mean(inputs, axis=axis) 113 | 114 | 115 | # !!! Deprecated 116 | def octree_upsample(data, octree, depth, channel, training): 117 | with tf.variable_scope('octree_upsample'): 118 | depad = octree_depad(data, octree, depth) 119 | up = upsample(depad, channel, training) 120 | return up 121 | 122 | 123 | def octree_upsample(data, octree, depth, channel, training): 124 | up = octree_deconv_bn_relu(data, octree, depth, channel, training, 125 | kernel_size=[2], stride=2, fast_mode=False) 126 | return up 127 | 128 | 129 | def octree_downsample(data, octree, depth, channel, training): 130 | with tf.variable_scope('octree_downsample'): 131 | down = downsample(data, channel, training) 132 | pad = octree_pad(down, octree, depth) 133 | return pad 134 | 135 | 136 | def octree_conv_bn(data, octree, depth, channel, training, kernel_size=[3], 137 | stride=1, fast_mode=False): 138 | if fast_mode == True: 139 | conv = octree_conv_fast(data, octree, depth, channel, kernel_size, stride) 140 | else: 141 | conv = octree_conv_memory( 142 | data, octree, depth, channel, kernel_size, stride) 143 | return tf.layers.batch_normalization(conv, axis=1, training=training) 144 | 145 | 146 | def octree_conv_bn_relu(data, octree, depth, channel, training, kernel_size=[3], 147 | stride=1, fast_mode=False): 148 | with tf.variable_scope('conv_bn_relu'): 149 | conv_bn = octree_conv_bn(data, octree, depth, channel, training, kernel_size, 150 | stride, fast_mode) 151 | rl = tf.nn.relu(conv_bn) 152 | return rl 153 | 154 | 155 | def octree_conv_bn_leakyrelu(data, octree, depth, channel, training): 156 | cb = octree_conv_bn(data, octree, depth, channel, training) 157 | return tf.nn.leaky_relu(cb, alpha=0.2) 158 | 159 | 160 | def octree_deconv_bn(data, octree, depth, channel, training, kernel_size=[3], 161 | stride=1, fast_mode=False): 162 | if fast_mode == True: 163 | conv = octree_deconv_fast( 164 | data, octree, depth, channel, kernel_size, stride) 165 | else: 166 | conv = octree_deconv_memory( 167 | data, octree, depth, channel, kernel_size, stride) 168 | return tf.layers.batch_normalization(conv, axis=1, training=training) 169 | 170 | 171 | def octree_deconv_bn_relu(data, octree, depth, channel, training, kernel_size=[3], 172 | stride=1, fast_mode=False): 173 | with tf.variable_scope('deconv_bn_relu'): 174 | conv_bn = octree_deconv_bn(data, octree, depth, channel, training, kernel_size, 175 | stride, fast_mode) 176 | rl = tf.nn.relu(conv_bn) 177 | return rl 178 | 179 | 180 | def octree_resblock(data, octree, depth, num_out, stride, training, bottleneck=4): 181 | num_in = int(data.shape[1]) 182 | channelb = int(num_out / bottleneck) 183 | if stride == 2: 184 | data, mask = octree_max_pool(data, octree, depth=depth) 185 | depth = depth - 1 186 | 187 | with tf.variable_scope("1x1x1_a"): 188 | block1 = octree_conv1x1_bn_relu(data, channelb, training=training) 189 | 190 | with tf.variable_scope("3x3x3"): 191 | block2 = octree_conv_bn_relu(block1, octree, depth, channelb, training) 192 | 193 | with tf.variable_scope("1x1x1_b"): 194 | block3 = octree_conv1x1_bn(block2, num_out, training=training) 195 | 196 | block4 = data 197 | if num_in != num_out: 198 | with tf.variable_scope("1x1x1_c"): 199 | block4 = octree_conv1x1_bn(data, num_out, training=training) 200 | 201 | return tf.nn.relu(block3 + block4) 202 | 203 | 204 | def octree_resblock2(data, octree, depth, num_out, training): 205 | num_in = int(data.shape[1]) 206 | with tf.variable_scope("conv_1"): 207 | conv = octree_conv_bn_relu(data, octree, depth, num_out/4, training) 208 | with tf.variable_scope("conv_2"): 209 | conv = octree_conv_bn(conv, octree, depth, num_out, training) 210 | 211 | link = data 212 | if num_in != num_out: 213 | with tf.variable_scope("conv_1x1"): 214 | link = octree_conv1x1_bn(data, num_out, training=training) 215 | 216 | out = tf.nn.relu(conv + link) 217 | return out 218 | 219 | 220 | def predict_module(data, num_output, num_hidden, training): 221 | # MLP with one hidden layer 222 | with tf.variable_scope('conv1'): 223 | conv = octree_conv1x1_bn_relu(data, num_hidden, training) 224 | with tf.variable_scope('conv2'): 225 | logit = octree_conv1x1(conv, num_output, use_bias=True) 226 | return logit 227 | 228 | 229 | def predict_label(data, num_output, num_hidden, training): 230 | logit = predict_module(data, num_output, num_hidden, training) 231 | # prob = tf.nn.softmax(logit, axis=1) # logit (1, num_output, ?, 1) 232 | label = tf.argmax(logit, axis=1, output_type=tf.int32) # predict (1, ?, 1) 233 | label = tf.reshape(label, [-1]) # flatten 234 | return logit, label 235 | 236 | 237 | def predict_signal(data, num_output, num_hidden, training): 238 | return tf.nn.tanh(predict_module(data, num_output, num_hidden, training)) 239 | 240 | 241 | def softmax_loss(logit, label_gt, num_class, label_smoothing=0.0): 242 | with tf.name_scope('softmax_loss'): 243 | label_gt = tf.cast(label_gt, tf.int32) 244 | onehot = tf.one_hot(label_gt, depth=num_class) 245 | loss = tf.losses.softmax_cross_entropy( 246 | onehot, logit, label_smoothing=label_smoothing) 247 | return loss 248 | 249 | 250 | def l2_regularizer(name, weight_decay): 251 | with tf.name_scope('l2_regularizer'): 252 | var = get_variables_with_name(name) 253 | regularizer = tf.add_n([tf.nn.l2_loss(v) for v in var]) * weight_decay 254 | return regularizer 255 | 256 | 257 | def label_accuracy(label, label_gt): 258 | label_gt = tf.cast(label_gt, tf.int32) 259 | accuracy = tf.reduce_mean(tf.to_float(tf.equal(label, label_gt))) 260 | return accuracy 261 | 262 | 263 | def softmax_accuracy(logit, label): 264 | with tf.name_scope('softmax_accuracy'): 265 | predict = tf.argmax(logit, axis=1, output_type=tf.int32) 266 | accu = label_accuracy(predict, tf.cast(label, tf.int32)) 267 | return accu 268 | 269 | 270 | def regress_loss(signal, signal_gt): 271 | return tf.reduce_mean(tf.reduce_sum(tf.square(signal-signal_gt), 1)) 272 | 273 | 274 | def normalize_signal(data): 275 | channel = data.shape[1] 276 | assert(channel == 3 or channel == 4) 277 | with tf.variable_scope("normalize"): 278 | if channel == 4: 279 | normals = tf.slice(data, [0, 0, 0, 0], [1, 3, -1, 1]) 280 | displacement = tf.slice(data, [0, 3, 0, 0], [1, 1, -1, 1]) 281 | normals = tf.nn.l2_normalize(normals, axis=1) 282 | output = tf.concat([normals, displacement], axis=1) 283 | else: 284 | output = tf.nn.l2_normalize(data, axis=1) 285 | return output 286 | 287 | 288 | def average_tensors(tower_tensors): 289 | avg_tensors = [] 290 | with tf.name_scope('avg_tensors'): 291 | for tensors in tower_tensors: 292 | tensors = [tf.expand_dims(tensor, 0) for tensor in tensors] 293 | avg_tensor = tf.concat(tensors, axis=0) 294 | avg_tensor = tf.reduce_mean(avg_tensor, 0) 295 | avg_tensors.append(avg_tensor) 296 | return avg_tensors 297 | 298 | 299 | def solver_single_gpu(total_loss, learning_rate_handle, gpu_num=1): 300 | with tf.variable_scope('solver'): 301 | update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) 302 | with tf.control_dependencies(update_ops): 303 | global_step = tf.Variable(0, trainable=False, name='global_step') 304 | lr = learning_rate_handle(global_step) 305 | solver = tf.train.MomentumOptimizer(lr, 0.9) \ 306 | .minimize(total_loss, global_step=global_step) 307 | return solver, lr 308 | 309 | 310 | def solver_multiple_gpus(total_loss, learning_rate_handle, gpu_num): 311 | tower_grads, variables = [], [] 312 | with tf.device('/cpu:0'): 313 | with tf.variable_scope('solver'): 314 | global_step = tf.Variable(0, trainable=False, name='global_step') 315 | lr = learning_rate_handle(global_step) 316 | opt = tf.train.MomentumOptimizer(lr, 0.9) 317 | 318 | for i in range(gpu_num): 319 | with tf.device('/gpu:%d' % i): 320 | with tf.name_scope('device_b%d' % i): 321 | grads_and_vars = opt.compute_gradients(total_loss[i]) 322 | grads, variables = zip(*grads_and_vars) 323 | tower_grads.append(grads) 324 | 325 | update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) 326 | # !!! Only get the update_ops defined on `device_0` to avoid the sync 327 | # between different GPUs to speed up the training process. !!! 328 | update_ops = [op for op in update_ops if 'device_0' in op.name] 329 | assert update_ops, 'The update ops of BN are empty, check the namescope \'device_0\'' 330 | with tf.device('/cpu:0'): 331 | with tf.name_scope('sync_and_apply_grad'): 332 | with tf.control_dependencies(update_ops): 333 | tower_grads = list(zip(*tower_grads)) 334 | avg_grads = average_tensors(tower_grads) 335 | grads_and_vars = list(zip(avg_grads, variables)) 336 | solver = opt.apply_gradients(grads_and_vars, global_step=global_step) 337 | return solver, lr 338 | 339 | 340 | def build_solver(total_loss, learning_rate_handle, gpu_num=1): 341 | assert (gpu_num > 0) 342 | the_solver = solver_single_gpu if gpu_num == 1 else solver_multiple_gpus 343 | return the_solver(total_loss, learning_rate_handle, gpu_num) 344 | 345 | 346 | def summary_train(names, tensors): 347 | with tf.name_scope('summary_train'): 348 | summaries = [] 349 | for it in zip(names, tensors): 350 | summaries.append(tf.summary.scalar(it[0], it[1])) 351 | summ = tf.summary.merge(summaries) 352 | return summ 353 | 354 | 355 | def summary_test(names): 356 | with tf.name_scope('summary_test'): 357 | summaries = [] 358 | summ_placeholder = [] 359 | for name in names: 360 | summ_placeholder.append(tf.placeholder(tf.float32)) 361 | summaries.append(tf.summary.scalar(name, summ_placeholder[-1])) 362 | summ = tf.summary.merge(summaries) 363 | return summ, summ_placeholder 364 | 365 | 366 | def loss_functions(logit, label_gt, num_class, weight_decay, var_name, label_smoothing=0.0): 367 | with tf.name_scope('loss'): 368 | loss = softmax_loss(logit, label_gt, num_class, label_smoothing) 369 | accu = softmax_accuracy(logit, label_gt) 370 | regularizer = l2_regularizer(var_name, weight_decay) 371 | return [loss, accu, regularizer] 372 | 373 | 374 | def loss_functions_seg(logit, label_gt, num_class, weight_decay, var_name, mask=-1): 375 | with tf.name_scope('loss_seg'): 376 | label_mask = label_gt > mask # filter label -1 377 | masked_logit = tf.boolean_mask(logit, label_mask) 378 | masked_label = tf.boolean_mask(label_gt, label_mask) 379 | loss = softmax_loss(masked_logit, masked_label, num_class) 380 | 381 | accu = softmax_accuracy(masked_logit, masked_label) 382 | regularizer = l2_regularizer(var_name, weight_decay) 383 | return [loss, accu, regularizer] 384 | 385 | 386 | def get_seg_label(octree, depth): 387 | with tf.name_scope('seg_label'): 388 | label = octree_property(octree, property_name='label', dtype=tf.float32, 389 | depth=depth, channel=1) 390 | label = tf.reshape(tf.cast(label, tf.int32), [-1]) 391 | return label 392 | 393 | 394 | def run_k_iterations(sess, k, tensors): 395 | num = len(tensors) 396 | avg_results = [0] * num 397 | for _ in range(k): 398 | iter_results = sess.run(tensors) 399 | for j in range(num): 400 | avg_results[j] += iter_results[j] 401 | 402 | for j in range(num): 403 | avg_results[j] /= k 404 | return avg_results 405 | 406 | 407 | def tf_IoU_per_shape(pred, label, class_num, mask=-1): 408 | with tf.name_scope('IoU'): 409 | label_mask = label > mask # filter label -1 410 | pred = tf.boolean_mask(pred, label_mask) 411 | label = tf.boolean_mask(label, label_mask) 412 | pred = tf.argmax(pred, axis=1, output_type=tf.int32) 413 | IoU, valid_part_num, esp = 0.0, 0.0, 1.0e-10 414 | for k in range(class_num): 415 | pk, lk = tf.equal(pred, k), tf.equal(label, k) 416 | # pk, lk = pred == k, label == k # why can this not output the right results? 417 | intsc = tf.reduce_sum(tf.cast(pk & lk, dtype=tf.float32)) 418 | union = tf.reduce_sum(tf.cast(pk | lk, dtype=tf.float32)) 419 | valid = tf.cast(tf.reduce_any(lk), dtype=tf.float32) 420 | valid_part_num += valid 421 | IoU += valid * intsc / (union + esp) 422 | IoU /= valid_part_num + esp 423 | return IoU, valid_part_num 424 | 425 | 426 | class Optimizer: 427 | def __init__(self, stype='SGD', var_list=None, mul=1.0): 428 | self.stype = stype # TODO: support more optimizers 429 | self.mul = mul # used to modulate the global learning rate 430 | self.var_list = var_list 431 | 432 | def __call__(self, total_loss, learning_rate): 433 | with tf.name_scope('solver'): 434 | update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) 435 | with tf.control_dependencies(update_ops): 436 | global_step = tf.Variable(0, trainable=False, name='global_step') 437 | lr = learning_rate(global_step) * self.mul 438 | solver = tf.train.MomentumOptimizer(lr, 0.9) \ 439 | .minimize(total_loss, global_step=global_step, 440 | var_list=self.var_list) 441 | return solver, lr 442 | 443 | 444 | 445 | def octree2points(octree, depth, pts_channel=4, output_normal=False): 446 | with tf.name_scope('octree2points'): 447 | signal = octree_signal(octree, depth, 4) # normal and displacement 448 | signal = tf.transpose(tf.squeeze(signal, [0, 3])) # (1, C, H, 1) -> (H, C) 449 | xyz = octree_xyz(octree, depth) 450 | xyz = tf.cast(xyz, dtype=tf.float32) 451 | 452 | mask = octree_child(octree, depth) > -1 453 | signal = tf.boolean_mask(signal, mask) 454 | xyz = tf.boolean_mask(xyz, mask) 455 | 456 | c = 3.0 ** 0.5 / 2.0 457 | normal, dis = tf.split(signal, [3, 1], axis=1) 458 | pts, idx = tf.split(xyz, [3, 1], axis=1) 459 | pts = (pts + 0.5) + normal * (dis * c) 460 | if pts_channel == 4: 461 | pts = tf.concat([pts, idx], axis=1) 462 | output = pts if not output_normal else (pts, normal) 463 | return output 464 | 465 | -------------------------------------------------------------------------------- /Pretrained/Config_d6_1p_pretrained.json: -------------------------------------------------------------------------------- 1 | { 2 | "_comment":"=========================================", 3 | "_comment":" Parameters for point-based shape generation r=e", 4 | "_comment":"=========================================", 5 | "exp_name":"Pretrained/ShapeNet_bs32_g2_all_1p_d6_pretrained", 6 | "train_data":"USE_YOUR_DATA", 7 | "train_batch_size":32, 8 | "lr_decay_epochs":10, 9 | "learning_rate":0.001, 10 | "max_training_epochs_d6":200, 11 | "max_training_epochs_d7":200, 12 | "ckpt":"Pretrained/d6_1p_pretrained", 13 | "gpu":"0", 14 | "num_of_input_points":3000, 15 | "sdf_data_sources":6, 16 | "input_normal_signals":false, 17 | 18 | "_comment":"=========================================", 19 | "_comment":" Parameters for network architecture", 20 | "_comment":"=========================================", 21 | "octree_depth":7, 22 | "points_per_node":1, 23 | "weight_decay":0.00005, 24 | "decoder_octree_depth":6, 25 | 26 | "_comment":"=========================================", 27 | "_comment":" Parameters for loss design", 28 | "_comment":"=========================================", 29 | "sdf_loss_weight":200.0, 30 | "sdf_grad_loss_weight":0.05, 31 | "geo_reg":10.0, 32 | "repulsion_weight":0.05, 33 | "normal_norm_reg_weight":0.0001, 34 | "num_neighbors_to_search":10, 35 | "stop_regularization_weight_gradient":true, 36 | "sdf_samples_each_iter":20000, 37 | "octree_split_loss_weighting":0.1, 38 | "node_receptive_field":2.0, 39 | 40 | "_comment":"=========================================", 41 | "_comment":" Radius parameters", 42 | "_comment":"=========================================", 43 | "constant_radius":false, 44 | "radius_range":2.0, 45 | "patch_radius_smoothness":10.0 46 | } -------------------------------------------------------------------------------- /Pretrained/Config_d7_1p_pretrained.json: -------------------------------------------------------------------------------- 1 | { 2 | "_comment":"=========================================", 3 | "_comment":" Parameters for point-based shape generation r=e", 4 | "_comment":"=========================================", 5 | "exp_name":"Pretrained/ShapeNet_bs32_g2_all_1p_d7_pretrained", 6 | "train_data":"USE_YOUR_DATA", 7 | "train_batch_size":32, 8 | "lr_decay_epochs":10, 9 | "learning_rate":0.001, 10 | "max_training_epochs_d6":200, 11 | "max_training_epochs_d7":200, 12 | "ckpt":"Pretrained/d7_1p_pretrained", 13 | "gpu":"0", 14 | "num_of_input_points":3000, 15 | "sdf_data_sources":6, 16 | "input_normal_signals":false, 17 | 18 | "_comment":"=========================================", 19 | "_comment":" Parameters for network architecture", 20 | "_comment":"=========================================", 21 | "octree_depth":7, 22 | "points_per_node":1, 23 | "weight_decay":0.00005, 24 | "decoder_octree_depth":7, 25 | 26 | "_comment":"=========================================", 27 | "_comment":" Parameters for loss design", 28 | "_comment":"=========================================", 29 | "sdf_loss_weight":200.0, 30 | "sdf_grad_loss_weight":0.05, 31 | "geo_reg":10.0, 32 | "repulsion_weight":0.05, 33 | "normal_norm_reg_weight":0.0001, 34 | "num_neighbors_to_search":10, 35 | "stop_regularization_weight_gradient":true, 36 | "sdf_samples_each_iter":20000, 37 | "octree_split_loss_weighting":0.1, 38 | "node_receptive_field":2.0, 39 | 40 | "_comment":"=========================================", 41 | "_comment":" Radius parameters", 42 | "_comment":"=========================================", 43 | "constant_radius":false, 44 | "radius_range":2.0, 45 | "patch_radius_smoothness":10.0 46 | } -------------------------------------------------------------------------------- /Pretrained/d6_1p_pretrained.zip: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Andy97/DeepMLS/551c51606a586f7ab3633621b9e4f57bb80c2345/Pretrained/d6_1p_pretrained.zip -------------------------------------------------------------------------------- /Pretrained/d7_1p_pretrained.zip: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Andy97/DeepMLS/551c51606a586f7ab3633621b9e4f57bb80c2345/Pretrained/d7_1p_pretrained.zip -------------------------------------------------------------------------------- /Pretrained/download_models.py: -------------------------------------------------------------------------------- 1 | import wget 2 | import zipfile 3 | import os 4 | 5 | if __name__ == '__main__': 6 | print("Downloading Pretrained Models...") 7 | if(not os.path.exists("d7_1p_pretrained.zip")): 8 | print("Downloading d7_1p_pretrained model...") 9 | wget.download("http://home.ustc.edu.cn/~freelin/DeepMLS/Pretrained_Models/d7_1p_pretrained.zip", 'd7_1p_pretrained.zip') 10 | print("\n") 11 | if(not os.path.exists("d6_1p_pretrained.zip")): 12 | print("Downloading d6_1p_pretrained model...") 13 | wget.download("http://home.ustc.edu.cn/~freelin/DeepMLS/Pretrained_Models/d6_1p_pretrained.zip", 'd6_1p_pretrained.zip') 14 | print("\n") 15 | 16 | if(not os.path.exists("d7_1p_pretrained")): 17 | with zipfile.ZipFile("d7_1p_pretrained.zip", 'r') as zip_ref: 18 | zip_ref.extractall("") 19 | 20 | if(not os.path.exists("d6_1p_pretrained")): 21 | with zipfile.ZipFile("d6_1p_pretrained.zip", 'r') as zip_ref: 22 | zip_ref.extractall("") -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # DeepMLS: Deep Implicit Moving Least-Squares Functions for 3D Reconstruction 2 | 3 | This repository contains the implementation of the paper: 4 | 5 | **Deep Implicit Moving Least-Squares Functions for 3D Reconstruction** [\[arXiv\]](https://arxiv.org/abs/2103.12266) 6 | Shi-Lin Liu, [Hao-Xiang Guo](https://haoxiangguo.cn/), [Hao Pan](https://haopan.github.io/), [Pengshuai Wang](https://wang-ps.github.io/), [Xin Tong](http://www.xtong.info/), [Yang Liu](https://xueyuhanlang.github.io/). 7 | 8 |
9 | 10 |
11 | 12 | 13 | If you find our code or paper useful, please consider citing 14 | 15 | ```bibtex 16 | @inproceedings{Liu2021MLS, 17 | author = {Shi-Lin Liu, Hao-Xiang Guo, Hao Pan, Pengshuai Wang, Xin Tong, Yang Liu}, 18 | title = {Deep Implicit Moving Least-Squares Functions for 3D Reconstruction}, 19 | booktitle = {IEEE/CVF Conference on Computer Vision and Pattern Recognition}, 20 | year = {2021}} 21 | ``` 22 | 23 | ## Installation 24 | First you have to make sure that you have all dependencies in place. 25 | The simplest way to do so, is to use [anaconda](https://www.anaconda.com/). 26 | 27 | You can create an anaconda environment called `deep_mls` using 28 | ``` 29 | conda env create -f environment.yml 30 | conda activate deep_mls 31 | ``` 32 | Next, a few customized tensorflow modules should be installed: 33 | #### [O-CNN](https://github.com/microsoft/O-CNN) Module 34 | [O-CNN](https://github.com/microsoft/O-CNN) is an octree-based convolution module, please take the following steps to install: 35 | ``` 36 | cd Octree && git clone https://github.com/microsoft/O-CNN/ 37 | cd O-CNN/octree/external && git clone --recursive https://github.com/wang-ps/octree-ext.git 38 | cd .. && mkdir build && cd build 39 | cmake .. && cmake --build . --config Release 40 | export PATH=`pwd`:$PATH 41 | cd ../../tensorflow/libs && python build.py --cuda /usr/local/cuda-10.0 42 | cp libocnn.so ../../../ocnn-tf/libs 43 | ``` 44 | 45 | #### Efficient Neighbor Searching Ops 46 | Neighbor searching is intensively used in DeepMLS. For efficiency reasons, we provide several customized neighbor searching ops: 47 | ``` 48 | cd points3d-tf/points3d 49 | bash build.sh 50 | ``` 51 | In this step, some errors like this may occur: 52 | ``` 53 | tensorflow_core/include/tensorflow/core/util/gpu_kernel_helper.h:22:10: fatal error: third_party/gpus/cuda/include/cuda_fp16.h: No such file or directory 54 | #include "third_party/gpus/cuda/include/cuda_fp16.h" 55 | ``` 56 | For solving this, please refer to [issue](https://github.com/tensorflow/tensorflow/issues/31349). 57 | Basically, We need to edit the codes in tensorflow framework, please modify 58 | ``` C 59 | #include "third_party/gpus/cuda/include/cuda_fp16.h" 60 | ``` 61 | in "site-packages/tensorflow_core/include/tensorflow/core/util/gpu_kernel_helper.h" to 62 | ``` C 63 | #include "cuda_fp16.h" 64 | ``` 65 | and 66 | ``` C 67 | #include "third_party/gpus/cuda/include/cuComplex.h" 68 | #include "third_party/gpus/cuda/include/cuda.h" 69 | ``` 70 | in "site-packages/tensorflow_core/include/tensorflow/core/util/gpu_device_functions.h" to 71 | ``` C 72 | #include "cuComplex.h" 73 | #include "cuda.h" 74 | ``` 75 | 76 | #### Modified Marching Cubes Module 77 | We have modified the [PyMCubes](https://github.com/pmneila/PyMCubes) to get a more efficient marching cubes method for extract 0-isosurface defined by mls points. 78 | To install: 79 | ``` 80 | git clone https://github.com/Andy97/PyMCubes 81 | cd PyMCubes && python setup.py install 82 | ``` 83 | 84 | ## Datasets 85 | 86 | ### Preprocessed ShapeNet Dataset 87 | We have provided the processed tfrecords file. This can be used directly. 88 | 89 | Our training data is **available** now! (total 130G+) 90 | Please download all zip files for extraction. 91 | [ShapeNet_points_all_train.zip.001](https://drive.google.com/file/d/1tP2y_lx8QIbxP9emRejpwQSkZ1x-MW5S/view?usp=sharing) 92 | [ShapeNet_points_all_train.zip.002](https://drive.google.com/file/d/1NrgLarnDAgvZ0aipKGYXmgQHFK887-XM/view?usp=sharing) 93 | [ShapeNet_points_all_train.zip.003](https://drive.google.com/file/d/1D49XjrkoH7ftuWKdXAKef1kRGRNxsU3I/view?usp=sharing) 94 | After extraction, please modify the "train_data" field in experiment config json file with this tfrecords name. 95 | 96 | #### Update(20220228) 97 | Since the training data in google drive is down, we provide download links in baidu netdisk: 98 | https://pan.baidu.com/s/1V70-VphJcN1fZaBdYNUN0g 99 | 100 | Please feel free to email freelin AT mail.ustc.edu.cn for access code. 101 | 102 | 103 | ### Build the Dataset 104 | If you want to build the dataset from your own data, please follow: 105 | 106 | #### Step 1: Get Watertight Meshes 107 | To acquire a watertight mesh, we first preprocess each mesh follow the [preprocess steps](https://github.com/autonomousvision/occupancy_networks#Building-the-dataset) of [Occupancy Networks](https://github.com/autonomousvision/occupancy_networks). 108 | 109 | #### Step 2: Get the groundtruth sdf pair 110 | From step 1, we have already gotten the watertight version of each model. Then, we utilize [OpenVDB](https://www.openvdb.org/) library to get the sdf values and gradients for training. 111 | For details, please refer to [here](vdb_tsdf). 112 | 113 | #### Update (2021-09-30): The script for generating training tfrecords is released, for details please refer to [Readme.md](https://github.com/Andy97/DeepMLS/tree/master/scripts). 114 | 115 | ## Usage 116 | 117 | #### Inference using pre-trained model 118 | We have provided pretrained models which can be used to inference: 119 | ``` 120 | #first download the pretrained models 121 | cd Pretrained && python download_models.py 122 | #then we can use either of the pretrained model to do the inference 123 | cd .. && python DeepMLS_Generation.py Pretrained/Config_d7_1p_pretrained.json --test 124 | ``` 125 | 126 | The input for the inference is defined in [here](https://github.com/Andy97/DeepMLS/blob/master/DeepMLS_Generation.py#L533). 127 | Your can replace it with other point cloud files in [examples](https://github.com/Andy97/DeepMLS/tree/master/examples) or your own data. 128 | 129 | #### Extract Isosurface from MLS Points 130 | After inference, now we have network predicted mls points. The next step is to extract the surface: 131 | ``` 132 | python mls_marching_cubes.py --i examples/d0fa70e45dee680fa45b742ddc5add59.ply.xyz --o examples/d0fa70e45dee680fa45b742ddc5add59_mc.obj --scale 133 | ``` 134 | 135 | #### Training 136 | Our code supports single and multiple gpu training. For details, please refer to the config json file. 137 | ``` 138 | python DeepMLS_Generation.py examples/Config_g2_bs32_1p_d6.json 139 | ``` 140 | #### Evaluation 141 | For evaluation of results, [ConvONet](https://github.com/autonomousvision/convolutional_occupancy_networks) has provided a great script. Please refer to [here](https://github.com/autonomousvision/convolutional_occupancy_networks#evaluation). 142 | -------------------------------------------------------------------------------- /environment.yml: -------------------------------------------------------------------------------- 1 | name: deep_mls 2 | dependencies: 3 | - _libgcc_mutex=0.1=main 4 | - ca-certificates=2021.1.19=h06a4308_1 5 | - certifi=2020.12.5=py37h06a4308_0 6 | - ld_impl_linux-64=2.33.1=h53a641e_7 7 | - libffi=3.3=he6710b0_2 8 | - libgcc-ng=9.1.0=hdf63c60_0 9 | - libstdcxx-ng=9.1.0=hdf63c60_0 10 | - ncurses=6.2=he6710b0_1 11 | - openssl=1.1.1j=h27cfd23_0 12 | - pip=21.0.1=py37h06a4308_0 13 | - python=3.7.10=hdb3f193_0 14 | - readline=8.1=h27cfd23_0 15 | - setuptools=52.0.0=py37h06a4308_0 16 | - sqlite=3.35.1=hdfb4753_0 17 | - tk=8.6.10=hbc83047_0 18 | - wheel=0.36.2=pyhd3eb1b0_0 19 | - xz=5.2.5=h7b6447c_0 20 | - zlib=1.2.11=h7b6447c_3 21 | - pip: 22 | - absl-py==0.12.0 23 | - astor==0.8.1 24 | - cached-property==1.5.2 25 | - gast==0.2.2 26 | - google-pasta==0.2.0 27 | - grpcio==1.36.1 28 | - h5py==3.2.1 29 | - importlib-metadata==3.7.3 30 | - keras-applications==1.0.8 31 | - keras-preprocessing==1.1.2 32 | - markdown==3.3.4 33 | - numpy==1.17.2 34 | - opt-einsum==3.3.0 35 | - pandas==1.2.3 36 | - plydata==0.4.3 37 | - plyfile==0.7.3 38 | - protobuf==3.15.6 39 | - python-dateutil==2.8.1 40 | - pytz==2021.1 41 | - six==1.15.0 42 | - tensorboard==1.15.0 43 | - tensorflow-estimator==1.15.1 44 | - tensorflow-gpu==1.15.0 45 | - termcolor==1.1.0 46 | - tqdm==4.59.0 47 | - typing-extensions==3.7.4.3 48 | - werkzeug==1.0.1 49 | - wget==3.2 50 | - wrapt==1.12.1 51 | - zipp==3.4.1 52 | 53 | -------------------------------------------------------------------------------- /examples/Config_g2_bs32_1p_d6.json: -------------------------------------------------------------------------------- 1 | { 2 | "_comment":"=========================================", 3 | "_comment":" Parameters for point-based shape generation r=e", 4 | "_comment":"=========================================", 5 | "exp_name":"ShapeNet_local_bs32_g2_all_1p_d6", 6 | "train_data":"../PointShape/data/ShapeNet_points_all_train_w_octree_grid_occupancy_compress.tfrecords", 7 | "train_batch_size":32, 8 | "lr_decay_epochs":10, 9 | "learning_rate":0.001, 10 | "max_training_epochs_d6":200, 11 | "max_training_epochs_d7":200, 12 | "ckpt":"ShapeNet_local_bs32_g2_all_1p_d6/model", 13 | "gpu":"0,1", 14 | "num_of_input_points":3000, 15 | "sdf_data_sources":6, 16 | "input_normal_signals":false, 17 | 18 | "_comment":"=========================================", 19 | "_comment":" Parameters for network architecture", 20 | "_comment":"=========================================", 21 | "octree_depth":7, 22 | "points_per_node":1, 23 | "weight_decay":0.00005, 24 | "decoder_octree_depth":6, 25 | 26 | "_comment":"=========================================", 27 | "_comment":" Parameters for loss design", 28 | "_comment":"=========================================", 29 | "sdf_loss_weight":200.0, 30 | "sdf_grad_loss_weight":0.05, 31 | "geo_reg":10.0, 32 | "repulsion_weight":0.05, 33 | "normal_norm_reg_weight":0.0001, 34 | "num_neighbors_to_search":10, 35 | "sdf_samples_each_iter":20000, 36 | "octree_split_loss_weighting":0.1, 37 | "node_receptive_field":2.0, 38 | 39 | "_comment":"=========================================", 40 | "_comment":" Radius parameters", 41 | "_comment":"=========================================", 42 | "constant_radius":false, 43 | "radius_range":2.0, 44 | "patch_radius_smoothness":10.0 45 | } -------------------------------------------------------------------------------- /examples/fcfc935c2ff7c66c858699aaad4acee4.ply: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Andy97/DeepMLS/551c51606a586f7ab3633621b9e4f57bb80c2345/examples/fcfc935c2ff7c66c858699aaad4acee4.ply -------------------------------------------------------------------------------- /media/teaser.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Andy97/DeepMLS/551c51606a586f7ab3633621b9e4f57bb80c2345/media/teaser.png -------------------------------------------------------------------------------- /mls_marching_cubes.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import numpy as np 4 | import tensorflow as tf 5 | from scipy import ndimage 6 | import mcubes 7 | import math 8 | import time 9 | 10 | sys.path.append("points3d-tf") 11 | 12 | os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' 13 | os.environ['CUDA_VISIBLE_DEVICES'] = '0' 14 | KNN_PREFERRED = 20 15 | print("KNN={}".format(KNN_PREFERRED)) 16 | assert(KNN_PREFERRED >= 2) 17 | print('====') 18 | sys.stdout.flush() 19 | 20 | struct2 = ndimage.generate_binary_structure(3, 3) 21 | 22 | import argparse 23 | parser = argparse.ArgumentParser(description='Marching Cube Input Output Arguments') 24 | parser.add_argument('--i', type=str, help='input grid sdf file or points file or directory') #if args.i is dir, we will process all points file under the folder 25 | parser.add_argument('--o', type=str, help='output marching cube file name') 26 | parser.add_argument('--res', type=int, default=7, help='grid resolution') 27 | parser.add_argument('--gpu', type=str, help='used gpu for program') 28 | parser.add_argument('--scale', action='store_true', help='scale mesh back') 29 | parser.add_argument('--overwrite', action='store_true', help='overwrite results') 30 | 31 | args = parser.parse_args() 32 | if(args.gpu is not None): 33 | os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu 34 | 35 | mls_points_ph = tf.placeholder(tf.float32) 36 | mls_radius_ph = tf.placeholder(tf.float32) 37 | query_points_ph = tf.placeholder(tf.float32) 38 | 39 | from points3d import get_neighbor_spatial_grid_radius_v_voting 40 | 41 | def rowwise_l2_norm_squared(feature): 42 | #assum input with size[n,f] out shape = [n,1] 43 | return tf.reduce_sum(tf.math.square(feature), axis=-1) 44 | 45 | def write_obj(obj_filename, vertices, triangles, output_max_component=True, scale=None): 46 | #normalize to -1,1 47 | resolution = 2**args.res 48 | print("mesh resolution in {}".format(resolution)) 49 | 50 | vertices /= (resolution / 2.0) 51 | vertices -= 1 52 | 53 | if(scale is not None): 54 | vertices *= scale 55 | 56 | with open(obj_filename,"w") as wf: 57 | for i in range(vertices.shape[0]): 58 | wf.write("v %lf %lf %lf\n" %(vertices[i][0], vertices[i][1], vertices[i][2])) 59 | for i in range(triangles.shape[0]): 60 | wf.write("f %d %d %d\n" %(triangles[i][0]+1, triangles[i][1]+1, triangles[i][2]+1)) 61 | 62 | def eval_sdf_from_mls(cur_projection, target_position, target_normal_normalized, src_points_num, target_points_num, \ 63 | per_point_squared_ball_radius, s2t_neighbor=None): 64 | #here we assert batch size=1 65 | if(s2t_neighbor == None): 66 | projection_points = tf.reshape(tf.transpose(tf.reshape(cur_projection, [1, -1, 3]), perm=[0,2,1]), [1, -1]) 67 | target_points = tf.concat([tf.reshape(target_position, [-1, 3]), tf.reshape(target_normal_normalized, [-1, 3])], axis=1) 68 | #fetch corresponding patch, now we get indices matrix [batch_size*n, KNN] 69 | s2t_neighbor = tf.reshape(get_neighbor_spatial_grid_radius_v_voting(projection_points, target_points, [target_points_num],\ 70 | tf.reshape(per_point_squared_ball_radius, [-1]), knn=KNN_PREFERRED), [-1,KNN_PREFERRED]) 71 | 72 | invalid_index_mask = tf.cast(tf.less(s2t_neighbor, 0), dtype=s2t_neighbor.dtype) 73 | #make -1 to 0 74 | s2t_neighbor += invalid_index_mask 75 | s2t_neighbor = tf.expand_dims(s2t_neighbor, axis=-1) 76 | 77 | #get per vertex patch vertices position & normal [batch_size*n, KNN, 3]+[batch_size*n, KNN, 3] 78 | s2t_patch_pos = tf.gather_nd(target_position, s2t_neighbor) 79 | s2t_patch_normal = tf.gather_nd(target_normal_normalized, s2t_neighbor) 80 | s2t_patch_squared_radius = tf.gather_nd(tf.reshape(per_point_squared_ball_radius, [-1]), s2t_neighbor) 81 | #compute mls weights 82 | s2t_patch_pos_diff = tf.tile(tf.reshape(cur_projection, [-1,1,3]),multiples=[1,KNN_PREFERRED,1]) - s2t_patch_pos 83 | s2t_patch_distance = rowwise_l2_norm_squared(s2t_patch_pos_diff) 84 | 85 | valid_neighbor_mask = tf.cast(1 - invalid_index_mask, dtype=s2t_patch_distance.dtype) 86 | #mls weight is stored in matrix with shape=[batch_size*n, KNN] 87 | s2t_mls_weight = -s2t_patch_distance / s2t_patch_squared_radius 88 | s2t_mls_weight -= tf.tile(tf.expand_dims(s2t_mls_weight[:,0], axis=-1), multiples=[1, KNN_PREFERRED]) 89 | s2t_mls_weight = tf.math.exp(s2t_mls_weight) 90 | 91 | #mask the patch verts out of ball 92 | s2t_mls_weight = s2t_mls_weight*valid_neighbor_mask 93 | #select projection points which are far away from surface 94 | s2t_mls_weight += tf.pad(tf.expand_dims(1 - valid_neighbor_mask[:,0], axis=-1), [[0,0],[0,KNN_PREFERRED-1]]) 95 | #mls weights should also be normalized 96 | s2t_mls_weight = s2t_mls_weight / tf.tile(tf.reshape(tf.reduce_sum(s2t_mls_weight,axis=-1),[-1,1]),multiples=[1,KNN_PREFERRED]) 97 | s2t_sdf = tf.reduce_sum(s2t_mls_weight*tf.reduce_sum(s2t_patch_pos_diff*s2t_patch_normal, axis=-1), axis=-1) 98 | return s2t_sdf 99 | 100 | def get_sdf(mls_points, mls_radius, query_points): 101 | #assume mls_points in shape [n,6] and query_points in shape[q,3] 102 | n_mls_points = tf.shape(mls_points)[0] 103 | n_query_points = query_points.shape[0] 104 | mls_points = tf.cast(mls_points, tf.float64) 105 | query_points = tf.cast(query_points, tf.float64) 106 | assert(mls_radius is not None) 107 | mls_radius = tf.cast(mls_radius, tf.float64) 108 | return eval_sdf_from_mls(query_points, mls_points[:,:3], mls_points[:,3:], \ 109 | n_query_points, n_mls_points, mls_radius) 110 | 111 | class MLS_Surface: 112 | def __init__(self, mls_points, mls_radius, sess=None, sdf_node=None, scale=None): 113 | self.mls_points = mls_points 114 | self.mls_radius = mls_radius 115 | if(sess is None): 116 | config = tf.ConfigProto() 117 | config.gpu_options.allow_growth = True 118 | self.sess = tf.Session(config=config) 119 | else: 120 | self.sess = sess 121 | 122 | self.sdf_node = sdf_node 123 | self.scale = scale 124 | 125 | def marching_cubes(self, output_obj_filename, voxel_resolution): 126 | voxel_dim = 2**voxel_resolution 127 | voxel_length = 2.0 / voxel_dim 128 | voxel_dim += 1 129 | 130 | grids = np.ones([voxel_dim, voxel_dim, voxel_dim])*1000000 #which means initial values are all not available (will be considered in our implemented marching_cubes_partial) 131 | assert(self.mls_radius is not None) 132 | 133 | mls_points_id = np.round((self.mls_points[:,:3] + 1) / voxel_length) 134 | 135 | np.clip(mls_points_id, 0, voxel_dim - 1, out = mls_points_id) 136 | mls_points_id = mls_points_id.astype(int) 137 | active_grids = np.zeros([voxel_dim, voxel_dim, voxel_dim]) 138 | active_grids[mls_points_id[:,0], mls_points_id[:, 1], mls_points_id[:, 2]] = 1.0; 139 | 140 | #might use larger dilation_layer_num 141 | max_radius = np.sqrt(np.max(self.mls_radius)) * 2 142 | dilation_layer_num = (int)(round(max_radius / voxel_length)) + 1 143 | print ("dilation layer number: ", dilation_layer_num) 144 | active_grids_large = ndimage.binary_dilation(active_grids, structure=struct2, iterations = dilation_layer_num) 145 | 146 | nonzeros = np.nonzero(active_grids_large) 147 | evaluated_points = np.stack(nonzeros, axis=1)*voxel_length - 1 148 | print("active number: ", nonzeros[0].shape) 149 | 150 | max_num_per_iter = 200000 151 | num_iter = math.ceil(nonzeros[0].shape[0] / max_num_per_iter) 152 | for i in range(0, num_iter): 153 | startid = i * max_num_per_iter 154 | endid = (i + 1) * max_num_per_iter 155 | if i == num_iter-1: 156 | endid = nonzeros[0].shape[0] 157 | if(self.sdf_node == None): 158 | sdf = self.sess.run(get_sdf(self.mls_points, self.mls_radius, evaluated_points[startid:endid])) 159 | else: 160 | sdf = self.sess.run(self.sdf_node, feed_dict={mls_points_ph:self.mls_points, mls_radius_ph:self.mls_radius, query_points_ph:evaluated_points[startid:endid]}) 161 | grids[nonzeros[0][startid:endid], nonzeros[1][startid:endid], nonzeros[2][startid:endid]] = sdf 162 | 163 | print('outputfilename: ', output_obj_filename) 164 | vertices, triangles = mcubes.marching_cubes_partial(-grids, 0.0) 165 | write_obj(output_obj_filename, vertices, triangles, scale=self.scale) 166 | 167 | def get_mls_points_radius_from_file(points_filename): 168 | mls_points = np.loadtxt(points_filename) 169 | if(mls_points.shape[1] != 6): 170 | #the xyz file does not contain normals: indicate this is not a mls file 171 | if(args.scale): 172 | return mls_points, None, None 173 | else: 174 | return mls_points, None 175 | if(os.path.exists(points_filename.replace(".xyz", "_radius.txt"))): 176 | mls_radius = np.loadtxt(points_filename.replace(".xyz", "_radius.txt")) 177 | else: 178 | print("radius of mls points not found") 179 | raise NotImplementedError 180 | mls_radius = None 181 | 182 | if(mls_radius.shape == ()): 183 | mls_radius = mls_radius*np.ones([mls_points.shape[0]]) 184 | assert(mls_points.shape[0] == mls_radius.shape[0]) 185 | 186 | if(args.scale): 187 | scale_file = points_filename.replace(".xyz", "_scale.txt") 188 | assert(os.path.exists(scale_file)) 189 | scale = 1.0 / np.loadtxt(scale_file) 190 | return mls_points, mls_radius, scale 191 | 192 | #normal of mls points should be normalized 193 | mls_postion = mls_points[:,:3] 194 | mls_normal = mls_points[:,3:] / np.linalg.norm(mls_points[:,3:], axis=1, keepdims=True) 195 | mls_points = np.concatenate([mls_postion, mls_normal], axis=1) 196 | 197 | return mls_points, mls_radius 198 | 199 | def process_folder(sess): 200 | #build graph in advance 201 | sdf_node = get_sdf(mls_points_ph, mls_radius_ph, query_points_ph) 202 | list = os.listdir(args.i) 203 | for file in list: 204 | if(file.endswith(".xyz")): 205 | points_filename = os.path.join(args.i, file) 206 | output_file_name = points_filename.replace(".xyz", "_mc.obj") 207 | if(not args.overwrite and os.path.exists(output_file_name)): 208 | continue 209 | if(args.scale): 210 | mls_points, mls_radius, scale = get_mls_points_radius_from_file(points_filename) 211 | else: 212 | mls_points, mls_radius = get_mls_points_radius_from_file(points_filename) 213 | scale = None 214 | if(mls_points.shape[1] != 6): 215 | continue 216 | mls_object = MLS_Surface(mls_points, mls_radius, sess=sess, sdf_node=sdf_node, scale=scale) 217 | 218 | if(args.overwrite or not os.path.exists(output_file_name)): 219 | mls_object.marching_cubes(output_file_name, args.res) 220 | 221 | def main(): 222 | config = tf.ConfigProto() 223 | config.gpu_options.allow_growth = True 224 | sess = tf.Session(config=config) 225 | 226 | start = time.time() 227 | if(args.i is not None and os.path.isdir(args.i)): 228 | #marching cube for all files under directory 229 | return process_folder(sess) 230 | 231 | #process single input file 232 | assert(args.i is not None and args.i.endswith(".xyz")) 233 | points_filename = args.i 234 | 235 | if(args.scale): 236 | mls_points, mls_radius, scale = get_mls_points_radius_from_file(points_filename) 237 | else: 238 | mls_points, mls_radius = get_mls_points_radius_from_file(points_filename) 239 | scale = None 240 | assert(mls_points.shape[1] == 6) 241 | mls_object = MLS_Surface(mls_points, mls_radius, sess=sess, scale=scale) 242 | 243 | #marching cube from implicit distance fields defined by mls control points 244 | mls_object.marching_cubes("marching_cube_128.obj" if args.o is None else args.o, args.res) 245 | 246 | end = time.time() 247 | print("{}s elapsed".format(end - start)) 248 | 249 | if __name__ == '__main__': 250 | print('Program running...') 251 | main() -------------------------------------------------------------------------------- /network_architecture.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import sys 3 | sys.path.append("Octree") 4 | from ocnn import * 5 | 6 | def preprocess_octree_data_from_unoriented_points(data, octree, depth): 7 | #subtract octree node mean position 8 | #print("=======================\ninput octree feature with shape:{}\n==========================".format(data.shape)) 9 | data = tf.reshape(data, [3, -1]) 10 | octree_node_xyz = tf.cast(octree_xyz(octree, depth), tf.int32) #we get nx4 11 | octree_node_center = tf.cast(octree_node_xyz[:,:3], dtype=tf.float32)*(0.5**(depth-1)) - 1 + (0.5**depth) 12 | #normalize to [-1, 1] 13 | data = (data - tf.transpose(octree_node_center))*(2**depth) 14 | 15 | #concat node occupancy indicator 16 | data = tf.concat([data, tf.ones([1, tf.shape(data)[1]], dtype=data.dtype)], axis=0) 17 | print("=======================\ninput octree feature with shape:{}\n==========================".format(data.shape)) 18 | 19 | #mask empty nodes to zero 20 | input_node_mask = tf.reshape(tf.greater(octree_property(octree, property_name="split", dtype=tf.float32, depth=depth, channel=1), 0), [1, -1]) 21 | input_node_mask = tf.cast(tf.tile(input_node_mask, multiples = [4,1]), data.dtype) 22 | data = tf.stop_gradient(data*input_node_mask) 23 | 24 | return tf.reshape(data, [1, 4, -1, 1]) 25 | 26 | def octree_network_unet(octree, input_octree_depth, output_octree_depth, octree_feature_channel, points_per_node, training, reuse=False, node_receptive_field=1.0, predict_radius=False, radius_range=2.0, gt_octree=None): 27 | #if gt_octree is None, this is ShapeAE Architecture else Shape Completion 28 | channels = [4, 64, 128, 128, 128, 64, 32, 16, 8] 29 | resblock_num = 3 30 | depth = input_octree_depth 31 | with tf.variable_scope('ocnn_encoder', reuse=reuse): 32 | with tf.variable_scope('signal_gt'): 33 | data = octree_property(octree, property_name="feature", dtype=tf.float32, 34 | depth=depth, channel=octree_feature_channel) 35 | if(octree_feature_channel == 3): 36 | #which means input signal does not include normal 37 | data = preprocess_octree_data_from_unoriented_points(data, octree, depth) 38 | else: 39 | data = tf.reshape(data, [1, octree_feature_channel, -1, 1]) 40 | 41 | with tf.variable_scope("front"): 42 | data = octree_conv_bn_relu(data, octree, depth, channels[depth], training) 43 | 44 | convd = [None]*10 45 | for d in range(depth, 1, -1): 46 | for i in range(0, resblock_num): 47 | with tf.variable_scope('resblock_%d_%d' % (d, i)): 48 | data = octree_resblock(data, octree, d, channels[d], 1, training) 49 | convd[d] = data #for skip connections 50 | if(d != 2): 51 | with tf.variable_scope('down_%d' % d): 52 | data = octree_conv_bn_relu(data, octree, d, channels[d-1], training, 53 | stride=2, kernel_size=[2]) 54 | 55 | code = data 56 | 57 | #decoder 58 | depth = output_octree_depth 59 | with tf.variable_scope('ocnn_decoder', reuse=reuse): 60 | data = code 61 | loss, accu = [], [] 62 | for d in range(2, depth + 1): 63 | for i in range(0, resblock_num): 64 | with tf.variable_scope('resblock_%d_%d' % (d, i)): 65 | data = octree_resblock(data, octree if gt_octree is None else gt_octree, d, channels[d], 1, training) 66 | 67 | with tf.variable_scope('predict_%d' % d): 68 | logit, label = predict_label(data, 2, 32, training) 69 | logit = tf.transpose(tf.squeeze(logit, [0,3])) # (1, C, H, 1) -> (H, C) 70 | 71 | with tf.variable_scope('octree_loss_%d' % d): 72 | with tf.variable_scope('label_gt'): 73 | label_gt = octree_property(octree if gt_octree is None else gt_octree, property_name="split", 74 | dtype=tf.float32, depth=d, channel=1) 75 | label_gt = tf.reshape(tf.cast(label_gt, dtype=tf.int32), [-1]) 76 | loss.append(softmax_loss(logit, label_gt, num_class=2)) 77 | accu.append(label_accuracy(label, label_gt)) 78 | 79 | signals_per_point = 6 80 | if(predict_radius): 81 | signals_per_point += 1 82 | 83 | if d == depth: 84 | with tf.variable_scope('regress_%d' % d): 85 | mls_points_local = predict_signal(data, points_per_node*signals_per_point, 128, training) #axis-angle 86 | #from signal to mls points, merge points 87 | mls_points_local = tf.reshape(tf.transpose(tf.squeeze(mls_points_local, [0,3])), [-1, signals_per_point]) # (1, C, H, 1) -> (H, C) 88 | position = tf.nn.tanh(mls_points_local[:,:3])*(0.5**depth)*node_receptive_field 89 | if(predict_radius): 90 | normal = mls_points_local[:,3:6] 91 | radius = tf.expand_dims(tf.math.pow(radius_range, tf.nn.tanh(mls_points_local[:,6])), axis=-1) 92 | normal = tf.concat([normal, radius], axis=-1) 93 | else: 94 | normal = mls_points_local[:,3:] 95 | 96 | octree_node_xyz = tf.cast(octree_xyz(octree if gt_octree is None else gt_octree, depth), tf.int32) #we get nx4 97 | octree_node_center = tf.cast(octree_node_xyz[:,:3], dtype=position.dtype)*(0.5**(depth-1)) - 1 + (0.5**depth) 98 | position += tf.reshape(tf.tile(tf.expand_dims(octree_node_center, axis=1), multiples = [1, points_per_node, 1]), [-1,3]) 99 | 100 | mls_points = tf.concat([position, normal], axis=-1) 101 | #points_nums = tf.segment_sum(tf.ones_like(point_segment, dtype=tf.int32), point_segment) #tf.cumsum() 102 | 103 | #mask empty octree node 104 | if(True): 105 | #using groundtruth supervision in last depth 106 | node_mask = tf.greater(label_gt, 0) 107 | else: 108 | print("with real split in unet decoder output") 109 | #using the real split supervision 110 | node_mask = tf.greater(label, 0) 111 | octree_node_xyz_valid = tf.boolean_mask(octree_node_xyz, node_mask) 112 | point_segment = octree_node_xyz_valid[:,3] 113 | 114 | #point_segment = tf.boolean_mask(octree_node_xyz[:,3], node_mask) 115 | point_segment = tf.cast(tf.reshape(tf.tile(tf.expand_dims(point_segment, axis=1), multiples=[1, points_per_node]), [-1]), tf.int32) 116 | points_mask = tf.reshape(tf.tile(tf.expand_dims(node_mask, axis=-1), multiples=[1, points_per_node]), [-1]) 117 | mls_points = tf.boolean_mask(mls_points, points_mask) 118 | 119 | if d < depth: 120 | with tf.variable_scope('up_%d' % d): 121 | data = octree_deconv_bn_relu(data, octree if gt_octree is None else gt_octree, d, channels[d-1], training, 122 | stride=2, kernel_size=[2]) 123 | #skip connections 124 | if(gt_octree is None): 125 | data = tf.concat([data, convd[d+1]], axis=1) 126 | else: 127 | skip, _ = octree_align(convd[d+1], octree, gt_octree, d+1) 128 | data = tf.concat([data, skip], axis=1) 129 | if(gt_octree is None): 130 | return mls_points, point_segment, octree_node_xyz_valid 131 | else: 132 | return loss, accu, mls_points, point_segment, octree_node_xyz_valid 133 | 134 | def octree_network_unet_completion_decode_shape(octree, input_octree_depth, output_octree_depth, octree_feature_channel, points_per_node, shape_batch_size, training, reuse=False, node_receptive_field=1.0, predict_radius=False, radius_range=2.0): 135 | assert(not training) 136 | #Shape Completion decode shape at test phase 137 | channels = [4, 64, 128, 128, 128, 64, 32, 16, 8] 138 | resblock_num = 3 139 | depth = input_octree_depth 140 | with tf.variable_scope('ocnn_encoder', reuse=reuse): 141 | with tf.variable_scope('signal_gt'): 142 | data = octree_property(octree, property_name="feature", dtype=tf.float32, 143 | depth=depth, channel=octree_feature_channel) 144 | if(octree_feature_channel == 3): 145 | #which means input signal does not include normal 146 | data = preprocess_octree_data_from_unoriented_points(data, octree, depth) 147 | else: 148 | data = tf.reshape(data, [1, octree_feature_channel, -1, 1]) 149 | 150 | with tf.variable_scope("front"): 151 | data = octree_conv_bn_relu(data, octree, depth, channels[depth], training) 152 | 153 | convd = [None]*10 154 | input_split_label = [None]*10 155 | 156 | for d in range(depth, 1, -1): 157 | input_split_label[d] = octree_property(octree, property_name="split", dtype=tf.float32, depth=d, channel=1) 158 | print("input split label with shape in depth {}: {}".format(d, input_split_label[d].shape)) 159 | for i in range(0, resblock_num): 160 | with tf.variable_scope('resblock_%d_%d' % (d, i)): 161 | data = octree_resblock(data, octree, d, channels[d], 1, training) 162 | convd[d] = data #for skip connections 163 | if(d != 2): 164 | with tf.variable_scope('down_%d' % d): 165 | data = octree_conv_bn_relu(data, octree, d, channels[d-1], training, 166 | stride=2, kernel_size=[2]) 167 | 168 | code = data 169 | 170 | #decoder 171 | depth = output_octree_depth 172 | with tf.variable_scope('ocnn_decoder', reuse=reuse): 173 | # init the octree 174 | with tf.variable_scope('octree_0'): 175 | #dis = False if flags.channel < 4 else True 176 | octree_prediction = octree_new(batch_size=shape_batch_size, channel=4, has_displace=False) 177 | with tf.variable_scope('octree_1'): 178 | octree_prediction = octree_grow(octree_prediction, target_depth=1, full_octree=True) 179 | with tf.variable_scope('octree_2'): 180 | octree_prediction = octree_grow(octree_prediction, target_depth=2, full_octree=True) 181 | 182 | data = code 183 | for d in range(2, depth + 1): 184 | for i in range(0, resblock_num): 185 | with tf.variable_scope('resblock_%d_%d' % (d, i)): 186 | data = octree_resblock(data, octree_prediction, d, channels[d], 1, training) 187 | 188 | with tf.variable_scope('predict_%d' % d): 189 | logit, label = predict_label(data, 2, 32, training) 190 | logit = tf.transpose(tf.squeeze(logit, [0,3])) # (1, C, H, 1) -> (H, C) 191 | 192 | #utilize input octree info 193 | #skip_label, _ = octree_align(tf.reshape(input_split_label[d], [1,1,-1,1]), octree, octree_prediction, d) 194 | #skip_label = tf.reshape(skip_label, [-1]) 195 | #print("label shape {}".format(label.shape)) 196 | #label = tf.cast(tf.greater(tf.cast(label, tf.float32)+skip_label, 0), tf.int32) 197 | 198 | with tf.variable_scope('octree_%d' % d): 199 | octree_prediction = octree_update(octree_prediction, label, depth=d, mask=1) 200 | 201 | if d == depth: 202 | signals_per_point = 6 203 | if(predict_radius): 204 | signals_per_point += 1 205 | with tf.variable_scope('regress_%d' % d): 206 | mls_points_local = predict_signal(data, points_per_node*signals_per_point, 128, training) #axis-angle 207 | #from signal to mls points, merge points 208 | mls_points_local = tf.reshape(tf.transpose(tf.squeeze(mls_points_local, [0,3])), [-1, signals_per_point]) # (1, C, H, 1) -> (H, C) 209 | position = tf.nn.tanh(mls_points_local[:,:3])*(0.5**depth)*node_receptive_field 210 | if(predict_radius): 211 | normal = mls_points_local[:,3:6] 212 | radius = tf.expand_dims(tf.math.pow(radius_range, tf.nn.tanh(mls_points_local[:,6])), axis=-1) 213 | normal = tf.concat([normal, radius], axis=-1) 214 | else: 215 | normal = mls_points_local[:,3:] 216 | 217 | octree_node_xyz = tf.cast(octree_xyz(octree_prediction, depth), tf.int32) #we get nx4 218 | octree_node_center = tf.cast(octree_node_xyz[:,:3], dtype=position.dtype)*(0.5**(depth-1)) - 1 + (0.5**depth) 219 | position += tf.reshape(tf.tile(tf.expand_dims(octree_node_center, axis=1), multiples = [1, points_per_node, 1]), [-1,3]) 220 | 221 | mls_points = tf.concat([position, normal], axis=-1) 222 | #points_nums = tf.segment_sum(tf.ones_like(point_segment, dtype=tf.int32), point_segment) #tf.cumsum() 223 | 224 | #mask empty octree node 225 | node_mask = tf.greater(label, 0) 226 | octree_node_xyz_valid = tf.boolean_mask(octree_node_xyz, node_mask) 227 | point_segment = octree_node_xyz_valid[:,3] 228 | 229 | #point_segment = tf.boolean_mask(octree_node_xyz[:,3], node_mask) 230 | point_segment = tf.cast(tf.reshape(tf.tile(tf.expand_dims(point_segment, axis=1), multiples=[1, points_per_node]), [-1]), tf.int32) 231 | points_mask = tf.reshape(tf.tile(tf.expand_dims(node_mask, axis=-1), multiples=[1, points_per_node]), [-1]) 232 | mls_points = tf.boolean_mask(mls_points, points_mask) 233 | 234 | if d < depth: 235 | with tf.variable_scope('octree_%d' % (d+1)): 236 | octree_prediction = octree_grow(octree_prediction, target_depth=d+1, full_octree=False) 237 | with tf.variable_scope('up_%d' % d): 238 | data = octree_deconv_bn_relu(data, octree_prediction, d, channels[d-1], training, 239 | stride=2, kernel_size=[2]) 240 | #skip connections 241 | skip, _ = octree_align(convd[d+1], octree, octree_prediction, d+1) 242 | data = tf.concat([data, skip], axis=1) 243 | 244 | return mls_points, point_segment, octree_node_xyz_valid 245 | -------------------------------------------------------------------------------- /points3d-tf/points3d/__init__.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import os 3 | from tensorflow.python.framework import ops 4 | 5 | _current_path = os.path.dirname(os.path.realpath(__file__)) 6 | _tf_points_module = tf.load_op_library(os.path.join(_current_path, 'libpoints3d.so')) 7 | 8 | get_neighbor_spatial_grid_local_self = _tf_points_module.get_neighbor_spatial_grid_local_self 9 | get_neighbor_spatial_grid_radius_v_voting = _tf_points_module.get_neighbor_spatial_grid_radius_v_voting 10 | 11 | ops.NotDifferentiable("GetNeighborSpatialGridRadiusVVoting") 12 | ops.NotDifferentiable("GetNeighborSpatialGridLocalSelf") -------------------------------------------------------------------------------- /points3d-tf/points3d/build.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | TF_CFLAGS=$(python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_compile_flags()))') 3 | TF_LFLAGS=$(python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_link_flags()))') 4 | 5 | nvcc -std=c++11 -c get_neighbor_points_spatial_grid_local_self.cu.cc $TF_CFLAGS -I /usr/local/cuda-10.0/include -x cu -Xcompiler -fPIC -D GOOGLE_CUDA=1 -I /usr/local -expt-relaxed-constexpr -DNDEBUG --expt-relaxed-constexpr --ptxas-options=--verbose -maxrregcount=55 6 | 7 | nvcc -std=c++11 -c get_neighbor_points_spatial_grid_radius_weighting_voting.cu.cc $TF_CFLAGS -I /usr/local/cuda-10.0/include -x cu -Xcompiler -fPIC -D GOOGLE_CUDA=1 -I /usr/local -expt-relaxed-constexpr -DNDEBUG --expt-relaxed-constexpr --ptxas-options=--verbose -maxrregcount=55 8 | 9 | g++ -std=c++11 -shared -o libpoints3d.so get_neighbor_points_spatial_grid_local_self.cc get_neighbor_points_spatial_grid_radius_weighting_voting.cc get_neighbor_points_spatial_grid_local_self.cu.o get_neighbor_points_spatial_grid_radius_weighting_voting.cu.o $TF_CFLAGS -I /usr/local/cuda-10.0/include -L /usr/local/cuda-10.0/lib64 -D GOOGLE_CUDA=1 $TF_LFLAGS -fPIC -lcudart 10 | 11 | rm *.o -------------------------------------------------------------------------------- /points3d-tf/points3d/get_neighbor_points_spatial_grid_local_self.cc: -------------------------------------------------------------------------------- 1 | #include "tensorflow/core/framework/op.h" 2 | #include "tensorflow/core/framework/op_kernel.h" 3 | #include "tensorflow/core/framework/shape_inference.h" 4 | #include "tensorflow/core/framework/common_shape_fns.h" 5 | 6 | namespace tensorflow { 7 | 8 | template 9 | void neighbor_points_spatial_grid_local_self(OpKernelContext* context, 10 | const int* n_source_points, const int batch_size, const T* in_source_points, 11 | const int knn, const int batch_points_total, int* neighbor_points_id_ptr); 12 | 13 | REGISTER_OP("GetNeighborSpatialGridLocalSelf") 14 | .Attr("T: {float, double}") 15 | .Input("source_points: T") 16 | .Input("source_points_num: int32") 17 | .Attr("knn: int = 10") 18 | .Output("neighbor_points: int32") 19 | .SetShapeFn([](shape_inference::InferenceContext* c) { 20 | c->set_output(0, c->MakeShape({c->Dim(c->input(0),0), c->UnknownDim()})); 21 | return Status::OK(); 22 | }) 23 | .Doc(R"doc( 24 | Find the neighbor points for each points using spatial partition and efficient insertion sort. 25 | )doc"); 26 | 27 | template 28 | class GetNeighborSpatialGridLocalSelfOp : public OpKernel { 29 | public: 30 | explicit GetNeighborSpatialGridLocalSelfOp(OpKernelConstruction* context) 31 | : OpKernel(context) 32 | { 33 | OP_REQUIRES_OK(context, context->GetAttr("knn", 34 | &this->knn_)); 35 | } 36 | 37 | void Compute(OpKernelContext* context) override { 38 | // in source points [\sum(n_source_points), 6] 39 | const Tensor& in_source_points = context->input(0); 40 | auto in_source_points_ptr = in_source_points.flat().data(); 41 | batch_points_total = in_source_points.dim_size(0); 42 | 43 | const Tensor& in_source_num = context->input(1); 44 | auto in_source_num_ptr = in_source_num.flat().data(); 45 | batch_size_ = in_source_num.dim_size(0); 46 | 47 | // out loss 48 | Tensor* out_vec = nullptr; 49 | TensorShape out_shape({batch_points_total, knn_}); 50 | OP_REQUIRES_OK(context, context->allocate_output("neighbor_points", out_shape, &out_vec)); 51 | auto out_ptr = out_vec->flat().data(); 52 | 53 | // find neighbor point 54 | neighbor_points_spatial_grid_local_self(context, in_source_num_ptr, batch_size_, in_source_points_ptr, knn_, batch_points_total, out_ptr); 55 | } 56 | 57 | private: 58 | int batch_points_total; 59 | int batch_size_; 60 | int knn_; 61 | }; 62 | 63 | REGISTER_KERNEL_BUILDER(Name("GetNeighborSpatialGridLocalSelf").Device(DEVICE_GPU).TypeConstraint("T"), 64 | GetNeighborSpatialGridLocalSelfOp); 65 | 66 | REGISTER_KERNEL_BUILDER(Name("GetNeighborSpatialGridLocalSelf").Device(DEVICE_GPU).TypeConstraint("T"), 67 | GetNeighborSpatialGridLocalSelfOp); 68 | } // namespace tensorflow 69 | -------------------------------------------------------------------------------- /points3d-tf/points3d/get_neighbor_points_spatial_grid_local_self.cu.cc: -------------------------------------------------------------------------------- 1 | #define EIGEN_USE_THREADS 2 | #define EIGEN_USE_GPU 3 | 4 | #include "tensorflow/core/util/gpu_kernel_helper.h" 5 | #include 6 | #include 7 | #include 8 | #include 9 | 10 | #include 11 | #include 12 | #include 13 | #include 14 | 15 | // Macro definition 16 | #define CudaAssert( X ) if ( !(X) ) {printf( "Thread %d:%d failed assert at %s:%d!\n", blockIdx.x, threadIdx.x, __FILE__, __LINE__ ); return;} 17 | 18 | namespace tensorflow { 19 | 20 | typedef Eigen::GpuDevice GPUDevice; 21 | 22 | template 23 | static __global__ void find_knn_grids_local_self( 24 | const int nthreads, const int* n_source_points, const T* in_source_points, const int batch_size, const int knn, const int CUBES_TOTAL, const int MAX_POINTS_IN_CUBE, 25 | const int *cube_points_num_ptr, const int *cube_points_indices_ptr, int *neighbor_points_id_ptr, T *neighbor_points_dist_ptr) { 26 | // 0 1 2 3 4 5 6 7 10 14 18 19 21 23 25 27 | __shared__ int disp_x[27];// = {0, -1, 0, 0, 0, 0, 1, -1, -1, -1, -1, 0, 0, 0, 0, 1, 1, 1, 1, -1, -1, -1, -1, 1, 1, 1, 1}; 28 | __shared__ int disp_y[27];// = {0, 0, -1, 0, 0, 1, 0, -1, 0, 0, 1, -1, -1, 1, 1, 0, 0, -1, 1, 1, 1, -1, -1, -1, -1, 1, 1}; 29 | __shared__ int disp_z[27];// = {0, 0, 0, -1, 1, 0, 0, 0, -1, 1, 0, -1, 1, -1, 1, -1, 1, 0, 0, -1, 1, -1, 1, -1, 1, -1, 1}; 30 | 31 | //fill the shared memory 32 | if(threadIdx.x < 27) 33 | { 34 | int i = threadIdx.x; 35 | 36 | if(i < 9) 37 | disp_x[i] = 0; 38 | else if (i < 18) 39 | disp_x[i] = -1; 40 | else 41 | disp_x[i] = 1; 42 | 43 | if ((i/3) % 3 == 0) 44 | disp_y[i] = 0; 45 | else if ((i/3) % 3 == 1) 46 | disp_y[i] = -1; 47 | else 48 | disp_y[i] = 1; 49 | 50 | if(i % 3 == 0) 51 | disp_z[i] = 0; 52 | else if(i % 3 == 1) 53 | disp_z[i] = -1; 54 | else 55 | disp_z[i] = 1; 56 | } 57 | __syncthreads(); 58 | 59 | int batch_index = 0; 60 | 61 | CUDA_1D_KERNEL_LOOP(index, nthreads) { 62 | 63 | for(int ii = batch_index; ii < batch_size; ii++) 64 | if(index < n_source_points[ii]) 65 | { 66 | batch_index = ii; 67 | break; 68 | } 69 | 70 | int* cur_neighbor_points_id_ptr = neighbor_points_id_ptr + index*knn; 71 | T *cur_neighbor_points_dist_ptr = neighbor_points_dist_ptr + index*knn; 72 | 73 | for(int ii = 0; ii < knn; ii++) 74 | cur_neighbor_points_id_ptr[ii] = -1; 75 | 76 | T sx = in_source_points[6 * index]; 77 | T sy = in_source_points[6 * index + 1]; 78 | T sz = in_source_points[6 * index + 2]; 79 | 80 | //find current cube and nearby cubes 81 | int scx = int((sx + 1) * 5); 82 | int scy = int((sy + 1) * 5); 83 | int scz = int((sz + 1) * 5); 84 | 85 | if(scx < 0) 86 | scx = 0; 87 | else if(scx > 9) 88 | scx = 9; 89 | 90 | if(scy < 0) 91 | scy = 0; 92 | else if (scy > 9) 93 | scy = 9; 94 | 95 | if(scz < 0) 96 | scz = 0; 97 | else if(scz > 9) 98 | scz = 9; 99 | 100 | CudaAssert(scx>=0 && scy>=0 && scz>=0 && scx<10 && scy<10 && scz<10); 101 | 102 | int k = 0; 103 | //T tnx,tny,tnz,tn_norm; 104 | T dx,dy,dz,cur_distance; 105 | //T dot_product; 106 | 107 | for (int j = 0; j < 27; j++) 108 | { 109 | int tcx = scx + disp_x[j]; 110 | int tcy = scy + disp_y[j]; 111 | int tcz = scz + disp_z[j]; 112 | if(tcx < 0 || tcx >=10) 113 | continue; 114 | if(tcy < 0 || tcy >=10) 115 | continue; 116 | if(tcz < 0 || tcz >=10) 117 | continue; 118 | int cube_index = 100*tcx + 10*tcy + tcz; 119 | const int points_in_cube = cube_points_num_ptr[batch_index*CUBES_TOTAL + cube_index]; 120 | const int *cube_indices = cube_points_indices_ptr + batch_index*CUBES_TOTAL*MAX_POINTS_IN_CUBE + cube_index * MAX_POINTS_IN_CUBE; 121 | 122 | for (int pt = 0; pt < points_in_cube; pt++) 123 | { 124 | int tid = cube_indices[pt]; 125 | if(tid == index) 126 | continue; 127 | 128 | dx = sx - in_source_points[6 * tid]; 129 | dy = sy - in_source_points[6 * tid + 1]; 130 | dz = sz - in_source_points[6 * tid + 2]; 131 | 132 | cur_distance = dx * dx + dy * dy + dz * dz; 133 | 134 | int iii; 135 | 136 | if(k < knn) 137 | { 138 | //do a insertion less than knn 139 | iii = k; 140 | } 141 | else if(cur_distance < cur_neighbor_points_dist_ptr[knn-1]) 142 | { 143 | //do the insertion 144 | iii = knn-1; 145 | } 146 | else 147 | continue; 148 | 149 | //the actual comparison 150 | for(; iii > 0; iii--) 151 | { 152 | if(cur_distance < cur_neighbor_points_dist_ptr[iii-1]) 153 | { 154 | cur_neighbor_points_dist_ptr[iii] = cur_neighbor_points_dist_ptr[iii-1]; 155 | cur_neighbor_points_id_ptr[iii] = cur_neighbor_points_id_ptr[iii-1]; 156 | } 157 | else 158 | break; 159 | } 160 | cur_neighbor_points_dist_ptr[iii] = cur_distance; 161 | cur_neighbor_points_id_ptr[iii] = tid; 162 | k++; 163 | } 164 | } 165 | 166 | //CudaAssert(k>=knn); 167 | if(k == 0) 168 | { 169 | cur_neighbor_points_dist_ptr[0] = 1e20; 170 | cur_neighbor_points_id_ptr[0] = -1; 171 | 172 | int start_index = 0; 173 | if(batch_index != 0) 174 | start_index = n_source_points[batch_index - 1]; 175 | 176 | //continue search until find nearest point at least 177 | for (int pt = start_index; pt < n_source_points[batch_index]; pt++) 178 | { 179 | if(pt == index) 180 | continue; 181 | 182 | dx = sx - in_source_points[6 * pt]; 183 | dy = sy - in_source_points[6 * pt + 1]; 184 | dz = sz - in_source_points[6 * pt + 2]; 185 | 186 | cur_distance = dx * dx + dy * dy + dz * dz; 187 | 188 | if(cur_distance < cur_neighbor_points_dist_ptr[0]) 189 | { 190 | cur_neighbor_points_dist_ptr[0] = cur_distance; 191 | cur_neighbor_points_id_ptr[0] = pt; 192 | } 193 | } 194 | } 195 | CudaAssert(cur_neighbor_points_id_ptr[0] >= 0); 196 | } 197 | } 198 | 199 | template 200 | static __global__ void gather_cube_pointsb_local( 201 | const int nthreads, const int* n_target_points, const T* in_target_points, const int MAX_POINTS_IN_CUBE, 202 | const int CUBES_TOTAL, int *cube_points_num_ptr, int *cube_points_indices_ptr) { 203 | CUDA_1D_KERNEL_LOOP(batch_index, nthreads) { 204 | int *cube_indices = cube_points_indices_ptr + batch_index * MAX_POINTS_IN_CUBE * CUBES_TOTAL; 205 | int *batch_cube_points_num_ptr = cube_points_num_ptr + batch_index*CUBES_TOTAL; 206 | 207 | int start_index = 0; 208 | if(batch_index != 0) 209 | start_index = n_target_points[batch_index - 1]; 210 | 211 | for (int n_target_index = start_index; n_target_index < n_target_points[batch_index]; n_target_index++) 212 | { 213 | T tx = in_target_points[6*n_target_index]; 214 | T ty = in_target_points[6*n_target_index + 1]; 215 | T tz = in_target_points[6*n_target_index + 2]; 216 | //find if point is inside cube 217 | int x = int((tx + 1) * 5); 218 | int y = int((ty + 1) * 5); 219 | int z = int((tz + 1) * 5); 220 | if(x < 0) 221 | x = 0; 222 | else if(x > 9) 223 | x = 9; 224 | 225 | if(y < 0) 226 | y = 0; 227 | else if (y > 9) 228 | y = 9; 229 | 230 | if(z < 0) 231 | z = 0; 232 | else if(z > 9) 233 | z = 9; 234 | 235 | int cube_index = 100*x+10*y+z; 236 | cube_indices[cube_index*MAX_POINTS_IN_CUBE + batch_cube_points_num_ptr[cube_index]] = n_target_index; 237 | batch_cube_points_num_ptr[cube_index]++; 238 | CudaAssert(batch_cube_points_num_ptr[cube_index] <= MAX_POINTS_IN_CUBE); 239 | } 240 | } 241 | } 242 | 243 | template 244 | void neighbor_points_spatial_grid_local_self(OpKernelContext* context, 245 | const int* n_source_points, const int batch_size, const T* in_source_points, 246 | const int knn, const int batch_points_total, int* neighbor_points_id_ptr) { 247 | // get GPU device 248 | GPUDevice d = context->eigen_device(); 249 | CudaLaunchConfig config; 250 | int nthreads; 251 | 252 | const int CUBES_TOTAL = 1000; 253 | const int MAX_POINTS_IN_CUBE = 4096; 254 | 255 | //assume input points are bounded by a unit sphere, and we partition it to 10x10x10 grids 256 | const TensorShape cube_index_shape({batch_size, CUBES_TOTAL, MAX_POINTS_IN_CUBE}); 257 | Tensor cube_points_indices; 258 | OP_REQUIRES_OK(context, context->allocate_temp(DT_INT32, cube_index_shape, &cube_points_indices)); 259 | auto cube_points_indices_ptr = cube_points_indices.flat().data(); 260 | 261 | const TensorShape cube_points_num_shape({batch_size, CUBES_TOTAL}); 262 | Tensor cube_points_num; 263 | OP_REQUIRES_OK(context, context->allocate_temp(DT_INT32, cube_points_num_shape, &cube_points_num)); 264 | auto cube_points_num_ptr = cube_points_num.flat().data(); 265 | cudaMemset(cube_points_num_ptr, 0, sizeof(int)*batch_size*CUBES_TOTAL); 266 | 267 | const TensorShape neighbor_points_distance_shape({batch_points_total, knn}); 268 | Tensor neighbor_points_distance; 269 | OP_REQUIRES_OK(context, context->allocate_temp(sizeof(T)==4?DT_FLOAT:DT_DOUBLE, neighbor_points_distance_shape, &neighbor_points_distance)); 270 | auto neighbor_points_distance_ptr = neighbor_points_distance.flat().data(); 271 | 272 | nthreads = batch_size; 273 | config = GetCudaLaunchConfig(nthreads, d); 274 | gather_cube_pointsb_local 275 | <<>>( 276 | nthreads, n_source_points, in_source_points, MAX_POINTS_IN_CUBE, CUBES_TOTAL, cube_points_num_ptr, cube_points_indices_ptr); 277 | 278 | nthreads = batch_points_total; 279 | config = GetCudaLaunchConfig(nthreads, d); 280 | //printf("%d block and %d threads per block\n", config.block_count, config.thread_per_block); 281 | find_knn_grids_local_self 282 | <<>>( 283 | nthreads, n_source_points, in_source_points, batch_size, knn, CUBES_TOTAL, MAX_POINTS_IN_CUBE, 284 | cube_points_num_ptr, cube_points_indices_ptr, neighbor_points_id_ptr, neighbor_points_distance_ptr); 285 | cudaDeviceSynchronize(); 286 | 287 | cudaError_t err2 = cudaGetLastError(); 288 | if (err2 != cudaSuccess) 289 | printf("Cuda error: %s\n", cudaGetErrorString(err2)); 290 | 291 | } 292 | 293 | //explicit instantiation 294 | template void neighbor_points_spatial_grid_local_self(OpKernelContext* context, 295 | const int* n_source_points, const int batch_size, 296 | const float* in_source_points, const int knn, const int batch_points_total, int* neighbor_points_id_ptr); 297 | 298 | template void neighbor_points_spatial_grid_local_self(OpKernelContext* context, 299 | const int* n_source_points, const int batch_size, 300 | const double* in_source_points, const int knn, const int batch_points_total, int* neighbor_points_id_ptr); 301 | 302 | } // namespace tensorflow 303 | -------------------------------------------------------------------------------- /points3d-tf/points3d/get_neighbor_points_spatial_grid_radius_weighting_voting.cc: -------------------------------------------------------------------------------- 1 | #include "tensorflow/core/framework/op.h" 2 | #include "tensorflow/core/framework/op_kernel.h" 3 | #include "tensorflow/core/framework/shape_inference.h" 4 | #include "tensorflow/core/framework/common_shape_fns.h" 5 | 6 | namespace tensorflow { 7 | 8 | template 9 | void neighbor_points_spatial_grid_no_normal_weight_local_varying_radius_voting(OpKernelContext* context, 10 | const int n_source_points, const int* n_target_points, const int batch_size, 11 | const T* in_source_points, const T* in_target_points, const T* squared_radius, const int knn, 12 | int* neighbor_points_id_ptr); 13 | 14 | REGISTER_OP("GetNeighborSpatialGridRadiusVVoting") 15 | .Attr("T: {float, double}") 16 | .Input("source_points: T") 17 | .Input("target_points: T") 18 | .Input("target_point_num: int32") 19 | .Input("squared_radius: T") 20 | .Attr("knn: int = 10") 21 | .Output("neighbor_points: int32") 22 | .SetShapeFn([](shape_inference::InferenceContext* c) { 23 | c->set_output(0, c->MakeShape({c->Dim(c->input(0),0),c->UnknownDim(),c->UnknownDim()})); 24 | return Status::OK(); 25 | }) 26 | .Doc(R"doc( 27 | Find the neighbor points for each source points in target pointset using spatial partition and efficient insertion sort. 28 | )doc"); 29 | 30 | template 31 | class GetNeighborSpatialGridRadiusVVotingOp : public OpKernel { 32 | public: 33 | explicit GetNeighborSpatialGridRadiusVVotingOp(OpKernelConstruction* context) 34 | : OpKernel(context) 35 | { 36 | OP_REQUIRES_OK(context, context->GetAttr("knn", 37 | &this->knn_)); 38 | } 39 | 40 | void Compute(OpKernelContext* context) override { 41 | // in source points [bs, 3*n_source_points] 42 | const Tensor& in_source_points = context->input(0); 43 | auto in_source_points_ptr = in_source_points.flat().data(); 44 | batch_size_ = in_source_points.dim_size(0); 45 | n_source_points_ = in_source_points.dim_size(1)/3; 46 | 47 | // in target points [\sum(n_target_points), 6] 48 | const Tensor& in_target_points = context->input(1); 49 | auto in_target_points_ptr = in_target_points.flat().data(); 50 | CHECK_EQ(in_target_points.dim_size(1), 6); 51 | 52 | const Tensor& in_target_num = context->input(2); 53 | auto in_target_num_ptr = in_target_num.flat().data(); 54 | CHECK_EQ(in_target_num.dim_size(0), batch_size_); 55 | 56 | //in shape [\sum(n_target_points)] 57 | const Tensor& in_squared_radius = context->input(3); 58 | auto in_squared_radius_ptr = in_squared_radius.flat().data(); 59 | CHECK_EQ(in_squared_radius.dim_size(0), in_target_points.dim_size(0)); 60 | 61 | // out tensor 62 | Tensor* out_vec = nullptr; 63 | TensorShape out_shape({batch_size_, n_source_points_, knn_}); 64 | OP_REQUIRES_OK(context, context->allocate_output("neighbor_points", out_shape, &out_vec)); 65 | auto out_ptr = out_vec->flat().data(); 66 | 67 | // find neighbor points 68 | neighbor_points_spatial_grid_no_normal_weight_local_varying_radius_voting(context, n_source_points_, in_target_num_ptr, 69 | batch_size_, in_source_points_ptr, in_target_points_ptr, in_squared_radius_ptr, knn_, out_ptr); 70 | } 71 | 72 | private: 73 | int n_source_points_; 74 | int batch_size_; 75 | int knn_; 76 | }; 77 | 78 | REGISTER_KERNEL_BUILDER(Name("GetNeighborSpatialGridRadiusVVoting").Device(DEVICE_GPU).TypeConstraint("T"), 79 | GetNeighborSpatialGridRadiusVVotingOp); 80 | 81 | REGISTER_KERNEL_BUILDER(Name("GetNeighborSpatialGridRadiusVVoting").Device(DEVICE_GPU).TypeConstraint("T"), 82 | GetNeighborSpatialGridRadiusVVotingOp); 83 | } // namespace tensorflow 84 | -------------------------------------------------------------------------------- /points3d-tf/points3d/get_neighbor_points_spatial_grid_radius_weighting_voting.cu.cc: -------------------------------------------------------------------------------- 1 | #define EIGEN_USE_THREADS 2 | #define EIGEN_USE_GPU 3 | 4 | #include "tensorflow/core/util/gpu_kernel_helper.h" 5 | #include 6 | #include 7 | #include 8 | #include 9 | 10 | #include 11 | #include 12 | #include 13 | #include 14 | 15 | // Macro definition 16 | #define CudaAssert( X ) if ( !(X) ) {printf( "Thread %d:%d failed assert at %s:%d!\n", blockIdx.x, threadIdx.x, __FILE__, __LINE__ ); return;} 17 | 18 | //assume mls weighting function is exp(-3d^2/r^2) 19 | 20 | namespace tensorflow { 21 | 22 | typedef Eigen::GpuDevice GPUDevice; 23 | 24 | template 25 | static __global__ void find_knn_grids_no_normalw_local_varying_radius_voting( 26 | const int nthreads, const int n_source_points, const int* n_target_points, const T* in_source_points, const T* in_target_points, 27 | const T* in_squared_radius, const int knn, const int CUBES_TOTAL, const int MAX_POINTS_IN_CUBE, 28 | const int *cube_points_num_ptr, const int *cube_points_indices_ptr, int *neighbor_points_id_ptr, T *neighbor_points_dist_ptr) { 29 | // 0 1 2 3 4 5 6 7 10 14 18 19 21 23 25 30 | __shared__ int disp_x[27];// = {0, -1, 0, 0, 0, 0, 1, -1, -1, -1, -1, 0, 0, 0, 0, 1, 1, 1, 1, -1, -1, -1, -1, 1, 1, 1, 1}; 31 | __shared__ int disp_y[27];// = {0, 0, -1, 0, 0, 1, 0, -1, 0, 0, 1, -1, -1, 1, 1, 0, 0, -1, 1, 1, 1, -1, -1, -1, -1, 1, 1}; 32 | __shared__ int disp_z[27];// = {0, 0, 0, -1, 1, 0, 0, 0, -1, 1, 0, -1, 1, -1, 1, -1, 1, 0, 0, -1, 1, -1, 1, -1, 1, -1, 1}; 33 | 34 | 35 | 36 | //fill the shared memory 37 | if(threadIdx.x < 27) 38 | { 39 | int i = threadIdx.x; 40 | 41 | if(i < 9) 42 | disp_x[i] = 0; 43 | else if (i < 18) 44 | disp_x[i] = -1; 45 | else 46 | disp_x[i] = 1; 47 | 48 | if ((i/3) % 3 == 0) 49 | disp_y[i] = 0; 50 | else if ((i/3) % 3 == 1) 51 | disp_y[i] = -1; 52 | else 53 | disp_y[i] = 1; 54 | 55 | if(i % 3 == 0) 56 | disp_z[i] = 0; 57 | else if(i % 3 == 1) 58 | disp_z[i] = -1; 59 | else 60 | disp_z[i] = 1; 61 | 62 | //printf("%d %d %d\n", threadIdx.x, blockIdx.x, gridDim.x); 63 | } 64 | __syncthreads(); 65 | 66 | CUDA_1D_KERNEL_LOOP(index, nthreads) { 67 | int batch_index = index / n_source_points; 68 | int n_source_index = index % n_source_points; 69 | 70 | int* cur_neighbor_points_id_ptr = neighbor_points_id_ptr + index*knn; 71 | T *cur_neighbor_points_dist_ptr = neighbor_points_dist_ptr + index*knn; 72 | 73 | for(int ii = 0; ii < knn; ii++) 74 | cur_neighbor_points_id_ptr[ii] = -1; 75 | 76 | T sx = in_source_points[(batch_index*3 + 0)*n_source_points + n_source_index]; 77 | T sy = in_source_points[(batch_index*3 + 1)*n_source_points + n_source_index]; 78 | T sz = in_source_points[(batch_index*3 + 2)*n_source_points + n_source_index]; 79 | 80 | //filter out padded data 81 | if(sx < -10) 82 | { 83 | cur_neighbor_points_id_ptr[0] = 0; 84 | continue; 85 | } 86 | 87 | //find current cube and nearby cubes 88 | int scx = int((sx + 1) * 5); 89 | int scy = int((sy + 1) * 5); 90 | int scz = int((sz + 1) * 5); 91 | 92 | if(scx < 0) 93 | scx = 0; 94 | else if(scx > 9) 95 | scx = 9; 96 | 97 | if(scy < 0) 98 | scy = 0; 99 | else if (scy > 9) 100 | scy = 9; 101 | 102 | if(scz < 0) 103 | scz = 0; 104 | else if(scz > 9) 105 | scz = 9; 106 | 107 | CudaAssert(scx>=0 && scy>=0 && scz>=0 && scx<10 && scy<10 && scz<10); 108 | 109 | int k = 0; 110 | 111 | T dx, dy, dz, cur_distance; 112 | for (int j = 0; j < 27; j++) 113 | { 114 | int tcx = scx + disp_x[j]; 115 | int tcy = scy + disp_y[j]; 116 | int tcz = scz + disp_z[j]; 117 | if(tcx < 0 || tcx >=10) 118 | continue; 119 | if(tcy < 0 || tcy >=10) 120 | continue; 121 | if(tcz < 0 || tcz >=10) 122 | continue; 123 | int cube_index = 100*tcx + 10*tcy + tcz; 124 | const int points_in_cube = cube_points_num_ptr[batch_index*CUBES_TOTAL + cube_index]; 125 | const int *cube_indices = cube_points_indices_ptr + batch_index*CUBES_TOTAL*MAX_POINTS_IN_CUBE + cube_index * MAX_POINTS_IN_CUBE; 126 | 127 | for (int pt = 0; pt < points_in_cube; pt++) 128 | { 129 | int tid = cube_indices[pt]; 130 | 131 | dx = sx - in_target_points[6*tid]; 132 | dy = sy - in_target_points[6*tid + 1]; 133 | dz = sz - in_target_points[6*tid + 2]; 134 | 135 | T squared_radius = in_squared_radius[tid]; 136 | cur_distance = dx * dx + dy * dy + dz * dz; 137 | //compute mls weights 138 | cur_distance = cur_distance / squared_radius; //sort in ascent order 139 | 140 | int iii; 141 | 142 | if(k < knn) 143 | { 144 | //do a insertion less than knn 145 | iii = k; 146 | } 147 | else if(cur_distance < cur_neighbor_points_dist_ptr[knn-1]) 148 | { 149 | //do the insertion 150 | iii = knn-1; 151 | } 152 | else 153 | continue; 154 | 155 | //the actual comparison 156 | for(; iii > 0; iii--) 157 | { 158 | if(cur_distance < cur_neighbor_points_dist_ptr[iii-1]) 159 | { 160 | cur_neighbor_points_dist_ptr[iii] = cur_neighbor_points_dist_ptr[iii-1]; 161 | cur_neighbor_points_id_ptr[iii] = cur_neighbor_points_id_ptr[iii-1]; 162 | } 163 | else 164 | break; 165 | } 166 | cur_neighbor_points_dist_ptr[iii] = cur_distance; 167 | cur_neighbor_points_id_ptr[iii] = tid; 168 | k++; 169 | } 170 | } 171 | 172 | if(k <= 1) 173 | { 174 | //first try another ring 175 | for(int dx1 = -2; dx1 <= 2; dx1++) 176 | for(int dx2 = -2; dx2 <= 2; dx2++) 177 | for(int dx3 = -2; dx3 <= 2; dx3++) 178 | { 179 | if(abs(dx1) != 2 && abs(dx2) != 2 && abs(dx3) != 2) 180 | continue; 181 | int tcx = scx + dx1; 182 | int tcy = scy + dx2; 183 | int tcz = scz + dx3; 184 | if(tcx < 0 || tcx >=10) 185 | continue; 186 | if(tcy < 0 || tcy >=10) 187 | continue; 188 | if(tcz < 0 || tcz >=10) 189 | continue; 190 | int cube_index = 100*tcx + 10*tcy + tcz; 191 | const int points_in_cube = cube_points_num_ptr[batch_index*CUBES_TOTAL + cube_index]; 192 | const int *cube_indices = cube_points_indices_ptr + batch_index*CUBES_TOTAL*MAX_POINTS_IN_CUBE + cube_index * MAX_POINTS_IN_CUBE; 193 | 194 | for (int pt = 0; pt < points_in_cube; pt++) 195 | { 196 | int tid = cube_indices[pt]; 197 | 198 | dx = sx - in_target_points[6*tid]; 199 | dy = sy - in_target_points[6*tid + 1]; 200 | dz = sz - in_target_points[6*tid + 2]; 201 | 202 | T squared_radius = in_squared_radius[tid];//[batch_index*n_target_points + tid]; 203 | cur_distance = dx * dx + dy * dy + dz * dz; 204 | //compute mls weights 205 | cur_distance = cur_distance / squared_radius; //sort in ascent order 206 | 207 | int iii; 208 | 209 | if(k < knn) 210 | { 211 | //do a insertion less than knn 212 | iii = k; 213 | } 214 | else if(cur_distance < cur_neighbor_points_dist_ptr[knn-1]) 215 | { 216 | //do the insertion 217 | iii = knn-1; 218 | } 219 | else 220 | continue; 221 | 222 | //the actual comparison 223 | for(; iii > 0; iii--) 224 | { 225 | if(cur_distance < cur_neighbor_points_dist_ptr[iii-1]) 226 | { 227 | cur_neighbor_points_dist_ptr[iii] = cur_neighbor_points_dist_ptr[iii-1]; 228 | cur_neighbor_points_id_ptr[iii] = cur_neighbor_points_id_ptr[iii-1]; 229 | } 230 | else 231 | break; 232 | } 233 | cur_neighbor_points_dist_ptr[iii] = cur_distance; 234 | cur_neighbor_points_id_ptr[iii] = tid; 235 | k++; 236 | } 237 | } 238 | } 239 | 240 | //CudaAssert(k>=knn); 241 | if(k <= 1) 242 | { 243 | //rare case 244 | k = 0; 245 | cur_neighbor_points_dist_ptr[0] = 1e20; 246 | cur_neighbor_points_id_ptr[0] = -1; 247 | 248 | int start_index = 0; 249 | if(batch_index != 0) 250 | start_index = n_target_points[batch_index - 1]; 251 | 252 | //continue search until find nearest point at least 253 | for (int pt = start_index; pt < n_target_points[batch_index]; pt++) 254 | { 255 | dx = sx - in_target_points[6 * pt]; 256 | dy = sy - in_target_points[6 * pt + 1]; 257 | dz = sz - in_target_points[6 * pt + 2]; 258 | 259 | cur_distance = dx * dx + dy * dy + dz * dz; 260 | T squared_radius = in_squared_radius[pt];//[batch_index*n_target_points + pt]; 261 | //compute mls weights 262 | cur_distance = cur_distance / squared_radius; //sort in ascent order 263 | 264 | int iii; 265 | 266 | if(k < knn) 267 | { 268 | //do a insertion less than knn 269 | iii = k; 270 | } 271 | else if(cur_distance < cur_neighbor_points_dist_ptr[knn-1]) 272 | { 273 | //do the insertion 274 | iii = knn-1; 275 | } 276 | else 277 | continue; 278 | 279 | //the actual comparison 280 | for(; iii > 0; iii--) 281 | { 282 | if(cur_distance < cur_neighbor_points_dist_ptr[iii-1]) 283 | { 284 | cur_neighbor_points_dist_ptr[iii] = cur_neighbor_points_dist_ptr[iii-1]; 285 | cur_neighbor_points_id_ptr[iii] = cur_neighbor_points_id_ptr[iii-1]; 286 | } 287 | else 288 | break; 289 | } 290 | cur_neighbor_points_dist_ptr[iii] = cur_distance; 291 | cur_neighbor_points_id_ptr[iii] = pt; 292 | k++; 293 | } 294 | } 295 | 296 | CudaAssert(cur_neighbor_points_id_ptr[0] >= 0); 297 | //voting according to normal 298 | if(k > 1) 299 | { 300 | if(k > knn) 301 | k = knn; 302 | T tnx,tny,tnz, tn_norm; 303 | T sn_norm; 304 | //filter according to normal 305 | tnx = 0; 306 | tny = 0; 307 | tnz = 0; 308 | //get average normal 309 | for(int iii = 0; iii < k; iii++) 310 | { 311 | int tid = cur_neighbor_points_id_ptr[iii]; 312 | T mls_weight = exp(cur_neighbor_points_dist_ptr[0] - cur_neighbor_points_dist_ptr[iii]); 313 | tnx += mls_weight * in_target_points[6 * tid + 3]; 314 | tny += mls_weight * in_target_points[6 * tid + 4]; 315 | tnz += mls_weight * in_target_points[6 * tid + 5]; 316 | } 317 | //normalize normal 318 | tn_norm = sqrt(tnx*tnx+tny*tny+tnz*tnz); 319 | CudaAssert(tn_norm > 1e-20); 320 | tnx /= tn_norm; 321 | tny /= tn_norm; 322 | tnz /= tn_norm; 323 | 324 | //remove neighbors with opposite normal direction 325 | for(int iii = k - 1; iii >= 0; iii--) 326 | { 327 | int tid = cur_neighbor_points_id_ptr[iii]; 328 | T dot_product = tnx*in_target_points[6 * tid + 3]; 329 | dot_product += tny*in_target_points[6 * tid + 4]; 330 | dot_product += tnz*in_target_points[6 * tid + 5]; 331 | sn_norm = in_target_points[6 * tid + 3]*in_target_points[6 * tid + 3] + in_target_points[6 * tid + 4]*in_target_points[6 * tid + 4] + in_target_points[6 * tid + 5]*in_target_points[6 * tid + 5]; 332 | 333 | if(dot_product < -0.5*sqrt(sn_norm)) 334 | { 335 | //remove point from neighbor list 336 | //cur_neighbor_points_id_ptr[iii] = cur_neighbor_points_id_ptr[k - 1]; 337 | for(int jjj = iii; jjj < k - 1; jjj++) 338 | cur_neighbor_points_id_ptr[jjj] = cur_neighbor_points_id_ptr[jjj + 1]; 339 | cur_neighbor_points_id_ptr[k - 1] = -1; 340 | k--; 341 | } 342 | } 343 | } 344 | 345 | } 346 | } 347 | 348 | 349 | template 350 | static __global__ void gather_cube_pointsb_local( 351 | const int nthreads, const int* n_target_points, const T* in_target_points, const int MAX_POINTS_IN_CUBE, 352 | const int CUBES_TOTAL, int *cube_points_num_ptr, int *cube_points_indices_ptr) { 353 | CUDA_1D_KERNEL_LOOP(batch_index, nthreads) { 354 | int *cube_indices = cube_points_indices_ptr + batch_index * MAX_POINTS_IN_CUBE * CUBES_TOTAL; 355 | int *batch_cube_points_num_ptr = cube_points_num_ptr + batch_index*CUBES_TOTAL; 356 | 357 | int start_index = 0; 358 | if(batch_index != 0) 359 | start_index = n_target_points[batch_index - 1]; 360 | 361 | for (int n_target_index = start_index; n_target_index < n_target_points[batch_index]; n_target_index++) 362 | { 363 | T tx = in_target_points[6*n_target_index]; 364 | T ty = in_target_points[6*n_target_index + 1]; 365 | T tz = in_target_points[6*n_target_index + 2]; 366 | //find if point is inside cube 367 | int x = int((tx + 1) * 5); 368 | int y = int((ty + 1) * 5); 369 | int z = int((tz + 1) * 5); 370 | if(x < 0) 371 | x = 0; 372 | else if(x > 9) 373 | x = 9; 374 | 375 | if(y < 0) 376 | y = 0; 377 | else if (y > 9) 378 | y = 9; 379 | 380 | if(z < 0) 381 | z = 0; 382 | else if(z > 9) 383 | z = 9; 384 | 385 | int cube_index = 100*x+10*y+z; 386 | cube_indices[cube_index*MAX_POINTS_IN_CUBE + batch_cube_points_num_ptr[cube_index]] = n_target_index; 387 | batch_cube_points_num_ptr[cube_index]++; 388 | CudaAssert(batch_cube_points_num_ptr[cube_index] <= MAX_POINTS_IN_CUBE); 389 | } 390 | } 391 | } 392 | 393 | template 394 | void neighbor_points_spatial_grid_no_normal_weight_local_varying_radius_voting(OpKernelContext* context, 395 | const int n_source_points, const int* n_target_points, const int batch_size, 396 | const T* in_source_points, const T* in_target_points, 397 | const T* in_squared_radius, const int knn, 398 | int* neighbor_points_id_ptr) { 399 | 400 | //time_t start_time = clock(); 401 | // get GPU device 402 | GPUDevice d = context->eigen_device(); 403 | CudaLaunchConfig config; 404 | int nthreads; 405 | 406 | const int CUBES_TOTAL = 1000; 407 | const int MAX_POINTS_IN_CUBE = 4096; 408 | 409 | //assume input points are bounded by a unit sphere, and we partition it to 10x10x10 grids 410 | const TensorShape cube_index_shape({batch_size, CUBES_TOTAL, MAX_POINTS_IN_CUBE}); 411 | Tensor cube_points_indices; 412 | OP_REQUIRES_OK(context, context->allocate_temp(DT_INT32, cube_index_shape, &cube_points_indices)); 413 | auto cube_points_indices_ptr = cube_points_indices.flat().data(); 414 | 415 | const TensorShape cube_points_num_shape({batch_size, CUBES_TOTAL}); 416 | Tensor cube_points_num; 417 | OP_REQUIRES_OK(context, context->allocate_temp(DT_INT32, cube_points_num_shape, &cube_points_num)); 418 | auto cube_points_num_ptr = cube_points_num.flat().data(); 419 | cudaMemset(cube_points_num_ptr, 0, sizeof(int)*batch_size*CUBES_TOTAL); 420 | 421 | const TensorShape neighbor_points_distance_shape({batch_size, n_source_points, knn}); 422 | Tensor neighbor_points_distance; 423 | OP_REQUIRES_OK(context, context->allocate_temp(sizeof(T)==4?DT_FLOAT:DT_DOUBLE, neighbor_points_distance_shape, &neighbor_points_distance)); 424 | auto neighbor_points_distance_ptr = neighbor_points_distance.flat().data(); 425 | 426 | nthreads = batch_size; 427 | config = GetCudaLaunchConfig(nthreads, d); 428 | gather_cube_pointsb_local 429 | <<>>( 430 | nthreads, n_target_points, in_target_points, MAX_POINTS_IN_CUBE, CUBES_TOTAL, cube_points_num_ptr, cube_points_indices_ptr); 431 | 432 | nthreads = batch_size * n_source_points; 433 | config = GetCudaLaunchConfig(nthreads, d); 434 | //printf("%d block and %d threads per block\n", config.block_count, config.thread_per_block); 435 | find_knn_grids_no_normalw_local_varying_radius_voting 436 | <<>>( 437 | nthreads, n_source_points, n_target_points, in_source_points, in_target_points, in_squared_radius, knn, CUBES_TOTAL, MAX_POINTS_IN_CUBE, 438 | cube_points_num_ptr, cube_points_indices_ptr, neighbor_points_id_ptr, neighbor_points_distance_ptr); 439 | cudaDeviceSynchronize(); 440 | 441 | cudaError_t err2 = cudaGetLastError(); 442 | if (err2 != cudaSuccess) 443 | printf("Cuda error: %s\n", cudaGetErrorString(err2)); 444 | 445 | } 446 | 447 | //explicit instantiation 448 | template void neighbor_points_spatial_grid_no_normal_weight_local_varying_radius_voting(OpKernelContext* context, 449 | const int n_source_points, const int* n_target_points, const int batch_size, 450 | const float* in_source_points, const float* in_target_points, 451 | const float* in_squared_radius, const int knn, 452 | int* neighbor_points_id_ptr); 453 | 454 | template void neighbor_points_spatial_grid_no_normal_weight_local_varying_radius_voting(OpKernelContext* context, 455 | const int n_source_points, const int* n_target_points, const int batch_size, 456 | const double* in_source_points, const double* in_target_points, 457 | const double* in_squared_radius, const int knn, 458 | int* neighbor_points_id_ptr); 459 | 460 | } // namespace tensorflow 461 | -------------------------------------------------------------------------------- /points3d-tf/pointsdataset/points_dataset_octreed7.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import tensorflow as tf 3 | import numpy as np 4 | sys.path.append("Octree/ocnn-tf") 5 | from libs import * 6 | 7 | def IntList_to_Bytes(int_list): 8 | #list_bytes = struct.pack('i'*len(int_list), *int_list) 9 | x = np.array(int_list, dtype=np.int32) 10 | list_bytes = x.tobytes() 11 | return list_bytes 12 | 13 | def DoubleList_to_Bytes(float_list): 14 | #list_bytes = struct.pack('d'*len(float_list), *float_list) 15 | x = np.array(float_list, dtype=np.float64) 16 | list_bytes = x.tobytes() 17 | return list_bytes 18 | 19 | def Float32List_to_Bytes(float_list): 20 | #list_bytes = struct.pack('f'*len(float_list), *float_list) 21 | x = np.array(float_list, dtype=np.float32) 22 | list_bytes = x.tobytes() 23 | return list_bytes 24 | 25 | class PointsPreprocessor: 26 | def __init__(self, depth, points_num, full_depth=2, sample_grid_points_number=None, data_sources=-1, return_gt_points=False, noise_stddev=0.0095): 27 | self._depth = depth 28 | self._full_depth = full_depth 29 | self._points_num = points_num 30 | self._sample_grid_points_number = sample_grid_points_number 31 | self._data_sources = data_sources 32 | self._return_gt_points = return_gt_points 33 | self._noise_stddev = noise_stddev 34 | 35 | 36 | def __call__(self, record): 37 | points_array_f16, filename, sdf_samples_depth6, sdf_samples_depth7, depth6_coord, depth7_coord, octree_d7_points_idx = self.parse_example_array(record) 38 | 39 | octree_points_d7_x = octree_d7_points_idx // 65536 40 | octree_points_d7_y = (octree_d7_points_idx // 256) % 256 41 | octree_points_d7_z = octree_d7_points_idx % 256 42 | octree_points_d7 = tf.cast(tf.stack([octree_points_d7_x, octree_points_d7_y, octree_points_d7_z], axis=1), tf.float32) / 128.0 - 1 43 | 44 | gt_points = tf.cast(tf.transpose(tf.reshape(points_array_f16, [6, 100000])), tf.float32) 45 | 46 | #sample 3000 points and add noise 47 | selected_points_idx = tf.random.uniform(shape=[3000], minval=0, maxval=100000, dtype=tf.int32) 48 | selected_points = tf.gather(gt_points, selected_points_idx) 49 | #add noise 50 | assert(self._points_num == 3000) 51 | if(self._noise_stddev > 1e-6): 52 | selected_points_position = selected_points[:,:3] + tf.random.normal([3000,3], stddev=self._noise_stddev)#0.0095 53 | else: 54 | selected_points_position = selected_points[:,:3] 55 | #selected_points_normal = selected_points[:,3:] 56 | #selected_points = tf.concat([selected_points_position, selected_points_normal], axis=-1) 57 | 58 | points = points_new(selected_points_position, [], selected_points_position, []) 59 | input_octree = points2octree(points, depth=self._depth, full_depth=self._full_depth, 60 | node_dis=False, split_label=True) #node_feature=True, 61 | 62 | gt_octree = points2octree(points_new(octree_points_d7, [], tf.ones(tf.shape(octree_points_d7_x)), []), depth=self._depth, full_depth=self._full_depth, 63 | node_dis=False, split_label=True) #node_feature=True, 64 | if(self._data_sources == 6): 65 | sdf_samples = sdf_samples_depth6 66 | sdf_samples_id = depth6_coord 67 | elif(self._data_sources == 7): 68 | sdf_samples = sdf_samples_depth7 69 | sdf_samples_id = depth7_coord 70 | else: 71 | raise NotImplementedError 72 | 73 | if(self._sample_grid_points_number is None): 74 | raise NotImplementedError 75 | else: 76 | padding_limit = self._sample_grid_points_number 77 | #random select sample_grid_points_number samples 78 | sdf_samples = tf.reshape(sdf_samples, [-1,4]) 79 | num_valid_samples = tf.reduce_sum(tf.cast(tf.greater_equal(sdf_samples_id, 0), tf.int32)) 80 | 81 | #randomly select points(uniformly) 82 | selected_idx = tf.random.shuffle(tf.range(num_valid_samples))[:padding_limit] 83 | selected_samples = tf.gather(sdf_samples, selected_idx) 84 | selected_samples_coord = tf.gather(sdf_samples_id, selected_idx) 85 | 86 | selected_samples_coordx = selected_samples_coord // 66049 87 | selected_samples_coordy = (selected_samples_coord // 257) % 257 88 | selected_samples_coordz = selected_samples_coord % 257 89 | 90 | selected_position = tf.cast(tf.stack([selected_samples_coordx, selected_samples_coordy, selected_samples_coordz], axis=1), tf.float32) / 128.0 - 1.0 91 | selected_sdf = selected_samples 92 | #pad if needed 93 | selected_position = tf.pad(selected_position, [[0, padding_limit - tf.shape(selected_position)[0]],[0,0]], constant_values=-100) 94 | selected_sdf = tf.pad(selected_sdf, [[0, padding_limit - tf.shape(selected_sdf)[0]],[0,0]], constant_values=0) 95 | if(self._return_gt_points): 96 | return input_octree, gt_octree, tf.transpose(selected_points), filename, selected_position, selected_sdf, gt_points 97 | else: 98 | return input_octree, gt_octree, tf.transpose(selected_points), filename, selected_position, selected_sdf 99 | 100 | def parse_example_array(self, record): 101 | features = { 'points': tf.FixedLenFeature([], tf.string), 102 | "filename": tf.FixedLenFeature([], tf.string)} 103 | SeqFeatures = {"sdf_samples_depth6": tf.FixedLenSequenceFeature([4000], tf.float32), 104 | "sdf_samples_depth7": tf.FixedLenSequenceFeature([4000], tf.float32), 105 | "depth6_coord": tf.FixedLenSequenceFeature([1000], tf.int64), 106 | "depth7_coord": tf.FixedLenSequenceFeature([1000], tf.int64), 107 | "octree_points_idx": tf.FixedLenSequenceFeature([2000], tf.int64)} 108 | parsed_all = tf.parse_single_sequence_example(record, context_features=features, sequence_features=SeqFeatures) 109 | parsed = parsed_all[0] 110 | feature_parsed = parsed_all[1] 111 | return tf.decode_raw(parsed['points'], tf.float16), parsed['filename'], feature_parsed['sdf_samples_depth6'], feature_parsed['sdf_samples_depth7'], tf.reshape(feature_parsed['depth6_coord'], [-1]), tf.reshape(feature_parsed['depth7_coord'], [-1]), tf.reshape(feature_parsed['octree_points_idx'], [-1]) 112 | 113 | def points_dataset_AE(record_name, batch_size, points_num, depth=5, full_depth=2, shuffle_data=True, sample_grid_points_number=None, data_sources=-1, return_gt_points=False, noise_stddev=0.0095): 114 | def merge_octrees(input_octrees, gt_octrees, points_arrays, filenames, grid_points, grid_fvalue): 115 | input_octree = octree_batch(input_octrees) 116 | gt_octree = octree_batch(gt_octrees) 117 | return input_octree, gt_octree, points_arrays, filenames, grid_points, grid_fvalue 118 | 119 | def merge_octrees_gtpoints(input_octrees, gt_octrees, points_arrays, filenames, grid_points, grid_fvalue, gt_points): 120 | input_octree = octree_batch(input_octrees) 121 | gt_octree = octree_batch(gt_octrees) 122 | return input_octree, gt_octree, points_arrays, filenames, grid_points, grid_fvalue, gt_points 123 | 124 | if("all_test_1000" in record_name): 125 | record_num = 1000 126 | else: 127 | record_num = sum(1 for _ in tf.python_io.tf_record_iterator(record_name)) 128 | shuffle_buffer_size = min(1000, 2*record_num) 129 | print("dataset using data sources:{}".format(data_sources)) 130 | print("{} records in {}".format(record_num, record_name)) 131 | if(shuffle_data): 132 | print("shuffle buffer size = {}".format(shuffle_buffer_size)) 133 | with tf.name_scope('points_dataset'): 134 | preprocess = PointsPreprocessor(depth, points_num, full_depth, sample_grid_points_number=sample_grid_points_number, data_sources=data_sources, return_gt_points=return_gt_points, noise_stddev=noise_stddev) 135 | if(return_gt_points): 136 | merge_octrees_op = merge_octrees_gtpoints 137 | else: 138 | merge_octrees_op = merge_octrees 139 | if(shuffle_data): 140 | return (record_num,) + tf.data.TFRecordDataset([record_name]).repeat().shuffle(shuffle_buffer_size) \ 141 | .map(preprocess, num_parallel_calls=8) \ 142 | .batch(batch_size).map(merge_octrees_op, num_parallel_calls=8) \ 143 | .prefetch(8).make_one_shot_iterator().get_next() 144 | else: 145 | return (record_num,) + tf.data.TFRecordDataset([record_name]).repeat() \ 146 | .map(preprocess, num_parallel_calls=8) \ 147 | .batch(batch_size).map(merge_octrees_op, num_parallel_calls=8) \ 148 | .prefetch(8).make_one_shot_iterator().get_next() 149 | 150 | def points_dataset_AE_multiple_GPU(record_name, batch_size, points_num, depth=5, full_depth=2, gpu_num=1, sample_grid_points_number=None, data_sources=-1, noise_stddev=0.0095): 151 | def merge_octrees(input_octrees, gt_octrees, points_arrays, filenames, grid_points, grid_fvalue): 152 | gpu_batch_size = batch_size // gpu_num 153 | return_list = [] 154 | for i in range(gpu_num): 155 | return_list.append(octree_batch(input_octrees[i*gpu_batch_size:(i+1)*gpu_batch_size])) 156 | for i in range(gpu_num): 157 | return_list.append(octree_batch(gt_octrees[i*gpu_batch_size:(i+1)*gpu_batch_size])) 158 | return_list += [points_arrays, filenames, grid_points, grid_fvalue] 159 | return return_list 160 | 161 | if("all_train" in record_name): 162 | record_num = 30661 163 | else: 164 | record_num = sum(1 for _ in tf.python_io.tf_record_iterator(record_name)) 165 | shuffle_buffer_size = min(1000, 2*record_num) 166 | print("multi-GPU dataset using data sources:{}".format(data_sources)) 167 | print("{} records in {}".format(record_num, record_name)) 168 | print("shuffle buffer size = {}".format(shuffle_buffer_size)) 169 | with tf.name_scope('points_dataset'): 170 | preprocess = PointsPreprocessor(depth, points_num, full_depth, sample_grid_points_number=sample_grid_points_number, data_sources=data_sources, noise_stddev=noise_stddev) 171 | return (record_num,) + tf.data.TFRecordDataset([record_name]).repeat().shuffle(shuffle_buffer_size) \ 172 | .map(preprocess, num_parallel_calls=8) \ 173 | .batch(batch_size).map(merge_octrees, num_parallel_calls=8) \ 174 | .prefetch(8).make_one_shot_iterator().get_next() -------------------------------------------------------------------------------- /scripts/Readme.md: -------------------------------------------------------------------------------- 1 | # Script for tfrecords generation 2 | 3 | Before using this script, please first compile [main.cpp](https://github.com/Andy97/DeepMLS/tree/master/vdb_tsdf) and get an executable file named "vdb_tsdf". 4 | 5 | The executable will be used to generate groudtruth sdf values and gradients as supervision for training, which will be used in script. ([here](https://github.com/Andy97/DeepMLS/blob/master/scripts/generate_tf_records.py#L79)) 6 | 7 | For each Shapenet model, there will be **2 input files** to this script: 8 | 9 | #### **1.Watertight mesh file**：the watertight version of each original shapenet model, generated using scripts from occupancy networks 10 | We assume that each mesh file has already been normalized to [-0.95, 0.95]^3 (which equals to [-1,1]^3 with 5% padding). 11 | 12 | This input is declared [here](https://github.com/Andy97/DeepMLS/blob/master/scripts/generate_tf_records.py#L53). 13 | 14 | #### **2.Sampled points file**: sampled 100k points with normal from input #1 15 | This input is the groundtruth points sampled from surface, we will draw 3k points from these points plus gaussian noise as network input. 16 | 17 | Please also note that since these points are sampled from input #1, all these points should also be bounded by bounding box [-0.95, 0.95]^3. 18 | 19 | This input is declared [here](https://github.com/Andy97/DeepMLS/blob/master/scripts/generate_tf_records.py#L42). 20 | -------------------------------------------------------------------------------- /scripts/generate_tf_records.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import struct 4 | 5 | import numpy as np 6 | import tensorflow as tf 7 | import trimesh 8 | 9 | def _float_feature(value): 10 | return tf.train.Feature(float_list=tf.train.FloatList(value=value)) 11 | 12 | def _bytes_feature(value): 13 | return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value])) 14 | 15 | def _int_feature(value): 16 | return tf.train.Feature(int64_list=tf.train.Int64List(value=value)) 17 | 18 | def get_octree_from_points(points, octree_depth=8): 19 | #get representative points for building the groundtruth octree 20 | assert(points.shape[1] == 6 or points.shape[1] == 3) 21 | points_pos = points[:,:3] 22 | octree_node_length = 2.0 / (2**octree_depth) 23 | octree_node_index = np.unique(np.clip(np.floor((points_pos + 1) / octree_node_length).astype(np.int32), a_min = 0, a_max = (2**octree_depth) - 1), axis=0) 24 | 25 | return octree_node_index 26 | 27 | def write_data_to_tfrecords(sample_list, tfrecords_filename): 28 | #options = tf.python_io.TFRecordOptions(compression_type=tf.python_io.TFRecordCompressionType.ZLIB) 29 | writer = tf.python_io.TFRecordWriter(tfrecords_filename)#, options=options 30 | n_data = len(sample_list) 31 | for i in range(n_data): 32 | if i % 100 == 0 or i == n_data-1: 33 | print('=========================\n') 34 | print('data loaded: {} / {}\n'.format(i+1, n_data)) 35 | print('=========================\n') 36 | sample_category = sample_list[i][0] 37 | sample_id = sample_list[i][2] 38 | #print("{}:{}".format(sample_category, sample_id)) 39 | signature = "{}:{}".format(sample_category, sample_id) 40 | 41 | #points_files contains the 100k points sampled from groundtruth surface, 3k input points to the network will be drawn from this set 42 | points_file = os.path.join("occupancy_networks-master\\data\\ShapeNet.build", sample_category, "4_pointcloud", sample_id+".npz") 43 | assert(os.path.exists(points_file)) 44 | points_dict = np.load(points_file) 45 | points_position = points_dict['points'] 46 | points_normal = points_dict['normals'] 47 | assert(points_position.shape == points_normal.shape) 48 | print("points in bbox:{} {}".format(points_position.min(), points_position.max())) 49 | assert(points_position.min() >= -1 and points_position.max() <= 1) 50 | points = np.concatenate((points_position, points_normal), axis=-1).astype(np.float32) 51 | assert(points.shape[0] == 100000 and points.shape[1] ==6) 52 | 53 | mesh_file = sample_list[i][1] + ".obj" 54 | assert(os.path.exists(mesh_file)) 55 | #prepare ground truth octree for decoder supervision 56 | #sample 5M points to build octree as gt octree for decode shape 57 | mesh = trimesh.load(mesh_file) 58 | sampled_points = trimesh.sample.sample_surface(mesh, 5000000)[0] 59 | #scale to [-1,1] with 5% padding 60 | #sampled_points *= 1.9 61 | assert(sampled_points.min() >= -1 and sampled_points.max() <= 1) 62 | octree = get_octree_from_points(sampled_points, 8) 63 | assert(octree.min() >=0 and octree.max() <= 255) 64 | octree_points = (octree*np.array([256*256, 256, 1])).sum(axis=-1) 65 | assert(octree_points.dtype == np.int) 66 | assert(octree_points.min() >=0 and octree_points.max() <= 16777215) 67 | #octree_points_file = sample_list[i][1] + "_octreep.npz" 68 | #np.savez_compressed(octree_points_file, octree_points=octree_points) 69 | #print("{} octree points".format(octree_points.shape)) 70 | 71 | octree_points_seq_length = 2000 72 | #pad octree_points 73 | octree_points = np.reshape(octree_points, [-1]) 74 | octree_points_pad_elements = (octree_points_seq_length - (octree_points.shape[0] % octree_points_seq_length)) % octree_points_seq_length 75 | octree_points = np.reshape(np.pad(octree_points, (0, octree_points_pad_elements), mode='constant', constant_values=octree_points[0]), [-1, octree_points_seq_length]) 76 | 77 | sdf_samples_file = sample_list[i][1] + ".mat.bin" 78 | if(not os.path.exists(sdf_samples_file)): 79 | os.system("vdb_tsdf {} {}".format(mesh_file, sdf_samples_file)) 80 | assert(os.path.exists(sdf_samples_file)) 81 | 82 | sdf_samples = np.reshape(np.fromfile(sdf_samples_file, dtype=np.float32), [-1, 7]) 83 | assert(sdf_samples.shape[1] == 7) 84 | sdf_points = sdf_samples[:,:3] 85 | assert(sdf_points.min() >= -1 and sdf_points.max() <= 1) 86 | 87 | #first get samples on 64^3 grid 88 | x = sdf_samples[:,0] 89 | y = sdf_samples[:,1] 90 | z = sdf_samples[:,2] 91 | coordx = np.round((x+1)*128).astype(int) 92 | coordy = np.round((y+1)*128).astype(int) 93 | coordz = np.round((z+1)*128).astype(int) 94 | depth6_ind = np.where(np.logical_and(np.logical_and(coordx % 4 == 0, coordy % 4 == 0), coordz % 4 == 0)) 95 | depth6_sdf_samples = sdf_samples[depth6_ind] 96 | #print("{} samples in depth 6".format(depth6_sdf_samples.shape)) 97 | 98 | #filter sdf samples with sdf > 2/32 99 | sdf_samples = sdf_samples[np.where(np.abs(sdf_samples[:,3]) < 2.0/32)] 100 | x = sdf_samples[:,0] 101 | y = sdf_samples[:,1] 102 | z = sdf_samples[:,2] 103 | coordx = np.round((x+1)*128).astype(int) 104 | coordy = np.round((y+1)*128).astype(int) 105 | coordz = np.round((z+1)*128).astype(int) 106 | depth7_ind = np.where(np.logical_and(np.logical_and(coordx % 2 == 0, coordy % 2 == 0), coordz % 2 == 0)) 107 | depth7_sdf_samples = sdf_samples[depth7_ind] 108 | #print("{} samples in depth 7".format(depth7_sdf_samples.shape)) 109 | 110 | def efficient_padding(sdf_samples): 111 | xyz_coord = np.round((sdf_samples[:,:3] + 1)*128).astype(np.int32) 112 | assert(xyz_coord.min() >= 0 and xyz_coord.max() <= 256) 113 | xyz_id = (xyz_coord*np.array([257*257, 257, 1])).sum(axis=-1) 114 | sdf_pad_elements = (1000 - (xyz_id.shape[0] % 1000)) % 1000 115 | 116 | xyz_id = np.pad(xyz_id, (0, sdf_pad_elements), mode='constant', constant_values=-100) #should be carefully dealt when parsing 117 | sdf_values = np.pad(np.reshape(sdf_samples[:,3:], [-1]), (0, 4*sdf_pad_elements), mode='constant', constant_values=-100) 118 | assert(xyz_id.shape[0] % 1000 == 0 and xyz_id.dtype == np.int32) 119 | assert(sdf_values.shape[0] == xyz_id.shape[0]*4 and sdf_values.dtype==np.float32) 120 | return np.reshape(xyz_id, [-1,1000]), np.reshape(sdf_values, [-1, 4000]) 121 | 122 | depth7_coord, depth7_sdf_samples = efficient_padding(depth7_sdf_samples) 123 | depth6_coord, depth6_sdf_samples = efficient_padding(depth6_sdf_samples) 124 | assert(depth7_coord.dtype == np.int32 and depth6_coord.dtype == np.int32) 125 | assert(depth7_sdf_samples.dtype == np.float32 and depth6_sdf_samples.dtype == np.float32) 126 | 127 | sequence_feature = {"depth7_coord":tf.train.FeatureList(feature=[_int_feature(coord_row.tolist()) for coord_row in depth7_coord]), 128 | "depth6_coord":tf.train.FeatureList(feature=[_int_feature(coord_row.tolist()) for coord_row in depth6_coord]), 129 | "sdf_samples_depth7":tf.train.FeatureList(feature=[_float_feature(sdf_row.tolist()) for sdf_row in depth7_sdf_samples]), 130 | "sdf_samples_depth6":tf.train.FeatureList(feature=[_float_feature(sdf_row.tolist()) for sdf_row in depth6_sdf_samples]), 131 | "octree_points_idx":tf.train.FeatureList(feature=[_int_feature(idx_row.tolist()) for idx_row in octree_points])} 132 | feature = {'points': _bytes_feature(np.reshape(np.transpose(points), [-1]).astype(np.float16).tobytes()), 133 | "filename":_bytes_feature(signature.encode('utf8')), 134 | } 135 | example = tf.train.SequenceExample(context=tf.train.Features(feature=feature), 136 | feature_lists=tf.train.FeatureLists(feature_list=sequence_feature)) 137 | writer.write(example.SerializeToString()) 138 | writer.close() 139 | 140 | def main(): 141 | cat = ['02691156', '02828884', '02933112', '02958343', '03001627', '03211117', '03636649', '03691459', '04090263', '04256520', '04379243', '04401088', '04530566'] 142 | rootpc = "occupancy_networks-master\\data\\ShapeNet.build\\" 143 | 144 | list=[] 145 | for item in cat: 146 | print("working on {}".format(item)) 147 | dir_point = os.path.join(rootpc, item, 'watertight_normalized') 148 | assert(os.path.exists(dir_point)) 149 | 150 | #training data split file follow occupancy networks 151 | split_file = os.path.join("convolutional_occupancy_networks-master\\data\\ShapeNet", item, "train.lst") 152 | assert(os.path.exists(split_file)) 153 | split_file = open(split_file, "r") 154 | fns = [line.replace("\r", "").replace("\n", "") for line in split_file] 155 | split_file.close() 156 | 157 | assert(len(fns) != 0) 158 | for fn in fns: 159 | list.append((item, os.path.join(dir_point, fn), fn)) 160 | 161 | 162 | print("{} samples total".format(len(list))) 163 | random.shuffle(list) 164 | 165 | tfrecords_filename = 'ShapeNet_points_{}_w_octree_grid_occupancy_compress.tfrecords'.format("all_train") 166 | write_data_to_tfrecords(list, tfrecords_filename) 167 | 168 | if __name__ == '__main__': 169 | main() 170 | -------------------------------------------------------------------------------- /utils.py: -------------------------------------------------------------------------------- 1 | import operator 2 | import functools 3 | import tensorflow as tf 4 | import json 5 | 6 | def config_reader(config_path): 7 | json_data=open(config_path).read() 8 | config = json.loads(json_data) 9 | return config 10 | 11 | def get_num_params(): 12 | num_params = 0 13 | for variable in tf.trainable_variables(): 14 | shape = variable.get_shape() 15 | num_params += functools.reduce(operator.mul, [dim.value for dim in shape], 1) 16 | return num_params 17 | 18 | def average_gradient(tower_grads): 19 | average_grads = [] 20 | for grad_and_vars in zip(*tower_grads): 21 | grads = [] 22 | for g, var in grad_and_vars: 23 | if g is not None: 24 | expanded_g = tf.expand_dims(g, 0) 25 | grads.append(expanded_g) 26 | if(len(grads) == 0): 27 | continue 28 | grad = tf.concat(grads, axis=0) 29 | grad = tf.reduce_mean(grad, axis=0) 30 | 31 | v = grad_and_vars[0][1] 32 | grad_and_var = (grad, v) 33 | average_grads.append(grad_and_var) 34 | return average_grads 35 | 36 | def tf_summary_from_dict(loss_dict, is_training): 37 | if(is_training): 38 | summary_name = "train_summary" 39 | else: 40 | summary_name = "test_summary" 41 | #create summary 42 | with tf.name_scope(summary_name): 43 | summary_list = [] 44 | gpu0_loss_dict = loss_dict[0] 45 | for item in gpu0_loss_dict: 46 | scalar_acc = 0 47 | for i in range(len(loss_dict)): 48 | scalar_acc += loss_dict[i][item] 49 | scalar_acc /= len(loss_dict) 50 | summary_item = tf.summary.scalar(item, scalar_acc) 51 | summary_list.append(summary_item) 52 | return tf.summary.merge(summary_list) 53 | 54 | def rowwise_l2_norm_squared(feature): 55 | #assum input with size[n,f] out shape = [n] 56 | return tf.reduce_sum(tf.math.square(feature), axis=-1) -------------------------------------------------------------------------------- /vdb_tsdf/ReadMe.MD: -------------------------------------------------------------------------------- 1 | ## Groundtruth SDF Data Generation 2 | Given a watertight mesh, we use [OpenVDB](https://www.openvdb.org/) library to generate the tsdf field. 3 | As clarified in paper, more sdf samples near the surface will be chosen. 4 | For details, please refer to the source code. 5 | 6 | ### Dependencies 7 | ##### 1. [OpenVDB](https://www.openvdb.org/) 8 | ##### 2. [OpenMesh](https://www.openmesh.org/) 9 | -------------------------------------------------------------------------------- /vdb_tsdf/main.cpp: -------------------------------------------------------------------------------- 1 | #pragma warning (disable: 4146) 2 | #define _USE_MATH_DEFINES 3 | #include 4 | #define NOMINMAX 5 | #include 6 | #include 7 | #include 8 | #include 9 | 10 | #include "OpenMesh\Core\IO\MeshIO.hh" 11 | #include "OpenMesh/Core/Mesh/TriMesh_ArrayKernelT.hh" 12 | 13 | #include 14 | inline bool debug_output(const char *string) 15 | { 16 | std::cout << "=======================================================\n"; 17 | std::cout << "Assertion failed:" << string << "!!!\n"; 18 | std::cout << "=======================================================\n"; 19 | system("pause"); 20 | return false; 21 | } 22 | 23 | #define myassert(expression) (bool)( \ 24 | (!!(expression)) || \ 25 | (debug_output(#expression)) \ 26 | ) 27 | 28 | // Select mesh type (TriMesh) and kernel (ArrayKernel) 29 | // and define my personal mesh type (MyMesh) 30 | struct DPTraits : public OpenMesh::DefaultTraits 31 | { 32 | typedef OpenMesh::Vec3d Point; // use double-values points/normals 33 | typedef OpenMesh::Vec3d Normal; 34 | }; 35 | typedef OpenMesh::TriMesh_ArrayKernelT TriMesh; 36 | 37 | void fetch_sdf(const char *model_filename, const char *sdf_filename, int grid_resolution = 256) 38 | { 39 | auto mesh = std::make_shared(); 40 | if (!OpenMesh::IO::read_mesh(*mesh, model_filename)) 41 | { 42 | std::cerr << "Error: cannot read mesh file " << model_filename << std::endl; 43 | return; 44 | } 45 | std::cout << "Finished reading mesh.\n"; 46 | 47 | double bbox[6] = { 1,1,1,-1,-1,-1 }; 48 | std::vector verts; 49 | std::vector faces; 50 | for (auto vtx = mesh->vertices_begin(); vtx != mesh->vertices_end(); vtx++) 51 | { 52 | auto& pt = mesh->point(*vtx); 53 | verts.push_back(openvdb::Vec3f(pt[0], pt[1], pt[2])); 54 | 55 | //for (int i = 0; i < 3; i++) 56 | //{ 57 | // bbox[i] = std::max(pt[i], bbox[i]); 58 | // bbox[i + 3] = std::min(pt[i], bbox[i + 3]); 59 | //} 60 | } 61 | for (auto face = mesh->faces_begin(); face != mesh->faces_end(); face++) 62 | { 63 | openvdb::Vec3I face_idx; 64 | auto fvhandle = mesh->fv_begin(*face); 65 | for (int vitr = 0; vitr < 3; vitr++, fvhandle++) 66 | { 67 | face_idx[vitr] = fvhandle->idx(); 68 | } 69 | faces.push_back(face_idx); 70 | } 71 | 72 | double diag_len = (bbox[0] - bbox[3])*(bbox[0] - bbox[3]) + (bbox[1] - bbox[4])*(bbox[1] - bbox[4]) 73 | + (bbox[2] - bbox[5])*(bbox[2] - bbox[5]); 74 | diag_len = std::sqrt(diag_len); 75 | double center[3] = { bbox[0] + bbox[3], bbox[1] + bbox[4], bbox[2] + bbox[5] }; 76 | for (int i = 0; i < 3; i++) 77 | { 78 | center[i] /= 2.0; 79 | } 80 | double voxel_size = 2.0 / grid_resolution; 81 | 82 | printf("Model center: %f,%f,%f, diagonal length: %f, voxel size: %f\n", center[0], center[1], center[2], 83 | diag_len, voxel_size); 84 | 85 | auto grid = openvdb::tools::meshToSignedDistanceField( 86 | *openvdb::math::Transform::createLinearTransform(voxel_size), 87 | verts, faces, std::vector(), grid_resolution / 4.0, grid_resolution / 4.0 88 | ); 89 | openvdb::tools::signedFloodFill(grid->tree()); 90 | 91 | std::vector>> sdf_grids; 92 | std::vector>> sdf_grids_grad; 93 | sdf_grids.resize(grid_resolution + 1); 94 | sdf_grids_grad.resize(grid_resolution + 1); 95 | for (int i = 0; i < sdf_grids.size(); i++) 96 | { 97 | sdf_grids[i].resize(grid_resolution + 1); 98 | sdf_grids_grad[i].resize(grid_resolution + 1); 99 | for (int j = 0; j < sdf_grids.size(); j++) 100 | { 101 | sdf_grids[i][j].resize(grid_resolution + 1); 102 | sdf_grids_grad[i][j].resize(grid_resolution + 1); 103 | std::fill(sdf_grids[i][j].begin(), sdf_grids[i][j].end(), -10000); 104 | std::fill(sdf_grids_grad[i][j].begin(), sdf_grids_grad[i][j].end(), openvdb::math::Vec3s(-10000, -10000, -10000)); 105 | } 106 | } 107 | auto grad_field = openvdb::tools::gradient(*grid); 108 | //openvdb::tools::signedFloodFill(grad_field->tree()); 109 | myassert(grid_resolution % 2 == 0); 110 | int half_grid_resolution = grid_resolution / 2; 111 | 112 | int filled_value_count = 0; 113 | int maxx, maxy, maxz; 114 | int minx, miny, minz; 115 | maxx = maxy = maxz = -1; 116 | minx = miny = minz = 1; 117 | 118 | //extract sdf samples with fabs(sdf) < sdf_truncation 119 | const double sdf_truncation = 2.0 / 16; 120 | std::vector sampled_points; 121 | sampled_points.clear(); 122 | 123 | //iterate the active voxels of grid 124 | // Iterate over all active values but don't allow them to be changed. 125 | for (openvdb::FloatGrid::ValueOnCIter iter = grid->cbeginValueOn(); iter.test(); ++iter) { 126 | const float& value = *iter; 127 | if (iter.isVoxelValue()) { 128 | //if (value <= 0) 129 | // std::cout << iter.getCoord() << " value: " << value << std::endl; 130 | int x, y, z; 131 | x = iter.getCoord().x(); 132 | y = iter.getCoord().y(); 133 | z = iter.getCoord().z(); 134 | 135 | maxx = std::max(x, maxx); 136 | maxy = std::max(y, maxy); 137 | maxz = std::max(z, maxz); 138 | 139 | minx = std::min(x, minx); 140 | miny = std::min(y, miny); 141 | minz = std::min(z, minz); 142 | 143 | x += half_grid_resolution; 144 | y += half_grid_resolution; 145 | z += half_grid_resolution; 146 | 147 | if (x >= 0 && y >= 0 && z >= 0 && x <= grid_resolution && y <= grid_resolution && z <= grid_resolution) 148 | { 149 | filled_value_count++; 150 | sdf_grids[x][y][z] = value; 151 | if (fabs(value) < sdf_truncation) 152 | { 153 | //aggressive generate more samples near surface 154 | if(fabs(value) < sdf_truncation / 4) 155 | sampled_points.push_back(OpenMesh::Vec3i(x, y, z)); 156 | else if(x % 4 == 0 && y % 4 == 0 && z % 4 == 0) 157 | sampled_points.push_back(OpenMesh::Vec3i(x, y, z)); 158 | else if(fabs(value) < sdf_truncation / 2 && x % 2 == 0 && y % 2 == 0 && z % 2 == 0) 159 | sampled_points.push_back(OpenMesh::Vec3i(x, y, z)); 160 | } 161 | } 162 | } 163 | else 164 | { 165 | std::cout << "Not voxel node sdf value" << std::endl; 166 | //int junk; 167 | //std::cin >> junk; 168 | } 169 | } 170 | std::cout << maxx << "," << maxy << "," << maxz << "\n"; 171 | std::cout << minx << "," << miny << "," << minz << "\n"; 172 | std::cout << sampled_points.size() << " points with |sdf| < " << sdf_truncation << "\n"; 173 | 174 | int filled_grad_count = 0; 175 | for (openvdb::Vec3SGrid::ValueOnCIter iter = grad_field->cbeginValueOn(); iter.test(); ++iter) { 176 | const openvdb::math::Vec3s& value = *iter; 177 | if (iter.isVoxelValue()) { 178 | //if (value <= 0) 179 | // std::cout << iter.getCoord() << " value: " << value << std::endl; 180 | int x, y, z; 181 | x = iter.getCoord().x(); 182 | y = iter.getCoord().y(); 183 | z = iter.getCoord().z(); 184 | 185 | maxx = std::max(x, maxx); 186 | maxy = std::max(y, maxy); 187 | maxz = std::max(z, maxz); 188 | 189 | minx = std::min(x, minx); 190 | miny = std::min(y, miny); 191 | minz = std::min(z, minz); 192 | 193 | x += half_grid_resolution; 194 | y += half_grid_resolution; 195 | z += half_grid_resolution; 196 | 197 | if (x >= 0 && y >= 0 && z >= 0 && x <= grid_resolution && y <= grid_resolution && z <= grid_resolution) 198 | { 199 | filled_grad_count++; 200 | sdf_grids_grad[x][y][z] = value; 201 | } 202 | } 203 | else 204 | { 205 | //std::cout << "Not voxel node sdf grad" << std::endl; 206 | //int junk; 207 | //std::cin >> junk; 208 | } 209 | } 210 | 211 | auto valid_grid_coeff = [](int x, int y, int z, int max) 212 | { 213 | if (x < 0 || y < 0 || z < 0) 214 | return false; 215 | if (x > max || y > max || z > max) 216 | return false; 217 | return true; 218 | }; 219 | 220 | int coordx, coordy, coordz; 221 | FILE *wf = fopen(sdf_filename, "wb"); 222 | std::vector buffered_data; 223 | buffered_data.reserve(sampled_points.size() * 7); 224 | for (int i = 0; i < sampled_points.size(); i++) 225 | { 226 | OpenMesh::Vec3i point = sampled_points[i]; 227 | coordx = point[0]; 228 | coordy = point[1]; 229 | coordz = point[2]; 230 | myassert(fabs(sdf_grids[coordx][coordy][coordz]) < sdf_truncation); 231 | if (sdf_grids_grad[coordx][coordy][coordz].x() < -1000) 232 | { 233 | //openvdb does not generate gradient we want 234 | float gradx, grady, gradz; 235 | //center difference 236 | //x 237 | if (valid_grid_coeff(coordx - 1, coordy, coordz, grid_resolution)) 238 | { 239 | myassert(sdf_grids[coordx - 1][coordy][coordz] > -1000); 240 | if (valid_grid_coeff(coordx + 1, coordy, coordz, grid_resolution)) 241 | { 242 | myassert(sdf_grids[coordx + 1][coordy][coordz] > -1000); 243 | //center difference 244 | gradx = (sdf_grids[coordx + 1][coordy][coordz] - sdf_grids[coordx - 1][coordy][coordz]) / 2.0; 245 | } 246 | else 247 | gradx = sdf_grids[coordx][coordy][coordz] - sdf_grids[coordx - 1][coordy][coordz]; 248 | } 249 | else 250 | { 251 | gradx = sdf_grids[coordx + 1][coordy][coordz] - sdf_grids[coordx][coordy][coordz]; 252 | } 253 | 254 | //y 255 | if (valid_grid_coeff(coordx, coordy - 1, coordz, grid_resolution)) 256 | { 257 | myassert(sdf_grids[coordx][coordy - 1][coordz] > -1000); 258 | if (valid_grid_coeff(coordx, coordy + 1, coordz, grid_resolution)) 259 | { 260 | myassert(sdf_grids[coordx][coordy + 1][coordz] > -1000); 261 | //center difference 262 | grady = (sdf_grids[coordx][coordy + 1][coordz] - sdf_grids[coordx][coordy - 1][coordz]) / 2.0; 263 | } 264 | else 265 | grady = sdf_grids[coordx][coordy][coordz] - sdf_grids[coordx][coordy - 1][coordz]; 266 | } 267 | else 268 | { 269 | grady = sdf_grids[coordx][coordy + 1][coordz] - sdf_grids[coordx][coordy][coordz]; 270 | } 271 | 272 | //z 273 | if (valid_grid_coeff(coordx, coordy, coordz - 1, grid_resolution)) 274 | { 275 | myassert(sdf_grids[coordx][coordy][coordz - 1] > -1000); 276 | if (valid_grid_coeff(coordx, coordy, coordz + 1, grid_resolution)) 277 | { 278 | myassert(sdf_grids[coordx][coordy][coordz + 1] > -1000); 279 | //center difference 280 | gradz = (sdf_grids[coordx][coordy][coordz + 1] - sdf_grids[coordx][coordy][coordz - 1]) / 2.0; 281 | } 282 | else 283 | gradz = sdf_grids[coordx][coordy][coordz] - sdf_grids[coordx][coordy][coordz - 1]; 284 | } 285 | else 286 | { 287 | gradz = sdf_grids[coordx][coordy][coordz + 1] - sdf_grids[coordx][coordy][coordz]; 288 | } 289 | 290 | sdf_grids_grad[coordx][coordy][coordz] = openvdb::math::Vec3s(gradx, grady, gradz); 291 | sdf_grids_grad[coordx][coordy][coordz].normalize(); 292 | } 293 | myassert(sdf_grids_grad[coordx][coordy][coordz].x() > -1000); 294 | sdf_grids_grad[coordx][coordy][coordz].normalize(); 295 | buffered_data.push_back(2.0*coordx / grid_resolution - 1); 296 | buffered_data.push_back(2.0*coordy / grid_resolution - 1); 297 | buffered_data.push_back(2.0*coordz / grid_resolution - 1); 298 | buffered_data.push_back(sdf_grids[coordx][coordy][coordz]); 299 | buffered_data.push_back(sdf_grids_grad[coordx][coordy][coordz].x()); 300 | buffered_data.push_back(sdf_grids_grad[coordx][coordy][coordz].y()); 301 | buffered_data.push_back(sdf_grids_grad[coordx][coordy][coordz].z()); 302 | } 303 | fwrite(&(buffered_data[0]), sizeof(float), buffered_data.size(), wf); 304 | fclose(wf); 305 | 306 | printf("done\n"); 307 | } 308 | 309 | int main(int argc, char *argv[]) 310 | { 311 | if (argc != 3 && argc != 4) 312 | { 313 | printf("Usage: vdb_tsdf model_obj_filename output_sdf_filename [resolution=256]\n"); 314 | return 0; 315 | } 316 | int grid_resolution = 256; 317 | 318 | if (argc >= 4) 319 | { 320 | sscanf(argv[3], "%d", &grid_resolution); 321 | printf("grid resolution set to %d\n", grid_resolution); 322 | } 323 | 324 | openvdb::initialize(); 325 | fetch_sdf(argv[1], argv[2], grid_resolution); 326 | std::cout << "Finished converting mesh." << std::endl; 327 | return 0; 328 | } --------------------------------------------------------------------------------