Skip to content

Commit

Permalink
Merge pull request #67 from N-Dekker/ReadFromStore-variadic-template
Browse files Browse the repository at this point in the history 
Replace READ_ELEMENT_IF with variadic template, move functions into unnamed namespace
  • Loading branch information
@thewtex
thewtex authored Aug 21, 2024
2 parents 6bb50d2 + 77a2f45 commit ccf414f
Showing 1 changed file with 130 additions and 123 deletions.
253 changes: 130 additions & 123 deletions src/itkOMEZarrNGFFImageIO.cxx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,105 +36,6 @@ namespace itk
{
namespace
{
template <typename TPixel>
void
ReadFromStore(const tensorstore::TensorStore<> & store, const ImageIORegion & storeIORegion, TPixel * buffer)
{
if (store.domain().num_elements() == storeIORegion.GetNumberOfPixels())
{
// Read the entire available voxel region.
// Allow tensorstore to perform any axis permutations or other index operations
// to map from store axes to ITK image axes.
auto arr = tensorstore::Array(buffer, store.domain().shape(), tensorstore::c_order);
tensorstore::Read(store, tensorstore::UnownedToShared(arr)).value();
}
else
{
// Read a requested voxel subregion.
// We cannot infer axis permutations by matching requested axis sizes.
// Therefore we assume that tensorstore axes are in "C-style" order
// with the last index as the fastest moving axis, aka "z,y,x" order,
// and must inverted to match ITK's "Fortran-style" order of axis indices
// with the first index as the fastest moving axis, aka "x,y,z" order.
// "C-style" is generally the default layout for new tensorstore arrays.
// Refer to https://google.github.io/tensorstore/driver/zarr/index.html#json-driver/zarr.metadata.order
//
// In the future this may be extended to permute axes based on
// OME-Zarr NGFF axis labels.
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)
{
// 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();

auto arr = tensorstore::Array(buffer, indexDomain.shape(), tensorstore::c_order);
auto indexedStore = store | tensorstore::AllDims().SizedInterval(indices, sizes);
tensorstore::Read(indexedStore, tensorstore::UnownedToShared(arr)).value();
}
}

// Update an existing "read" specification for an "http" driver to retrieve remote files.
// Note that an "http" driver specification may operate on an HTTP or HTTPS connection.
void
MakeKVStoreHTTPDriverSpec(nlohmann::json & spec, const std::string & fullPath)
{
// Decompose path into a base URL and reference subpath according to TensorStore HTTP KVStore driver spec
// https://google.github.io/tensorstore/kvstore/http/index.html
spec["kvstore"] = { { "driver", "http" } };

// Naively decompose the URL into "base" and "resource" components.
// Generally assumes that the spec will only be used once to access a specific resource.
// For example, the URL "http://localhost/path/to/resource.json" will be split
// into components "http://localhost/path/to" and "resource.json".
//
// Could be revisited for a better root "base_url" at the top level allowing acces
// to multiple subpaths. For instance, decomposing the example above into
// "http://localhost/" and "path/to/resource.json" would allow for a given HTTP spec
// to be more easily reused with different subpaths.
//
spec["kvstore"]["base_url"] = fullPath.substr(0, fullPath.find_last_of("/"));
spec["kvstore"]["path"] = fullPath.substr(fullPath.find_last_of("/") + 1);
}

} // namespace

thread_local tensorstore::Context tsContext = tensorstore::Context::Default();

OMEZarrNGFFImageIO::OMEZarrNGFFImageIO()
{
this->AddSupportedWriteExtension(".zarr");
this->AddSupportedWriteExtension(".zr2");
this->AddSupportedWriteExtension(".zr3");
this->AddSupportedWriteExtension(".zip");
this->AddSupportedWriteExtension(".memory");

this->AddSupportedReadExtension(".zarr");
this->AddSupportedReadExtension(".zr2");
this->AddSupportedReadExtension(".zr3");
this->AddSupportedReadExtension(".zip");
this->AddSupportedWriteExtension(".memory");

this->Self::SetCompressor("");
this->Self::SetMaximumCompressionLevel(9);
this->Self::SetCompressionLevel(2);
}


void
OMEZarrNGFFImageIO::PrintSelf(std::ostream & os, Indent indent) const
{
Superclass::PrintSelf(os, indent);
os << indent << "DatasetIndex: " << m_DatasetIndex << std::endl;
os << indent << "TimeIndex: " << m_TimeIndex << std::endl;
os << indent << "ChannelIndex: " << m_ChannelIndex << std::endl;
}

IOComponentEnum
tensorstoreToITKComponentType(const tensorstore::DataType dtype)
{
Expand Down Expand Up @@ -246,6 +147,132 @@ itkToTensorstoreComponentType(const IOComponentEnum itkComponentType)
}
}

template <typename TPixel>
void
ReadFromStore(const tensorstore::TensorStore<> & store, const ImageIORegion & storeIORegion, TPixel * buffer)
{
if (store.domain().num_elements() == storeIORegion.GetNumberOfPixels())
{
// Read the entire available voxel region.
// Allow tensorstore to perform any axis permutations or other index operations
// to map from store axes to ITK image axes.
auto arr = tensorstore::Array(buffer, store.domain().shape(), tensorstore::c_order);
tensorstore::Read(store, tensorstore::UnownedToShared(arr)).value();
}
else
{
// Read a requested voxel subregion.
// We cannot infer axis permutations by matching requested axis sizes.
// Therefore we assume that tensorstore axes are in "C-style" order
// with the last index as the fastest moving axis, aka "z,y,x" order,
// and must inverted to match ITK's "Fortran-style" order of axis indices
// with the first index as the fastest moving axis, aka "x,y,z" order.
// "C-style" is generally the default layout for new tensorstore arrays.
// Refer to https://google.github.io/tensorstore/driver/zarr/index.html#json-driver/zarr.metadata.order
//
// In the future this may be extended to permute axes based on
// OME-Zarr NGFF axis labels.
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)
{
// 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();

auto arr = tensorstore::Array(buffer, indexDomain.shape(), tensorstore::c_order);
auto indexedStore = store | tensorstore::AllDims().SizedInterval(indices, sizes);
tensorstore::Read(indexedStore, tensorstore::UnownedToShared(arr)).value();
}
}

// Reads from the store if the specified pixel type and the ITK component type match.
template <typename TPixel>
bool
ReadFromStoreIfTypesMatch(const IOComponentEnum componentType,
const tensorstore::TensorStore<> & store,
const ImageIORegion & storeIORegion,
void * buffer)
{
if (tensorstoreToITKComponentType(tensorstore::dtype_v<TPixel>) == componentType)
{
ReadFromStore(store, storeIORegion, static_cast<TPixel *>(buffer));
return true;
}
return false;
}

// Tries to read from the specified store, trying any of the specified pixel types.
template <typename... TPixel>
bool
TryToReadFromStore(const IOComponentEnum componentType,
const tensorstore::TensorStore<> & store,
const ImageIORegion & storeIORegion,
void * buffer)
{
return (ReadFromStoreIfTypesMatch<TPixel>(componentType, store, storeIORegion, buffer) || ...);
}

// Update an existing "read" specification for an "http" driver to retrieve remote files.
// Note that an "http" driver specification may operate on an HTTP or HTTPS connection.
void
MakeKVStoreHTTPDriverSpec(nlohmann::json & spec, const std::string & fullPath)
{
// Decompose path into a base URL and reference subpath according to TensorStore HTTP KVStore driver spec
// https://google.github.io/tensorstore/kvstore/http/index.html
spec["kvstore"] = { { "driver", "http" } };

// Naively decompose the URL into "base" and "resource" components.
// Generally assumes that the spec will only be used once to access a specific resource.
// For example, the URL "http://localhost/path/to/resource.json" will be split
// into components "http://localhost/path/to" and "resource.json".
//
// Could be revisited for a better root "base_url" at the top level allowing acces
// to multiple subpaths. For instance, decomposing the example above into
// "http://localhost/" and "path/to/resource.json" would allow for a given HTTP spec
// to be more easily reused with different subpaths.
//
spec["kvstore"]["base_url"] = fullPath.substr(0, fullPath.find_last_of("/"));
spec["kvstore"]["path"] = fullPath.substr(fullPath.find_last_of("/") + 1);
}

} // namespace

thread_local tensorstore::Context tsContext = tensorstore::Context::Default();

OMEZarrNGFFImageIO::OMEZarrNGFFImageIO()
{
this->AddSupportedWriteExtension(".zarr");
this->AddSupportedWriteExtension(".zr2");
this->AddSupportedWriteExtension(".zr3");
this->AddSupportedWriteExtension(".zip");
this->AddSupportedWriteExtension(".memory");

this->AddSupportedReadExtension(".zarr");
this->AddSupportedReadExtension(".zr2");
this->AddSupportedReadExtension(".zr3");
this->AddSupportedReadExtension(".zip");
this->AddSupportedWriteExtension(".memory");

this->Self::SetCompressor("");
this->Self::SetMaximumCompressionLevel(9);
this->Self::SetCompressionLevel(2);
}


void
OMEZarrNGFFImageIO::PrintSelf(std::ostream & os, Indent indent) const
{
Superclass::PrintSelf(os, indent);
os << indent << "DatasetIndex: " << m_DatasetIndex << std::endl;
os << indent << "TimeIndex: " << m_TimeIndex << std::endl;
os << indent << "ChannelIndex: " << m_ChannelIndex << std::endl;
}

// Returns TensorStore KvStore driver name appropriate for this path.
// Options are file, zip. TODO: http, gcs (GoogleCouldStorage), etc.
std::string
Expand Down Expand Up @@ -601,14 +628,6 @@ OMEZarrNGFFImageIO::ReadImageInformation()
ReadArrayMetadata(std::string(this->GetFileName()) + "/" + path, driver);
}

// We call tensorstoreToITKComponentType for each type.
// Hopefully compiler will optimize it away via constant propagation and inlining.
#define READ_ELEMENT_IF(typeName) \
else if (tensorstoreToITKComponentType(tensorstore::dtype_v<typeName>) == this->GetComponentType()) \
{ \
ReadFromStore<typeName>(store, storeIORegion, reinterpret_cast<typeName *>(buffer)); \
}

void
OMEZarrNGFFImageIO::Read(void * buffer)
{
Expand All @@ -634,23 +653,11 @@ OMEZarrNGFFImageIO::Read(void * buffer)
<< storeIORegion;
}

if (false)
{
}
READ_ELEMENT_IF(float)
READ_ELEMENT_IF(int8_t)
READ_ELEMENT_IF(uint8_t)
READ_ELEMENT_IF(int16_t)
READ_ELEMENT_IF(uint16_t)
READ_ELEMENT_IF(int32_t)
READ_ELEMENT_IF(uint32_t)
READ_ELEMENT_IF(int64_t)
READ_ELEMENT_IF(uint64_t)
READ_ELEMENT_IF(float)
READ_ELEMENT_IF(double)
else
if (const IOComponentEnum componentType{ this->GetComponentType() };
!TryToReadFromStore<int8_t, uint8_t, int16_t, uint16_t, int32_t, uint32_t, int64_t, uint64_t, float, double>(
componentType, store, storeIORegion, buffer))
{
itkExceptionMacro("Unsupported component type: " << GetComponentTypeAsString(this->GetComponentType()));
itkExceptionMacro("Unsupported component type: " << GetComponentTypeAsString(componentType));
}
}

Expand Down

0 comments on commit ccf414f

Please sign in to comment.