ENH: Allow user-defined slicing over time or channel axes (#43)

* ENH: Allow user-defined slicing over time or channel axes Changes: - Adds axes metadata parsing and exposes result in API - Exposes optional parameter to allow slicing over channel or time axes Addresses issue where 5D t/c/z/y/x volume subregions could not be read with OMEZarrNGFFImageIO because Tensorstore was not provided with enough information to map from 5D axes to 3D output image. Could be extended in the future to allow the user to override the default mapping of ITK dimension 0 -> "x", 1 -> "y", 2 -> "z" for different input axes. * DOC: Remove whitespace
InsightSoftwareConsortium · Aug 9, 2023 · 08d3617 · 08d3617
1 parent ea7bd2c
commit 08d3617
Show file tree

Hide file tree

Showing 4 changed files with 236 additions and 8 deletions.
diff --git a/include/itkOMEZarrNGFFImageIO.h b/include/itkOMEZarrNGFFImageIO.h
@@ -22,10 +22,28 @@
 
 
 #include <fstream>
+#include <string>
+#include <vector>
 #include "itkImageIOBase.h"
 
 namespace itk
 {
+/** \class OMEZarrNGFFAxis
+ *
+ * \brief Represent an OME-Zarr NGFF axis
+ *
+ * Open Microscopy Environment Zarr Next Generation File Format
+ * specification can be found at https://github.com/ome/ngff
+ *
+ * \ingroup IOOMEZarrNGFF
+ */
+struct IOOMEZarrNGFF_EXPORT OMEZarrNGFFAxis
+{
+  std::string name;
+  std::string type;
+  std::string unit;
+};
+
 /** \class OMEZarrNGFFImageIO
  *
  * \brief Read and write OMEZarrNGFF images.
@@ -54,6 +72,8 @@ class IOOMEZarrNGFF_EXPORT OMEZarrNGFFImageIO : public ImageIOBase
   itkTypeMacro(OMEZarrNGFFImageIO, ImageIOBase);
 
   static constexpr unsigned MaximumDimension = 5; // OME-NGFF specifies up to 5D data
+  static constexpr int      INVALID_INDEX = -1;   // for specifying enumerated axis slice indices
+  using AxesCollectionType = std::vector<OMEZarrNGFFAxis>;
 
   /** The different types of ImageIO's can support data of varying
    * dimensionality. For example, some file formats are strictly 2D
@@ -126,6 +146,20 @@ class IOOMEZarrNGFF_EXPORT OMEZarrNGFFImageIO : public ImageIOBase
   itkGetConstMacro(DatasetIndex, int);
   itkSetMacro(DatasetIndex, int);
 
+  /** If there is a time axis, at what index should it be sliced? */
+  itkGetConstMacro(TimeIndex, int);
+  itkSetMacro(TimeIndex, int);
+
+  /** If there are multiple channels, which one should be read? */
+  itkGetConstMacro(ChannelIndex, int);
+  itkSetMacro(ChannelIndex, int);
+
+  /** Get the available axes in the OME-Zarr store in ITK (Fortran-style) order.
+   *  This is reversed from the default C-style order of
+   *  axes as used in the Zarr / NumPy / Tensorstore interface.
+   */
+  itkGetConstMacro(StoreAxes, const AxesCollectionType &);
+
   bool
   CanStreamRead() override
   {
@@ -149,6 +183,17 @@ class IOOMEZarrNGFF_EXPORT OMEZarrNGFFImageIO : public ImageIOBase
   void
   ReadArrayMetadata(std::string path, std::string driver);
 
+  /** Process requested store region for given configuration */
+  ImageIORegion
+  ConfigureTensorstoreIORegion(const ImageIORegion & ioRegion) const;
+
+  /** Helper method to get axes in tensorstore C-style order*/
+  AxesCollectionType
+  GetAxesInStoreOrder() const
+  {
+    return AxesCollectionType(m_StoreAxes.rbegin(), m_StoreAxes.rend());
+  }
+
   /** Sets the requested dimension, and initializes spatial metadata to identity. */
   void
   InitializeIdentityMetadata(unsigned nDims)
@@ -182,7 +227,10 @@ class IOOMEZarrNGFF_EXPORT OMEZarrNGFFImageIO : public ImageIOBase
   const std::vector<std::string> dimensionUnits = { "millimeter", "millimeter", "millimeter", "index", "second" };
 
 private:
-  int m_DatasetIndex = 0; // first, highest resolution scale by default
+  int                m_DatasetIndex = 0; // first, highest resolution scale by default
+  int                m_TimeIndex = INVALID_INDEX;
+  int                m_ChannelIndex = INVALID_INDEX;
+  AxesCollectionType m_StoreAxes;
 };
 } // end namespace itk
 

diff --git a/src/itkOMEZarrNGFFImageIO.cxx b/src/itkOMEZarrNGFFImageIO.cxx
@@ -38,9 +38,9 @@ namespace
 {
 template <typename TPixel>
 void
-ReadFromStore(const tensorstore::TensorStore<> & store, const ImageIORegion & ioRegion, TPixel * buffer)
+ReadFromStore(const tensorstore::TensorStore<> & store, const ImageIORegion & storeIORegion, TPixel * buffer)
 {
-  if (store.domain().num_elements() == ioRegion.GetNumberOfPixels())
+  if (store.domain().num_elements() == storeIORegion.GetNumberOfPixels())
   {
     // Read the entire available voxel region.
     // Allow tensorstore to perform any axis permutations or other index operations
@@ -61,14 +61,15 @@ ReadFromStore(const tensorstore::TensorStore<> & store, const ImageIORegion & io
     //
     // In the future this may be extended to permute axes based on
     // OME-Zarr NGFF axis labels.
-    const auto                      dimension = ioRegion.GetImageDimension();
+    const auto                      dimension = store.rank();
     std::vector<tensorstore::Index> indices(dimension);
     std::vector<tensorstore::Index> sizes(dimension);
     for (size_t dim = 0; dim < dimension; ++dim)
     {
-      // Reverse order of axes to match assumed C-style Zarr storage
-      indices[(dimension - 1) - dim] = ioRegion.GetIndex(dim);
-      sizes[(dimension - 1) - dim] = ioRegion.GetSize(dim);
+      // Input IO region is assumed to already be reversed from ITK requested region
+      // to match assumed C-style Zarr storage
+      indices[dim] = storeIORegion.GetIndex(dim);
+      sizes[dim] = storeIORegion.GetSize(dim);
     }
     auto indexDomain = tensorstore::IndexDomainBuilder(dimension).origin(indices).shape(sizes).Finalize().value();
 
@@ -401,6 +402,74 @@ OMEZarrNGFFImageIO::ReadArrayMetadata(std::string path, std::string driver)
   }
 }
 
+ImageIORegion
+OMEZarrNGFFImageIO::ConfigureTensorstoreIORegion(const ImageIORegion & ioRegion) const
+{
+  // Set up IO region to match known store dimensions
+  itkAssertOrThrowMacro(m_StoreAxes.size() == store.rank(), "Detected mismatch in axis count and store rank");
+  ImageIORegion storeRegion(store.rank());
+  itkAssertOrThrowMacro(storeRegion.GetImageDimension(), store.rank());
+  auto storeAxes = this->GetAxesInStoreOrder();
+
+  for (size_t storeIndex = 0; storeIndex < store.rank(); ++storeIndex)
+  {
+    auto axisName = storeAxes.at(storeIndex).name;
+
+    // Optionally slice time or channel indices
+    if (axisName == "t")
+    {
+      storeRegion.SetSize(storeIndex, 0);
+      if (m_TimeIndex == INVALID_INDEX)
+      {
+        itkWarningMacro(<< "The OME-Zarr store contains a time \"t\" axis but no time point has been specified. "
+                           "Reading along a time axis is not currently supported. Data will be read from the first "
+                           "available time point by default.");
+        storeRegion.SetIndex(storeIndex, 0);
+      }
+      else
+      {
+        storeRegion.SetIndex(storeIndex, m_TimeIndex);
+      }
+    }
+    else if (axisName == "c")
+    {
+      storeRegion.SetSize(storeIndex, 0);
+      if (m_ChannelIndex == INVALID_INDEX)
+      {
+        itkWarningMacro(<< "The OME-Zarr store contains a channel \"c\" axis but no channel index has been specified. "
+                           "Reading along a channel axis is not currently supported. Data will be read from the first "
+                           "available channel by default.");
+        storeRegion.SetIndex(storeIndex, 0);
+      }
+      else
+      {
+        storeRegion.SetIndex(storeIndex, m_ChannelIndex);
+      }
+    }
+    // Set requested region on X/Y/Z axes
+    else if (axisName == "x")
+    {
+      itkAssertOrThrowMacro(ioRegion.GetImageDimension() >= 1, "Failed to read from \"x\" axis into ITK axis \"0\"");
+      storeRegion.SetSize(storeIndex, ioRegion.GetSize(0));
+      storeRegion.SetIndex(storeIndex, ioRegion.GetIndex(0));
+    }
+    else if (axisName == "y")
+    {
+      itkAssertOrThrowMacro(ioRegion.GetImageDimension() >= 2, "Failed to read from \"y\" axis into ITK axis \"1\"");
+      storeRegion.SetSize(storeIndex, ioRegion.GetSize(1));
+      storeRegion.SetIndex(storeIndex, ioRegion.GetIndex(1));
+    }
+    else if (axisName == "z")
+    {
+      itkAssertOrThrowMacro(ioRegion.GetImageDimension() >= 3, "Failed to read from \"z\" axis into ITK axis \"2\"");
+      storeRegion.SetSize(storeIndex, ioRegion.GetSize(2));
+      storeRegion.SetIndex(storeIndex, ioRegion.GetIndex(2));
+    }
+  }
+
+  return storeRegion;
+}
+
 void
 addCoordinateTransformations(OMEZarrNGFFImageIO * io, nlohmann::json ct)
 {
@@ -478,9 +547,23 @@ OMEZarrNGFFImageIO::ReadImageInformation()
   if (json.contains("axes")) // optional before 0.3
   {
     this->InitializeIdentityMetadata(json.at("axes").size());
+
+    m_StoreAxes.resize(json.at("axes").size());
+    auto targetIt = m_StoreAxes.rbegin();
+    for (const auto & axis : json.at("axes"))
+    {
+      *targetIt = (OMEZarrNGFFAxis{ axis.at("name"), axis.at("type"), (axis.contains("unit") ? axis.at("unit") : "") });
+      ++targetIt;
+    }
+    itkAssertOrThrowMacro(targetIt == m_StoreAxes.rend(),
+                          "Internal error: failed to fully parse axes from OME-Zarr metadata");
   }
   else
   {
+    if (version == "0.4")
+    {
+      itkExceptionMacro(<< "\"axes\" field is missing from OME-Zarr image metadata at " << zattrsFilePath);
+    }
     this->SetNumberOfDimensions(0);
   }
 
@@ -521,7 +604,7 @@ OMEZarrNGFFImageIO::ReadImageInformation()
 #define READ_ELEMENT_IF(typeName)                                                                     \
   else if (tensorstoreToITKComponentType(tensorstore::dtype_v<typeName>) == this->GetComponentType()) \
   {                                                                                                   \
-    ReadFromStore<typeName>(store, m_IORegion, reinterpret_cast<typeName *>(buffer));                 \
+    ReadFromStore<typeName>(store, storeIORegion, reinterpret_cast<typeName *>(buffer));              \
   }
 
 void
@@ -541,6 +624,7 @@ OMEZarrNGFFImageIO::Read(void * buffer)
     itkAssertOrThrowMacro(this->GetNumberOfComponents() == 1,
                           "Reading an image subregion is currently supported only for single channel images");
   }
+  auto storeIORegion = this->ConfigureTensorstoreIORegion(m_IORegion);
 
   if (false)
   {}

diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt
@@ -118,3 +118,15 @@ itk_add_test(
     itkOMEZarrNGFFHTTPTest
     1
 )
+itk_add_test(
+  NAME IOOMEZarrNGFFHTTP_TimeSlice
+  COMMAND IOOMEZarrNGFFTestDriver
+    itkOMEZarrNGFFHTTPTest
+    2
+)
+itk_add_test(
+  NAME IOOMEZarrNGFFHTTP_TimeAndChannelSlice
+  COMMAND IOOMEZarrNGFFTestDriver
+    itkOMEZarrNGFFHTTPTest
+    3
+)
diff --git a/test/itkOMEZarrNGFFHTTPTest.cxx b/test/itkOMEZarrNGFFHTTPTest.cxx
@@ -155,6 +155,86 @@ test3DImage()
 
   return EXIT_SUCCESS;
 }
+
+bool
+testTimeSlice()
+{
+  // Read a subregion of an arbitrary time point from a 3D image buffer into a 2D image
+  using ImageType = itk::Image<unsigned char, 2>;
+  const std::string resourceURL = "https://s3.embl.de/i2k-2020/ngff-example-data/v0.4/tyx.ome.zarr";
+
+  auto imageIO = itk::OMEZarrNGFFImageIO::New();
+  imageIO->SetDatasetIndex(0);
+  imageIO->SetTimeIndex(2);
+
+  auto requestedRegion = itk::ImageRegion<2>();
+  requestedRegion.SetSize(itk::MakeSize(50, 50));
+  requestedRegion.SetIndex(itk::MakeIndex(100, 100));
+
+  // Set up reader
+  auto reader = itk::ImageFileReader<ImageType>::New();
+  reader->SetFileName(resourceURL);
+  reader->SetImageIO(imageIO);
+  reader->GetOutput()->SetRequestedRegion(requestedRegion);
+
+  // Read
+  reader->Update();
+
+  auto image = reader->GetOutput();
+  image->Print(std::cout);
+
+  ITK_TEST_EXPECT_EQUAL(image->GetBufferedRegion().GetSize(), requestedRegion.GetSize());
+  ITK_TEST_EXPECT_EQUAL(image->GetBufferedRegion().GetIndex(), requestedRegion.GetIndex());
+
+  typename ImageType::SpacingType expectedSpacing;
+  expectedSpacing.Fill(0.65);
+  ITK_TEST_EXPECT_EQUAL(image->GetSpacing(), expectedSpacing);
+  ITK_TEST_EXPECT_EQUAL(image->GetOrigin(), itk::MakePoint(0.0, 0.0));
+
+  return EXIT_SUCCESS;
+}
+
+bool
+testTimeAndChannelSlice()
+{
+  // Read a subregion of an arbitrary channel and time point from a 5D image buffer into a 3D image
+  using ImageType = itk::Image<unsigned char, 3>;
+  const std::string resourceURL = "https://s3.embl.de/i2k-2020/ngff-example-data/v0.4/tczyx.ome.zarr";
+
+  auto imageIO = itk::OMEZarrNGFFImageIO::New();
+  imageIO->SetDatasetIndex(2);
+  imageIO->SetTimeIndex(0);
+  imageIO->SetChannelIndex(0);
+
+  auto requestedRegion = itk::ImageRegion<3>();
+  requestedRegion.SetSize(itk::MakeSize(10, 20, 30));
+  requestedRegion.SetIndex(itk::MakeIndex(5, 10, 15));
+
+  // Set up reader
+  auto reader = itk::ImageFileReader<ImageType>::New();
+  reader->SetFileName(resourceURL);
+  reader->SetImageIO(imageIO);
+  reader->GetOutput()->SetRequestedRegion(requestedRegion);
+
+  // Read
+  reader->Update();
+
+  auto image = reader->GetOutput();
+  image->Print(std::cout);
+
+  ITK_TEST_EXPECT_EQUAL(image->GetBufferedRegion().GetSize(), requestedRegion.GetSize());
+  ITK_TEST_EXPECT_EQUAL(image->GetBufferedRegion().GetIndex(), requestedRegion.GetIndex());
+
+  typename ImageType::SpacingType expectedSpacing;
+  expectedSpacing.SetElement(0, 2.6);
+  expectedSpacing.SetElement(1, 2.6);
+  expectedSpacing.SetElement(2, 4.0);
+  ITK_TEST_EXPECT_EQUAL(image->GetSpacing(), expectedSpacing);
+  ITK_TEST_EXPECT_EQUAL(image->GetOrigin(), itk::MakePoint(0.0, 0.0, 0.0));
+
+  return EXIT_SUCCESS;
+}
+
 } // namespace
 
 int
@@ -177,6 +257,10 @@ itkOMEZarrNGFFHTTPTest(int argc, char * argv[])
     case 1:
       return test3DImage();
       break;
+    case 2:
+      return testTimeSlice();
+    case 3:
+      return testTimeAndChannelSlice();
     default:
       throw std::invalid_argument("Invalid test case ID: " + std::to_string(testCase));
   }