image_writer.cc

#include "drake/systems/sensors/image_writer.h"

#include <unistd.h>

#include <regex>
#include <string>
#include <utility>
#include <vector>

#include <vtkImageData.h>
#include <vtkNew.h>
#include <vtkPNGWriter.h>
#include <vtkSmartPointer.h>
#include <vtkTIFFWriter.h>
#include "fmt/ostream.h"

#include "drake/common/text_logging.h"

namespace drake {
namespace systems {
namespace sensors {

using std::move;

template <PixelType kPixelType>
void SaveToFileHelper(const std::string& file_path,
                      const Image<kPixelType>& image) {
  const int width = image.width();
  const int height = image.height();
  const int num_channels = Image<kPixelType>::kNumChannels;

  vtkSmartPointer<vtkImageWriter> writer;
  vtkNew<vtkImageData> vtk_image;
  vtk_image->SetDimensions(width, height, 1);

  // NOTE: This excludes *many* of the defined `PixelType` values.
  switch (kPixelType) {
    case PixelType::kRgba8U:
      vtk_image->AllocateScalars(VTK_UNSIGNED_CHAR, num_channels);
      writer = vtkSmartPointer<vtkPNGWriter>::New();
      break;
    case PixelType::kDepth32F:
      vtk_image->AllocateScalars(VTK_FLOAT, num_channels);
      writer = vtkSmartPointer<vtkTIFFWriter>::New();
      break;
    case PixelType::kLabel16I:
      vtk_image->AllocateScalars(VTK_UNSIGNED_SHORT, num_channels);
      writer = vtkSmartPointer<vtkPNGWriter>::New();
      break;
    default:
      throw std::logic_error(
          "Unsupported image type; cannot be written to file");
  }

  auto image_ptr = reinterpret_cast<typename Image<kPixelType>::T*>(
      vtk_image->GetScalarPointer());
  const int num_scalar_components = vtk_image->GetNumberOfScalarComponents();
  DRAKE_DEMAND(num_scalar_components == num_channels);

  for (int v = height - 1; v >= 0; --v) {
    for (int u = 0; u < width; ++u) {
      for (int c = 0; c < num_channels; ++c) {
        image_ptr[c] =
            static_cast<typename Image<kPixelType>::T>(image.at(u, v)[c]);
      }
      image_ptr += num_scalar_components;
    }
  }

  writer->SetFileName(file_path.c_str());
  writer->SetInputData(vtk_image.GetPointer());
  writer->Write();
}

void SaveToPng(const std::string& file_path, const ImageRgba8U& image) {
  SaveToFileHelper(file_path, image);
}

void SaveToTiff(const std::string& file_path, const ImageDepth32F& image) {
  SaveToFileHelper(file_path, image);
}

void SaveToPng(const std::string& file_path, const ImageLabel16I& image) {
  SaveToFileHelper(file_path, image);
}

ImageWriter::ImageWriter() {
  // NOTE: This excludes *many* of the defined `PixelType` values.
  labels_[PixelType::kRgba8U] = "color";
  extensions_[PixelType::kRgba8U] = ".png";
  labels_[PixelType::kDepth32F] = "depth";
  extensions_[PixelType::kLabel16I] = ".png";
  labels_[PixelType::kLabel16I] = "label";
  extensions_[PixelType::kDepth32F] = ".tiff";
}

template <PixelType kPixelType>
const InputPort<double>& ImageWriter::DeclareImageInputPort(
    std::string port_name, std::string file_name_format, double publish_period,
    double start_time) {
  // Test to confirm valid pixel type.
  static_assert(kPixelType == PixelType::kRgba8U ||
                    kPixelType == PixelType::kDepth32F ||
                    kPixelType == PixelType::kLabel16I,
                "ImageWriter::DeclareImageInputPort() the only supported "
                "pixel types are: kRgba8U, kDepth32F, and kLabel16I");

  if (publish_period <= 0) {
    throw std::logic_error("ImageWriter: publish period must be positive");
  }

  // Confirms the implied directory is valid.
  spruce::path test_dir =
      DirectoryFromFormat(file_name_format, port_name, kPixelType);
  if (!IsDirectoryValid(test_dir)) {
    throw std::logic_error(
        fmt::format("ImageWriter: The format string `{}` implied the invalid "
                    "directory: '{}'",
                    file_name_format, test_dir.getStr()));
  }

  // Confirms file has appropriate extension.
  const std::string& extension = extensions_[kPixelType];
  if (file_name_format.substr(file_name_format.size() - extension.size()) !=
      extension) {
    file_name_format += extension;
  }
  // TODO(SeanCurtis-TRI): Handle other issues that may arise with filename:
  //  - invalid symbols
  //  - invalid length
  //  - more?

  // Now configure the system for the valid port declaration.
  const auto& port =
      DeclareAbstractInputPort(port_name, systems::Value<Image<kPixelType>>());

  PublishEvent<double> event(
      Event<double>::TriggerType::kPeriodic,
      [this, port_index = port.get_index()](const Context<double>& context,
                                            const PublishEvent<double>&) {
        WriteImage<kPixelType>(context, port_index);
      });
  DeclarePeriodicEvent<PublishEvent<double>>(publish_period, start_time, event);
  port_data_.emplace_back(std::move(file_name_format), kPixelType);

  return port;
}

template <PixelType kPixelType>
void ImageWriter::WriteImage(const Context<double>& context, int index) const {
  const ImagePortData& data = port_data_[index];
  const Image<kPixelType>* image =
      this->EvalInputValue<Image<kPixelType>>(context, index);
  if (image) {
    const std::string& port_name = get_input_port(index).get_name();
    SaveToFileHelper(MakeFileName(data.format, data.pixel_type,
                                  context.get_time(), port_name, data.count++),
                     *image);
    return;
  }
  throw std::logic_error(
      fmt::format("ImageWriter: {} image input port {} is not connected",
                  labels_.at(data.pixel_type), index));
}

std::string ImageWriter::MakeFileName(const std::string& format,
                                      PixelType pixel_type, double time,
                                      const std::string& port_name,
                                      int count) const {
  DRAKE_DEMAND(labels_.count(pixel_type) > 0);

  int64_t u_time = static_cast<int64_t>(time * 1e6 + 0.5);
  int m_time = static_cast<int>(time * 1e3 + 0.5);
  return fmt::format(format, fmt::arg("port_name", port_name),
                     fmt::arg("image_type", labels_.at(pixel_type)),
                     fmt::arg("time_double", time),
                     fmt::arg("time_usec", u_time),
                     fmt::arg("time_msec", m_time), fmt::arg("count", count));
}

spruce::path ImageWriter::DirectoryFromFormat(const std::string& format,
                                              const std::string& port_name,
                                              PixelType pixel_type) const {
  // Extract the directory.
  size_t index = format.rfind('/');
  std::string dir_format = format.substr(0, index + 1);
  // NOTE: [bcdelmosu] are all the characters in: double, msec, and usec.
  // Technically, this will also key on '{time_mouse}', but if someone is
  // putting that in their file path, they deserve whatever they get.
  std::regex invalid_args("\\{count|time_[bcdelmosu]+\\}");
  std::smatch match;
  std::regex_search(dir_format, match, invalid_args);
  if (!match.empty()) {
    throw std::logic_error(
        "ImageWriter: The directory path cannot include time or image count");
  }
  std::string dir = MakeFileName(dir_format, pixel_type, 0, port_name, 0);
  return spruce::path(dir);
}

bool ImageWriter::IsDirectoryValid(const spruce::path& file_path) {
  if (file_path.exists()) {
    if (file_path.isDir()) {
      if (access(file_path.getStr().c_str(), W_OK) == 0) {
        return true;
      } else {
        drake::log()->error(
            "ImageWriter: no write permissions for the given folder: " +
            file_path.getStr());
      }
    } else {
      drake::log()->error(
          "ImageWriter: the provided folder isn't a directory: " +
          file_path.getStr());
    }
  } else {
    drake::log()->error("ImageWriter: the provided folder doesn't exist: " +
                        file_path.getStr());
  }
  return false;
}

template const InputPort<double>& ImageWriter::DeclareImageInputPort<
    PixelType::kRgba8U>(std::string port_name, std::string file_name_format,
                        double publish_period, double start_time);
template const InputPort<double>& ImageWriter::DeclareImageInputPort<
    PixelType::kDepth32F>(std::string port_name, std::string file_name_format,
                          double publish_period, double start_time);
template const InputPort<double>& ImageWriter::DeclareImageInputPort<
    PixelType::kLabel16I>(std::string port_name, std::string file_name_format,
                          double publish_period, double start_time);

template void ImageWriter::WriteImage<PixelType::kRgba8U>(
    const Context<double>& context, int index) const;
template void ImageWriter::WriteImage<PixelType::kDepth32F>(
    const Context<double>& context, int index) const;
template void ImageWriter::WriteImage<PixelType::kLabel16I>(
    const Context<double>& context, int index) const;

}  // namespace sensors
}  // namespace systems
}  // namespace drake