Skip to content
Snippets Groups Projects
Commit 46430997 authored by fima's avatar fima :beers:
Browse files

Merge branch 'ome_zarr_export' into 'main'

OME-Zarr export and import

See merge request !113
parents a8ed0f2e 0e8295e1
No related branches found
No related tags found
1 merge request!113OME-Zarr export and import
......@@ -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
......@@ -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
"""
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
......@@ -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):
......
......@@ -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
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment