SimToCalibration.cpp

/*
 * Copyright (C) 2024: Arizona Board of Regents on Behalf of the University of Arizona
 */

 /**
  * @file SimToCalibration.cpp
  * @brief Apache Strap-Down Pilotage program that reads the configuration, images, and poses from a Blender simulation
  * and converts them into noisy, distorted images to be fed into a calibration workflow.
  *
 * @author ReliaSolve.
 * @date December 12th, 2024.
 */

#include <iostream>
#include <string>
#include <vector>
#include <array>
#include <fstream>
#include <sstream>
#include <filesystem>
#include <ASDP_Core_API.h>
#include <nlohmann/json.hpp>
#include <glm/glm.hpp>
#include <glm/gtc/matrix_transform.hpp>
#define GLM_ENABLE_EXPERIMENTAL
#include <glm/gtx/matrix_decompose.hpp>
#include <glm/gtc/type_ptr.hpp>
#include "PNMImage.h"
#include "CameraConfig.h"
#include "ImageResampler.h"
#include "AddNoise.h"

using namespace asdp;
using json = nlohmann::json;

static std::string VERSION = "3.3.0";

std::string doConversions(std::string rootDirectory, std::string outDirectory,
  std::vector<CameraConfig> &cameraConfigs, int verbosity, std::vector< std::vector<int> > poseEntries,
  double sigma, bool fixEndianness = true)
{
  // Open the stream files in the output directory, creating the directory if it does not exist.

  // Make the subdirectory based on our serial number and make it if it does not exist.
  std::filesystem::create_directories(outDirectory);
  if (!std::filesystem::exists(outDirectory)) {
    return "Error: Could not create directory: " + outDirectory;
  }

  if (verbosity > 0) {
    std::cout << "Writing resampled images to: " << outDirectory << std::endl;
  }

  // Now we read the images based on the poseEntries values, resample, add noise, and write them to the output directory.
  // The poseEntries vector is a vector of vectors, where each inner vector has 3 values:
  // the frame number, the camera number, and the number of noisy frames to generate for that camera.
  std::string broken;
#pragma omp parallel for
  for (int f = 0; f < poseEntries.size(); f++) {
    std::vector<int> entry = poseEntries[f];
    size_t frameID = entry[0];
    size_t cameraID = entry[1];
    size_t noisyFrames = entry[2];

    // Look up the camera whose ID matches the cameraID.
    CameraConfig* which = nullptr;
    for (auto& cameraConfig : cameraConfigs) {
      if (cameraConfig.id == cameraID) {
        which = &cameraConfig;
        break;
      }
    }
    if (which == nullptr) {
#pragma omp critical
      {
        std::cerr << "Error: Could not find camera configuration for camera ID: " << cameraID << std::endl;
        broken = "Error: Could not find camera configuration for camera ID: " + std::to_string(cameraID);
      }
      continue;
    }
    auto const& cameraConfig = *which;

    // Skip each frame if we have already encountered an error.
    // We cannot break or return from within an OpenMP parallel for loop.
    if (!broken.empty()) {
      continue;
    }

    std::string directory = rootDirectory;
    if (verbosity >= 1000) {
#pragma omp critical
      {
        std::cout << "  Processing frame " << frameID << " in directory " << directory << std::endl;
      }
    }

    // Read an image for the specified camera.
    std::string imageName = directory + "/" + std::to_string(frameID) + "_" + std::to_string(cameraID) + ".pgm";
    std::shared_ptr<PNMImage> image = std::make_shared<PNMImage>(imageName, fixEndianness);
    if (!image->width()) {
#pragma omp critical
      {
        std::cerr << "Error: Could not read image file: " + imageName << std::endl;
        broken = "Error reading image file: " + imageName;
      }
      continue;
    }
    if (verbosity >= 2000) {
#pragma omp critical
      {
        std::cout << "    Image: " << imageName << ", size " << image->width() << "x" << image->height() << std::endl;
      }
    }

    // Resample each image and write it to the appropriate output file(s).
    {
      std::array<float, 3> zero = { 0.0f,0.0f,0.0f };
      std::array<float, 3> velocityInCameraFrame = zero;
      std::array<float, 3> angleVelInCameraFrame = zero;

      std::shared_ptr<PNMImage> resampledImage = Resample_Image(*image, cameraConfigs[cameraID-1],
        1 / 60.0f, angleVelInCameraFrame);
      if (!resampledImage) {
#pragma omp critical
        {
          std::cerr << "Error resampling image for camera " + std::to_string(cameraConfig.id);
          broken = "Error resampling image for camera " + std::to_string(cameraConfig.id);
        }
        continue;
      }
      if (verbosity >= 4000) {
#pragma omp critical
        {
          std::cout << "      Resampled image " << frameID << " to size " << resampledImage->width() << "x" << resampledImage->height() << std::endl;
        }
      }

      // Adjust the gain and offset if our values are not unity mapping.
      if (cameraConfig.offset != 0 || cameraConfig.gain != 1) {
        // Apply the inverse mapping by offset and gain to the image so that applying the offset and gain will return
        // to the original value.
        if (verbosity >= 1500) {
          std::cout << "    Applying offset " << -cameraConfig.offset << " and gain " << 1 / cameraConfig.gain << " to image " << imageName << std::endl;
        }
        size_t count = image->width() * image->height();
        for (size_t i = 0; i < count; ++i) {
          float value = static_cast<float>(image->data()[i]);
          value = value / cameraConfig.gain - cameraConfig.offset;
          if (value < 0) {
            value = 0;
          }
          if (value > 65535) {
            value = 65535;
          }
          image->data()[i] = static_cast<uint16_t>(value);
        }
      }

      // Add noise to make as many images as requested, writing each one to the output directory.
      for (int i = 1; i <= noisyFrames; i++) {
        std::shared_ptr<PNMImage> noisyImage = Add_Noise(*resampledImage, sigma);
        if (!noisyImage) {
          #pragma omp critical
          {
            std::cerr << "Error adding noise to resampled image for camera " + std::to_string(cameraConfig.id);
            broken = "Error adding noise to resampled image for camera " + std::to_string(cameraConfig.id);
          }
          continue;
        }
        std::string noisyName = outDirectory + "/" + std::to_string(frameID) + "_" + std::to_string(cameraConfig.id) + "_" + std::to_string(i) + ".pgm";
        // We always fix the endianness of the output files.
        if (!noisyImage->Write(noisyName)) {
#pragma omp critical
          {
            std::cerr << "Error writing resampled image to file: " + noisyName;
            broken = "Error writing resampled image to file: " + noisyName;
          }
          continue;
        }
        if (verbosity >= 4000) {
#pragma omp critical
          {
            std::cout << "      Wrote resampled image to " << noisyName << std::endl;
          }
        }
      }
    }
  };

  return broken;
}

void usage(const std::string& programName)
{
  std::cout << "Usage: " << programName << " [options] rootDirectory outDirectory\n"
    << "rootDirectory is the directory where the simulation files are stored\n"
    << "outDirectory is the directory where the output files will be stored\n"
    << "\n"
    << "Options:\n"
    << "  --verbosity <number>      Specify the verbosity (default 1)\n"
    << "  --configFile <filename>   Specify the configuration file name (default config.json)\n"
    << "  --poseFile <filename>     File listing the cameras to report at each frame (default poses.csv)\n"
    << "  --sigma <number>          Standard deviation of pixel-counts of noise to add (default 1000.0)\n"
    << "  --dontFixEndianness       Do not fix the endianness of the input files (output files always fixed)\n"
    << "  --help                    Show this help message and exit\n";
}

int main(int argc, char* argv[])
{
  // Default values for command line arguments
  std::string rootDirectory;
  std::string outDirectory;
  std::string configFileName = "config.json";
  std::string poseFileName = "poses.csv";
  int verbosity = 1;
  double sigma = 1000.0;
  bool fixEndianness = true;
  unsigned realParams = 0;

  // Parse the command line arguments
  for (int i = 1; i < argc; ++i) {
    std::string arg = argv[i];
    if (arg == "--help") {
      usage(argv[0]);
      return 0;
    } else if (arg == "--verbosity") {
      if (i + 1 < argc) {
        verbosity = atoi(argv[++i]);
      } else {
        std::cerr << "--verbosity option requires one argument." << std::endl;
        usage(argv[0]);
        return 1;
      }
    } else if (arg == "--configFile") {
      if (i + 1 < argc) {
        configFileName = argv[++i];
      } else {
        std::cerr << "--configFile option requires one argument." << std::endl;
        usage(argv[0]);
        return 1;
      }
    } else if (arg == "--poseFile") {
      if (i + 1 < argc) {
        poseFileName = argv[++i];
      } else {
        std::cerr << "--poseFile option requires one argument." << std::endl;
        usage(argv[0]);
        return 1;
      }
    } else if (arg == "--sigma") {
      if (i + 1 < argc) {
        sigma = atof(argv[++i]);
      } else {
        std::cerr << "--sigma option requires one argument." << std::endl;
        usage(argv[0]);
        return 1;
      }
    } else if (arg == "--dontFixEndianness") {
      fixEndianness = false;
    } else if (arg[0] == '-') {
      std::cerr << "Unrecognized flag argument: " << arg << std::endl;
      usage(argv[0]);
      return 1;
    } else switch (++realParams) {
      case 1:
        rootDirectory = arg;
        break;
      case 2:
        outDirectory = arg;
        break;
      default:
        std::cerr << "Unknown option: " << arg << std::endl;
        usage(argv[0]);
        return 1;
    }
  }
  if (realParams != 2) {
    usage(argv[0]);
    return 1;
  }

  if (verbosity > 0) {
    std::cout << "SimToCalibration version " << VERSION << std::endl;
  }

  // Read all needed values from the configuration file
  std::vector<CameraConfig> cameraConfigs;
  std::ifstream configFile(configFileName);
  if (!configFile) {
    std::cerr << "Error: Could not open configuration file: " << configFileName << std::endl;
    return 2;
  }
  json config = json::parse(configFile);
  try {
    for (json& cameraConfig : config["cameras"]) {
      cameraConfigs.push_back(CameraConfig(cameraConfig));
    }
    if (verbosity > 0) {
      std::cout << "Found " << cameraConfigs.size() << " cameras" << std::endl;
    }
  } catch (std::exception& e) {
    std::cerr << "Error: Could not read configuration file: " << e.what() << std::endl;
    return 3;
  }

  // Read all needed values from the pose file. Skip the header, then read each line, storing
  // the first, fourth, and fifth values into an integer vector.  Make a vector of these vectors.
  std::ifstream poseFile(poseFileName);
  if (!poseFile) {
    std::cerr << "Error: Could not open pose file: " << poseFileName << std::endl;
    return 4;
  }
  std::string line;
  std::getline(poseFile, line); // Skip the header line
  std::vector< std::vector<int> > poseEntries;
  while (std::getline(poseFile, line)) {
    std::istringstream lineStream(line);
    std::string value;
    std::vector<int> poseEntry;
    for (int i = 0; i < 5; ++i) {
      if (!std::getline(lineStream, value, ',')) {
        std::cerr << "Error: Expected 5 values in pose file: " << poseFileName << std::endl;
        return 5;
      }
      if (i == 0 || i == 3 || i == 4) { poseEntry.push_back(std::stoi(value)); }
    }
    poseEntries.push_back(poseEntry);
  }

  if (verbosity > 0) {
    std::cout << "Root Directory: " << rootDirectory << std::endl;
    std::cout << "Output Directory: " << outDirectory << std::endl;
  }

  // Run the conversions.
  std::string message = doConversions(rootDirectory, outDirectory, cameraConfigs, verbosity, poseEntries, sigma,
    fixEndianness);
  if (!message.empty()) {
    std::cerr << "Error converting: " << message << std::endl;
    return 5;
  }

  return 0;
}