void ScanOctree::trim_octree()

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;
        }
      }
    }
  }
}