Example: rgbd_dataset2rawlog

C++ example source code:

/* +------------------------------------------------------------------------+
   |                     Mobile Robot Programming Toolkit (MRPT)            |
   |                          https://www.mrpt.org/                         |
   |                                                                        |
   | Copyright (c) 2005-2023, Individual contributors, see AUTHORS file     |
   | See: https://www.mrpt.org/Authors - All rights reserved.               |
   | Released under BSD License. See: https://www.mrpt.org/License          |
   +------------------------------------------------------------------------+ */

/* ------------------------------------------------------
                  rgbd2rawlog
    A small tool to translate RGB+D datasets from TUM:
    http://cvpr.in.tum.de/data/datasets/rgbd-dataset/

    into MRPT rawlog binary format, ready to be parsed
    by RawLogViewer or user programs.

Usage:
      rgbd_dataset2rawlog  [PATH_TO_RGBD_DATASET] [OUTPUT_NAME]

Where expected input files are:
    - PATH_TO_RGBD_DATASET/accelerometer.txt
    - PATH_TO_RGBD_DATASET/depth.txt
    - PATH_TO_RGBD_DATASET/rgb.txt
    - PATH_TO_RGBD_DATASET/rgb/... (Images, 3x8u)
    - PATH_TO_RGBD_DATASET/depth/...  (Images, 1x16u)

Output files:
    - OUTPUT_NAME.rawlog: The output rawlog file.
    - OUTPUT_NAME_Images/...: External RGB images.

   ------------------------------------------------------ */

#include <mrpt/io/CFileGZOutputStream.h>
#include <mrpt/io/CTextFileLinesParser.h>
#include <mrpt/obs/CObservation3DRangeScan.h>
#include <mrpt/obs/CObservationIMU.h>
#include <mrpt/serialization/CArchive.h>
#include <mrpt/system/filesystem.h>

#include <iostream>

using namespace std;
using namespace mrpt;
using namespace mrpt::obs;
using namespace mrpt::serialization;

const double KINECT_FPS = 30.0;

// Based on:
// http://stackoverflow.com/questions/218488/finding-best-matching-key-for-a-given-key-in-a-sorted-stl-container
template <class CONTAINER>
typename CONTAINER::const_iterator find_closest(
    const CONTAINER& data, const typename CONTAINER::key_type& searchkey)
{
    typename CONTAINER::const_iterator upper = data.lower_bound(searchkey);
    if (upper == data.begin() || upper->first == searchkey) return upper;
    typename CONTAINER::const_iterator lower = upper;
    --lower;
    if (upper == data.end() ||
        (searchkey - lower->first) < (upper->first - searchkey))
        return lower;
    return upper;
}

void rgbd2rawlog(const string& src_path, const string& out_name)
{
    const string in_fil_acc = src_path + string("/accelerometer.txt");
    const string in_fil_depth =
        mrpt::system::fileExists(src_path + string("/ir.txt"))
        ? src_path + string("/ir.txt")
        : src_path + string("/depth.txt");
    const string in_fil_rgb = src_path + string("/rgb.txt");

    // ASSERT_FILE_EXISTS_(in_fil_acc)
    // ASSERT_FILE_EXISTS_(in_fil_depth)
    // ASSERT_FILE_EXISTS_(in_fil_rgb)

    // make a list with RGB & DEPTH files ------------------------
    map<double, string> list_rgb, list_depth;
    map<double, vector<double>> list_acc;
    std::istringstream line;
    if (mrpt::system::fileExists(in_fil_rgb))
    {
        mrpt::io::CTextFileLinesParser fparser(in_fil_rgb);
        while (fparser.getNextLine(line))
        {
            double tim;
            string f;
            if (line >> tim >> f) list_rgb[tim] = f;
        }
    }
    if (mrpt::system::fileExists(in_fil_depth))
    {
        mrpt::io::CTextFileLinesParser fparser(in_fil_depth);
        while (fparser.getNextLine(line))
        {
            double tim;
            string f;
            if (line >> tim >> f) list_depth[tim] = f;
        }
    }
    if (mrpt::system::fileExists(in_fil_acc))
    {
        mrpt::io::CTextFileLinesParser fparser(in_fil_acc);
        while (fparser.getNextLine(line))
        {
            double tim;
            double ax, ay, az;
            if (line >> tim >> ax >> ay >> az)
            {
                vector<double>& v = list_acc[tim];
                v.resize(3);
                v[0] = ax;
                v[1] = ay;
                v[2] = az;
            }
        }
    }
    cout << "Parsed: " << list_depth.size() << " / " << list_rgb.size() << " / "
         << list_acc.size() << " depth/rgb/acc entries.\n";

    const bool only_ir = list_depth.size() > 10 && list_rgb.empty();
    const bool only_rgb = list_depth.empty() && list_rgb.size() > 10;

    // Create output directory for images ------------------------------
    const string out_img_dir = out_name + string("_Images");

    cout << "Creating images directory: " << out_img_dir << endl;
    mrpt::system::createDirectory(out_img_dir);

    // Create rawlog file ----------------------------------------------
    const string out_rawlog_fil = out_name + string(".rawlog");
    cout << "Creating rawlog: " << out_rawlog_fil << endl;
    mrpt::io::CFileGZOutputStream f_out(out_rawlog_fil);

    // Fill out the common field to all entries:
    CObservation3DRangeScan obs;
    obs.sensorLabel = "KINECT";
    obs.range_is_depth = true;  // Kinect style: ranges are actually depth
    // values, not Euclidean distances.

    // Range & RGB images are already registered and warped to remove
    // distortion:
    const double FOCAL = 525.0;
    obs.cameraParams.nrows = 480;
    obs.cameraParams.ncols = 640;
    obs.cameraParams.fx(FOCAL);
    obs.cameraParams.fy(FOCAL);
    obs.cameraParams.cx(319.5);
    obs.cameraParams.cy(239.5);
    obs.cameraParamsIntensity = obs.cameraParams;
    obs.relativePoseIntensityWRTDepth = mrpt::poses::CPose3D(
        0, 0, 0, mrpt::DEG2RAD(-90), 0,
        mrpt::DEG2RAD(-90));  // No translation between rgb & range
    // cameras, and rotation only due to XYZ
    // axes conventions.

    CObservationIMU obs_imu;
    obs_imu.sensorLabel = "KINECT_ACC";
    obs_imu.dataIsPresent[IMU_X_ACC] = true;
    obs_imu.dataIsPresent[IMU_Y_ACC] = true;
    obs_imu.dataIsPresent[IMU_Z_ACC] = true;

    // Go thru the data:
    unsigned int counter = 0;
    if (!only_ir && !only_rgb)
    {
        for (map<double, string>::const_iterator it_list_rgb = list_rgb.begin();
             it_list_rgb != list_rgb.end(); ++it_list_rgb)
        {
            cout << "Processing image " << counter << "\r";
            cout.flush();
            counter++;

            // This is not the most efficient solution in the world, but...come
            // on! it's just a rawlog converter tool!
            map<double, string>::const_iterator it_list_depth =
                find_closest(list_depth, it_list_rgb->first);

            const double At =
                std::abs(it_list_rgb->first - it_list_depth->first);
            if (At > (1. / KINECT_FPS) * .5)
            {
                cout << "\nWarning: Discarding observation for too separated "
                        "RGB/D timestamps: "
                     << At * 1e3 << " ms\n";
            }
            else
            {
                // OK, we accept this RGB-DEPTH pair:
                const double avrg_time =
                    .5 * (it_list_rgb->first + it_list_depth->first);
                obs.timestamp = mrpt::system::time_tToTimestamp(avrg_time);

                // RGB img:
                obs.hasIntensityImage = true;
                obs.intensityImage.loadFromFile(
                    src_path + string("/") + it_list_rgb->second);
                const string sRGBfile =
                    mrpt::format("%.06f_rgb.png", avrg_time);
                obs.intensityImage.saveToFile(
                    out_img_dir + string("/") + sRGBfile);
                obs.intensityImage.setExternalStorage(sRGBfile);

                // Depth:
                obs.hasRangeImage = true;
                obs.rangeImage_forceResetExternalStorage();
                mrpt::img::CImage depth_img;
                if (!depth_img.loadFromFile(
                        src_path + string("/") + it_list_depth->second))
                    throw std::runtime_error(
                        string("Error loading depth image!: ") +
                        it_list_depth->second);

                const unsigned int w = depth_img.getWidth();
                const unsigned int h = depth_img.getHeight();
                obs.rangeImage_setSize(h, w);

                for (unsigned int row = 0; row < h; row++)
                {
                    const uint16_t* ptr = depth_img.ptrLine<uint16_t>(row);
                    for (unsigned int col = 0; col < w; col++)
                        obs.rangeImage(row, col) = (*ptr++) * (1.f / 5000);
                }

                const string sDepthfile =
                    mrpt::format("%.06f_depth.bin", avrg_time);
                obs.rangeImage_convertToExternalStorage(
                    sDepthfile, out_img_dir + string("/"));

                // save:
                archiveFrom(f_out) << obs;

                // Search for acc data:
                map<double, vector<double>>::const_iterator it_list_acc =
                    find_closest(list_acc, avrg_time);
                const double At_acc =
                    std::abs(it_list_rgb->first - it_list_acc->first);
                if (At_acc < 1e-2)
                {
                    obs_imu.timestamp =
                        mrpt::system::time_tToTimestamp(avrg_time);

                    obs_imu.rawMeasurements[IMU_X_ACC] = it_list_acc->second[0];
                    obs_imu.rawMeasurements[IMU_Y_ACC] = it_list_acc->second[1];
                    obs_imu.rawMeasurements[IMU_Z_ACC] = it_list_acc->second[2];

                    archiveFrom(f_out) << obs_imu;
                }
            }
        }
    }
    else if (only_ir)
    {
        for (map<double, string>::const_iterator it_list_depth =
                 list_depth.begin();
             it_list_depth != list_depth.end(); ++it_list_depth)
        {
            cout << "Processing image " << counter << "\r";
            cout.flush();
            counter++;

            const double avrg_time = it_list_depth->first;
            obs.timestamp =
                mrpt::system::time_tToTimestamp(it_list_depth->first);

            // Depth:
            obs.hasRangeImage = true;
            obs.rangeImage_forceResetExternalStorage();
            mrpt::img::CImage depth_img;
            if (!depth_img.loadFromFile(
                    src_path + string("/") + it_list_depth->second))
                throw std::runtime_error(
                    string("Error loading depth image!: ") +
                    it_list_depth->second);

            const unsigned int w = depth_img.getWidth();
            const unsigned int h = depth_img.getHeight();
            obs.rangeImage_setSize(h, w);

            for (unsigned int row = 0; row < h; row++)
            {
                const uint16_t* ptr = depth_img.ptrLine<uint16_t>(row);
                for (unsigned int col = 0; col < w; col++)
                    obs.rangeImage(row, col) = (*ptr++) * (1.f / 5000);
            }

            const string sDepthfile =
                mrpt::format("%.06f_depth.bin", avrg_time);
            obs.rangeImage_convertToExternalStorage(
                sDepthfile, out_img_dir + string("/"));

            // save:
            archiveFrom(f_out) << obs;
        }
    }
    else if (only_rgb)
    {
        for (map<double, string>::const_iterator it_list_rgb = list_rgb.begin();
             it_list_rgb != list_rgb.end(); ++it_list_rgb)
        {
            cout << "Processing image " << counter << "\r";
            cout.flush();
            counter++;

            const double avrg_time = it_list_rgb->first;
            obs.timestamp = mrpt::system::time_tToTimestamp(it_list_rgb->first);

            // RGB img:
            obs.hasIntensityImage = true;
            obs.intensityImage.loadFromFile(
                src_path + string("/") + it_list_rgb->second);
            const string sRGBfile = mrpt::format("%.06f_rgb.png", avrg_time);
            obs.intensityImage.saveToFile(out_img_dir + string("/") + sRGBfile);
            obs.intensityImage.setExternalStorage(sRGBfile);

            // save:
            archiveFrom(f_out) << obs;
        }
    }

    cout << "\nAll done!\n";
}

// ------------------------------------------------------
//                      MAIN
// ------------------------------------------------------
int main(int argc, char** argv)
{
    try
    {
        if (argc != 3)
        {
            cerr << "Usage: " << argv[0]
                 << " [PATH_TO_RGBD_DATASET] [OUTPUT_NAME]\n";
            return 1;
        }

        rgbd2rawlog(argv[1], argv[2]);

        return 0;
    }
    catch (const std::exception& e)
    {
        std::cerr << "MRPT error: " << mrpt::exception_to_str(e) << std::endl;
        return -1;
    }
}