Voxelize KITTI point clouds with given poses

voxelize.cpp

#include <rv/Laserscan.h>
#include <fstream>
#include <iostream>
#include "io/KITTIReader.h"
#include "util/kitti_utils.h"

using namespace rv;

class Voxelgrid {
 public:
  class Voxel {
   public:
    Eigen::Vector4f point{Eigen::Vector4f::Zero()};
    uint32_t count{0};
  };

  void initialize(float resolution, const Eigen::Vector3f& min, const Eigen::Vector3f& max) {
    resolution_ = resolution;
    sizex_ = std::ceil((max.x() - min.x()) / resolution_);
    sizey_ = std::ceil((max.y() - min.y()) / resolution_);
    sizez_ = std::ceil((max.z() - min.z()) / resolution_);

    //    voxels_.resize(sizex_ * sizey_ * sizez_);

    center_.head(3) = 0.5 * (max - min) + min;
    center_[3] = 0;
  }

  void insert(const Eigen::Vector4f& p) {
    Eigen::Vector4f tp = p - center_;
    uint32_t i = std::floor((tp.x() + 0.5 * resolution_ * sizex_) / resolution_);
    uint32_t j = std::floor((tp.y() + 0.5 * resolution_ * sizey_) / resolution_);
    uint32_t k = std::floor((tp.z() + 0.5 * resolution_ * sizez_) / resolution_);

    if ((i >= sizex_) || (j >= sizey_) || (k >= sizez_)) {
      //      std::cerr << "should not happend.(1)" << std::endl;
      //      return;
    }
    int32_t gidx = i + j * sizex_ + k * sizex_ * sizey_;
    if (gidx < 0) {
      //      std::cerr << "should not happend.(2)" << std::endl;
      //      return;
    }

    if (voxels_.find(gidx) == voxels_.end()) {
      occupied_.push_back(gidx);
      voxels_[gidx] = Voxel();
    }
    float n = voxels_[gidx].count;
    voxels_[gidx].point = (1. / (n + 1)) * (n * voxels_[gidx].point + p);
    voxels_[gidx].count += 1;
  }

  std::vector<Eigen::Vector4f> points() const {
    std::vector<Eigen::Vector4f> pts;

    for (uint32_t i = 0; i < occupied_.size(); ++i) {
      pts.push_back(voxels_.find(occupied_[i])->second.point);
    }

    return pts;
  }

 protected:
  float resolution_;

  uint32_t sizex_, sizey_, sizez_;
  std::map<int32_t, Voxel> voxels_;
  std::vector<uint32_t> occupied_;

  Eigen::Vector4f center_;
};

int main(int argc, char** argv) {
  if (argc < 4) {
    std::cerr << "Usage: ./voxelize [scan directory] [pose file] [result filename] [resolution (optional)]"
              << std::endl;

    return 1;
  }

  float resolution = 0.2;

  if (argc > 4) resolution = std::stof(argv[4]);

  std::string scan_dir = argv[1];
  std::string poses_filename = argv[2];
  std::string result_filename = argv[3];

  std::cout << "1. Determining extent of complete registered point cloud ... " << std::flush;

  Laserscan scan;

  KITTIReader reader(scan_dir);
  std::vector<Eigen::Matrix4f> poses = KITTI::Odometry::loadPoses(poses_filename);

  Eigen::Vector3f min, max;
  bool firsttime = true;

  for (uint32_t t = 0; t < poses.size(); ++t) {
    const Eigen::Matrix4f& pose = poses[t];
    reader.read(scan);

    for (uint32_t i = 0; i < scan.size(); ++i) {
      const Eigen::Vector4f p = pose * scan.point(i).vec;

      if (firsttime) {
        min = max = p.head(3);
        firsttime = false;
      } else {
        min.x() = std::min(min.x(), p.x());
        min.y() = std::min(min.y(), p.y());
        min.z() = std::min(min.z(), p.z());

        max.x() = std::max(max.x(), p.x());
        max.y() = std::max(max.y(), p.y());
        max.z() = std::max(max.z(), p.z());
      }
    }
  }

  std::cout << "finished." << std::endl;

  std::cout << "min = " << min.transpose() << ", max = " << max.transpose() << std::endl;

  Voxelgrid grid;
  grid.initialize(resolution, min, max);

  reader.reset();

  std::cout << "2. Adding points to voxel grid ... " << std::flush;

  for (uint32_t t = 0; t < poses.size(); ++t) {
    const Eigen::Matrix4f& pose = poses[t];
    reader.read(scan);

    for (uint32_t i = 0; i < scan.size(); ++i) {
      if (scan.point(i).vec.head(3).norm() < 2.0) continue;
      Eigen::Vector4f p = pose * scan.point(i).vec;
      p[3] = scan.remission(i);

      grid.insert(p);
    }
  }

  std::cout << "finished." << std::endl;

  std::cout << "3. Saving point cloud to file: " << result_filename << std::flush;

  std::ofstream out(result_filename.c_str(), std::ofstream::out | std::ofstream::binary);

  std::vector<Eigen::Vector4f> pts = grid.points();
  std::cout << ". Writing " << pts.size() << " points ..." << std::flush;

  for (uint32_t i = 0; i < pts.size(); ++i) {
    out.write((const char*)pts[i].data(), 4 * sizeof(float));
  }

  out.close();

  std::cout << "finished." << std::endl;

  return 0;
}