Skip to content

Commit

Permalink
Update tensorstore zarr writer to writer sequences
Browse files Browse the repository at this point in the history
  • Loading branch information
@JesseMckinzie
JesseMckinzie committed Jul 2, 2024
1 parent e24c7f9 commit 97bc1c7
Show file tree
Hide file tree
Showing 4 changed files with 140 additions and 106 deletions.
207 changes: 112 additions & 95 deletions src/cpp/writer/tswriter.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

#include "tensorstore/array.h"
#include "tensorstore/open.h"
#include "tensorstore/index_space/dim_expression.h"

#include "tswriter.h"
#include "../utilities/utilities.h"
Expand All @@ -18,118 +19,134 @@ TsWriterCPP::TsWriterCPP(
_image_shape(image_shape),
_chunk_shape(chunk_shape),
_dtype_code(GetDataTypeCode(dtype_str)) {

TENSORSTORE_CHECK_OK_AND_ASSIGN(_source, tensorstore::Open(
GetZarrSpecToWrite(_filename, _image_shape, _chunk_shape, GetEncodedType(_dtype_code)),
tensorstore::OpenMode::create |
tensorstore::OpenMode::delete_existing,
tensorstore::ReadWriteMode::write).result());


TENSORSTORE_CHECK_OK_AND_ASSIGN(_source, tensorstore::Open(
GetZarrSpecToWrite(_filename, _image_shape, _chunk_shape, GetEncodedType(_dtype_code)),
tensorstore::OpenMode::create |
tensorstore::OpenMode::delete_existing,
tensorstore::ReadWriteMode::write).result()
);
}


void TsWriterCPP::WriteImageData(py::array& py_image) {
void TsWriterCPP::WriteImageData(
py::array& py_image,
const Seq& rows,
const Seq& cols,
const Seq& layers,
const Seq& channels,
const Seq& tsteps) {

// use switch instead of template to avoid creating functions for each datatype
switch(_dtype_code)
{
case (1): {
auto data_array = tensorstore::Array(py_image.mutable_unchecked<std::uint8_t, 1>().data(0), _image_shape, tensorstore::c_order);
tensorstore::IndexTransform<> output_transform = tensorstore::IdentityTransform(_source.domain());

// Write data array to TensorStore
auto write_result = tensorstore::Write(tensorstore::UnownedToShared(data_array), _source).result();
output_transform = (std::move(output_transform) | tensorstore::Dims(_t_index).ClosedInterval(tsteps.Start(), tsteps.Stop())).value();
output_transform = (std::move(output_transform) | tensorstore::Dims(_c_index).ClosedInterval(channels.Start(), channels.Stop())).value();
output_transform = (std::move(output_transform) | tensorstore::Dims(_z_index).ClosedInterval(layers.Start(), layers.Stop())).value();
output_transform = (std::move(output_transform) | tensorstore::Dims(_y_index).ClosedInterval(rows.Start(), rows.Stop()) |
tensorstore::Dims(_x_index).ClosedInterval(cols.Start(), cols.Stop())).value();

if (!write_result.ok()) {
std::cerr << "Error writing image: " << write_result.status() << std::endl;
}
// use switch instead of template to avoid creating functions for each datatype
switch(_dtype_code)
{
case (1): {
auto data_array = tensorstore::Array(py_image.mutable_unchecked<std::uint8_t, 1>().data(0), _image_shape, tensorstore::c_order);

break;
}
case (2): {
auto data_array = tensorstore::Array(py_image.mutable_unchecked<std::uint16_t, 1>().data(0), _image_shape, tensorstore::c_order);
// Write data array to TensorStore
auto write_result = tensorstore::Write(tensorstore::UnownedToShared(data_array), _source | output_transform).result();

if (!write_result.ok()) {
std::cerr << "Error writing image: " << write_result.status() << std::endl;
}

auto write_result = tensorstore::Write(tensorstore::UnownedToShared(data_array), _source).result();
if (!write_result.ok()) {
std::cerr << "Error writing image: " << write_result.status() << std::endl;
break;
}
break;
}
case (4): {
auto data_array = tensorstore::Array(py_image.mutable_unchecked<std::uint32_t, 1>().data(0), _image_shape, tensorstore::c_order);

auto write_result = tensorstore::Write(tensorstore::UnownedToShared(data_array), _source).result();
if (!write_result.ok()) {
std::cerr << "Error writing image: " << write_result.status() << std::endl;
case (2): {
auto data_array = tensorstore::Array(py_image.mutable_unchecked<std::uint16_t, 1>().data(0), _image_shape, tensorstore::c_order);

auto write_result = tensorstore::Write(tensorstore::UnownedToShared(data_array), _source | output_transform).result();
if (!write_result.ok()) {
std::cerr << "Error writing image: " << write_result.status() << std::endl;
}
break;
}
case (4): {
auto data_array = tensorstore::Array(py_image.mutable_unchecked<std::uint32_t, 1>().data(0), _image_shape, tensorstore::c_order);

auto write_result = tensorstore::Write(tensorstore::UnownedToShared(data_array), _source | output_transform).result();
if (!write_result.ok()) {
std::cerr << "Error writing image: " << write_result.status() << std::endl;
}
break;
}
break;
}
case (8): {
auto data_array = tensorstore::Array(py_image.mutable_unchecked<std::uint64_t, 1>().data(0), _image_shape, tensorstore::c_order);

auto write_result = tensorstore::Write(tensorstore::UnownedToShared(data_array), _source).result();
if (!write_result.ok()) {
std::cerr << "Error writing image: " << write_result.status() << std::endl;
case (8): {
auto data_array = tensorstore::Array(py_image.mutable_unchecked<std::uint64_t, 1>().data(0), _image_shape, tensorstore::c_order);

auto write_result = tensorstore::Write(tensorstore::UnownedToShared(data_array), _source | output_transform).result();
if (!write_result.ok()) {
std::cerr << "Error writing image: " << write_result.status() << std::endl;
}
break;
}
break;
}
case (16): {
auto data_array = tensorstore::Array(py_image.mutable_unchecked<std::int8_t, 1>().data(0), _image_shape, tensorstore::c_order);

auto write_result = tensorstore::Write(tensorstore::UnownedToShared(data_array), _source).result();
if (!write_result.ok()) {
std::cerr << "Error writing image: " << write_result.status() << std::endl;
case (16): {
auto data_array = tensorstore::Array(py_image.mutable_unchecked<std::int8_t, 1>().data(0), _image_shape, tensorstore::c_order);

auto write_result = tensorstore::Write(tensorstore::UnownedToShared(data_array), _source | output_transform).result();
if (!write_result.ok()) {
std::cerr << "Error writing image: " << write_result.status() << std::endl;
}
break;
}
break;
}
case (32): {
auto data_array = tensorstore::Array(py_image.mutable_unchecked<std::int16_t, 1>().data(0), _image_shape, tensorstore::c_order);

auto write_result = tensorstore::Write(tensorstore::UnownedToShared(data_array), _source).result();
if (!write_result.ok()) {
std::cerr << "Error writing image: " << write_result.status() << std::endl;
case (32): {
auto data_array = tensorstore::Array(py_image.mutable_unchecked<std::int16_t, 1>().data(0), _image_shape, tensorstore::c_order);

auto write_result = tensorstore::Write(tensorstore::UnownedToShared(data_array), _source | output_transform).result();
if (!write_result.ok()) {
std::cerr << "Error writing image: " << write_result.status() << std::endl;
}
break;
}
case (64): {
auto data_array = tensorstore::Array(py_image.mutable_unchecked<std::int32_t, 1>().data(0), _image_shape, tensorstore::c_order);

auto write_result = tensorstore::Write(tensorstore::UnownedToShared(data_array), _source | output_transform).result();
if (!write_result.ok()) {
std::cerr << "Error writing image: " << write_result.status() << std::endl;
}
break;
}
break;
}
case (64): {
auto data_array = tensorstore::Array(py_image.mutable_unchecked<std::int32_t, 1>().data(0), _image_shape, tensorstore::c_order);

auto write_result = tensorstore::Write(tensorstore::UnownedToShared(data_array), _source).result();
if (!write_result.ok()) {
std::cerr << "Error writing image: " << write_result.status() << std::endl;
case (128): {
auto data_array = tensorstore::Array(py_image.mutable_unchecked<std::int64_t, 1>().data(0), _image_shape, tensorstore::c_order);

// Write data array to TensorStore
auto write_result = tensorstore::Write(tensorstore::UnownedToShared(data_array), _source | output_transform).result();
if (!write_result.ok()) {
std::cerr << "Error writing image: " << write_result.status() << std::endl;
}
break;
}
break;
}
case (128): {
auto data_array = tensorstore::Array(py_image.mutable_unchecked<std::int64_t, 1>().data(0), _image_shape, tensorstore::c_order);

// Write data array to TensorStore
auto write_result = tensorstore::Write(tensorstore::UnownedToShared(data_array), _source).result();
if (!write_result.ok()) {
std::cerr << "Error writing image: " << write_result.status() << std::endl;
case (256): {
auto data_array = tensorstore::Array(py_image.mutable_unchecked<float, 1>().data(0), _image_shape, tensorstore::c_order);

auto write_result = tensorstore::Write(tensorstore::UnownedToShared(data_array), _source | output_transform).result();
if (!write_result.ok()) {
std::cerr << "Error writing image: " << write_result.status() << std::endl;
}
break;
}
break;
}
case (256): {
auto data_array = tensorstore::Array(py_image.mutable_unchecked<float, 1>().data(0), _image_shape, tensorstore::c_order);

auto write_result = tensorstore::Write(tensorstore::UnownedToShared(data_array), _source).result();
if (!write_result.ok()) {
std::cerr << "Error writing image: " << write_result.status() << std::endl;
case (512): {
auto data_array = tensorstore::Array(py_image.mutable_unchecked<double, 1>().data(0), _image_shape, tensorstore::c_order);

auto write_result = tensorstore::Write(tensorstore::UnownedToShared(data_array), _source | output_transform).result();
if (!write_result.ok()) {
std::cerr << "Error writing image: " << write_result.status() << std::endl;
}
break;
}
break;
}
case (512): {
auto data_array = tensorstore::Array(py_image.mutable_unchecked<double, 1>().data(0), _image_shape, tensorstore::c_order);

auto write_result = tensorstore::Write(tensorstore::UnownedToShared(data_array), _source).result();
if (!write_result.ok()) {
std::cerr << "Error writing image: " << write_result.status() << std::endl;
default: {
// should not be reached
std::cerr << "Error writing image: unsupported data type" << std::endl;
}
break;
}
default: {
// should not be reached
std::cerr << "Error writing image: unsupported data type" << std::endl;
}
}
}
}
22 changes: 19 additions & 3 deletions src/cpp/writer/tswriter.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,14 +12,30 @@ namespace bfiocpp{

class TsWriterCPP{
public:
TsWriterCPP(const std::string& fname, const std::vector<std::int64_t>& image_shape, const std::vector<std::int64_t>& chunk_shape, const std::string& dtype);

void WriteImageData(py::array& py_image);
TsWriterCPP(
const std::string& fname,
const std::vector<std::int64_t>& image_shape,
const std::vector<std::int64_t>& chunk_shape,
const std::string& dtype);

void WriteImageData(
py::array& py_image,
const Seq& rows,
const Seq& cols,
const Seq& layers,
const Seq& channels,
const Seq& tsteps);

private:
std::string _filename;

std::vector<std::int64_t> _image_shape, _chunk_shape;

static constexpr int _t_index = 0,
_c_index = 1,
_z_index = 2,
_y_index = 3,
_x_index = 4;

uint16_t _dtype_code;

Expand Down
4 changes: 2 additions & 2 deletions src/python/bfiocpp/tswriter.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ def __init__(
file_name, image_shape, chunk_shape, str(dtype)
)

def write_image_data(self, image_data: np.ndarray):
def write_image_data(self, image_data: np.ndarray, rows: int, cols: int, layers: int, channels: int, tsteps: int):

Check warning on line 19 in src/python/bfiocpp/tswriter.py

View workflow job for this annotation

GitHub Actions / flake8 

E241: multiple spaces after ','

Check warning on line 19 in src/python/bfiocpp/tswriter.py

View workflow job for this annotation

GitHub Actions / flake8 

E501: line too long (119 > 88 characters)
"""Write image data to file
image_data: 5d numpy array containing image data
Expand All @@ -27,7 +27,7 @@ def write_image_data(self, image_data: np.ndarray):
raise ValueError("Image data must be a 5d numpy array")

try:
self._image_writer.write_image_data(image_data.flatten())
self._image_writer.write_image_data(image_data.flatten(), rows, cols, layers, channels, tsteps)

Check warning on line 30 in src/python/bfiocpp/tswriter.py

View workflow job for this annotation

GitHub Actions / flake8 

E501: line too long (107 > 88 characters)

except Exception as e:
raise RuntimeError(f"Error writing image data: {e.what}")
Expand Down
13 changes: 7 additions & 6 deletions tests/test_write.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -77,8 +77,14 @@ def test_write_zarr_2d(self):
# Use the temporary directory
test_file_path = os.path.join(dir, 'out/test.ome.zarr')

rows = Seq(0, br._Y - 1, 1)
cols = Seq(0, br._X - 1, 1)
layers = Seq(0, 0, 1)
channels = Seq(0, 0, 1)
tsteps = Seq(0, 0, 1)

bw = TSWriter(test_file_path, tmp.shape, tmp.shape, str(tmp.dtype))
bw.write_image_data(tmp)
bw.write_image_data(tmp, rows, cols, layers, channels, tsteps)
bw.close()

br = TSReader(
Expand All @@ -87,11 +93,6 @@ def test_write_zarr_2d(self):
"",
)

rows = Seq(0, br._Y - 1, 1)
cols = Seq(0, br._X - 1, 1)
layers = Seq(0, 0, 1)
channels = Seq(0, 0, 1)
tsteps = Seq(0, 0, 1)
tmp = br.data(rows, cols, layers, channels, tsteps)

assert tmp.dtype == np.uint8
Expand Down

0 comments on commit 97bc1c7

Please sign in to comment.