Complete record dataset

tools/dataset/CMakeLists.txt

@@ -6,13 +6,9 @@ set_outdir(
   "${OUT_DIR}/bin/${DIR_NAME}"
 )
 
-include_directories(
-  ${PRO_DIR}/src
-)
-
 ## record
 
-add_executable(record record.cc)
+add_executable(record record.cc dataset.cc)
 target_link_libraries(record mynteye ${OpenCV_LIBS})
 
 if(OS_WIN)

tools/dataset/dataset.cc

@@ -0,0 +1,127 @@
+#include "dataset/dataset.h"
+
+#ifdef USE_OPENCV2
+#include <opencv2/highgui/highgui.hpp>
+#else
+#include <opencv2/imgcodecs/imgcodecs.hpp>
+#endif
+
+#include <glog/logging.h>
+
+#include <iomanip>
+#include <limits>
+#include <stdexcept>
+#include <utility>
+
+#include "internal/files.h"
+
+#define FULL_PRECISION \
+  std::fixed << std::setprecision(std::numeric_limits<double>::max_digits10)
+
+#define IMAGE_FILENAME_WIDTH 6
+
+MYNTEYE_BEGIN_NAMESPACE
+
+namespace tools {
+
+Dataset::Dataset(std::string outdir) : outdir_(std::move(outdir)) {
+  VLOG(2) << __func__;
+  if (!files::mkdir(outdir_)) {
+    LOG(FATAL) << "Create directory failed: " << outdir_;
+  }
+}
+
+Dataset::~Dataset() {
+  VLOG(2) << __func__;
+  for (auto &&it = stream_writers_.begin(); it != stream_writers_.end(); it++) {
+    if (it->second) {
+      it->second->ofs.flush();
+      it->second->ofs.close();
+    }
+  }
+  if (motion_writer_) {
+    motion_writer_->ofs.flush();
+    motion_writer_->ofs.close();
+  }
+}
+
+void Dataset::SaveStreamData(
+    const Stream &stream, const device::StreamData &data) {
+  auto &&writer = GetStreamWriter(stream);
+  auto seq = stream_counts_[stream];
+  writer->ofs << seq << ", " << data.img->frame_id << ", "
+              << data.img->timestamp << ", " << data.img->exposure_time
+              << std::endl;
+  if (data.frame) {
+    std::stringstream ss;
+    ss << writer->outdir << OS_SEP << std::dec
+       << std::setw(IMAGE_FILENAME_WIDTH) << std::setfill('0') << seq << ".png";
+    cv::Mat img(
+        data.frame->height(), data.frame->width(), CV_8UC1, data.frame->data());
+    cv::imwrite(ss.str(), img);
+  }
+  ++stream_counts_[stream];
+}
+
+void Dataset::SaveMotionData(const device::MotionData &data) {
+  auto &&writer = GetMotionWriter();
+  auto seq = motion_count_;
+  writer->ofs << seq << ", " << data.imu->frame_id << ", "
+              << data.imu->timestamp << ", " << data.imu->accel[0] << ", "
+              << data.imu->accel[1] << ", " << data.imu->accel[2] << ", "
+              << data.imu->gyro[0] << ", " << data.imu->gyro[1] << ", "
+              << data.imu->gyro[2] << ", " << data.imu->temperature
+              << std::endl;
+  ++motion_count_;
+}
+
+Dataset::writer_t Dataset::GetStreamWriter(const Stream &stream) {
+  try {
+    return stream_writers_.at(stream);
+  } catch (const std::out_of_range &e) {
+    writer_t writer = std::make_shared<Writer>();
+    switch (stream) {
+      case Stream::LEFT: {
+        writer->outdir = outdir_ + OS_SEP "left";
+      } break;
+      case Stream::RIGHT: {
+        writer->outdir = outdir_ + OS_SEP "right";
+      } break;
+      default:
+        LOG(FATAL) << "Unsupported stream: " << stream;
+    }
+    writer->outfile = writer->outdir + OS_SEP "stream.txt";
+
+    files::mkdir(writer->outdir);
+    writer->ofs.open(writer->outfile, std::ofstream::out);
+    writer->ofs << "seq, frame_id, timestamp, exposure_time" << std::endl;
+    writer->ofs << FULL_PRECISION;
+
+    stream_writers_[stream] = writer;
+    stream_counts_[stream] = 0;
+    return writer;
+  }
+}
+
+Dataset::writer_t Dataset::GetMotionWriter() {
+  if (motion_writer_ == nullptr) {
+    writer_t writer = std::make_shared<Writer>();
+    writer->outdir = outdir_;
+    writer->outfile = writer->outdir + OS_SEP "motion.txt";
+
+    files::mkdir(writer->outdir);
+    writer->ofs.open(writer->outfile, std::ofstream::out);
+    writer->ofs << "seq, frame_id, timestamp, accel_x, accel_y, accel_z, "
+                   "gyro_x, gyro_y, gyro_z, temperature"
+                << std::endl;
+    writer->ofs << FULL_PRECISION;
+
+    motion_writer_ = writer;
+    motion_count_ = 0;
+  }
+  return motion_writer_;
+}
+
+}  // namespace tools
+
+MYNTEYE_END_NAMESPACE

tools/dataset/dataset.h

@@ -0,0 +1,50 @@
+#ifndef MYNTEYE_TOOLS_DATASET_H_  // NOLINT
+#define MYNTEYE_TOOLS_DATASET_H_
+#pragma once
+
+#include <fstream>
+#include <map>
+#include <memory>
+#include <string>
+
+#include "internal/callbacks.h"
+#include "mynteye/mynteye.h"
+
+MYNTEYE_BEGIN_NAMESPACE
+
+namespace tools {
+
+class Dataset {
+ public:
+  struct Writer {
+    std::ofstream ofs;
+    std::string outdir;
+    std::string outfile;
+  };
+
+  using writer_t = std::shared_ptr<Writer>;
+
+  explicit Dataset(std::string outdir);
+  ~Dataset();
+
+  void SaveStreamData(const Stream &stream, const device::StreamData &data);
+  void SaveMotionData(const device::MotionData &data);
+
+ private:
+  writer_t GetStreamWriter(const Stream &stream);
+  writer_t GetMotionWriter();
+
+  std::string outdir_;
+
+  std::map<Stream, writer_t> stream_writers_;
+  writer_t motion_writer_;
+
+  std::map<Stream, std::size_t> stream_counts_;
+  std::size_t motion_count_;
+};
+
+}  // namespace tools
+
+MYNTEYE_END_NAMESPACE
+
+#endif  // MYNTEYE_TOOLS_DATASET_H_ NOLINT

tools/dataset/record.cc

@@ -8,6 +8,8 @@
 
 #include "internal/times.h"
 
+#include "dataset/dataset.h"
+
 MYNTEYE_USE_NAMESPACE
 
 int main(int argc, char *argv[]) {
@@ -68,6 +70,14 @@ int main(int argc, char *argv[]) {
   device->EnableMotionDatas();
   device->Start(Source::ALL);
 
+  const char *outdir;
+  if (argc >= 2) {
+    outdir = argv[1];
+  } else {
+    outdir = "./dataset";
+  }
+  tools::Dataset dataset(outdir);
+
   cv::namedWindow("frame");
 
   std::size_t img_count = 0;
@@ -95,11 +105,28 @@ int main(int argc, char *argv[]) {
     cv::hconcat(left_img, right_img, img);
     cv::imshow("frame", img);
 
+    {  // save
+      for (auto &&left : left_datas) {
+        dataset.SaveStreamData(Stream::LEFT, left);
+      }
+      for (auto &&right : right_datas) {
+        dataset.SaveStreamData(Stream::RIGHT, right);
+      }
+
+      for (auto &&motion : motion_datas) {
+        dataset.SaveMotionData(motion);
+      }
+
+      std::cout << "\rSaved " << img_count << " imgs"
+                << ", " << imu_count << " imus" << std::flush;
+    }
+
     char key = static_cast<char>(cv::waitKey(1));
     if (key == 27 || key == 'q' || key == 'Q') {  // ESC/Q
       break;
     }
   }
+  std::cout << " to " << outdir << std::endl;
   auto &&time_end = times::now();
 
   device->Stop(Source::ALL);