Merge pull request #91 from bioio-devs/zarr-writer-single-time

allow zarr writer to do incremental writes of single timepoints

Author: toloudis

bioio/tests/conftest.py

@@ -57,7 +57,8 @@ def da_random_from_shape(
 
 
 array_constructor = pytest.mark.parametrize(
-    "array_constructor", [np_random_from_shape, da_random_from_shape]
+    "array_constructor",
+    [np_random_from_shape, da_random_from_shape],
 )
 
 DUMMY_PLUGIN_NAME = "dummy-plugin"

bioio/tests/writers/test_ome_zarr_writer_2.py

@@ -193,3 +193,93 @@ def test_write_ome_zarr(
     axes = node.metadata["axes"]
     dims = "".join([a["name"] for a in axes]).upper()
     assert dims == "TCZYX"
+
+
+@array_constructor
+@pytest.mark.parametrize(
+    "shape, num_levels, scaling, expected_shapes",
+    [
+        (
+            (4, 2, 2, 64, 32),  # easy, powers of two
+            3,
+            (1, 1, 1, 2, 2),  # downscale xy by two
+            [(4, 2, 2, 64, 32), (4, 2, 2, 32, 16), (4, 2, 2, 16, 8)],
+        ),
+        (
+            (4, 2, 2, 8, 6),
+            1,  # no downscaling
+            (1, 1, 1, 1, 1),
+            [(4, 2, 2, 8, 6)],
+        ),
+    ],
+)
+@pytest.mark.parametrize("filename", ["e.zarr"])
+def test_write_ome_zarr_iterative(
+    array_constructor: Callable,
+    filename: str,
+    shape: DimTuple,
+    num_levels: int,
+    scaling: Tuple[float, float, float, float, float],
+    expected_shapes: List[DimTuple],
+    tmp_path: pathlib.Path,
+) -> None:
+    # TCZYX order, downsampling x and y only
+    im = array_constructor(shape, dtype=np.uint8)
+    C = shape[1]
+
+    shapes = compute_level_shapes(shape, scaling, num_levels)
+    chunk_sizes = compute_level_chunk_sizes_zslice(shapes)
+
+    # Create an OmeZarrWriter object
+    writer = OmeZarrWriter()
+
+    # Initialize the store. Use s3 url or local directory path!
+    save_uri = tmp_path / filename
+    writer.init_store(str(save_uri), shapes, chunk_sizes, im.dtype)
+
+    # Write the image iteratively as if we only have one timepoint at a time
+    for t in range(shape[0]):
+        t4d = im[t]
+        t5d = np.expand_dims(t4d, axis=0)
+        writer.write_t_batches_array(t5d, channels=[], tbatch=1, toffset=t)
+
+    # TODO: get this from source image
+    physical_scale = {
+        "c": 1.0,  # default value for channel
+        "t": 1.0,
+        "z": 1.0,
+        "y": 1.0,
+        "x": 1.0,
+    }
+    physical_units = {
+        "x": "micrometer",
+        "y": "micrometer",
+        "z": "micrometer",
+        "t": "minute",
+    }
+    meta = writer.generate_metadata(
+        image_name="TEST",
+        channel_names=[f"c{i}" for i in range(C)],
+        physical_dims=physical_scale,
+        physical_units=physical_units,
+        channel_colors=[0xFFFFFF for i in range(C)],
+    )
+    writer.write_metadata(meta)
+
+    # Read written result and check basics
+    reader = Reader(parse_url(save_uri))
+    node = list(reader())[0]
+    num_levels_read = len(node.data)
+    assert num_levels_read == num_levels
+    for level, shape in zip(range(num_levels), expected_shapes):
+        read_shape = node.data[level].shape
+        assert read_shape == shape
+    axes = node.metadata["axes"]
+    dims = "".join([a["name"] for a in axes]).upper()
+    assert dims == "TCZYX"
+
+    # check lvl 0 data values got written in order
+    for t in range(shape[0]):
+        t4d = im[t]
+        read_t4d = node.data[0][t]
+        np.testing.assert_array_equal(t4d, read_t4d)

bioio/writers/ome_zarr_writer_2.py

@@ -220,7 +220,7 @@ def build_ome(
 class ZarrLevel:
     shape: DimTuple
     chunk_size: DimTuple
-    dtype: str
+    dtype: np.dtype
     zarray: zarr.core.Array
 
 
@@ -440,7 +440,7 @@ def _create_levels(
             self.levels.append(level)
 
     def _downsample_and_write_batch_t(
-        self, data_tczyx: da.Array, start_t: int, end_t: int
+        self, data_tczyx: da.Array, start_t: int, end_t: int, toffset: int = 0
     ) -> None:
         dtype = data_tczyx.dtype
         if len(data_tczyx.shape) != 5:
@@ -451,7 +451,11 @@ def _downsample_and_write_batch_t(
         # write level 0 first
         for k in range(start_t, end_t):
             subset = data_tczyx[[k - start_t]]
-            da.to_zarr(subset, self.levels[0].zarray, region=(slice(k, k + 1),))
+            da.to_zarr(
+                subset,
+                self.levels[0].zarray,
+                region=(slice(k + toffset, k + toffset + 1),),
+            )
 
         # downsample to next level then write
         for j in range(1, len(self.levels)):
@@ -463,7 +467,11 @@ def _downsample_and_write_batch_t(
             # write ti to zarr
             for k in range(start_t, end_t):
                 subset = data_tczyx[[k - start_t]]
-                da.to_zarr(subset, self.levels[j].zarray, region=(slice(k, k + 1),))
+                da.to_zarr(
+                    subset,
+                    self.levels[j].zarray,
+                    region=(slice(k + toffset, k + toffset + 1),),
+                )
 
         log.info(f"Completed {start_t} to {end_t}")
 
@@ -540,6 +548,7 @@ def write_t_batches_array(
         im: Union[da.Array, np.ndarray],
         channels: List[int] = [],
         tbatch: int = 4,
+        toffset: int = 0,
         debug: bool = False,
     ) -> None:
         """
@@ -551,6 +560,8 @@ def write_t_batches_array(
             An ArrayLike object. Should be 5D TCZYX.
         tbatch:
             The number of T to write at a time.
+        toffset:
+            The offset to start writing T from. All T in the input array will be written
         """
         # if isinstance(im, (np.ndarray)):
         #     im_da = da.from_array(im)
@@ -571,7 +582,9 @@ def write_t_batches_array(
                 if channels:
                     for t in range(len(ti)):
                         ti[t] = [ti[t][c] for c in channels]
-                self._downsample_and_write_batch_t(da.asarray(ti), start_t, end_t)
+                self._downsample_and_write_batch_t(
+                    da.asarray(ti), start_t, end_t, toffset
+                )
         log.info("Finished loop over T")
 
     def _get_scale_ratio(self, level: int) -> Tuple[float, float, float, float, float]: