Merge pull request #87 from JesseMckinzie/tswriter

Add tensorstore zarr writer

src/bfio/base_classes.py

@@ -887,3 +887,34 @@ def read_metadata(self):
     def read_image(self, *args):
         """Abstract read image executor."""
         pass
+
+
+class TSAbstractWriter(TSAbstractBackend):
+    """Base class for file writers.
+
+    All writer objects must be a subclass of AbstractWriter.
+    """
+
+    _metadata: ome_types.OME = None
+
+    @abc.abstractmethod
+    def __init__(self, frontend: BioBase):
+        """Initialize the reader object.
+
+        Args:
+            frontend (BioBase): The BioBase object associated with the backend.
+        """
+        super().__init__(frontend)
+
+    @abc.abstractmethod
+    def write_metadata(self):
+        """Read OME metadata from the file.
+
+        Subclasses must override this to properly retrieve and format the data.
+        """
+        pass
+
+    @abc.abstractmethod
+    def write_image(self, *args):
+        """Abstract read image executor."""
+        pass

src/bfio/bfio.py

@@ -12,7 +12,7 @@
 
 from bfio import backends
 from bfio.base_classes import BioBase
-from bfio.ts_backends import TensorstoreReader
+from bfio.ts_backends import TensorstoreReader, TensorstoreWriter
 
 
 class BioReader(BioBase):
@@ -1062,6 +1062,8 @@ class if specified. *Defaults to None.*
         # Ensure backend is supported
         if self._backend_name == "python":
             self._backend = backends.PythonWriter(self)
+        elif self._backend_name == "tensorstore":
+            self._backend = TensorstoreWriter(self)
         elif self._backend_name == "bioformats":
             try:
                 self._backend = backends.JavaWriter(self)
@@ -1121,9 +1123,11 @@ def set_backend(self, backend: typing.Optional[str] = None) -> None:
             "python",
             "bioformats",
             "zarr",
+            "tensorstore",
         ]:
             raise ValueError(
-                'Keyword argument backend must be one of ["python","bioformats","zarr"]'
+                "Keyword argument backend must be one of "
+                + '["python","bioformats","zarr","tensorstore"]'
             )
         if backend == "python":
             extension = "".join(self._file_path.suffixes)
@@ -1344,6 +1348,11 @@ def write(
         X_tile_end = numpy.ceil(X[1] / self._TILE_SIZE).astype(int) * self._TILE_SIZE
         Y_tile_end = numpy.ceil(Y[1] / self._TILE_SIZE).astype(int) * self._TILE_SIZE
 
+        # Ensure end is not out of bounds to avoid tensorstore errors
+        if self._backend_name == "tensorstore":
+            X_tile_end = min(X_tile_end, self._DIMS["X"])
+            Y_tile_end = min(Y_tile_end, self._DIMS["Y"])
+
         # Read the image
         self._backend.write_image(
             [X_tile_start, X_tile_end], [Y_tile_start, Y_tile_end], Z, C, T, image

src/bfio/ts_backends.py

@@ -3,14 +3,15 @@
 import logging
 from pathlib import Path
 from typing import Dict
+import shutil
 
 
 # Third party packages
 import ome_types
 from xml.etree import ElementTree as ET
 
 
-from bfiocpp import TSReader, Seq, FileType, get_ome_xml
+from bfiocpp import TSReader, TSWriter, Seq, FileType, get_ome_xml
 import bfio.base_classes
 from bfio.utils import clean_ome_xml_for_known_issues
 import zarr
@@ -276,3 +277,88 @@ def read_zarr_metadata(self):
 
                 self._metadata = omexml
                 return self._metadata
+
+
+class TensorstoreWriter(bfio.base_classes.TSAbstractWriter):
+    logger = logging.getLogger("bfio.backends.TensorstoreWriter")
+
+    def __init__(self, frontend):
+
+        super().__init__(frontend)
+
+        self._init_writer()
+
+    def _init_writer(self):
+        """_init_writer Initializes file writing.
+
+        This method is called exactly once per object. Once it is called,
+        all other methods of setting metadata will throw an error.
+
+        NOTE: For Zarr, it is not explicitly necessary to make the file
+              read-only once writing has begun. Thus functionality is mainly
+              incorporated to remain consistent with the OME TIFF formats.
+              In the future, it may be reasonable to not enforce read-only
+
+        """
+
+        if self.frontend._file_path.exists():
+            shutil.rmtree(self.frontend._file_path)
+
+        # Tensorstore writer currently only supports zarr
+        if not self.frontend._file_path.name.endswith(".zarr"):
+            raise ValueError("File type must be zarr to use tensorstore writer.")
+
+        shape = (
+            self.frontend.T,
+            self.frontend.C,
+            self.frontend.Z,
+            self.frontend.Y,
+            self.frontend.X,
+        )
+
+        self._writer = TSWriter(
+            str(self.frontend._file_path.resolve()),
+            shape,
+            (1, 1, 1, self.frontend._TILE_SIZE, self.frontend._TILE_SIZE),
+            self.frontend.dtype,
+            "TCZYX",
+        )
+
+    def write_metadata(self):
+
+        # Create the metadata
+        metadata_path = (
+            Path(self.frontend._file_path).joinpath("OME").joinpath("METADATA.ome.xml")
+        )
+
+        metadata_path.parent.mkdir(parents=True, exist_ok=True)
+
+        with open(metadata_path, "w") as fw:
+            fw.write(str(self.frontend._metadata.to_xml()))
+
+        # This is recommended to do for cloud storage to increase read/write
+        # speed, but it also increases write speed locally when threading.
+        zarr.consolidate_metadata(str(self.frontend._file_path.resolve()))
+
+    def write_image(self, X, Y, Z, C, T, image):
+
+        cols = Seq(X[0], X[-1] - 1, 1)
+        rows = Seq(Y[0], Y[-1] - 1, 1)
+        layers = Seq(Z[0], Z[-1] - 1, 1)
+
+        if len(C) == 1:
+            channels = Seq(C[0], C[0], 1)
+        else:
+            channels = Seq(C[0], C[-1], 1)
+
+        if len(T) == 1:
+            tsteps = Seq(T[0], T[0], 1)
+        else:
+            tsteps = Seq(T[0], T[-1], 1)
+
+        self._writer.write_image_data(
+            image.flatten(), rows, cols, layers, channels, tsteps
+        )
+
+    def close(self):
+        pass

tests/test_write.py

@@ -1,12 +1,14 @@
 # -*- coding: utf-8 -*-
 import unittest
-import requests, pathlib, shutil, logging, sys
+import requests, pathlib, shutil, logging, sys, os
 import bfio
 import numpy as np
+import tempfile
 from ome_zarr.utils import download as zarr_download
 
 TEST_IMAGES = {
     "Plate1-Blue-A-12-Scene-3-P3-F2-03.czi": "https://downloads.openmicroscopy.org/images/Zeiss-CZI/idr0011/Plate1-Blue-A_TS-Stinger/Plate1-Blue-A-12-Scene-3-P3-F2-03.czi",
+    "5025551.zarr": "https://uk1s3.embassy.ebi.ac.uk/idr/zarr/v0.4/idr0054A/5025551.zarr",
 }
 
 TEST_DIR = pathlib.Path(__file__).with_name("data")
@@ -28,7 +30,7 @@ def setUpModule():
     for file, url in TEST_IMAGES.items():
         logger.info(f"setup - Downloading: {file}")
 
-        if not file.endswith(".ome.zarr"):
+        if not file.endswith(".zarr"):
             if TEST_DIR.joinpath(file).exists():
                 continue
 
@@ -41,6 +43,11 @@ def setUpModule():
                 shutil.rmtree(TEST_DIR.joinpath(file))
             zarr_download(url, str(TEST_DIR))
 
+def tearDownModule():
+    """Remove test images"""
+
+    logger.info("teardown - Removing test images...")
+    shutil.rmtree(TEST_DIR)
 
 class TestOmeTiffWrite(unittest.TestCase):
     @classmethod
@@ -93,3 +100,35 @@ def test_write_java(self):
         with bfio.BioReader("4d_array_bf.ome.tif") as br:
             assert image.shape == br.shape
             assert np.array_equal(image[:], br[:])
+
+
+class TestOmeZarrWriter(unittest.TestCase):
+
+    def test_write_zarr_tensorstore(self):
+
+        with bfio.BioReader(str(TEST_DIR.joinpath("5025551.zarr"))) as br:
+
+            actual_shape = br.shape
+            actual_dtype = br.dtype
+
+            actual_image = br[:]
+
+            actual_mdata = br.metadata
+
+        with tempfile.TemporaryDirectory() as dir:
+
+            # Use the temporary directory
+            test_file_path = os.path.join(dir, 'out/test.ome.zarr')
+
+            with bfio.BioWriter(test_file_path, metadata=actual_mdata, backend="tensorstore") as bw:
+
+                expanded = np.expand_dims(actual_image, axis=-1)
+                bw[:] = expanded
+
+            with bfio.BioReader(test_file_path) as br:
+
+
+                assert br.shape == actual_shape
+                assert br.dtype == actual_dtype
+
+                assert br[:].sum() == actual_image.sum()