def search(base, aim):
'''
both base and aim are sorted
'''
n = TfUtil.get_tensor_shape(base)[0]
m = TfUtil.get_tensor_shape(aim)[0]
def body(i, j, idx):
catch = tf.logical_or(tf.greater_equal(base[i, 0], aim[j]),
tf.equal(i, n - 1))
def found(i, j, idx):
match = tf.cast(tf.equal(base[i, 0], aim[j]), tf.int32)
tmp = (i + 1) * match - 1
idx = tf.concat([idx, tf.reshape(tmp, [1, 1])], 0)
return i + 1, j + 1, idx
i, j, idx = tf.cond(catch, lambda: found(i, j, idx), lambda:
(i + 1, j, idx))
return i, j, idx
i = tf.constant(0)
j = tf.constant(0)
idx = tf.zeros([0, 1], tf.int32)
cond = lambda i, j, idx: tf.logical_and(tf.less(i, n), tf.less(j, m))
i, j, idx = tf.while_loop(cond,
body, [i, j, idx],
shape_invariants=[
tf.TensorShape([]),
tf.TensorShape([]),
tf.TensorShape([None, 1])
])
return tf.squeeze(idx, 1)
def voxel_sampling_grouping(bsg, features, dset_shape_idx):
xyz = get_ele(features, 'xyz', dset_shape_idx)
xyz = tf.expand_dims(xyz, 0)
dsb = {}
dsb['grouped_pindex'], dsb['vox_index'], dsb['grouped_bot_cen_top'], dsb['empty_mask'],\
dsb['out_bot_cen_top'], dsb['flatting_idx'], dsb['flat_valid_mask'], dsb['nb_enough_p'], others =\
bsg.grouping_multi_scale(xyz)
# **********
for s in range(1, bsg._num_scale):
for item in [
'grouped_pindex', 'grouped_bot_cen_top', 'empty_mask',
'nb_enough_p'
]:
features[item + '_%d' % (s)] = tf.squeeze(dsb[item][s], [0, 1])
for item in ['flatting_idx', 'flat_valid_mask']:
features[item + '_%d' % (s)] = dsb[item][s]
# **********
vertices = tf.concat([features['vertex_f'], \
tf.cast(features['vertex_uint8'], tf.float32)], -1)
if TfUtil.t_shape(dsb['grouped_pindex'][0])[-1] == 0:
vertices_gped = vertices
else:
vertices = TfUtil.gather(vertices, dsb['grouped_pindex'][0], 0)
raise NotImplementedError
vertices_gped = tf.squeeze(vertices_gped, [0, 1, 2])
tmp = dset_shape_idx['shape']['vertex_f'][-1]
features['vertex_f'] = vertices_gped[..., 0:tmp]
features['vertex_uint8'] = tf.cast(vertices_gped[..., tmp:], tf.uint8)
return features
def search_idx_scope(self, flat_vidx, aim_scale):
with tf.variable_scope('search_idx_scope_Aimscale_%d' % (aim_scale)):
flatvidx_rawidxscope = self.flatvidx[aim_scale]
aim_idx = OctreeTf.search(flatvidx_rawidxscope, flat_vidx)
check_all_searched = tf.assert_equal(TfUtil.tshape0(aim_idx),
TfUtil.tshape0(flat_vidx),
message="search failed")
with tf.control_dependencies([check_all_searched]):
return aim_idx
def sampling_grouping(self, scale, points, grouped_pindex, xyz):
with tf.variable_scope('SG_S%d' % (scale)):
if grouped_pindex is None:
# global scale
new_points = points
new_xyz = None
else:
new_points = TfUtil.gather_third_d(points, grouped_pindex)
new_xyz = TfUtil.gather_second_d(xyz, grouped_pindex)
new_xyz = tf.reduce_mean(new_xyz, 2)
new_points = self.normalize_feature(new_points, 2)
self.log_tensor_p(new_points, '', 'saming grouping norm')
return new_points, new_xyz
def add_point_cloud(self,
xyzs,
sortedidxs_base=None,
record=False,
rawfn=None):
'''
xyzs input is not sorted for all scales
'''
assert TfUtil.tsize(xyzs) == 2
vidx_1d = self.get_vidx1d(xyzs)
sorted_idxs_base, end_idxs_lastscale = self.get_sort_idx(
xyzs, vidx_1d, sortedidxs_base)
if self._scale == 0:
vidx_rawscope_lsc = OctreeTf.get_idx_scope(end_idxs_lastscale)
flatvidx_rawidxscopes, vidx_rawscopes = self.get_upper_scale_scope(
vidx_rawscope_lsc)
flat_idx_lsc = OctreeTf.get_flat_idx(vidx_rawscope_lsc[:, 0:-2])
flatvidx_rawscope_lsc = tf.concat(
[tf.expand_dims(flat_idx_lsc, -1), vidx_rawscope_lsc[:, -2:]],
-1)
flatvidx_rawidxscopes.append(flatvidx_rawscope_lsc)
self.sorted_idxs_base = sorted_idxs_base
self.flatvidx = [d[:, 0:-2] for d in flatvidx_rawidxscopes]
self.vidx = [d[:, 0:-2] for d in vidx_rawscopes]
self.idxscope_pervs = [d[:, -2:] for d in vidx_rawscopes]
self.show_summary()
#self.merge_neighbour()
#self.check_sort(xyzs)
if record:
self.recording(rawfn)
#self.test_search_inv(xyzs)
return sorted_idxs_base, end_idxs_lastscale
def test_search_inv(self, xyzs, test_n, aim_scale, ply=False):
from utils import ply_util
assert aim_scale < self._scale_num
xyzs = check_xyz_normed(xyzs)
xyzs = tf.gather(xyzs, self.sorted_idxs_base, 0)
vn0 = TfUtil.tshape0(xyzs)
sp_idx = tf.random_shuffle(tf.range(vn0))[0:test_n]
sp_idx = tf.contrib.framework.sort(sp_idx)
xyzs_sp = tf.gather(xyzs, sp_idx, 0)
idx_scopes = self.search_neighbours(xyzs_sp, aim_scale)
if not ply:
return idx_scopes
for i in range(5):
ply_dir = '/tmp/octree_plys'
fn = os.path.join(ply_dir, 'base_%d.ply' % (i))
xyz_i = xyzs_sp[i:i + 1, :]
ply_util.create_ply(xyz_i, fn)
fn = os.path.join(ply_dir, 'neighbors_%d.ply' % (i))
xyzs_neig_i = xyzs[idx_scopes[i, 0]:idx_scopes[i, 1], :]
min_i = tf.reduce_min(xyzs_neig_i, 0)
max_i = tf.reduce_max(xyzs_neig_i, 0)
scope_i = max_i - min_i
c0 = tf.reduce_all(tf.greater_equal(xyz_i - min_i, 0))
c1 = tf.reduce_all(tf.greater_equal(max_i - xyz_i, 0))
check = tf.assert_equal(tf.logical_and(c0, c1), True)
ply_util.create_ply(xyzs_neig_i, fn)
return idx_scopes
def get_vidx1d_all_scales(xyzs, _min, _res, _scale_num):
with tf.variable_scope('get_vidx1d_all_scales_Scale%d' % (_scale_num)):
xyzmin = tf.reduce_min(xyzs)
with tf.control_dependencies([tf.greater_equal(xyzmin, 0)]):
xyzs = tf.identity(xyzs)
def body(vidx_1ds, i, min_i):
res_i = _res / (tf.cast(tf.pow(2, i), tf.float32))
vidx_3d = OctreeTf.get_vidx3d(xyzs, min_i, res_i)
vidx_1d = idx3d_to_1d(vidx_3d)
min_i = min_i + tf.expand_dims(res_i, 0) * tf.cast(
vidx_3d, tf.float32)
i += 1
vidx_1ds = tf.concat(
[vidx_1ds, tf.reshape(vidx_1d, [-1, 1])], 1)
return vidx_1ds, i, min_i
i = tf.constant(0)
min_i = tf.expand_dims(_min, 0)
vertex_num1 = TfUtil.tshape0(xyzs)
vidx_1ds = tf.zeros([vertex_num1, 0], tf.int32)
cond = lambda vidx_1ds, i, min_i: tf.less(i, _scale_num)
vidx_1ds, i, min_i = tf.while_loop(cond, body, [vidx_1ds, i, min_i],\
shape_invariants=[tf.TensorShape([vertex_num1,None]),\
tf.TensorShape([]), tf.TensorShape([None,3])])
vidx_1ds.set_shape([vertex_num1, _scale_num])
return vidx_1ds
def get_flat_idx(vidxs):
sn = TfUtil.get_tensor_shape(vidxs)[1]
i = 0
flat_vidx = vidxs[:, i] * pow(8, sn - 1)
for i in range(1, sn):
flat_vidx += vidxs[:, i] * pow(8, sn - i - 1)
return flat_vidx
def show_summary(self):
# vns: [4, 18, 125, 630, 2592, 11050]
# pn_pervs: [50431, 11207, 1613, 320, 77, 18]
vns = [TfUtil.tshape0(v) for v in self.flatvidx]
pn_pervs = []
for s in range(self._scale_num):
idxs = self.idxscope_pervs[s]
np = idxs[:, 1] - idxs[:, 0]
pn_pervs.append(tf.reduce_mean(np).numpy())
print(vns)
print(pn_pervs)
def gather_sg_params(self, features):
if not hasattr(self, 'sg_settings'):
return {}
sg_scale_num = self.sg_settings['num_sg_scale']
assert sg_scale_num >= self.model_paras.scale_num - 1, "not enough sg scales"
sg_params = {}
for s in range(sg_scale_num):
for item in ['grouped_pindex', 'grouped_bot_cen_top', 'empty_mask', \
'nb_enough_p', 'flatting_idx', 'flat_valid_mask']:
if s == 0:
sg_params[item] = []
else:
item_v = features[item + '_%d' % (s)]
sg_params[item].append(item_v)
if s > 0:
assert TfUtil.tsize(sg_params['flatting_idx'][-1]) == 3
assert TfUtil.tsize(sg_params['grouped_pindex'][-1]) == 3
assert TfUtil.tsize(sg_params['grouped_bot_cen_top'][-1]) == 4
return sg_params
def parse_inputs(self, features):
inputs = {}
points = []
for e in self.data_configs['feed_data']:
ele = self.get_ele(features, e)
if e == 'xyz':
xyz = ele = self.normalize_xyz(ele)
if e == 'color':
ele = self.normalize_color(ele)
points.append(ele)
inputs['points'] = points = tf.concat(points, -1)
assert TfUtil.tsize(points) == 3
sg_params = self.gather_sg_params(features)
inputs.update(sg_params)
vshape = TfUtil.get_tensor_shape(points)
#self.batch_size = vshape[0]
self.num_vertex0 = vshape[1]
self.log_tensor_p(points, 'points', 'raw_input')
return inputs, xyz
def interp_uppooling(self, scale, points, flatting_idx, last_points):
with tf.variable_scope('UpPool_S%d' % (scale)):
if flatting_idx is None:
voexl_num = TfUtil.get_tensor_shape(last_points)[1]
points = tf.tile(points, [1, voexl_num, 1, 1])
else:
points = TfUtil.gather_third_d(points, flatting_idx)
points = tf.reduce_mean(points, 2, keepdims=True)
self.log_tensor_p(points, 'uppooling', 'scale%d' % (scale))
if scale == 0:
last_points = tf.transpose(last_points, [0, 2, 1, 3])
if self.sg_settings['num_sg_scale'] > 1:
fnl = TfUtil.get_tensor_shape(last_points)[2]
points = tf.tile(points, [1, 1, fnl, 1])
else:
fnl = TfUtil.get_tensor_shape(last_points)[1]
points = tf.tile(points, [1, fnl, 1, 1])
points = tf.concat([last_points, points], -1)
self.log_tensor_p(points, 'mean & interperation',
'scale%d' % (scale))
return points
def check_sort(self, xyzs):
xyzs = tf.gather(xyzs, self.sorted_idxs_base)
scale = self._scale_num - 1
idxscope_perv = self.idxscope_pervs[scale]
flatvidx = self.flatvidx[scale]
vn = TfUtil.get_tensor_shape(idxscope_perv)[0]
#for i in np.random.choice(vn, 10):
for i in range(0, 10):
s, e = idxscope_perv[i]
xyzs_i = xyzs[s:e, :]
flatvidx_i = flatvidx[i]
min_i = tf.reduce_min(xyzs_i, 0)
max_i = tf.reduce_max(xyzs_i, 0)
import pdb
pdb.set_trace() # XXX BREAKPOINT
pass
def update_octree_idx_scope(self, vidx_rawscope_lsc):
# too complicate, give up for now
if self._scale == self._scale_num - 1:
#************************************************************
vn, c = TfUtil.get_tensor_shape(vidx_rawscope_lsc)
with tf.control_dependencies([tf.assert_equal(c, 3)]):
vn = tf.identity(vn)
import pdb
pdb.set_trace() # XXX BREAKPOINT
def body(j):
vidx_j = vidx_rawscope_lsc[j, 0]
self._nodes[vidx_j]._idx_scope = vidx_rawscope_lsc[j, 1:3]
return j + 1
i = tf.constant(0)
cond = tf.less(i, vn)
i = tf.while_loop(cond, body, [i])
else:
#************************************************************
hist = tf.histogram_fixed_width(vidx_rawscope_lsc[:, 0], [0, 8], 8)
start_idx_i = 0
for i in range(8):
cond = lambda: tf.greater(hist[i], 0)
def do_update(start_idx):
end_idx = start_idx + hist[i]
sub_vidx_scopes = vidx_rawscope_lsc[start_idx:end_idx, 1:]
self._nodes[i]._idx_scope = tf.stack(
[sub_vidx_scopes[0, -2], sub_vidx_scopes[-1, -1]])
self._nodes[i].update_octree_idx_scope(sub_vidx_scopes)
return end_idx
def no_op(start_idx):
return start_idx
start_idx_i = tf.cond(cond, lambda: do_update(start_idx_i),
lambda: no_op(start_idx_i))
import pdb
pdb.set_trace() # XXX BREAKPOINT
pass
def normalize_feature(self, points, dim):
assert TfUtil.tsize(points) == 4
mean_p = tf.reduce_mean(points, dim, keepdims=True)
points = points - mean_p
return points
def get_sort_idx(self, xyzs, vidx_1d, sortedidxs_base):
'''
xyzs: [N,3]
vidx_1d: [N] 0~7
sorted_idxs_base: [N] sorted idx of input xyzs
end_idxs_lastscale: [M, self._scale_num-self._scale] the voxel idx of each scale,
and the end idx of sorted data idx at the last scale.
'''
sort_idx = tf.contrib.framework.argsort(vidx_1d, axis=-1)
if self._scale != 0:
sortedidxs_base = tf.gather(sortedidxs_base, sort_idx)
else:
sortedidxs_base = sort_idx
hist = tf.histogram_fixed_width(vidx_1d, [0, 8], nbins=8)
def add_one_voxel(start_idx, vidx):
#print('scale {} vid {} hist: {}'.format(self._scale, vidx, hist[vidx]))
i3d = tf.constant(idx1d_to_3d_np(vidx), tf.float32)
min_i = self._min + self._resolution * i3d
self._nodes[vidx] = OctreeTf(resolution=self._resolution / 2,
scale=self._scale + 1,
min_xyz=min_i)
end_idx = start_idx + hist[vidx]
indices_i = sort_idx[start_idx:end_idx]
if self._scale == self._scale_num - 1:
#sortedidxs_base_i = sortedidxs_base[start_idx:end_idx]
sortedidxs_base_i = tf.constant([], tf.int32)
end_idx_lastscale = tf.reshape(end_idx, [1, 1])
#check
if self._check_optial:
xyzs_i = tf.gather(xyzs, indices_i, axis=0)
mini = tf.reduce_min(xyzs_i, 0)
maxi = tf.reduce_max(xyzs_i, 0)
scope_i = maxi - mini
check_i = tf.reduce_all(
tf.less_equal(scope_i, self._resolution))
#print('scale {} vid {} hist: {}'.format(self._scale, vidx, hist[vidx]))
else:
# After sorted, the vertices belong to this sub voxel are together
# already. xyzs_i belong to vidx voxel only, but it is not sorted inside
# vidx voxel
xyzs_i = tf.gather(xyzs, indices_i, axis=0)
sortedidxs_base_i = sortedidxs_base[start_idx:end_idx]
sortedidxs_base_i, end_idx_lastscale = self._nodes[
vidx].add_point_cloud(xyzs_i, sortedidxs_base_i)
m = self._scale_num - self._scale - 1
tmp = tf.concat([
tf.zeros([1, m], tf.int32),
tf.reshape(start_idx, [1, 1])
], -1)
end_idx_lastscale += tmp
return end_idx, sortedidxs_base_i, end_idx_lastscale
def no_op(start_idx, vidx):
#print('scale {} vid {} empty'.format(self._scale, vidx))
return start_idx, tf.constant([], tf.int32), \
tf.zeros([0,self._scale_num-self._scale], dtype=tf.int32)
sorted_idxs_base_new = []
end_idxs_lastscale = []
start_idx_i = 0
for i in range(8):
not_empty = tf.greater(hist[i], 0)
start_idx_i, sorted_idx_i, end_idx_lastscale_i = tf.cond(
not_empty, lambda: add_one_voxel(start_idx_i, i),
lambda: no_op(start_idx_i, i))
sorted_idxs_base_new.append(sorted_idx_i)
ni = TfUtil.get_tensor_shape(end_idx_lastscale_i)[0]
cur_vidx = tf.ones([ni, 1], tf.int32) * i
end_idx_lastscale_i = tf.concat([cur_vidx, end_idx_lastscale_i],
-1)
end_idxs_lastscale.append(end_idx_lastscale_i)
sorted_idxs_base_new = tf.concat(sorted_idxs_base_new, 0)
if self._scale == self._scale_num - 1:
sorted_idxs_base_new = sortedidxs_base
end_idxs_lastscale = tf.concat(end_idxs_lastscale, 0)
return sorted_idxs_base_new, end_idxs_lastscale
def gather_second_d(inputs, indices): return TfUtil.gather_second_d(inputs, indices)
def get_tensor_shape(tensor): return TfUtil.get_tensor_shape(tensor)
def parse_record(tfrecord_serialized, is_training, dset_shape_idx, \
bsg=None, gen_ply=False):
with tf.variable_scope('parse_record'):
from utils.aug_data_tf import aug_main, aug_views
#if dset_shape_idx!=None:
# from aug_data_tf import aug_data, tf_Rz
# R = tf_Rz(1)
# import pdb; pdb.set_trace() # XXX BREAKPOINT
eles = dset_shape_idx['indices'].keys()
features_map = {
'vertex_f': tf.FixedLenFeature([], tf.string),
'vertex_uint8': tf.FixedLenFeature([], tf.string),
}
if 'vertex_i' in eles:
features['vertex_i'] = tf.FixedLenFeature([], tf.string)
if 'face_i' in eles:
features['face_i'] = tf.FixedLenFeature([], tf.string)
features['valid_num_face'] = tf.FixedLenFeature([1],
tf.int64,
default_value=-1)
tfrecord_features = tf.parse_single_example(tfrecord_serialized,
features=features_map,
name='pl_features')
#*************
if 'vertex_i' in eles:
vertex_i = tf.decode_raw(tfrecord_features['vertex_i'], tf.int32)
vertex_i = tf.reshape(vertex_i,
dset_shape_idx['shape']['vertex_i'])
vertex_f = tf.decode_raw(tfrecord_features['vertex_f'], tf.float32)
vertex_f = tf.reshape(vertex_f, dset_shape_idx['shape']['vertex_f'])
vertex_uint8 = tf.decode_raw(tfrecord_features['vertex_uint8'],
tf.uint8)
vertex_uint8 = tf.reshape(vertex_uint8,
dset_shape_idx['shape']['vertex_uint8'])
if 'face_i' in eles:
face_i = tf.decode_raw(tfrecord_features['face_i'], tf.int32)
face_i = tf.reshape(face_i, dset_shape_idx['shape']['face_i'])
valid_num_face = tfrecord_features['valid_num_face']
#*************
features = { "vertex_f": vertex_f, \
"vertex_uint8": vertex_uint8}
if 'vertex_i' in tfrecord_features:
features['vertex_i'] = vertex_i
if 'face_i' in tfrecord_features:
features['face_i'] = face_i
features['valid_num_face'] = valid_num_face
#*************
if bsg is not None:
features = voxel_sampling_grouping(bsg, features, dset_shape_idx)
#*************
label_category = get_ele(features, 'label_category', dset_shape_idx)
labels = tf.squeeze(label_category, 1)
for item in features.keys():
if 'nb_enough_p' in item:
assert TfUtil.tsize(features[item]) == 0, item
elif 'label' in item:
assert TfUtil.tsize(features[item]) == 1, item
elif 'grouped_bot_cen_top' in item:
assert TfUtil.tsize(features[item]) == 3, item
else:
assert TfUtil.tsize(features[item]) == 2, item
return features, labels
def get_shape0(t):
return TfUtil.tshape0(t)