Merge pull request #10 from janelia-cellmap/refactor

added test for
tiff and mrc
chunk writing for different dim of input array data
modified code to make sure chunk dim aligns with shape of the input data

Author: yuriyzubov

.gitignore

@@ -6,4 +6,5 @@ dist/
 .zarrify_venv
 .python-version
 *.egg-info/
-build/
+build/
+.vscode/

pyproject.toml

@@ -29,7 +29,7 @@ dependencies = [
     "numcodecs>=0.13.0,<0.16",
     "pydantic-zarr>=0.4.0,<0.5.0",
     "pint>=0.20.0,<1.0.0",
-    "tifffile>=2025.1.10,<2026.0.0",
+    "tifffile>=2024.1.0,<2025.5.10",
     "zarr==2.18.3",
     "bokeh>=3.1.0"
 ]

src/zarrify/formats/mrc.py

@@ -11,6 +11,7 @@
 from numcodecs import Zstd
 import logging
 import copy
+import numpy as np
 
 class Mrc3D(Volume):
 
@@ -31,8 +32,14 @@ def __init__(
 
         self.memmap = mrcfile.mmap(self.src_path, mode="r")
         self.ndim = self.memmap.data.ndim
-        self.shape = self.memmap.data.shape
+        self.shape = np.squeeze(self.memmap.data.shape)
         self.dtype = self.memmap.data.dtype
+        
+        # Scale metadata parameters to match data dimensionality
+        self.metadata["axes"] = self.metadata["axes"][-self.ndim:]
+        self.metadata["scale"] = self.metadata["scale"][-self.ndim:]
+        self.metadata["translation"] = self.metadata["translation"][-self.ndim:]
+        self.metadata["units"] = self.metadata["units"][-self.ndim:]
 
 
     def write_to_zarr(
@@ -53,6 +60,10 @@ def write_to_zarr(
                          format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 
         src_path = copy.copy(self.src_path)
+        
+        if len(zarr_chunks) != len(self.shape):
+           zarr_chunks = self.reshape_to_arr_shape(zarr_chunks, self.shape)
+        
         z_arr = self.get_output_array(dest, zarr_chunks, comp)
 
         out_slices = slices_from_chunks(

src/zarrify/formats/tiff.py

@@ -34,6 +34,13 @@ def __init__(
 
         self.shape = self.zarr_arr.shape
         self.dtype = self.zarr_arr.dtype
+        self.ndim = self.zarr_arr.ndim
+        
+        # Scale metadata parameters to match data dimensionality
+        self.metadata["axes"] = self.metadata["axes"][-self.ndim:]
+        self.metadata["scale"] = self.metadata["scale"][-self.ndim:]
+        self.metadata["translation"] = self.metadata["translation"][-self.ndim:]
+        self.metadata["units"] = self.metadata["units"][-self.ndim:]
 
     def write_to_zarr(self,
         dest: str,
@@ -42,6 +49,10 @@ def write_to_zarr(self,
         comp : ABCMeta = Zstd(level=6),
         ):
 
+        # reshape chunk shape to align with arr shape
+        if len(zarr_chunks) != self.shape:
+           zarr_chunks = self.reshape_to_arr_shape(zarr_chunks, self.shape)
+             
         z_arr = self.get_output_array(dest, zarr_chunks, comp)
         chunks_list = np.arange(0, z_arr.shape[0], z_arr.chunks[0])
 

src/zarrify/formats/tiff_stack.py

@@ -37,7 +37,14 @@ def __init__(
         probe_image_arr = da.from_zarr(probe_image_store)
 
         self.dtype = probe_image_arr.dtype
-        self.shape = [len(self.stack_list)] + list(probe_image_arr.shape)
+        self.shape = np.squeeze([len(self.stack_list)] + list(probe_image_arr.shape))
+        self.ndim = len(self.shape)
+        
+        # Scale metadata parameters to match data dimensionality
+        self.metadata["axes"] = self.metadata["axes"][-self.ndim:]
+        self.metadata["scale"] = self.metadata["scale"][-self.ndim:]
+        self.metadata["translation"] = self.metadata["translation"][-self.ndim:]
+        self.metadata["units"] = self.metadata["units"][-self.ndim:]
 
     def write_tile_slab_to_zarr(
         self, chunk_num: int, zarray: zarr.Array, src_volume: list
@@ -73,6 +80,10 @@ def write_to_zarr(self,
         comp : ABCMeta = Zstd(level=6),
         ):
 
+        # reshape chunk shape to align with arr shape
+        if len(zarr_chunks) != len(self.shape):
+           zarr_chunks = self.reshape_to_arr_shape(zarr_chunks, self.shape)
+        
         z_arr = self.get_output_array(dest, zarr_chunks, comp)
         chunks_list = np.arange(0, z_arr.shape[0], z_arr.chunks[0])
 

src/zarrify/utils/volume.py

@@ -31,6 +31,10 @@ def get_output_array(self, dest: str, chunks: list[int], comp: ABCMeta) -> zarr.
             chunks=chunks,
             compressor=comp,
         )
+        
+    def reshape_to_arr_shape(self, param_arr, ref_arr):
+        from itertools import cycle, islice
+        return list(islice(cycle(param_arr), len(ref_arr)))
 
     def add_ome_metadata(self, dest: str):
         """Add selected tiff metadata to zarr attributes file (.zattrs).

tests/test_to_zarr.py

@@ -0,0 +1,59 @@
+from pathlib import Path
+
+import pytest
+
+import mrcfile
+import tifffile
+import zarr
+import numpy as np
+import os
+
+from zarrify.utils.dask_utils import initialize_dask_client
+from zarrify import to_zarr
+
+@pytest.fixture
+def create_test_file(tmp_path):
+    """Factory fixture to create test files with specified format and dimensions"""
+    def _create_file(ext: str, shape: tuple):
+        data = np.random.rand(*shape).astype(np.uint8)
+        file_path = tmp_path / f"test_image_{len(shape)}d.{ext}"
+        
+        if ext in ["tiff", "tif"]:
+            tifffile.imwrite(file_path, data)
+        elif ext == "mrc":
+            with mrcfile.new(file_path, overwrite=True) as mrc:
+                mrc.set_data(data.astype(data.dtype))
+        else:
+            raise ValueError("Unsupported file format")
+        
+        return file_path, data
+    return _create_file
+
+import itertools
+
+# Test parameters
+FORMATS = ['tif', 'tiff', 'mrc']
+SHAPES = [(40, 50), (1, 30, 50)]  # 2D and 3D
+
+@pytest.mark.parametrize("ext,shape", list(itertools.product(FORMATS, SHAPES)))
+def test_to_zarr(create_test_file, ext, shape):
+    
+    src_path, expected_data = create_test_file(ext, shape)
+    dest_path = Path(f"{src_path.with_suffix('')}.zarr")
+    
+    dask_client = initialize_dask_client('local')
+    
+    # convert to zarr
+    to_zarr(src_path, dest_path, dask_client)
+    
+    if src_path.suffix.lstrip('.') in ['tif', 'tiff']:
+        src_data = tifffile.imread(src_path)
+    elif src_path.suffix.lstrip('.') == 'mrc':
+        with mrcfile.open(src_path, permissive=True) as mrc:
+            src_data = mrc.data
+    
+    # store array in s0 by convention
+    dest_data = zarr.open(f'{dest_path}/s0', mode='r')
+    assert np.array_equal(dest_data[:], src_data)
+    assert np.array_equal(dest_data[:], expected_data)
+