[sensors] Encapsulate VTK image readers and writers

The file format is now a first-class citizen, passed as an argument, and reading or writing from files or a memory buffer are now equally well supported. Switch LCM image receiver system to use the VTK PNG decoder instead of custom code. Solve the VTK TIFF upside down dilemma noted in a prior glTF render engine TODO comment. Add acceptance test for write-and-readback.
RobotLocomotion · Aug 21, 2023 · 287327c · 287327c
1 parent b8644b8
commit 287327c
Show file tree

Hide file tree

Showing 11 changed files with 537 additions and 138 deletions.
diff --git a/geometry/render_gltf_client/BUILD.bazel b/geometry/render_gltf_client/BUILD.bazel
@@ -98,6 +98,7 @@ drake_cc_library(
         "//geometry/render:render_camera",
         "//geometry/render:render_engine",
         "//systems/sensors:image",
+        "//systems/sensors:vtk_image_reader_writer",
         "@picosha2_internal//:picosha2",
         "@vtk//:vtkCommonCore",
         "@vtk//:vtkCommonDataModel",

diff --git a/geometry/render_gltf_client/internal_render_client.cc b/geometry/render_gltf_client/internal_render_client.cc
@@ -21,6 +21,7 @@
 #include "drake/common/temp_directory.h"
 #include "drake/common/text_logging.h"
 #include "drake/geometry/render_gltf_client/internal_http_service_curl.h"
+#include "drake/systems/sensors/vtk_image_reader_writer.h"
 
 namespace drake {
 namespace geometry {
@@ -37,8 +38,10 @@ using render::RenderCameraCore;
 using systems::sensors::CameraInfo;
 using systems::sensors::ImageDepth16U;
 using systems::sensors::ImageDepth32F;
+using systems::sensors::ImageFileFormat;
 using systems::sensors::ImageLabel16I;
 using systems::sensors::ImageRgba8U;
+using systems::sensors::internal::MakeReader;
 
 /* Adds field_name = field_data to the map, assumes data_map does **not**
  already have the key `field_name`. */
@@ -88,30 +91,17 @@ void AddField<RenderImageType>(DataFieldsMap* data_map,
   }
 }
 
-void ReadPngFile(const std::string& path, vtkImageExport* image_exporter) {
-  // Load the PNG file from disk if possible.
-  vtkNew<vtkPNGReader> png_reader;
-  if (!png_reader->CanReadFile(path.c_str())) {
+void ReadImageFile(ImageFileFormat format, const std::string& path,
+                   vtkImageExport* image_exporter) {
+  // Load the image file from disk if possible.
+  auto reader = MakeReader(format, path);
+  if (!reader->CanReadFile(path.c_str())) {
     throw std::runtime_error(
-        fmt::format("RenderClient: cannot load '{}' as PNG.", path));
+        fmt::format("RenderClient: cannot load '{}' as {}.", path, format));
   }
-  png_reader->SetFileName(path.c_str());
-  image_exporter->SetInputConnection(png_reader->GetOutputPort());
+  image_exporter->SetInputConnection(reader->GetOutputPort());
   image_exporter->ImageLowerLeftOff();
-  png_reader->Update();  // Loads the image.
-}
-
-void ReadTiffFile(const std::string& path, vtkImageExport* image_exporter) {
-  // Load the TIFF file from disk if possible.
-  vtkNew<vtkTIFFReader> tiff_reader;
-  if (!tiff_reader->CanReadFile(path.c_str())) {
-    throw std::runtime_error(
-        fmt::format("RenderClient: cannot load '{}' as TIFF.", path));
-  }
-  tiff_reader->SetFileName(path.c_str());
-  image_exporter->SetInputConnection(tiff_reader->GetOutputPort());
-  // TODO(svenevs): why is ImageLowerLeftOff() not required for this exporter?
-  tiff_reader->Update();
+  reader->Update();  // Loads the image.
 }
 
 /* Verifies the loaded image has the correct dimensions.  This includes
@@ -269,7 +259,10 @@ std::string RenderClient::RenderOnServer(
 
    NOTE: Do not rely on or trust the server to (correctly) report a valid mime
    type for the sent image.  VTK image readers' `CanReadFile` methods check
-   if the file *content* can actually be loaded (regardless of extension). */
+   if the file *content* can actually be loaded (regardless of extension).
+
+   TODO(jwnimmer-tri) Add image_reader_writer.h helper function(s) to subsume
+   this manual file type detection. */
   vtkNew<vtkPNGReader> png_reader;
   if (png_reader->CanReadFile(bin_out_path.c_str())) {
     return RenameHttpServiceResponse(bin_out_path, scene_path, ".png");
@@ -322,7 +315,7 @@ void RenderClient::LoadColorImage(const std::string& path,
   DRAKE_DEMAND(color_image_out != nullptr);
 
   vtkNew<vtkImageExport> image_exporter;
-  ReadPngFile(path, image_exporter);
+  ReadImageFile(ImageFileFormat::kPng, path, image_exporter);
   const int width = color_image_out->width();
   const int height = color_image_out->height();
   VerifyImportedImageDimensions(width, height, image_exporter, path);
@@ -391,9 +384,9 @@ void RenderClient::LoadDepthImage(const std::string& path,
   vtkNew<vtkImageExport> image_exporter;
   const std::string ext{fs::path{path}.extension()};
   if (ext == ".png") {
-    ReadPngFile(path, image_exporter);
+    ReadImageFile(ImageFileFormat::kPng, path, image_exporter);
   } else if (ext == ".tiff") {
-    ReadTiffFile(path, image_exporter);
+    ReadImageFile(ImageFileFormat::kTiff, path, image_exporter);
   } else {
     throw std::runtime_error("RenderClient: unsupported file extension");
   }

diff --git a/systems/sensors/BUILD.bazel b/systems/sensors/BUILD.bazel
@@ -31,6 +31,7 @@ drake_cc_package_library(
         ":color_palette",
         ":gyroscope",
         ":image",
+        ":image_file_format",
         ":image_to_lcm_image_array_t",
         ":image_writer",
         ":lcm_image_array_to_images",
@@ -158,6 +159,15 @@ drake_cc_library(
     ],
 )
 
+drake_cc_library(
+    name = "image_file_format",
+    srcs = ["image_file_format.cc"],
+    hdrs = ["image_file_format.h"],
+    deps = [
+        "//common:essential",
+    ],
+)
+
 drake_cc_library(
     name = "lcm_image_traits",
     srcs = [
@@ -204,8 +214,8 @@ drake_cc_library(
     ],
     deps = [
         ":lcm_image_traits",
+        ":vtk_image_reader_writer",
         "//lcmtypes:image_array",
-        "@libpng",
         "@vtk//:vtkCommonCore",
         "@vtk//:vtkIOImage",
         "@zlib",
@@ -307,9 +317,12 @@ drake_cc_library(
     interface_deps = [
         ":image",
         "//common:essential",
-        "//systems/framework",
+        "//systems/framework:leaf_system",
     ],
     deps = [
+        ":vtk_image_reader_writer",
+        "@vtk//:vtkCommonCore",
+        "@vtk//:vtkCommonDataModel",
         "@vtk//:vtkIOImage",
     ],
 )
@@ -328,6 +341,20 @@ drake_cc_library(
     ],
 )
 
+drake_cc_library(
+    name = "vtk_image_reader_writer",
+    srcs = ["vtk_image_reader_writer.cc"],
+    hdrs = ["vtk_image_reader_writer.h"],
+    internal = True,
+    visibility = ["//:__subpackages__"],
+    deps = [
+        ":image_file_format",
+        "//common:unused",
+        "@vtk//:vtkCommonCore",
+        "@vtk//:vtkIOImage",
+    ],
+)
+
 # === test/ ===
 
 drake_cc_googletest(
@@ -396,6 +423,13 @@ drake_cc_googletest(
     ],
 )
 
+drake_cc_googletest(
+    name = "image_file_format_test",
+    deps = [
+        ":image_file_format",
+    ],
+)
+
 drake_cc_googletest(
     name = "gyroscope_test",
     data = ["//examples/pendulum:models"],
@@ -567,4 +601,16 @@ drake_cc_googletest(
     ],
 )
 
+drake_cc_googletest(
+    name = "vtk_image_reader_writer_test",
+    deps = [
+        ":vtk_image_reader_writer",
+        "//common:temp_directory",
+        "//common/test_utilities:expect_throws_message",
+        "@vtk//:vtkCommonCore",
+        "@vtk//:vtkCommonDataModel",
+        "@vtk//:vtkIOImage",
+    ],
+)
+
 add_lint_tests()
diff --git a/systems/sensors/image_file_format.cc b/systems/sensors/image_file_format.cc
@@ -0,0 +1,23 @@
+#include "drake/systems/sensors/image_file_format.h"
+
+#include "drake/common/drake_assert.h"
+
+namespace drake {
+namespace systems {
+namespace sensors {
+
+std::string to_string(ImageFileFormat format) {
+  switch (format) {
+    case ImageFileFormat::kJpeg:
+      return "jpeg";
+    case ImageFileFormat::kPng:
+      return "png";
+    case ImageFileFormat::kTiff:
+      return "tiff";
+  }
+  DRAKE_UNREACHABLE();
+}
+
+}  // namespace sensors
+}  // namespace systems
+}  // namespace drake
diff --git a/systems/sensors/image_file_format.h b/systems/sensors/image_file_format.h
@@ -0,0 +1,28 @@
+#pragma once
+
+#include <string>
+
+#include "drake/common/fmt.h"
+
+namespace drake {
+namespace systems {
+namespace sensors {
+
+/** The image file formats known to Drake. */
+enum class ImageFileFormat {
+  /** mime-type: image/jpeg. */
+  kJpeg,
+  /** mime-type: image/png. */
+  kPng,
+  /** mime-type: image/tiff. */
+  kTiff,
+};
+
+std::string to_string(ImageFileFormat);
+
+}  // namespace sensors
+}  // namespace systems
+}  // namespace drake
+
+DRAKE_FORMATTER_AS(, drake::systems::sensors, ImageFileFormat, x,
+                   drake::systems::sensors::to_string(x))
diff --git a/systems/sensors/image_writer.cc b/systems/sensors/image_writer.cc
@@ -12,14 +12,16 @@
 
 // To ease build system upkeep, we annotate VTK includes with their deps.
 #include <vtkImageData.h>     // vtkCommonDataModel
+#include <vtkImageWriter.h>   // vtkIOImage
 #include <vtkNew.h>           // vtkCommonCore
-#include <vtkPNGWriter.h>     // vtkIOImage
 #include <vtkSmartPointer.h>  // vtkCommonCore
-#include <vtkTIFFWriter.h>    // vtkIOImage
+
+#include "drake/systems/sensors/vtk_image_reader_writer.h"
 
 namespace drake {
 namespace systems {
 namespace sensors {
+namespace {
 
 template <PixelType kPixelType>
 void SaveToFileHelper(const Image<kPixelType>& image,
@@ -37,19 +39,19 @@ void SaveToFileHelper(const Image<kPixelType>& image,
     case PixelType::kRgba8U:
     case PixelType::kGrey8U:
       vtk_image->AllocateScalars(VTK_UNSIGNED_CHAR, num_channels);
-      writer = vtkSmartPointer<vtkPNGWriter>::New();
+      writer = internal::MakeWriter(ImageFileFormat::kPng, file_path);
       break;
     case PixelType::kDepth16U:
       vtk_image->AllocateScalars(VTK_UNSIGNED_SHORT, num_channels);
-      writer = vtkSmartPointer<vtkPNGWriter>::New();
+      writer = internal::MakeWriter(ImageFileFormat::kPng, file_path);
       break;
     case PixelType::kDepth32F:
       vtk_image->AllocateScalars(VTK_FLOAT, num_channels);
-      writer = vtkSmartPointer<vtkTIFFWriter>::New();
+      writer = internal::MakeWriter(ImageFileFormat::kTiff, file_path);
       break;
     case PixelType::kLabel16I:
       vtk_image->AllocateScalars(VTK_UNSIGNED_SHORT, num_channels);
-      writer = vtkSmartPointer<vtkPNGWriter>::New();
+      writer = internal::MakeWriter(ImageFileFormat::kPng, file_path);
       break;
     default:
       throw std::logic_error(
@@ -71,11 +73,12 @@ void SaveToFileHelper(const Image<kPixelType>& image,
     }
   }
 
-  writer->SetFileName(file_path.c_str());
   writer->SetInputData(vtk_image.GetPointer());
   writer->Write();
 }
 
+}  // namespace
+
 void SaveToPng(const ImageRgba8U& image, const std::string& file_path) {
   SaveToFileHelper(image, file_path);
 }