def test(self, show=False, use_normal=False): data_folder = '../../PointSR_data/CAD/mesh_MC16k' phase = data_folder.split('/')[-2] + data_folder.split('/')[-1] save_path = os.path.join(MODEL_DIR, 'result/' + phase) self.saver = tf.train.Saver() _, restore_model_path = model_utils.pre_load_checkpoint(MODEL_DIR) print restore_model_path config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True with tf.Session(config=config) as sess: self.saver.restore(sess, restore_model_path) samples = glob(data_folder + "/.xyz") samples.sort() total_time = 0 #input, dist, edge, data_radius, name = data_provider.load_patch_data(NUM_POINT, True, 30) #edge = np.reshape(edge,[-1,NUM_EDGE,6]) for i, item in tqdm(enumerate(samples)): input = np.loadtxt(item) edge = np.loadtxt( item.replace('mesh_MC16k', 'mesh_edge').replace('.xyz', '_edge.xyz')) idx = np.all(edge[:, 0:3] == edge[:, 3:6], axis=-1) edge = edge[idx == False] l = len(edge) idx = range(l) * (1300 / l) + list( np.random.permutation(l)[:1300 % l]) edge = edge[idx] # # coord = input[:, 0:3] # # centroid = np.mean(coord, axis=0, keepdims=True) # # coord = coord - centroid # # furthest_distance = np.amax(np.sqrt(np.sum(abs(coord) ** 2, axis=-1))) # # coord = coord / furthest_distance # # input[:, 0:3] = coord input = np.expand_dims(input, axis=0) # input = data_provider.jitter_perturbation_point_cloud(input, sigma=0.01, clip=0.02) start_time = time.time() edge_pl = tf.placeholder(tf.float32, [1, edge.shape[0], 6]) dist_gt_pl = tf.sqrt( tf.reduce_min(model_utils.distance_point2edge( self.pred, edge_pl), axis=-1)) pred, pred_dist, dist_gt = sess.run( [self.pred, self.pred_dist, dist_gt_pl], feed_dict={ self.pointclouds_input: input[:, :, 0:3], self.pointclouds_radius: np.ones(BATCH_SIZE), edge_pl: np.expand_dims(edge, axis=0) }) total_time += time.time() - start_time norm_pl = np.zeros_like(pred) ##--------------visualize predicted point cloud---------------------- if show: f, axis = plt.subplots(3) axis[0].imshow( pc_util.point_cloud_three_views(input[:, 0:3], diameter=5)) axis[1].imshow( pc_util.point_cloud_three_views(pred[0, :, :], diameter=5)) axis[2].imshow( pc_util.point_cloud_three_views(gt[:, 0:3], diameter=5)) plt.show() path = os.path.join(save_path, item.split('/')[-1][:-4] + ".ply") # rgba =data_provider.convert_dist2rgba(pred_dist2,scale=10) # data_provider.save_ply(path, np.hstack((pred[0, ...],rgba,pred_dist2.reshape(NUM_ADDPOINT,1)))) path = os.path.join(save_path, item.split('/')[-1][:-4] + "_gt.ply") rgba = data_provider.convert_dist2rgba(dist_gt[0], scale=5) data_provider.save_ply( path, np.hstack( (pred[0, ...], rgba, dist_gt.reshape(NUM_ADDPOINT, 1)))) path = path.replace(phase, phase + "_input") path = path.replace('xyz', 'ply') rgba = data_provider.convert_dist2rgba(pred_dist[0], scale=5) data_provider.save_ply( path, np.hstack((input[0], rgba, pred_dist.reshape(NUM_POINT, 1)))) print total_time / len(samples)
def build_graph(self, is_training=True, scope='generator'): bn_decay = 0.95 self.step = tf.Variable(0, trainable=False) self.pointclouds_input = tf.placeholder(tf.float32, shape=(BATCH_SIZE, NUM_POINT, 3)) self.pointclouds_radius = tf.placeholder(tf.float32, shape=(BATCH_SIZE)) self.pointclouds_poisson = tf.placeholder(tf.float32, shape=(BATCH_SIZE, NUM_POINT, 3)) # self.pointclouds_dist = tf.placeholder(tf.float32, shape=(BATCH_SIZE, NUM_POINT)) self.pointclouds_idx = tf.placeholder(tf.int32, shape=(BATCH_SIZE, NUM_POINT, 2)) self.pointclouds_edge = tf.placeholder(tf.float32, shape=(BATCH_SIZE, NUM_EDGE, 6)) self.pointclouds_plane = tf.placeholder(tf.float32, shape=(BATCH_SIZE, NUM_FACE, 9)) self.pointclouds_plane_normal = tf.placeholder(tf.float32, shape=(BATCH_SIZE, NUM_FACE, 3)) # create the generator model self.pred_dist, self.pred_coord, self.idx, self.transform = MODEL_GEN.get_gen_model( self.pointclouds_input, is_training, scope=scope, bradius=1.0, num_addpoint=NUM_ADDPOINT, reuse=None, use_normal=False, use_bn=False, use_ibn=False, bn_decay=bn_decay, up_ratio=UP_RATIO, idx=self.pointclouds_idx, is_crop=IS_CROP) ###calculate the distance ground truth of upsample_point self.pointclouds_dist = model_utils.distance_point2edge( self.pred_coord, self.pointclouds_edge) self.pointclouds_dist = tf.sqrt( tf.reduce_min(self.pointclouds_dist, axis=-1)) self.pointclouds_dist_truncated = tf.minimum(0.5, self.pointclouds_dist) self.pred_dist = tf.minimum(0.5, tf.maximum(0.0, self.pred_dist)) # gather the edge self.pred_edgecoord = tf.gather_nd(self.pred_coord, self.idx) self.pred_edgedist = tf.gather_nd(self.pred_dist, self.idx) self.edgedist = tf.gather_nd(self.pointclouds_dist_truncated, self.idx) # ## The following code is okay when the batch size is 1 self.edge_threshold = tf.constant( 0.05, tf.float32, [1]) # select a small value when use 1k points if True: indics = tf.where( tf.less_equal(self.pred_edgedist, self.edge_threshold)) #(?,2) self.select_pred_edgecoord = tf.gather_nd(self.pred_edgecoord, indics) #(?,3) self.select_pred_edgedist = tf.gather_nd(self.pred_edgedist, indics) #(?,3) else: indics = tf.where( tf.less_equal(self.pointclouds_dist_truncated, self.edge_threshold)) # (?,2) self.select_pred_edgecoord = tf.gather_nd(self.pred_coord, indics) # (?,3) self.select_pred_edgedist = tf.gather_nd(self.pred_dist, indics) # (?,3) # if is_training is False: # input_dist = model_utils.distance_point2edge(self.pointclouds_input,self.pointclouds_edge) # input_dist = tf.sqrt(tf.reduce_min(input_dist,axis=-1)) # indics = tf.where(tf.less_equal(input_dist,self.edge_threshold)) # self.select_input_edge = tf.gather_nd(self.pointclouds_input, indics) # self.select_input_edgedist = tf.gather_nd(input_dist,indics) if is_training is False: return self.dist_mseloss = 1.0 / (0.4 + self.pointclouds_dist_truncated) * ( self.pointclouds_dist_truncated - self.pred_dist)**2 self.dist_mseloss = 5 * tf.reduce_mean( self.dist_mseloss / tf.expand_dims(self.pointclouds_radius**2, axis=-1)) tf.summary.scalar('loss/dist_loss', self.dist_mseloss) tf.summary.histogram('dist/gt', self.pointclouds_dist_truncated) tf.summary.histogram('dist/edge_dist', self.edgedist) tf.summary.histogram('dist/pred', self.pred_dist) # weight = tf.pow(0.98, tf.to_float(tf.div(self.step,200))) weight = tf.maximum(0.5 - tf.to_float(self.step) / 20000.0, 0.0) # weight = tf.maximum(0.5 - tf.to_float(self.step) / 200000.0, 0.0) self.edgemask = tf.to_float( tf.less_equal( weight * self.edgedist + (1 - weight) * self.pred_edgedist, 0.15)) # self.edgemask = tf.to_float(tf.less_equal(self.edgedist,0.45)) self.edge_loss = 50 * tf.reduce_sum( self.edgemask * self.edgedist**2 / tf.expand_dims(self.pointclouds_radius**2, axis=-1)) / ( tf.reduce_sum(self.edgemask) + 1.0) tf.summary.scalar('weight', weight) tf.summary.histogram('loss/edge_mask', self.edgemask) tf.summary.scalar('loss/edge_loss', self.edge_loss) with tf.device('/gpu:0'): self.plane_dist = model_utils.distance_point2mesh( self.pred_coord, self.pointclouds_plane) self.plane_dist = tf.reduce_min(self.plane_dist, axis=2) # idx = tf.argmin(self.plane_dist, axis=2,output_type=tf.int32) # idx0 = tf.tile(tf.reshape(tf.range(BATCH_SIZE), (BATCH_SIZE, 1)), (1, NUM_POINT*UP_RATIO/2)) # face_normal = tf.gather_nd(self.pointclouds_plane_normal,tf.stack([idx0,idx],axis=-1)) # dist = tf.where(tf.is_nan(dist),tf.zeros_like(dist),dist) self.plane_loss = 500 * tf.reduce_mean( self.plane_dist / tf.expand_dims(self.pointclouds_radius**2, axis=-1)) tf.summary.scalar('loss/plane_loss', self.plane_loss) #self.perulsionloss = 10*model_utils.get_uniform_loss1_orthdistance(self.pred_coord,face_normal, numpoint=NUM_POINT*UP_RATIO) self.perulsionloss = 500 * model_utils.get_perulsion_loss1( self.pred_coord, numpoint=NUM_POINT * UP_RATIO) tf.summary.scalar('loss/perulsion_loss', self.perulsionloss) # # Enforce the transformation as orthogonal matrix # K = transform.get_shape()[1].value # BxKxK # mat_diff = tf.matmul(transform, tf.transpose(transform, perm=[0, 2, 1])) # mat_diff -= tf.constant(np.eye(K), dtype=tf.float32) # self.mat_diff_loss = 0.01*tf.nn.l2_loss(mat_diff) # tf.summary.scalar('loss/mat_loss', self.mat_diff_loss) self.total_loss = self.dist_mseloss + self.plane_loss + self.edge_loss + self.perulsionloss + tf.losses.get_regularization_loss( ) gen_update_ops = [ op for op in tf.get_collection(tf.GraphKeys.UPDATE_OPS) if op.name.startswith(scope) ] gen_tvars = [ var for var in tf.trainable_variables() if var.name.startswith(scope) ] with tf.control_dependencies(gen_update_ops): self.pre_gen_train = tf.train.AdamOptimizer( BASE_LEARNING_RATE, beta1=0.9).minimize(self.total_loss, var_list=gen_tvars, colocate_gradients_with_ops=False, global_step=self.step) # merge summary and add pointclouds summary tf.summary.scalar('loss/regularation', tf.losses.get_regularization_loss()) tf.summary.scalar('loss/total_loss', self.total_loss) self.merged = tf.summary.merge_all() self.pointclouds_image_input = tf.placeholder( tf.float32, shape=[None, 500, 1500, 1]) pointclouds_input_summary = tf.summary.image( '1_input', self.pointclouds_image_input, max_outputs=1) self.pointclouds_image_pred = tf.placeholder( tf.float32, shape=[None, 500, 1500, 1]) pointclouds_pred_summary = tf.summary.image( '2_pred', self.pointclouds_image_pred, max_outputs=1) self.pointclouds_image_gt = tf.placeholder(tf.float32, shape=[None, 500, 1500, 1]) pointclouds_gt_summary = tf.summary.image('3_edge', self.pointclouds_image_gt, max_outputs=1) self.image_merged = tf.summary.merge([ pointclouds_input_summary, pointclouds_pred_summary, pointclouds_gt_summary ])
def crop_patch_boxes(self, save_root_path, use_dijkstra=True, id=0, boxes = None, scale_ratio=1,random_ratio=0.7): if save_root_path[-1]=='/': save_root_path = save_root_path[:-1] if not os.path.exists(save_root_path): os.makedirs(save_root_path) os.makedirs(save_root_path + "_dist") os.makedirs(save_root_path+"_edge") os.makedirs(save_root_path + "_edgepoint") os.makedirs(save_root_path + "_face") os.makedirs(save_root_path + "_facepoint") ## the first part if boxes is not None: centroid = [] for box in boxes: cen = np.mean(np.reshape(box,[-1,3]),axis=0) centroid.append(cen) centroid = np.asarray(centroid) idx = self.nbrs.query_ball_point(centroid, r=0.03) total_idx = [] for item in idx: total_idx.extend(item) total_idx = np.asarray(total_idx) corner_num = int(self.patch_num*random_ratio) idx = np.random.permutation(len(total_idx))[:corner_num] seeds1 = total_idx[idx] config = tf.ConfigProto() config.gpu_options.visible_device_list = str(id) config.gpu_options.allow_growth = True config.allow_soft_placement = True self.sess = tf.Session(config=config) tf_point = tf.placeholder(tf.float32, [1, self.patch_size, 3]) tf_edge = tf.placeholder(tf.float32, [1, self.edge.shape[0], 6]) tf_face = tf.placeholder(tf.float32, [1, self.face.shape[0], 9]) tf_edge_dist = model_utils.distance_point2edge(tf_point, tf_edge) tf_edge_dist = tf.sqrt(tf.squeeze(tf_edge_dist, axis=0)) tf_face_dist = model_utils.distance_point2mesh(tf_point, tf_face) tf_face_dist = tf.sqrt(tf.squeeze(tf_face_dist, axis=0)) i = -1 for seed in tqdm(seeds1): t0 = time.time() i = i + 1 assert scale_ratio >= 1.0 patch_size = int(self.patch_size * np.random.uniform(1.0, scale_ratio)) try: if use_dijkstra: idx = self.geodesic_knn(seed, patch_size) else: idx = self.bfs_knn(seed, patch_size) except: print "has exception" continue idxx = np.random.permutation(patch_size)[:self.patch_size] idxx.sort() idx = idx[idxx] point = self.data[idx] clean_point = self.clean_data[idx] t1 = time.time() edge_dist, face_dist = self.sess.run([tf_edge_dist, tf_face_dist], feed_dict={tf_point: np.expand_dims(clean_point, axis=0), tf_edge: np.expand_dims(self.edge, axis=0), tf_face: np.expand_dims(self.face, axis=0)}) subedge = self.get_subedge(edge_dist) subface = self.get_subface(face_dist) t2 = time.time() # print t1-t0, t2-t1 print "patch:%d point:%d subedge:%d subface:%d" % (i, patch_size, len(subedge), len(subface)) dist_min = np.reshape(np.min(edge_dist, axis=-1), [-1, 1]) np.savetxt('%s/%s_%d.xyz' % (save_root_path, self.name, i), point, fmt='%0.6f') np.savetxt('%s_dist/%s_%d.xyz' % (save_root_path, self.name, i), np.concatenate([point, dist_min], axis=-1), fmt='%0.6f') np.savetxt('%s_edge/%s_%d.xyz' % (save_root_path, self.name, i), subedge, fmt='%0.6f') np.savetxt('%s_edgepoint/%s_%d.xyz' % (save_root_path, self.name, i), sampling_from_edge(subedge), fmt='%0.6f') np.savetxt('%s_face/%s_%d.xyz' % (save_root_path, self.name, i), subface, fmt='%0.6f') np.savetxt('%s_facepoint/%s_%d.xyz' % (save_root_path, self.name, i), sampling_from_face(subface), fmt='%0.6f') self.sess.close() else: corner_num = 0 ## second part seeds2 = self.get_idx((self.patch_num-corner_num)*4, random_ratio=0.0) seeds2 = np.asarray(seeds2).astype(np.int32) ##remove those seed near the corner if boxes is not None: d1 = self.data[seeds2] d2 = centroid dist = spatial.distance_matrix(d1,d2) min_distance = np.amin(dist,axis=1) new_seed = [] for item1, item2 in zip(seeds2, min_distance): if item2 >0.05: new_seed.append(item1) new_seed = np.asarray(new_seed) seeds2 = new_seed[:(self.patch_num-corner_num)] else: seeds2 = seeds2[:(self.patch_num-corner_num)] self.edge = np.asarray([[10, 10, 10, 20, 20, 20]]) config = tf.ConfigProto() config.gpu_options.visible_device_list=str(id) config.gpu_options.allow_growth = True config.allow_soft_placement = True self.sess = tf.Session(config=config) tf_point = tf.placeholder(tf.float32, [1, self.patch_size, 3]) tf_edge = tf.placeholder(tf.float32, [1, self.edge.shape[0], 6]) tf_face = tf.placeholder(tf.float32, [1, self.face.shape[0], 9]) tf_edge_dist = model_utils.distance_point2edge(tf_point, tf_edge) tf_edge_dist = tf.sqrt(tf.squeeze(tf_edge_dist, axis=0)) tf_face_dist = model_utils.distance_point2mesh(tf_point, tf_face) tf_face_dist = tf.sqrt(tf.squeeze(tf_face_dist, axis=0)) i = corner_num-1 for seed in tqdm(seeds2): t0 = time.time() i = i+1 # patch_size = self.patch_size*np.random.randint(1,scale_ratio+1) assert scale_ratio>=1.0 patch_size = int(self.patch_size*np.random.uniform(1.0, scale_ratio)) try: if use_dijkstra: idx = self.geodesic_knn(seed, patch_size) else: idx = self.bfs_knn(seed, patch_size) except: print "has exception" continue idxx = np.random.permutation(patch_size)[:self.patch_size] idxx.sort() idx = idx[idxx] point = self.data[idx] clean_point = self.clean_data[idx] t1 = time.time() edge_dist,face_dist = self.sess.run([tf_edge_dist,tf_face_dist],feed_dict={tf_point: np.expand_dims(clean_point, axis=0), tf_edge: np.expand_dims(self.edge, axis=0), tf_face: np.expand_dims(self.face, axis=0)}) subedge = self.get_subedge(edge_dist) subface = self.get_subface(face_dist) t2 = time.time() # print t1-t0, t2-t1 print "patch:%d point:%d subedge:%d subface:%d" % (i, patch_size, len(subedge),len(subface)) dist_min = np.reshape(np.min(edge_dist, axis=-1),[-1,1]) np.savetxt('%s/%s_%d.xyz' % (save_root_path, self.name, i), point, fmt='%0.6f') np.savetxt('%s_dist/%s_%d.xyz' % (save_root_path,self.name, i), np.concatenate([point,dist_min],axis=-1), fmt='%0.6f') np.savetxt('%s_edge/%s_%d.xyz' % (save_root_path, self.name, i), subedge,fmt='%0.6f') np.savetxt('%s_edgepoint/%s_%d.xyz' % (save_root_path, self.name, i), sampling_from_edge(subedge),fmt='%0.6f') np.savetxt('%s_face/%s_%d.xyz' % (save_root_path, self.name, i), subface, fmt='%0.6f') np.savetxt('%s_facepoint/%s_%d.xyz' % (save_root_path, self.name, i), sampling_from_face(subface), fmt='%0.6f') self.sess.close()
def crop_patch(self, save_root_path, use_dijkstra=True, id=0, scale_ratio=1,random_ratio=0.7,cal_edge=False): if save_root_path[-1]=='/': save_root_path = save_root_path[:-1] if not os.path.exists(save_root_path): os.makedirs(save_root_path) os.makedirs(save_root_path + "_dist") os.makedirs(save_root_path+"_edge") os.makedirs(save_root_path + "_edgepoint") os.makedirs(save_root_path + "_face") os.makedirs(save_root_path + "_facepoint") seeds = self.get_idx(self.patch_num, random_ratio=random_ratio) if self.edge.shape[0] == 0: self.edge = np.asarray([[10, 10, 10, 20, 20, 20]]) config = tf.ConfigProto() config.gpu_options.visible_device_list=str(id) config.gpu_options.allow_growth = True config.allow_soft_placement = True self.sess = tf.Session(config=config) tf_point = tf.placeholder(tf.float32, [1, self.patch_size, 3]) tf_edge = tf.placeholder(tf.float32, [1, self.edge.shape[0], 6]) tf_face = tf.placeholder(tf.float32, [1, self.face.shape[0], 9]) tf_edge_dist = model_utils.distance_point2edge(tf_point, tf_edge) tf_edge_dist = tf.sqrt(tf.squeeze(tf_edge_dist, axis=0)) tf_face_dist = model_utils.distance_point2mesh(tf_point, tf_face) tf_face_dist = tf.sqrt(tf.squeeze(tf_face_dist, axis=0)) i = -1 for seed in tqdm(seeds): t0 = time.time() i = i+1 # patch_size = self.patch_size*np.random.randint(1,scale_ratio+1) assert scale_ratio>=1.0 patch_size = int(self.patch_size*np.random.uniform(1.0, scale_ratio)) try: if use_dijkstra: idx = self.geodesic_knn(seed, patch_size) else: idx = self.bfs_knn(seed, patch_size) except: print "has exception" continue idxx = np.random.permutation(patch_size)[:self.patch_size] idxx.sort() idx = idx[idxx] point = self.data[idx] clean_point = self.clean_data[idx] if cal_edge: t1 = time.time() edge_dist,face_dist = self.sess.run([tf_edge_dist,tf_face_dist],feed_dict={tf_point: np.expand_dims(clean_point, axis=0), tf_edge: np.expand_dims(self.edge, axis=0), tf_face: np.expand_dims(self.face, axis=0)}) subedge = self.get_subedge(edge_dist) subface = self.get_subface(face_dist) dist_min = np.reshape(np.min(edge_dist, axis=-1), [-1, 1]) else: subedge = np.asarray([[10,10,10,20,20,20]]) subface = np.asarray([[10,10,10,20,20,20,30,30,30]]) dist_min=np.zeros((point.shape[0],1)) t2 = time.time() # print t1-t0, t2-t1 print "patch:%d point:%d subedge:%d subface:%d" % (i, patch_size, len(subedge),len(subface)) np.savetxt('%s/%s_%d.xyz' % (save_root_path, self.name, i), point, fmt='%0.6f') np.savetxt('%s_dist/%s_%d.xyz' % (save_root_path,self.name, i), np.concatenate([point,dist_min],axis=-1), fmt='%0.6f') np.savetxt('%s_edge/%s_%d.xyz' % (save_root_path, self.name, i), subedge,fmt='%0.6f') np.savetxt('%s_edgepoint/%s_%d.xyz' % (save_root_path, self.name, i), sampling_from_edge(subedge),fmt='%0.6f') np.savetxt('%s_face/%s_%d.xyz' % (save_root_path, self.name, i), subface, fmt='%0.6f') np.savetxt('%s_facepoint/%s_%d.xyz' % (save_root_path, self.name, i), sampling_from_face(subface), fmt='%0.6f') # np.savetxt('%s/%s_%d.xyz' % (save_root_path, self.name, i), point, fmt='%0.6f') # np.savetxt('%s/%s_%d_dist.xyz' % (save_root_path, self.name, i), np.concatenate([point, dist_min], axis=-1), # fmt='%0.6f') # np.savetxt('%s/%s_%d_edge.xyz' % (save_root_path, self.name, i), subedge, fmt='%0.6f') # np.savetxt('%s/%s_%d_edgepoint.xyz' % (save_root_path, self.name, i), sampling_from_edge(subedge), # fmt='%0.6f') # np.savetxt('%s/%s_%d_face.xyz' % (save_root_path, self.name, i), subface, fmt='%0.6f') # np.savetxt('%s/%s_%d_facepoint.xyz' % (save_root_path, self.name, i), sampling_from_face(subface), # fmt='%0.6f') self.sess.close()
def build_graph(self,is_training=True,scope='generator'): bn_decay = 0.95 self.step = tf.Variable(0, trainable=False) self.pointclouds_input = tf.placeholder(tf.float32, shape=(BATCH_SIZE, NUM_POINT, 3)) self.pointclouds_radius = tf.placeholder(tf.float32, shape=(BATCH_SIZE)) self.pointclouds_idx = tf.placeholder(tf.int32,shape=(BATCH_SIZE,NUM_POINT,2)) self.pointclouds_edge = tf.placeholder(tf.float32, shape=(BATCH_SIZE, NUM_EDGE, 6)) self.pointclouds_surface = tf.placeholder(tf.float32, shape=(BATCH_SIZE, NUM_FACE, 9)) # create the generator model self.pred_dist, self.pred_coord,self.idx = MODEL_GEN.get_gen_model(self.pointclouds_input, is_training, scope=scope, bradius=1.0, num_addpoint=NUM_ADDPOINT,reuse=None, use_normal=False, use_bn=False,use_ibn=False, bn_decay=bn_decay, up_ratio=UP_RATIO,idx=self.pointclouds_idx,is_crop=IS_CROP) # calculate the ground truth point-to-edge distances for upsampled_points self.gt_dist = model_utils.distance_point2edge(self.pred_coord, self.pointclouds_edge) self.gt_dist = tf.sqrt(tf.reduce_min(self.gt_dist, axis=-1)) self.gt_dist_truncated = tf.minimum(0.5, self.gt_dist) self.pred_dist = tf.minimum(0.5,tf.maximum(0.0,self.pred_dist)) if is_training is False: # identify the edge points in the inference phase # We first filter the potential edge points with the top-k (NUM_ADDPOINT) minimum predicted point-to-edge distances. self.pred_edgecoord = tf.gather_nd(self.pred_coord, self.idx) self.pred_edgedist = tf.gather_nd(self.pred_dist, self.idx) # The following code is okay when the batch size is 1. self.edge_threshold = tf.constant(0.05, tf.float32, [1]) indics = tf.where(tf.less_equal(self.pred_edgedist, self.edge_threshold)) # (?,2) self.select_pred_edgecoord = tf.gather_nd(self.pred_edgecoord, indics) # (?,3) self.select_pred_edgedist = tf.gather_nd(self.pred_edgedist, indics) # (?,3) return # Loss1: Edge distance regressionn loss self.dist_mseloss = (self.gt_dist_truncated - self.pred_dist) ** 2 self.dist_mseloss = 5 * tf.reduce_mean(self.dist_mseloss / tf.expand_dims(self.pointclouds_radius ** 2, axis=-1)) tf.summary.histogram('dist/gt', self.gt_dist_truncated) tf.summary.histogram('dist/pred', self.pred_dist) tf.summary.scalar('loss/dist_loss', self.dist_mseloss) # Loss2: Edge loss # We first filter the potential edge points with the top-k (NUM_ADDPOINT) minimum predicted point-to-edge distances. self.pred_edgecoord = tf.gather_nd(self.pred_coord, self.idx) self.pred_edgedist = tf.gather_nd(self.pred_dist, self.idx) self.gt_edgedist_truncated = tf.gather_nd(self.gt_dist_truncated, self.idx) # At the beginning of the training, the predicted point-to-edge distance is not accurate. # When identifying the edge points, we use the weighted sum of the predicted and ground truth point-to-edge distances. weight = tf.maximum(0.5 - tf.to_float(self.step) / 20000.0, 0.0) self.edgemask = tf.to_float(tf.less_equal(weight * self.gt_edgedist_truncated + (1 - weight) * self.pred_edgedist, 0.15)) self.edge_loss = 50*tf.reduce_sum(self.edgemask * self.gt_edgedist_truncated ** 2 / tf.expand_dims(self.pointclouds_radius ** 2, axis=-1)) / (tf.reduce_sum(self.edgemask) + 1.0) tf.summary.scalar('weight',weight) tf.summary.histogram('dist/edge_dist', self.gt_edgedist_truncated) tf.summary.histogram('loss/edge_mask', self.edgemask) tf.summary.scalar('loss/edge_loss', self.edge_loss) # Loss3: Surface loss self.surface_dist = model_utils.distance_point2mesh(self.pred_coord, self.pointclouds_surface) self.surface_dist = tf.reduce_min(self.surface_dist, axis=2) self.surface_loss = 500*tf.reduce_mean(self.surface_dist / tf.expand_dims(self.pointclouds_radius ** 2, axis=-1)) tf.summary.scalar('loss/plane_loss', self.surface_loss) # Loss4: Repulsion loss self.perulsionloss = 500*model_utils.get_perulsion_loss(self.pred_coord, numpoint=NUM_POINT * UP_RATIO) tf.summary.scalar('loss/perulsion_loss', self.perulsionloss) # Total loss self.total_loss = self.surface_loss + self.edge_loss + self.perulsionloss + self.dist_mseloss + tf.losses.get_regularization_loss() # make optimizer gen_update_ops = [op for op in tf.get_collection(tf.GraphKeys.UPDATE_OPS) if op.name.startswith(scope)] gen_tvars = [var for var in tf.trainable_variables() if var.name.startswith(scope)] with tf.control_dependencies(gen_update_ops): self.gen_train = tf.train.AdamOptimizer(BASE_LEARNING_RATE, beta1=0.9).minimize(self.total_loss, var_list=gen_tvars, colocate_gradients_with_ops=False, global_step=self.step) # merge summary and add pointclouds summary tf.summary.scalar('loss/regularation', tf.losses.get_regularization_loss()) tf.summary.scalar('loss/total_loss', self.total_loss) self.merged = tf.summary.merge_all() self.pointclouds_image_input = tf.placeholder(tf.float32, shape=[None, 500, 1500, 1]) pointclouds_input_summary = tf.summary.image('1_input', self.pointclouds_image_input, max_outputs=1) self.pointclouds_image_pred = tf.placeholder(tf.float32, shape=[None, 500, 1500, 1]) pointclouds_pred_summary = tf.summary.image('2_pred', self.pointclouds_image_pred, max_outputs=1) self.pointclouds_image_gt = tf.placeholder(tf.float32, shape=[None, 500, 1500, 1]) pointclouds_gt_summary = tf.summary.image('3_edge', self.pointclouds_image_gt, max_outputs=1) self.image_merged = tf.summary.merge([pointclouds_input_summary, pointclouds_pred_summary, pointclouds_gt_summary])
def build_graph(self, is_training=True, scope='generator'): bn_decay = 0.95 self.step = tf.Variable(0, trainable=False) self.pointclouds_input = tf.placeholder(tf.float32, shape=(BATCH_SIZE, NUM_POINT, 3)) self.pointclouds_radius = tf.placeholder(tf.float32, shape=(BATCH_SIZE)) self.pointclouds_poisson = tf.placeholder(tf.float32, shape=(BATCH_SIZE, NUM_POINT, 3)) # self.pointclouds_dist = tf.placeholder(tf.float32, shape=(BATCH_SIZE, NUM_POINT)) self.pointclouds_idx = tf.placeholder(tf.int32, shape=(BATCH_SIZE, NUM_POINT, 2)) self.pointclouds_edge = tf.placeholder(tf.float32, shape=(BATCH_SIZE, NUM_EDGE, 6)) self.pointclouds_plane = tf.placeholder(tf.float32, shape=(BATCH_SIZE, NUM_FACE, 9)) # create the generator model self.pred_dist, self.pred_coord, self.idx = MODEL_GEN.get_gen_model( self.pointclouds_input, is_training, scope=scope, bradius=1.0, num_addpoint=NUM_ADDPOINT, reuse=None, use_normal=False, use_bn=False, use_ibn=False, bn_decay=bn_decay, up_ratio=UP_RATIO, idx=self.pointclouds_idx, is_crop=True) ###calculate the distance ground truth of upsample_point self.pointclouds_dist = model_utils.distance_point2edge( self.pred_coord, self.pointclouds_edge) self.pointclouds_dist = tf.sqrt( tf.reduce_min(self.pointclouds_dist, axis=-1)) self.pointclouds_dist_truncated = tf.minimum(0.5, self.pointclouds_dist) self.pred_dist = tf.minimum(0.5, tf.maximum(0.0, self.pred_dist)) ##### # project_point = model_utils.projected_point(self.pred_coord,self.pointclouds_edge) # offset = project_point-self.pred_coord # dist = tf.reduce_sum(offset*offset,axis=-1) ##### # gather the edge self.pred_edgecoord = tf.gather_nd(self.pred_coord, self.idx) self.pred_edgedist = tf.gather_nd(self.pred_dist, self.idx) self.edgedist = tf.gather_nd(self.pointclouds_dist_truncated, self.idx) # ## The following code is okay when the batch size is 1 self.edge_threshold = tf.constant(0.1, tf.float32, [1]) if True: indics = tf.where( tf.less_equal(self.pred_edgedist, self.edge_threshold)) #(?,2) self.select_pred_edgecoord = tf.gather_nd(self.pred_edgecoord, indics) #(?,3) self.select_pred_edgedist = tf.gather_nd(self.pred_edgedist, indics) #(?,3) else: indics = tf.where( tf.less_equal(self.pointclouds_dist_truncated, self.edge_threshold)) # (?,2) self.select_pred_edgecoord = tf.gather_nd(self.pred_coord, indics) # (?,3) self.select_pred_edgedist = tf.gather_nd(self.pred_dist, indics) # (?,3) if is_training is False: input_dist = model_utils.distance_point2edge( self.pointclouds_input, self.pointclouds_edge) input_dist = tf.sqrt(tf.reduce_min(input_dist, axis=-1)) indics = tf.where(tf.less_equal(input_dist, self.edge_threshold)) self.select_input_edge = tf.gather_nd(self.pointclouds_input, indics) self.select_input_edgedist = tf.gather_nd(input_dist, indics) if is_training is False: return self.dist_mseloss = 1.0 / (0.2 + self.pointclouds_dist_truncated) * ( self.pointclouds_dist_truncated - self.pred_dist)**2 self.dist_mseloss = 10 * tf.reduce_mean( self.dist_mseloss / tf.expand_dims(self.pointclouds_radius**2, axis=-1)) tf.summary.scalar('loss/dist_loss', self.dist_mseloss) tf.summary.histogram('dist/gt', self.pointclouds_dist_truncated) tf.summary.histogram('dist/pred', self.pred_dist) self.plane_loss = model_utils.distance_point2mesh( self.pred_coord, self.pointclouds_plane) self.plane_loss = 500 * tf.reduce_mean( self.plane_loss / tf.expand_dims(self.pointclouds_radius**2, axis=-1)) tf.summary.scalar('loss/plane_loss', self.plane_loss) weight = tf.pow(0.97, tf.to_float(tf.div(self.step, 100))) self.edgemask = tf.to_float( tf.less_equal( weight * self.edgedist + (1 - weight) * self.pred_edgedist, 0.15)) self.edge_loss = 10 * tf.reduce_sum( self.edgemask * self.edgedist / tf.expand_dims(self.pointclouds_radius**2, axis=-1)) / ( tf.reduce_sum(self.edgemask) + 1.0) tf.summary.scalar('weight', weight) tf.summary.histogram('loss/edge_mask', self.edgemask) tf.summary.scalar('loss/edge_loss', self.edge_loss) self.perulsionloss = 10 * model_utils.get_perulsion_loss1( self.pred_coord, numpoint=NUM_POINT * UP_RATIO) tf.summary.scalar('loss/perulsion_loss', self.perulsionloss) self.total_loss = self.dist_mseloss + self.plane_loss + self.edge_loss + self.perulsionloss + tf.losses.get_regularization_loss( ) gen_update_ops = [ op for op in tf.get_collection(tf.GraphKeys.UPDATE_OPS) if op.name.startswith(scope) ] gen_tvars = [ var for var in tf.trainable_variables() if var.name.startswith(scope) ] with tf.control_dependencies(gen_update_ops): self.pre_gen_train = tf.train.AdamOptimizer( BASE_LEARNING_RATE, beta1=0.9).minimize(self.total_loss, var_list=gen_tvars, colocate_gradients_with_ops=True, global_step=self.step) # merge summary and add pointclouds summary tf.summary.scalar('loss/regularation', tf.losses.get_regularization_loss()) tf.summary.scalar('loss/total_loss', self.total_loss) self.merged = tf.summary.merge_all() self.pointclouds_image_input = tf.placeholder( tf.float32, shape=[None, 500, 1500, 1]) pointclouds_input_summary = tf.summary.image( '1_input', self.pointclouds_image_input, max_outputs=1) self.pointclouds_image_pred = tf.placeholder( tf.float32, shape=[None, 500, 1500, 1]) pointclouds_pred_summary = tf.summary.image( '2_pred', self.pointclouds_image_pred, max_outputs=1) self.pointclouds_image_gt = tf.placeholder(tf.float32, shape=[None, 500, 1500, 1]) pointclouds_gt_summary = tf.summary.image('3_edge', self.pointclouds_image_gt, max_outputs=1) self.image_merged = tf.summary.merge([ pointclouds_input_summary, pointclouds_pred_summary, pointclouds_gt_summary ])
def test_hierarical_prediction(self): data_folder = '../../PointSR_data/CAD_imperfect/simu_noise' phase = data_folder.split('/')[-2] + "_" + data_folder.split('/')[-1] save_path = os.path.join(MODEL_DIR, 'result/' + phase) self.saver = tf.train.Saver() _, restore_model_path = model_utils.pre_load_checkpoint(MODEL_DIR) print restore_model_path config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True with tf.Session(config=config) as sess: self.saver.restore(sess, restore_model_path) total_time = 0 samples = glob(data_folder + "/*.xyz") samples.sort() for point_path in samples: print point_path edge_path = point_path.replace('simu_noise', 'mesh_edge') edge_path = edge_path.replace('_noise.xyz', '_edge.xyz') gm = GKNN(point_path, edge_path, patch_size=NUM_POINT, patch_num=30, add_noise=False) l = gm.edge.shape[0] idxx = range(l) * (NUM_EDGE / l) + list( np.random.permutation(l)[:NUM_EDGE % l]) edge = gm.edge[idxx] ## get patch seed from farthestsampling points = tf.convert_to_tensor(np.expand_dims(gm.data, axis=0), dtype=tf.float32) start = time.time() seed = farthest_point_sample(gm.data.shape[0] / 2, points).eval()[0] print "aaaaa", time.time() - start seed1_num = int(gm.data.shape[0] / NUM_POINT * 30) seed_list1 = seed[:seed1_num] seed_left = seed[seed1_num:] # seed2_num = int(gm.data.shape[0] / NUM_POINT * 1) # seed_list2 = gm.get_seed_fromdensity(seed2_num) # seed_list = np.concatenate([seed_list1, seed_list2]) seed_list = np.unique(seed_list1) inputs = [] up_point_list = [] up_edge_list = [] up_edgedist_list = [] input_edge_list = [] input_edgedist_list = [] fail = 0 for seed in tqdm(seed_list): try: patch_size = NUM_POINT * np.random.randint(1, 5) point = gm.bfs_knn(seed, patch_size) idx = np.random.permutation(patch_size)[:NUM_POINT] idx.sort() point = point[idx] except: fail = fail + 1 continue #get the idx idx1 = np.reshape(np.arange(NUM_POINT), [1, NUM_POINT]) idx0 = np.zeros((1, NUM_POINT)) idx = np.stack((idx0, idx1), axis=-1) up_point, up_edgepoint, up_edgedist, input_edge, input_edgedist = self.patch_prediction( point, edge, idx, sess) inputs.append(point) up_point_list.append(up_point) up_edge_list.append(up_edgepoint) up_edgedist_list.append(up_edgedist) input_edge_list.append(input_edge) input_edgedist_list.append(input_edgedist) print "total %d fails" % fail input = np.concatenate(inputs, axis=0) path = os.path.join( save_path, point_path.split('/')[-1][:-4] + "_input.xyz") data_provider.save_xyz(path, gm.data) input_edge = np.concatenate(input_edge_list, axis=0) input_edgedist = np.concatenate(input_edgedist_list, axis=0) rgba = data_provider.convert_dist2rgba(input_edgedist, scale=10) path = os.path.join( save_path, point_path.split('/')[-1][:-4] + "_inputedge.ply") data_provider.save_ply( path, np.hstack((input_edge, rgba, input_edgedist.reshape(-1, 1)))) pred = np.concatenate(up_point_list, axis=0) path = os.path.join( save_path, point_path.split('/')[-1][:-4] + "_output.xyz") data_provider.save_xyz(path, pred) pred_edge = np.concatenate(up_edge_list, axis=0) t1 = time.time() print "total %d edgepoint" % pred_edge.shape[0] edge_dist = np.zeros(pred_edge.shape[0]) for sid in range(0, pred_edge.shape[0], 20000): eid = np.minimum(pred_edge.shape[0], sid + 20000) tf_point = tf.placeholder(tf.float32, [1, eid - sid, 3]) tf_edge = tf.placeholder(tf.float32, [1, gm.edge.shape[0], 6]) pred_edge_dist_tf = model_utils.distance_point2edge( tf_point, tf_edge) pred_edge_dist_tf = tf.sqrt( tf.reduce_min(pred_edge_dist_tf, axis=-1)) edge_dist[sid:eid] = sess.run(pred_edge_dist_tf, feed_dict={ tf_point: np.expand_dims( pred_edge[sid:eid], axis=0), tf_edge: np.expand_dims(gm.edge, axis=0) }) t3 = time.time() print "tf time %f" % (t3 - t1) rgba = data_provider.convert_dist2rgba(edge_dist, scale=10) path = os.path.join( save_path, point_path.split('/')[-1][:-4] + "_outputedgeerror.ply") data_provider.save_ply( path, np.hstack((pred_edge, rgba, edge_dist.reshape(-1, 1)))) pred_edgedist = np.concatenate(up_edgedist_list, axis=0) rgba = data_provider.convert_dist2rgba(pred_edgedist, scale=10) path = os.path.join( save_path, point_path.split('/')[-1][:-4] + "_outputedge.ply") data_provider.save_ply( path, np.hstack((pred_edge, rgba, pred_edgedist.reshape(-1, 1)))) print total_time / len(samples)