in octree/octree/transform_octree.cpp [162:284]
void ScanOctree::trim_octree(vector<char>& octree_buffer,
const OctreeParser& octree_in,
vector<vector<int>>& drop_flags) {
// calculate the node number for the octree_out
int depth = octree_in.info().depth();
int depth_full = octree_in.info().full_layer();
vector<int> node_num_nempty(depth + 1, 0);
for (int d = 0; d < depth_full; ++d) {
node_num_nempty[d] = 1 << (3 * d);
}
for (int d = depth_full; d <= depth; ++d) {
int num = 0;
for (auto v : drop_flags[d]) {
if (v == 0) num++;
}
node_num_nempty[d] = num;
}
vector<int> node_num(depth + 1, 0);
node_num[0] = 1;
for (int d = 1; d <= depth; ++d) {
node_num[d] = 8 * node_num_nempty[d - 1];
}
// initialize
OctreeInfo info_out = octree_in.info();
info_out.set_nnum(node_num.data());
info_out.set_nempty(node_num_nempty.data());
info_out.set_nnum_cum();
info_out.set_ptr_dis();
octree_buffer.resize(info_out.sizeof_octree());
OctreeParser octree_out;
octree_out.set_cpu(octree_buffer.data(), &info_out);
// copy data
// !!! current channel_key = 1
int channel_feature = octree_in.info().channel(OctreeInfo::kFeature);
int location_feature = octree_in.info().locations(OctreeInfo::kFeature);
int channel_label = octree_in.info().channel(OctreeInfo::kLabel);
int location_label = octree_in.info().locations(OctreeInfo::kLabel);
int channel_split = octree_in.info().channel(OctreeInfo::kSplit);
int location_split = octree_in.info().locations(OctreeInfo::kSplit);
for (int d = 1; d <= depth; ++d) {
int num = octree_in.info().node_num(d - 1);
vector<int>& drop = drop_flags[d - 1];
vector<int>& drop_d = drop_flags[d];
// copy children and key
// !!! Caveat: currently channel_key is 1, and channel_child is 1
const int* child_in = octree_in.children_cpu(d - 1);
const uintk* key_in = octree_in.key_cpu(d);
int* child_out = octree_out.mutable_children_cpu(d);
uintk* key_out = octree_out.mutable_key_cpu(d);
for (int i = 0, j = 0, id = 0; i < num; ++i) {
// the node is dropped or empty
if (drop[i] == 1) continue;
int t = child_in[i];
for (int k = 8 * t; k < 8 * t + 8; ++k) {
key_out[j] = key_in[k];
// the node is non-empty and kept
int ch = drop_d[k] == 0 ? id++ : -1;
child_out[j] = ch;
j++;
}
}
// copy feature
if (location_feature == -1 || d == depth) {
int nnum_in = octree_in.info().node_num(d);
int nnum_out = octree_out.info().node_num(d);
const float* feature_in = octree_in.feature_cpu(d);
float* feature_out = octree_out.mutable_feature_cpu(d);
for (int i = 0, j = 0; i < num; ++i) {
// the node is dropped or empty
if (drop[i] == 1) continue;
int t = child_in[i];
for (int k = 8 * t; k < 8 * t + 8; ++k) {
for (int c = 0; c < channel_feature; ++c) {
feature_out[c * nnum_out + j] =
drop_d[k] == 0 ? feature_in[c * nnum_in + k] : 0;
}
j++;
}
}
}
// copy label
if ((location_label == -1 || d == depth) && channel_label != 0) {
const float* label_in = octree_in.label_cpu(d);
float* label_out = octree_out.mutable_label_cpu(d);
for (int i = 0, j = 0; i < num; ++i) {
// the node is dropped or empty
if (drop[i] == 1) continue;
int t = child_in[i];
for (int k = 8 * t; k < 8 * t + 8; ++k) {
label_out[j] = drop_d[k] == 0 ? label_in[k] : -1;
++j;
}
}
}
// copy split
if ((location_split == -1 || d == depth) && channel_split != 0) {
const float* split_in = octree_in.split_cpu(d);
float* split_out = octree_out.mutable_split_cpu(d);
for (int i = 0, j = 0; i < num; ++i) {
// the node is dropped or empty
if (drop[i] == 1) continue;
int t = child_in[i];
for (int k = 8 * t; k < 8 * t + 8; ++k) {
split_out[j] = drop_d[k] == 0 ? split_in[k] : 0;
++j;
}
}
}
}
}