OME-Zarr export and import

docs/io.md

@@ -9,3 +9,5 @@ Currently, it is possible to directly load `tiff`, `h5`, `nii`,`txm`, `vol` and
             - load
             - save
             - Downloader
+            - export_ome_zarr
+            - import_ome_zarr
\ No newline at end of file

qim3d/io/__init__.py

@@ -4,3 +4,4 @@ from .saving import DataSaver, save
 from .sync import Sync
 from .convert import convert
 from ..utils import logger
+from .ome_zarr import export_ome_zarr, import_ome_zarr

qim3d/io/ome_zarr.py

+"""
+Exporting data to different formats.
+"""
+
+import os
+import numpy as np
+import zarr
+from ome_zarr.io import parse_url
+from ome_zarr.writer import write_image
+from ome_zarr.reader import Reader
+from ome_zarr import scale
+import math
+import shutil
+from qim3d.utils.logger import log
+from scipy.ndimage import zoom
+
+
+class OMEScaler(
+    scale.Scaler,
+):
+    """Scaler in the style of OME-Zarr.
+    This is needed because their current zoom implementation is broken."""
+
+    def __init__(self, order=0, downscale=2, max_layer=5, method="scaleZYX"):
+        self.order = order
+        self.downscale = downscale
+        self.max_layer = max_layer
+        self.method = method
+
+    def scaleZYX(self, base):
+        """Downsample using :func:`scipy.ndimage.zoom`."""
+        rv = [base]
+        log.info(f"- Scale 0: {rv[-1].shape}")
+
+        for i in range(self.max_layer):
+            rv.append(zoom(base, zoom=1 / (self.downscale * (i + 1)), order=self.order))
+            log.info(f"- Scale {i+1}: {rv[-1].shape}")
+        return list(rv)
+
+
+def export_ome_zarr(
+    path, data, chunk_size=100, downsample_rate=2, order=0, replace=False
+):
+    """
+    Export image data to OME-Zarr format with pyramidal downsampling.
+
+    Automatically calculates the number of downsampled scales such that the smallest scale fits within the specified `chunk_size`. 
+    
+    Args:
+        path (str): The directory where the OME-Zarr data will be stored.
+        data (np.ndarray): The image data to be exported.
+        chunk_size (int, optional): The size of the chunks for storing data. Defaults to 100.
+        downsample_rate (int, optional): Factor by which to downsample the data for each scale. Must be greater than 1. Defaults to 2.
+        order (int, optional): Interpolation order to use when downsampling. Defaults to 0 (nearest-neighbor).
+        replace (bool, optional): Whether to replace the existing directory if it already exists. Defaults to False.
+
+    Raises:
+        ValueError: If the directory already exists and `replace` is False.
+        ValueError: If `downsample_rate` is less than or equal to 1.
+
+    Example:
+        ```python
+        import qim3d
+
+        downloader = qim3d.io.Downloader()
+        data = downloader.Snail.Escargot(load_file=True)
+
+        qim3d.io.export_ome_zarr("Escargot.zarr", data, chunk_size=100, downsample_rate=2)
+
+        ```
+
+    """
+
+    # Check if directory exists
+    if os.path.exists(path):
+        if replace:
+            shutil.rmtree(path)
+        else:
+            raise ValueError(
+                f"Directory {path} already exists. Use replace=True to overwrite."
+            )
+
+    # Check if downsample_rate is valid
+    if downsample_rate <= 1:
+        raise ValueError("Downsample rate must be greater than 1.")
+
+    log.info(f"Exporting data to OME-Zarr format at {path}")
+
+    # Get the number of scales
+    min_dim = np.max(np.shape(data))
+    nscales = math.ceil(math.log(min_dim / chunk_size) / math.log(downsample_rate)) + 1
+    log.info(f"Number of scales: {nscales + 1}")
+
+    # Create scaler
+    scaler = OMEScaler(
+        downscale=downsample_rate, max_layer=nscales, method="scaleZYX", order=order
+    )
+
+    # write the image data
+    os.mkdir(path)
+    store = parse_url(path, mode="w").store
+    root = zarr.group(store=store)
+    write_image(
+        image=data,
+        group=root,
+        axes="zyx",
+        storage_options=dict(chunks=(chunk_size, chunk_size, chunk_size)),
+        scaler=scaler,
+    )
+
+
+def import_ome_zarr(path, scale=0, load=True):
+    """
+    Import image data from an OME-Zarr file.
+
+    This function reads OME-Zarr formatted volumetric image data and returns the specified scale. 
+    The image data can be lazily loaded (as Dask arrays) or fully computed into memory.
+
+    Args:
+        path (str): The file path to the OME-Zarr data.
+        scale (int or str, optional): The scale level to load. 
+            If 'highest', loads the finest scale (scale 0). 
+            If 'lowest', loads the coarsest scale (last available scale). Defaults to 0.
+        load (bool, optional): Whether to compute the selected scale into memory. 
+            If False, returns a lazy Dask array. Defaults to True.
+
+    Returns:
+        np.ndarray or dask.array.Array: The requested image data, either as a NumPy array if `load=True`, 
+        or a Dask array if `load=False`.
+
+    Raises:
+        ValueError: If the requested `scale` does not exist in the data.
+
+    Example:
+        ```python
+        import qim3d
+
+        data = qim3d.io.import_ome_zarr("Escargot.zarr", scale=0, load=True)
+
+        ```
+
+    """
+
+    # read the image data
+    #store = parse_url(path, mode="r").store
+
+    reader = Reader(parse_url(path))
+    nodes = list(reader())
+    image_node = nodes[0]
+    dask_data = image_node.data
+
+    log.info(f"Data contains {len(dask_data)} scales:")
+    for i in np.arange(len(dask_data)):
+        log.info(f"- Scale {i}: {dask_data[i].shape}")
+
+    if scale == "highest":
+        scale = 0
+
+    if scale == "lowest":
+        scale = len(dask_data) - 1
+
+    if scale >= len(dask_data):
+        raise ValueError(f"Scale {scale} does not exist in the data. Please choose a scale between 0 and {len(dask_data)-1}.")
+
+    log.info(f"\nLoading scale {scale} with shape {dask_data[scale].shape}")
+
+    if load:
+        vol = dask_data[scale].compute()
+    else:
+        vol = dask_data[scale]
+
+    return vol

qim3d/utils/misc.py

@@ -223,7 +223,7 @@ def downscale_img(img, max_voxels=512**3):
     zoom_factor = (max_voxels / total_voxels) ** (1 / 3)
 
     # Downscale image
-    return zoom(img, zoom_factor)
+    return zoom(img, zoom_factor, order=0)
 
 
 def scale_to_float16(arr: np.ndarray):

qim3d/viz/k3d.py

@@ -22,7 +22,7 @@ def vol(
     grid_visible=False,
     cmap=None,
     samples="auto",
-    max_voxels=412**3,
+    max_voxels=512**3,
     data_type="scaled_float16",
     **kwargs,
 ):
@@ -97,7 +97,7 @@ def vol(
 
     if original_shape != new_shape:
         log.warning(
-            f"Downsampled image for visualization. From {original_shape} to {new_shape}"
+            f"Downsampled image for visualization, from {original_shape} to {new_shape}"
         )
 
     # Scale the image to float16 if needed