diff --git a/src/bfio/base_classes.py b/src/bfio/base_classes.py index 5bf86a5..a3dc8c2 100644 --- a/src/bfio/base_classes.py +++ b/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 diff --git a/src/bfio/bfio.py b/src/bfio/bfio.py index fbd335e..6feef4c 100644 --- a/src/bfio/bfio.py +++ b/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 diff --git a/src/bfio/ts_backends.py b/src/bfio/ts_backends.py index 8512f83..1e552c5 100644 --- a/src/bfio/ts_backends.py +++ b/src/bfio/ts_backends.py @@ -3,6 +3,7 @@ import logging from pathlib import Path from typing import Dict +import shutil # Third party packages @@ -10,7 +11,7 @@ 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 diff --git a/tests/test_write.py b/tests/test_write.py index 64c3f06..6cdcf60 100644 --- a/tests/test_write.py +++ b/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()