diff --git a/pyproject.toml b/pyproject.toml index 03ecbae3..c9254a6f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,17 +33,19 @@ classifiers = [# https://pypi.python.org/pypi?%3Aaction=list_classifiers dependencies = [ "adjustText", "anndata>=0.12.4", # https://github.com/scverse/anndata/issues/2166 - "bioio<2", + "bioio>=3", "bioio-nd2", + "bioio-ome-tiff", + "bioio-ome-zarr", "bioio-tifffile", "centrosome", "cp-measure>=0.1.16", "dask-image", - "dask<=2025.11.0", + "dask", "decorator", "filelock", "flox", - "fsspec!=2023.9.0", # 2023.9.0 causes ome-zarr write image to fail + "fsspec", "igraph", "itk-elastix", "itk", @@ -54,8 +56,8 @@ dependencies = [ "natsort", "numcodecs", "numpy", - "ome-zarr<0.12.0", - "pandas", + "ome-zarr", + "pandas<3", "pint", "psutil", "pyarrow", @@ -68,13 +70,13 @@ dependencies = [ "stardist", "statsmodels", "tensorflow", - "tifffile<=2025.5.10", + "tifffile", "xarray", - "zarr<3" + "zarr>=3" ] [project.optional-dependencies] -# Added pysam as an optional extra for Linux/macOS users + dialout = [ "pysam" ] @@ -86,7 +88,7 @@ dask-ml = [ "dask-ml" ] cellpose = [ - "cellpose<4" + "cellpose" ] ufish = [ "ufish" diff --git a/requirements.txt b/requirements.txt index 7b7f36d2..9bd56c95 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,13 +1,15 @@ -anndata==0.12.16 +anndata==0.13.2 adjustText==1.4.0 -bioio-nd2==1.3.0 +bioio-nd2==2.3.0 +bioio-ome-tiff==1.4.0 +bioio-ome-zarr==3.5.1 bioio-tifffile==1.3.0 -bioio==1.6.1 +bioio==3.4.0 centrosome==1.3.3 cp-measure==0.1.19 cython==3.2.4 dask-image==2026.5.0 -dask==2025.11.0 +dask===2026.7.1 decorator==5.3.1 filelock==3.29.4 flox==0.11.2 @@ -22,7 +24,7 @@ matplotlib==3.10.9 natsort==8.4.0 numcodecs==0.15.1 numpy==2.4.6 -ome-zarr==0.11.1 +ome-zarr==0.18.0 pandas==2.3.3 pint==0.25.3 psutil==7.2.2 @@ -36,6 +38,7 @@ shapely==2.1.2 stardist==0.9.2 statsmodels==0.14.6 tensorflow==2.21.0 -tifffile==2025.5.10 +tifffile==2026.7.14 +tifffile==2026.7.14 xarray==2026.2.0 -zarr==2.18.7 +zarr==3.2.1 diff --git a/scallops/_bioio_zarr_reader.py b/scallops/_bioio_zarr_reader.py deleted file mode 100644 index 6118c466..00000000 --- a/scallops/_bioio_zarr_reader.py +++ /dev/null @@ -1,263 +0,0 @@ -import logging -from pathlib import Path -from typing import Any, Dict, List, Optional, Tuple - -import xarray as xr -from bioio_base import constants, dimensions, exceptions, io, reader, types -from fsspec.spec import AbstractFileSystem -from ome_zarr.io import parse_url -from ome_zarr.reader import Reader as ZarrReader - -logger = logging.getLogger("scallops") - - -# Same as https://github.com/bioio-devs/bioio-ome-zarr/blob/main/bioio_ome_zarr/reader.py but fixes bug in channel names -# Also checks to see if zarr path is {zarr_path}/images/image1 with only 1 image -# See https://github.com/bioio-devs/bioio-ome-zarr/pull/22 -class ScallopsZarrReader(reader.Reader): - """The main class of each reader plugin. This class is subclass of the abstract class reader - (BaseReader) in bioio-base. - - Parameters - ---------- - image: types.PathLike - String or Path to the ZARR root - fs_kwargs: Dict[str, Any] - Ignored - """ - - _xarray_dask_data: Optional["xr.DataArray"] = None - _xarray_data: Optional["xr.DataArray"] = None - _mosaic_xarray_dask_data: Optional["xr.DataArray"] = None - _mosaic_xarray_data: Optional["xr.DataArray"] = None - _dims: Optional[dimensions.Dimensions] = None - _metadata: Optional[Any] = None - _scenes: Optional[Tuple[str, ...]] = None - _current_scene_index: int = 0 - # Do not provide default value because - # they may not need to be used by your reader (i.e. input param is an array) - _fs: "AbstractFileSystem" - _path: str - - # Required Methods - - def __init__( - self, - image: types.PathLike, - fs_kwargs: Dict[str, Any] = {}, - ): - # Expand details of provided image - self._fs, self._path = io.pathlike_to_fs( - image, - enforce_exists=False, - fs_kwargs=fs_kwargs, - ) - - # Enforce valid image - if not self._is_supported_image(self._fs, self._path): - raise exceptions.UnsupportedFileFormatError( - self.__class__.__name__, - self._path, - "Could not find a .zgroup or .zarray file at the provided path.", - ) - - self._zarr = get_zarr_reader(self._fs, self._path).zarr - self._physical_pixel_sizes: Optional[types.PhysicalPixelSizes] = None - self._channel_names: Optional[List[str]] = None - - @staticmethod - def _is_supported_image(fs: AbstractFileSystem, path: str, **kwargs: Any) -> bool: - try: - get_zarr_reader(fs, path) - return True - except AttributeError: - return False - - @classmethod - def is_supported_image( - cls, - image: types.ImageLike, - fs_kwargs: Dict[str, Any] = {}, - **kwargs: Any, - ) -> bool: - if isinstance(image, (str, Path)): - return cls._is_supported_image(None, str(image), **kwargs) - else: - return reader.Reader.is_supported_image( - cls, image, fs_kwargs=fs_kwargs, **kwargs - ) - - @property - def scenes(self) -> Tuple[str, ...]: - if self._scenes is None: - scenes = self._zarr.root_attrs["multiscales"] - - # if (each scene has a name) and (that name is unique) use name. - # otherwise generate scene names. - if all("name" in scene for scene in scenes) and ( - len({scene["name"] for scene in scenes}) == len(scenes) - ): - self._scenes = tuple(str(scene["name"]) for scene in scenes) - else: - self._scenes = tuple( - f"scene_{i}" - for i in range(len(self._zarr.root_attrs["multiscales"])) - ) - return self._scenes - - @property - def resolution_levels(self) -> Tuple[int, ...]: - """ - Returns - ------- - resolution_levels: Tuple[str, ...] - Return the available resolution levels for the current scene. - By default these are ordered from highest resolution to lowest - resolution. - """ - return tuple( - rl - for rl in range( - len( - self._zarr.root_attrs["multiscales"][self.current_scene_index][ - "datasets" - ] - ) - ) - ) - - def _read_delayed(self) -> xr.DataArray: - return self._xarr_format(delayed=True) - - def _read_immediate(self) -> xr.DataArray: - return self._xarr_format(delayed=False) - - def _xarr_format(self, delayed: bool) -> xr.DataArray: - data_path = self._zarr.root_attrs["multiscales"][self.current_scene_index][ - "datasets" - ][self.current_resolution_level]["path"] - image_data = self._zarr.load(data_path) - - axes = self._zarr.root_attrs["multiscales"][self.current_scene_index].get( - "axes" - ) - if axes: - dims = [sub["name"].upper() for sub in axes] - else: - dims = list(reader.Reader._guess_dim_order(image_data.shape)) - - if not delayed: - image_data = image_data.compute() - - coords = self._get_coords( - dims, - image_data.shape, - scene=self.current_scene, - channel_names=self.channel_names, - ) - - return xr.DataArray( - image_data, - dims=dims, - coords=coords, - attrs={constants.METADATA_UNPROCESSED: self._zarr.root_attrs}, - ) - - # Optional Methods - @property - def physical_pixel_sizes(self) -> types.PhysicalPixelSizes: - """Return the physical pixel sizes of the image.""" - if self._physical_pixel_sizes is None: - try: - z_size, y_size, x_size = self._get_pixel_size( - list(self.dims.order), - ) - except Exception as e: - logger.warning(f"Could not parse zarr pixel size: {e}") - z_size, y_size, x_size = None, None, None - - self._physical_pixel_sizes = types.PhysicalPixelSizes( - z_size, y_size, x_size - ) - return self._physical_pixel_sizes - - def _get_pixel_size( - self, - dims: List[str], - ) -> Tuple[Optional[float], Optional[float], Optional[float]]: - # OmeZarr file may contain an additional set of "coordinateTransformations" - # these coefficents are applied to all resolution levels. - if ( - "coordinateTransformations" - in self._zarr.root_attrs["multiscales"][self.current_scene_index] - ): - universal_res_consts = self._zarr.root_attrs["multiscales"][ - self.current_scene_index - ]["coordinateTransformations"][0]["scale"] - else: - universal_res_consts = [1.0 for _ in range(len(dims))] - - coord_transform = self._zarr.root_attrs["multiscales"][ - self.current_scene_index - ]["datasets"][self.current_resolution_level]["coordinateTransformations"] - - spatial_coeffs = {} - - for dim in [ - dimensions.DimensionNames.SpatialX, - dimensions.DimensionNames.SpatialY, - dimensions.DimensionNames.SpatialZ, - ]: - if dim in dims: - dim_index = dims.index(dim) - spatial_coeffs[dim] = ( - coord_transform[0]["scale"][dim_index] - * universal_res_consts[dim_index] - ) - else: - spatial_coeffs[dim] = None - - return ( - spatial_coeffs[dimensions.DimensionNames.SpatialZ], - spatial_coeffs[dimensions.DimensionNames.SpatialY], - spatial_coeffs[dimensions.DimensionNames.SpatialX], - ) - - @property - def channel_names(self) -> Optional[List[str]]: - if self._channel_names is None: - if "omero" in self._zarr.root_attrs: - self._channel_names = [ - str(channel["label"]) - for channel in self._zarr.root_attrs["omero"]["channels"] - ] - return self._channel_names - - @staticmethod - def _get_coords( - dims: List[str], - shape: Tuple[int, ...], - scene: str, - channel_names: Optional[List[str]], - ) -> Dict[str, Any]: - coords: Dict[str, Any] = {} - - # Use dims for coord determination - if dimensions.DimensionNames.Channel in dims: - # Generate channel names if no existing channel names - if channel_names is None: - coords[dimensions.DimensionNames.Channel] = [ - f"channel_{i}" - for i in range(shape[dims.index(dimensions.DimensionNames.Channel)]) - ] - else: - coords[dimensions.DimensionNames.Channel] = channel_names - - return coords - - -def get_zarr_reader(fs: AbstractFileSystem, path: str) -> ZarrReader: - if fs is not None: - path = fs.unstrip_protocol(path) - - return ZarrReader(parse_url(path, mode="r")) diff --git a/scallops/cli/features.py b/scallops/cli/features.py index c52040cd..cca59c2f 100644 --- a/scallops/cli/features.py +++ b/scallops/cli/features.py @@ -16,7 +16,6 @@ import dask.array import dask.array as da import fsspec -import numpy as np import pandas as pd import pyarrow as pa import pyarrow.parquet as pq @@ -52,7 +51,6 @@ pluralize, read_anndata_zarr, ) -from scallops.zarr_io import _read_ome_zarr_array logger = _get_cli_logger() @@ -240,44 +238,17 @@ def single_feature( output_fs, _ = fsspec.core.url_to_fs(output_dir) - zarr_inputs = True - - for f in file_list: - if not isinstance(f, (zarr.Group, zarr.Array)): - zarr_inputs = False - break - - if zarr_inputs and stacked_image_tuple is not None: - for f in stacked_file_list: - if not isinstance(f, (zarr.Group, zarr.Array)): - zarr_inputs = False - break - if not zarr_inputs: - image = _read_image(file_list, metadata) - else: - image = [] - for f in file_list: - array, _, _, _ = _read_ome_zarr_array(f) - image.append(array) + image = _read_image(file_list, metadata) n_channels1 = None if stacked_image_tuple is not None: - if not zarr_inputs: - stacked_image = _read_image(stacked_file_list, stacked_metadata) - n_channels1 = image.sizes["c"] - # clear coords to avoid issues with xr.concat - for c in list(image.coords.keys()): - del image.coords[c] - for c in list(stacked_image.coords.keys()): - del stacked_image.coords[c] - image = xr.concat((image, stacked_image), dim="c") - else: - n_channels1 = 0 - for img in image: - n_channels1 += np.prod(img.shape[:-2]) - n_channels1 = int(n_channels1) - for f in stacked_file_list: - array, _, _, _ = _read_ome_zarr_array(f) - image.append(array) + stacked_image = _read_image(stacked_file_list, stacked_metadata) + n_channels1 = image.sizes["c"] + # clear coords to avoid issues with xr.concat + for c in list(image.coords.keys()): + del image.coords[c] + for c in list(stacked_image.coords.keys()): + del stacked_image.coords[c] + image = xr.concat((image, stacked_image), dim="c") for label_name in label_name_to_features: features = label_name_to_features[label_name] @@ -297,6 +268,7 @@ def single_feature( if zarr_labels is None: logger.info(f"Unable to read {label_name} labels for {image_key}.") continue + label_image = da.from_zarr(zarr_labels) label_prefix = _label_name_to_prefix[label_name] merged_df = None if objects_dir is not None: @@ -310,7 +282,7 @@ def single_feature( if merged_df is None: logger.info(f"Find {label_name} objects for {image_key}.") - merged_df = find_objects(zarr_labels) + merged_df = find_objects(label_image) objects_path = f"{output_dir}{output_sep}{label_name}{output_sep}{image_key}-objects.parquet" merged_df.index.name = "label" merged_df.columns = f"{label_prefix}_" + merged_df.columns @@ -350,7 +322,7 @@ def single_feature( ): c = tokens[token_index] if c[0] == "s": - if c in channel_names: + if channel_names is not None and c in channel_names: channel_names[str(n_channels1 + int(c[1:]))] = ( channel_names.pop(c) ) @@ -398,8 +370,8 @@ def single_feature( merged_df[c] = merged_df[c].astype(int) df = label_features( objects_df=merged_df, - label_image=zarr_labels if zarr_inputs else da.from_zarr(zarr_labels), - intensity_image=image if zarr_inputs else image.data, + label_image=label_image, + intensity_image=image.data, features=features, normalize=normalize, bounding_box_columns=[ @@ -510,9 +482,18 @@ def run_pipeline_compute_features(arguments: argparse.Namespace) -> None: if features_plot is None: features_plot = [] if dask_server_url is None and arguments.dask_cluster is None: - dask_cluster_parameters = _dask_workers_threads( - threads_per_worker=4 if "sizeshape" in unique_features else 1 - ) + if dask_server_url is None and arguments.dask_cluster is None: + threads_per_worker = 1 + if "sizeshape" in unique_features: + threads_per_worker = 4 + else: + for feature in unique_features: + if feature.startswith("correlationpearsonbox"): + threads_per_worker = 2 + break + dask_cluster_parameters = _dask_workers_threads( + threads_per_worker=threads_per_worker + ) objects_dir_sep = None if objects_dir is not None: diff --git a/scallops/cli/find_objects.py b/scallops/cli/find_objects.py index 88bf71e0..04e7ed53 100644 --- a/scallops/cli/find_objects.py +++ b/scallops/cli/find_objects.py @@ -8,6 +8,7 @@ import argparse import json +import dask.array as da import fsspec import zarr from zarr import Group @@ -49,7 +50,7 @@ def _execute( logger.info(f"Skipping finding objects for {metadata['id']}.") return logger.info(f"Finding objects for {metadata['id']}.") - array = file_list[0][list(file_list[0].keys())[0]] + array = da.from_zarr(file_list[0][list(file_list[0].keys())[0]]) df = find_objects(array) df.index.name = "label" prefix = _label_name_to_prefix.get(label_name) diff --git a/scallops/cli/pooled_if_sbs.py b/scallops/cli/pooled_if_sbs.py index 69696e51..57d7464e 100644 --- a/scallops/cli/pooled_if_sbs.py +++ b/scallops/cli/pooled_if_sbs.py @@ -179,8 +179,7 @@ def _peaks_to_bases( def spot_detection_pipeline( image_tuple: tuple[tuple[str, ...], list[str], dict], iss_channels: list[int], - file_separator: str, - root: zarr.Group | str, + output: str, max_filter_width: int, sigma_log: float | list[float], z_index: int | str, @@ -226,17 +225,20 @@ def spot_detection_pipeline( """ _, file_list, metadata = image_tuple image_key = metadata["id"] + output_fs = fsspec.url_to_fs(output)[0] + output_sep = output_fs.sep + output = output.rstrip(output_sep) + points_path = f"{output}{output_sep}points" + points_protocol = _get_fs_protocol(output_fs) + if points_protocol != "file": + points_path = f"{points_protocol}://{points_path}" + peaks_path = f"{points_path}{output_sep}{image_key}-peaks.parquet" + if not force: - points_path = ( - f"{_get_store_path(root).rstrip(_get_sep(root))}{_get_sep(root)}points" - ) - points_protocol = _get_fs_protocol(_get_fs(root)) - if points_protocol != "file": - points_path = f"{points_protocol}://{points_path}" - peaks_path = f"{points_path}{_get_sep(root)}{image_key}-peaks.parquet" if is_parquet_file(peaks_path): logger.info(f"Skipping spot detection for {image_key}") return [] + root = open_ome_zarr(output, mode="a") image = _images2fov(file_list, metadata, dask=True) image = _z_projection(image, z_index) if expected_cycles is not None: @@ -325,7 +327,6 @@ def spot_detection_pipeline( root=root, image=loged, output_format=output_image_format, - file_separator=file_separator, zarr_format="zarr", compute=compute, ) @@ -340,7 +341,6 @@ def spot_detection_pipeline( root=root, image=std_arr, output_format=output_image_format, - file_separator=file_separator, metadata=dict(parent=image_key), compute=compute, ) @@ -355,7 +355,6 @@ def spot_detection_pipeline( root=root, image=maxed, output_format=output_image_format, - file_separator=file_separator, zarr_format="zarr", compute=compute, ) @@ -363,16 +362,10 @@ def spot_detection_pipeline( else: del maxed if "peaks" in save_keys: - points_path = ( - f"{_get_store_path(root).rstrip(_get_sep(root))}{_get_sep(root)}points" - ) - protocol = _get_fs_protocol(_get_fs(root)) - if protocol != "file": - points_path = f"{protocol}://{points_path}" - _get_fs(root).makedirs(points_path, exist_ok=True) - peaks_path = f"{points_path}{_get_sep(root)}{image_key}-peaks.parquet" - if _get_fs(root).exists(peaks_path): - _get_fs(root).rm(peaks_path, recursive=True) + output_fs.makedirs(points_path, exist_ok=True) + + if output_fs.exists(peaks_path): + output_fs.rm(peaks_path, recursive=True) dask_delayed.append( _to_parquet( @@ -863,7 +856,6 @@ def spot_detect_main(arguments: argparse.Namespace): chunks = (chunks, chunks) output = _add_suffix(output, ".zarr") - root = open_ome_zarr(output, mode="a") exp_gen = _set_up_experiment(images, image_pattern, group_by, subset=subset) with ( _create_default_dask_config(), @@ -874,8 +866,7 @@ def spot_detect_main(arguments: argparse.Namespace): delayed_results += spot_detection_pipeline( img, iss_channels=channels, - file_separator=None, - root=root, + output=output, z_index=z_index, output_image_format="zarr", max_filter_width=max_filter_width, @@ -1291,7 +1282,7 @@ def reads_main(arguments: argparse.Namespace): for key in image_keys: reads_pipeline( key, - spots_root=zarr.open(spots, "r"), + spots_root=zarr.open(spots, mode="r"), labels_root=zarr.open(labels + labels_fs.sep + "labels", mode="r"), barcodes_file=barcodes_file, file_separator=output_fs.sep, diff --git a/scallops/cli/register.py b/scallops/cli/register.py index 4bbb0c12..1f1a6e17 100644 --- a/scallops/cli/register.py +++ b/scallops/cli/register.py @@ -32,6 +32,7 @@ pluralize, ) from scallops.registration.itk import ( + _get_chunk_size, _itk_align_reference_time_zarr, _itk_transform_image_zarr, _load_itk_parameters, @@ -393,6 +394,7 @@ def single_registration( output_root=label_output_root, flip_y=flip_moving_y, flip_x=flip_moving_x, + chunksize=chunksize, ) else: # align to t=reference_timepoint @@ -439,8 +441,8 @@ def single_registration( parameter_object_across_channels=parameter_object_across_channels, ) moving_image_attrs = moving_image[0].attrs.copy() + chunk_size = _get_chunk_size(moving_image[reference_timepoint]) del moving_image - if len(moving_label_keys) > 0: _transform_labels_t( image_key=image_key, @@ -450,6 +452,7 @@ def single_registration( moving_image_attrs=moving_image_attrs, moving_image_spacing=moving_image_spacing, label_output_root=label_output_root, + chunk_size=chunk_size, ) return image_key @@ -463,6 +466,7 @@ def _transform_labels_t( moving_image_attrs, moving_image_spacing, label_output_root, + chunk_size, ): # transform_dest structure is image_key/t=1 # assume labels are named image_key-t-suffix @@ -495,6 +499,7 @@ def _transform_labels_t( output_names=output_names, moving_image_spacing=moving_image_spacing, output_root=label_output_root, + chunk_size=chunk_size, ) @@ -588,6 +593,7 @@ def _transform_labels( attrs: None | dict, flip_y: bool = False, flip_x: bool = False, + chunksize: tuple[int, int] | None = None, ): """Transform and save labels. @@ -622,12 +628,13 @@ def _transform_labels( image_spacing=moving_image_spacing, ) del array - + # match chunk size of fixed image _write_zarr_image( name=output_names[i], root=output_root, image=transformed_array, group="labels", + storage_options=dict(chunks=chunksize) if chunksize is not None else None, ) @@ -693,6 +700,7 @@ def single_transformix( output_root=output_root, moving_image_spacing=image_spacing, attrs=None, + chunksize=None, ) else: # see if transform dir has subdirectory describing channel transformation diff --git a/scallops/cli/util.py b/scallops/cli/util.py index 0e05501f..44fc06ec 100644 --- a/scallops/cli/util.py +++ b/scallops/cli/util.py @@ -26,6 +26,7 @@ import numpy as np import xarray as xr import zarr +from dask.delayed import Delayed from distributed import Client from scallops.io import save_ome_tiff @@ -199,11 +200,11 @@ def _write_image( root: zarr.Group | str, image: np.ndarray | xr.DataArray | da.Array, output_format: str, - file_separator: str, + file_separator: str = "/", metadata: dict | None = None, compute: bool = True, **kwargs, -) -> None: +) -> list[Delayed]: """Write image data to Zarr or TIFF format. :param name: Name of the image. @@ -215,7 +216,7 @@ def _write_image( :param compute: Whether to compute the Dask array before saving. """ if output_format == "zarr": - _write_zarr_image( + return _write_zarr_image( name=name, root=root, image=image, diff --git a/scallops/features/generate.py b/scallops/features/generate.py index 59b9234b..35b71ff0 100644 --- a/scallops/features/generate.py +++ b/scallops/features/generate.py @@ -121,6 +121,39 @@ def _create_dd_metadata( ).fillna(0) +def _rewrite_channels( + funcs, all_required_channels: np.ndarray, channel_names: list[str] +) -> np.ndarray: + channel_names_subset = [] + channel_map = dict() + for i in range(len(all_required_channels)): + original_channel_index = all_required_channels[i] + channel_map[original_channel_index] = i + channel_names_subset.append(channel_names[original_channel_index]) + for f in funcs: + for key in f.keywords: + if key in ["c", "c1", "c2"]: + val = f.keywords[key] + if isinstance(val, int): + f.keywords[key] = channel_map[val] + elif isinstance(val, (list, tuple)): + new_val = [] + + for v in val: + if isinstance(v, int): + new_val.append(channel_map[v]) + elif isinstance(v, tuple): + new_val.append((channel_map[v[0]], channel_map[v[1]])) + else: + raise ValueError() + if isinstance(val, tuple): + new_val = tuple(new_val) + f.keywords[key] = new_val + else: + raise ValueError(f"{val} is not a valid type.") + return channel_names_subset + + def label_features( objects_df: pd.DataFrame, label_image: da.Array | zarr.Array, @@ -199,10 +232,17 @@ def label_features( if channel_index < 0 or channel_index >= nchannels: raise ValueError("Channel index out of range") channel_names_[channel_index] = channel_name - funcs, requires_intensity_image = _create_funcs( + funcs, all_required_channels = _create_funcs( features=features, n_channels=len(channel_names_) ) - if not requires_intensity_image: # Don't pass intensity image if not needed + if is_dask_array and len(all_required_channels) != len(channel_names_): + channel_names_ = _rewrite_channels( + funcs, + all_required_channels=all_required_channels, + channel_names=channel_names_, + ) + intensity_image = intensity_image[..., all_required_channels] + if len(all_required_channels) == 0: # Don't pass intensity image if not needed intensity_image = None intensity_image_delayed = ( delayed(intensity_image) if not isinstance(intensity_image, da.Array) else None @@ -384,20 +424,18 @@ def normalize_features(features: Iterable[str]) -> set[str]: def _create_funcs( features: Iterable[str], n_channels: Sequence[str], -) -> tuple[list[Callable], bool]: +) -> tuple[list[Callable], np.ndarray[int]]: """Create feature functionss. :param features: Iterable of feature names to be processed. :param n_channels: Number of channels in image :return: A tuple containing: - list of partial functions corresponding to the provided features. - - boolean indicating whether any feature requires intensity image + - set of required intensity image channel indices. """ funcs = [] - requires_intensity_image = False - features_dict = _features.copy() func_tuple_to_params = defaultdict(lambda: []) # (func_name, non-channel parameter values) @@ -412,7 +450,7 @@ def _create_funcs( if param_name not in ("c", "c1", "c2"): key.append(params[param_name]) func_tuple_to_params[tuple(key)].append(params) - + all_required_channels = set() for func_tuple, params_list in func_tuple_to_params.items(): func_name = func_tuple[0] if "c" in params_list[0]: @@ -420,7 +458,7 @@ def _create_funcs( for params in params_list: for val in params["c"]: channels.add(val) - requires_intensity_image = True + all_required_channels.update(channels) new_params = dict(c=tuple(sorted(channels))) for p in params_list[0]: if p != "c": @@ -428,7 +466,6 @@ def _create_funcs( f = partial(features_dict[func_name], **new_params) funcs.append(f) elif "c1" in params_list[0]: - requires_intensity_image = True seen_symmetric = set() # c1 always < c2 for params in params_list: @@ -447,6 +484,9 @@ def _create_funcs( additional_params[p] = params_list[0][p] rewrite_func = _features_rewrite.get(func_name) if len(seen_symmetric) > 0: + for c1, c2 in seen_symmetric: + all_required_channels.add(c1) + all_required_channels.add(c2) if rewrite_func is not None: new_params = dict(c=list(seen_symmetric)) new_params.update(additional_params) @@ -462,8 +502,9 @@ def _create_funcs( for params in params_list: f = partial(features_dict[func_name], **params) funcs.append(f) - - return funcs, requires_intensity_image + all_required_channels = np.array(list(all_required_channels)) + all_required_channels.sort() + return funcs, all_required_channels def _get_params( diff --git a/scallops/features/util.py b/scallops/features/util.py index 05c68b7e..d10c64b1 100644 --- a/scallops/features/util.py +++ b/scallops/features/util.py @@ -99,22 +99,26 @@ def _slice_anndata( if var is not None: var_indices = _normalize_index(var, data.var.index) X = data.X - layers = dict(data.layers) - obsm = dict(data.obsm) - varm = dict(data.varm) + layers = dict() + obsm = dict() + varm = dict() + for key in data.layers.keys(): + if key is not None: + layers[key] = data.layers[key] if obs_indices is not None: X = X[obs_indices] - for key in layers: - layers[key] = layers[key][obs_indices] - for key in obsm: - obsm[key] = obsm[key][obs_indices] + for key in layers.keys(): + if key is not None: + layers[key] = layers[key][obs_indices] + for key in data.obsm.keys(): + obsm[key] = data.obsm[key][obs_indices] if var_indices is not None: X = X[:, var_indices] - for key in layers: + for key in layers.keys(): layers[key] = layers[key][:, var_indices] - for key in varm: - varm[key] = varm[key][var_indices] + for key in data.varm.keys(): + varm[key] = data.varm[key][var_indices] obs = data.obs.iloc[obs_indices] if obs_indices is not None else data.obs var = data.var.iloc[var_indices] if var_indices is not None else data.var return anndata.AnnData(X=X, obs=obs, var=var, layers=layers, obsm=obsm, varm=varm) diff --git a/scallops/io.py b/scallops/io.py index b9093da9..8d43e183 100644 --- a/scallops/io.py +++ b/scallops/io.py @@ -55,7 +55,6 @@ from xarray.core.utils import equivalent from zarr.storage import StoreLike -from scallops._bioio_zarr_reader import ScallopsZarrReader from scallops.experiment.elements import Experiment, _LazyLoadData from scallops.externals.tifffile2014 import imsave from scallops.utils import forceTCZYX, mlcs @@ -235,7 +234,9 @@ def _create_image(path: str, **kwargs) -> bioio.BioImage: base_path_lc, ext = os.path.splitext(path_lc) if "reader" not in img_args: if ext in ["", ".zarr", "/", ".zarr/"]: - img_args["reader"] = ScallopsZarrReader + import bioio_ome_zarr + + img_args["reader"] = bioio_ome_zarr.Reader elif ext in [".tiff", ".tif"] and os.path.splitext(base_path_lc)[1] != ".ome": img_args["reader"] = bioio_tifffile.Reader return bioio.BioImage(path, **img_args) diff --git a/scallops/registration/itk.py b/scallops/registration/itk.py index 7bad5abe..f976ddfd 100644 --- a/scallops/registration/itk.py +++ b/scallops/registration/itk.py @@ -31,9 +31,8 @@ from scallops.io import _download_file, _get_fs_protocol, get_image_spacing from scallops.registration.landmarks import _get_translation, find_landmarks -from scallops.utils import _dask_from_array_no_copy from scallops.xr import _get_dims -from scallops.zarr_io import open_ome_zarr, write_zarr +from scallops.zarr_io import _create_array_kwargs, open_ome_zarr, write_zarr logger = logging.getLogger("scallops") @@ -252,7 +251,10 @@ def _array_to_itk( ) itk_image = itk.image_view_from_array(data) - + if ( + len(spacing) == 3 + ): # assume image had z-axis that was max-projected and z-axis is no longer present + spacing = spacing[1], spacing[2] itk_image.SetSpacing(spacing) return itk_image @@ -324,23 +326,24 @@ def _init_callback(init_params: dict[str, Any]) -> dict[str, Any]: attrs = init_params["attrs"] dtype = init_params["dtype"] chunk_size = init_params["chunk_size"] - zarr_dataset = None + zarr_array = None group = None if image_root is not None: images_group = image_root.require_group("images", overwrite=False) group = images_group.create_group( image_name.replace("/", "-"), overwrite=True ) - zarr_dataset = group.create_dataset( - "0", + zarr_array = group.create_array( + "s0", shape=shape, chunks=(1,) * (len(shape) - 2) + chunk_size, dtype=dtype, overwrite=True, + **_create_array_kwargs(), ) return { - "data": zarr_dataset, + "data": zarr_array, "group": group, "dims": dims, "coords": coords, @@ -352,31 +355,33 @@ def done(d: dict[str, Any]): :param d: Dictionary containing dataset and metadata information. """ - data = d["data"] + zarr_array = d["data"] group = d["group"] dims = d["dims"] coords = d["coords"] image_attrs = d["attrs"] - if data is not None: + if zarr_array is not None: write_zarr( grp=group, - data=data, + data=zarr_array, image_attrs=image_attrs, coords=coords, dims=dims, zarr_format="zarr", ) - def _write_callback(x, idx, val): - if x is None: + def _write_callback(zarr_array, idx, val): + if zarr_array is None: # do not save image return - if isinstance(idx, int): - idx = (idx,) + if isinstance(val, xr.DataArray): val = val.data - if not isinstance(val, da.Array): - val = _dask_from_array_no_copy(val, chunks=x.chunks[-2:]) - da.store(val, x, regions=idx, compute=True) + if isinstance(val, da.Array): + if isinstance(idx, int): + idx = (idx,) + da.store(val, zarr_array, regions=idx, compute=True) + else: + zarr_array[idx] = val _itk_align_reference_time( moving_image=moving_image, @@ -676,7 +681,7 @@ def _write_value(x, idx, val): ) ) - result_data = init_dict["data"] + result_array = init_dict["data"] output_fs = fsspec.core.url_to_fs(output_dir)[0] if output_dir is not None else None unrolled_t_index = 0 @@ -837,7 +842,7 @@ def _write_value(x, idx, val): index = (i, j) if not unroll_channels else unrolled_t_index + j logger.info(f"Writing t={i}, c={j}.") - write_callback(result_data, index, image_i_j) + write_callback(result_array, index, image_i_j) del image_i_j del transform_parameter_object @@ -1182,12 +1187,13 @@ def _itk_transform_image_zarr( image_name.replace("/", "-"), overwrite=True ) chunks = (1,) * len(transform_dims) + (chunksize or (1024, 1024)) - data = group.create_dataset( - "0", + data = group.create_array( + "s0", shape=dim_sizes + output_size, chunks=chunks, dtype=image.dtype, overwrite=True, + **_create_array_kwargs(), ) _itk_transform_image( diff --git a/scallops/stitch/_align.py b/scallops/stitch/_align.py index c5ee0ed8..a37fe484 100644 --- a/scallops/stitch/_align.py +++ b/scallops/stitch/_align.py @@ -386,7 +386,7 @@ def stitch_align( original_tile_shape, crop_width, ) - logger.info(f"Crop for writing: {fuse_crop_width}.") + logger.info(f"Crop for writing: {fuse_crop_width[0], fuse_crop_width[1]}.") if evaluate_stitching: # Global evaluation logger.info( diff --git a/scallops/stitch/_stitch.py b/scallops/stitch/_stitch.py index fb3eefe8..1911798c 100644 --- a/scallops/stitch/_stitch.py +++ b/scallops/stitch/_stitch.py @@ -6,7 +6,6 @@ from collections.abc import Sequence from typing import Literal -import dask.array as da import fsspec import numpy as np import pandas as pd @@ -14,13 +13,12 @@ import pyarrow.parquet as pq import zarr from sklearn.cluster import AgglomerativeClustering -from zarr.errors import PathNotFoundError from scallops.cli.util import _get_cli_logger, cli_metadata from scallops.io import is_parquet_file, read_image from scallops.stitch._align import stitch_align from scallops.stitch._plots import _qc_report -from scallops.stitch.fuse import _create_label_ome_metadata, _create_ome_metadata, _fuse +from scallops.stitch.fuse import _fuse from scallops.stitch.shift_utils import _zncc, convert_stage_positions from scallops.stitch.utils import ( _download_path, @@ -31,8 +29,13 @@ tile_overlap_mask, tile_source_labels, ) -from scallops.utils import _dask_from_array_no_copy -from scallops.zarr_io import _da_to_zarr_kwargs, _omero_channels, is_ome_zarr_array +from scallops.zarr_io import ( + _attrs_axes_scales, + _create_array_kwargs, + _create_zarr_attrs, + _current_format, + is_ome_zarr_array, +) logger = _get_cli_logger() @@ -87,14 +90,14 @@ def _single_stitch( if is_ome_zarr_array(image_output_root.get(f"images/{image_key}")): logger.info(f"Skipping stitching for {image_key}.") return - except PathNotFoundError: + except: # noqa: E722 pass elif not no_save_labels: try: if is_ome_zarr_array(image_output_root.get(f"labels/{image_key}-mask")): logger.info(f"Skipping stitching for {image_key}.") return - except PathNotFoundError: + except: # noqa: E722 pass elif is_parquet_file(f"{other_output_path}{image_key}-positions.parquet"): logger.info(f"Skipping stitching for {image_key}.") @@ -164,7 +167,8 @@ def _single_stitch( if isinstance(image_spacing, np.ndarray) else image_spacing ) - + output_metadata["physical_pixel_sizes"] = output_metadata["image_spacing"] + output_metadata["physical_pixel_units"] = ["micrometer", "micrometer"] auto_radial_correction = radial_correction_k == "auto" stitch_result = stitch_align( filepaths=filepaths, @@ -378,11 +382,8 @@ def _single_stitch( blend, image_output_root, image_key, - fused_y_size, - fused_x_size, fused_tile_shape, chunk_size, - image_spacing, no_save_labels, no_save_image, ffp_path, @@ -408,11 +409,8 @@ def _write_arrays( blend, image_output_root, image_key, - fused_y_size, - fused_x_size, fused_tile_shape, chunk_size, - image_spacing, no_save_labels, no_save_image, ffp_path, @@ -429,60 +427,46 @@ def _write_arrays( channel_filter_percentiles, ): gc.collect() + fmt = _current_format() if not no_save_labels: labels_group = image_output_root.require_group("labels") group = labels_group.create_group(image_key + "-mask", overwrite=True) - zarr_kwargs = _da_to_zarr_kwargs() - array = group.create_dataset( - name="0", - shape=(fused_y_size, fused_x_size), + group.create_array( + name="s0", + data=tile_overlap_mask( + stitch_positions_df, + fill=blend != "none", + tile_shape=fused_tile_shape, + ), chunks=chunk_size, - dtype=np.uint8, overwrite=True, - **zarr_kwargs, + **_create_array_kwargs(fmt=fmt), ) - - da.store( - _dask_from_array_no_copy( - tile_overlap_mask( - stitch_positions_df, - fill=blend != "none", - tile_shape=fused_tile_shape, - ), - chunks=chunk_size, - ), - array, - compute=True, + image_attrs, axes, scale_dict = _attrs_axes_scales( + metadata, None, ["y", "x"], np.uint8 ) - group.attrs.update( - _create_label_ome_metadata(image_spacing, image_key + "-mask") + zarr_attrs = _create_zarr_attrs( + fmt, group, ["y", "x"], image_attrs, axes, scale_dict ) + group.attrs.update(zarr_attrs) + if blend == "none": group = labels_group.create_group(image_key + "-tile", overwrite=True) - array = group.create_dataset( - name="0", - shape=(fused_y_size, fused_x_size), + group.create_array( + name="s0", + data=tile_source_labels(stitch_positions_df, fused_tile_shape), chunks=chunk_size, - dtype=np.uint16, overwrite=True, - **zarr_kwargs, + **_create_array_kwargs(fmt=fmt), ) - da.store( - _dask_from_array_no_copy( - tile_source_labels(stitch_positions_df, fused_tile_shape), - chunks=chunk_size, - ), - array, - compute=True, + image_attrs, axes, scale_dict = _attrs_axes_scales( + metadata, None, ["y", "x"], np.uint16 ) - label_metadata = _create_label_ome_metadata( - image_spacing, image_key + "-tile" + zarr_attrs = _create_zarr_attrs( + fmt, group, ["y", "x"], image_attrs, axes, scale_dict ) - label_metadata["multiscales"][0]["metadata"] = { - "source": f"../../images/{image_key}" - } - group.attrs.update(label_metadata) + group.attrs.update(zarr_attrs) cleanup_paths = [] if not no_save_image: group = image_output_root.require_group("images").require_group( @@ -536,12 +520,14 @@ def _write_arrays( dfp = dfp.values - tile_channel_names = _fuse( + _fuse( df=stitch_positions_df_local, group=group, z_index=z_index, blend=blend, output_channels=output_channels, + channel_names=channel_names, + output_metadata=metadata, ffp=ffp, dfp=dfp, crop_width=fuse_crop_width, @@ -553,24 +539,6 @@ def _write_arrays( channel_filter_percentiles=channel_filter_percentiles, ) - ome_metadata = _create_ome_metadata( - image_key=image_key, - stitch_coords=stitch_positions_df, - **metadata, - ) - # Fall back to the channel names read from the source tiles when the user - # did not explicitly pass --channel-name. - if channel_names is None: - channel_names = tile_channel_names - if channel_names is not None: - if output_channels is not None: - channel_names = channel_names[output_channels] - channel_names = [ - name.replace("_", "-").replace(" ", "-") for name in channel_names - ] - ome_metadata["omero"] = _omero_channels(channel_names) - group.attrs.update(ome_metadata) - # cleanup for path in cleanup_paths: if os.path.exists(path): diff --git a/scallops/stitch/fuse.py b/scallops/stitch/fuse.py index 80e101df..c516039c 100644 --- a/scallops/stitch/fuse.py +++ b/scallops/stitch/fuse.py @@ -4,10 +4,9 @@ import os import threading from collections.abc import Sequence -from typing import Literal +from typing import Any, Literal import dask -import dask.array as da import numpy as np import pandas as pd import psutil @@ -24,89 +23,25 @@ from scallops.stitch._radial import radial_correct from scallops.stitch.shift_utils import calc_best_shift from scallops.stitch.utils import _crop_image, dtype_convert -from scallops.utils import _cpu_count, _dask_from_array_no_copy -from scallops.zarr_io import _da_to_zarr_kwargs +from scallops.utils import _cpu_count +from scallops.zarr_io import ( + _attrs_axes_scales, + _create_array_kwargs, + _create_zarr_attrs, + _current_format, +) logger = logging.getLogger("scallops") -def _create_label_ome_metadata(image_spacing: tuple[float, float], label_name: str): - return { - "multiscales": [ - { - "axes": [ - {"name": "y", "type": "space", "unit": "micrometer"}, - {"name": "x", "type": "space", "unit": "micrometer"}, - ], - "datasets": [ - { - "coordinateTransformations": [ - { - "scale": [ - float(image_spacing[0]), - float(image_spacing[1]), - ], - "type": "scale", - } - ], - "path": "0", - } - ], - "name": f"/labels/{label_name}", - "version": "0.4", - } - ] - } - - -def _create_ome_metadata( - image_spacing: tuple[float, float], - stitch_coords: pd.DataFrame, - image_key: str, - **kwargs, -): - metadata = {} - metadata.update(**kwargs) - metadata["stitch_coords"] = dict() - for c in stitch_coords: # convert to dict - metadata["stitch_coords"][c] = stitch_coords[c].to_list() - return { - "multiscales": [ - { - "metadata": metadata, - "axes": [ - {"name": "c", "type": "channel"}, - {"name": "y", "type": "space", "unit": "micrometer"}, - {"name": "x", "type": "space", "unit": "micrometer"}, - ], - "datasets": [ - { - "coordinateTransformations": [ - { - "scale": [ - 1.0, - float(image_spacing[0]), - float(image_spacing[1]), - ], - "type": "scale", - } - ], - "path": "0", - } - ], - "name": f"/images/{image_key}", - "version": "0.4", - } - ] - } - - def _fuse( df: pd.DataFrame, group: zarr.Group, z_index: int | Literal["max"] | None = None, blend: Literal["none", "linear"] = "none", output_channels: Sequence[int] | None = None, + channel_names: Sequence[str] | None = None, + output_metadata: dict[str, Any] | None = None, ffp: np.ndarray | None = None, dfp: np.ndarray | None = None, crop_width: tuple[int, int] | None = None, @@ -117,7 +52,7 @@ def _fuse( channel_window: int = 2, channel_cross_correlation_upsample: int = 2, channel_filter_percentiles: tuple[float, float] | None = None, -) -> list[str] | None: +): """Use stitching coordinates to fuse tiles. :param df: Stitch output dataframe. @@ -125,6 +60,8 @@ def _fuse( :param z_index: z-index or 'max'. Ignored when df contains `z_index` column. :param blend: Blending mode :param output_channels: Optional output channels to include + :param channel_names: Optional channel names to output, overwriting channel names derived from image tiles + :param output_metadata: Optional output metadata :param ffp: ffp image for illumination correction :param dfp: dfp image for illumination correction :param crop_width: Image crop width @@ -139,8 +76,6 @@ def _fuse( :param channel_cross_correlation_upsample: Perform subpixel alignment for registration across channels if greater than one :param channel_filter_percentiles: Replace data outside of percentile range [q1, q2] with uniform noise over the range [q1,q2] for registration across channels. - :return: Channel names read from the source tiles (e.g. from nd2 metadata), or - None when the tiles carry no named channel coordinate. """ assert blend in ["none", "linear"] if crop_width is not None and crop_width[0] <= 0 and crop_width[1] <= 0: @@ -243,8 +178,8 @@ def _fuse( locks.append(threading.Lock()) locks = np.array(locks) partition_tree = shapely.STRtree(partition_boxes) - - result = group.create_dataset( + fmt = _current_format() + result_zarr_array = group.create_array( shape=( len(output_channels), # c fused_y_size, @@ -252,9 +187,9 @@ def _fuse( ), dtype=target_dtype, chunks=(1,) + chunk_size, - name="0", + name="s0", overwrite=True, - **_da_to_zarr_kwargs(), + **_create_array_kwargs(fmt), ) _fuse_image_delayed = delayed(_fuse_image) @@ -392,17 +327,30 @@ def _fuse( target = target / weights_sum target = np.round(target).astype(target_dtype) - with ProgressBar(): - logger.info("Writing to disk.") - target = _dask_from_array_no_copy(target, chunks=(1,) + chunk_size) - da.store( - target, - result, - regions=(slice(channel_batch, channel_batch + channels_per_batch),), - compute=True, - ) + sl = slice(channel_batch, channel_batch + channels_per_batch) + + result_zarr_array[sl] = target - return tile_channel_names + # Fall back to the channel names read from the source tiles when the user + # did not explicitly pass --channel-name. + if channel_names is None: + channel_names = tile_channel_names + if channel_names is not None: + channel_names = [channel_names[index] for index in output_channels] + channel_names = [ + name.replace("_", "-").replace(" ", "-") for name in channel_names + ] + + image_attrs, axes, scale_dict = _attrs_axes_scales( + output_metadata, + {"c": channel_names}, + ["c", "y", "x"], + np.uint16, + ) + zarr_attrs = _create_zarr_attrs( + fmt, group, ["c", "y", "x"], image_attrs, axes, scale_dict + ) + group.attrs.update(zarr_attrs) def _register_across_channels( diff --git a/scallops/stitch/utils.py b/scallops/stitch/utils.py index 44a97d4f..e5bf4374 100644 --- a/scallops/stitch/utils.py +++ b/scallops/stitch/utils.py @@ -328,19 +328,18 @@ def get_tile_position(image: bioio.BioImage, image_index: int = 0): ome_metadata = _get_ome(image) physical_size_y_unit = None physical_size_x_unit = None - if ome_metadata is not None: - values = [ - ome_metadata.images[image_index].pixels.planes[0].position_y, - ome_metadata.images[image_index].pixels.planes[0].position_x, - ] - physical_size_y_unit = ( - ome_metadata.images[image_index].pixels.planes[0].position_y_unit.value - ) - physical_size_x_unit = ( - ome_metadata.images[image_index].pixels.planes[0].position_x_unit.value - ) - elif "multiscales" in image.metadata: - metadata = image.metadata["multiscales"][0]["metadata"] + values = None + if ome_metadata is not None and image_index < len(ome_metadata.images): + img = ome_metadata.images[image_index] + if len(img.pixels.planes) > 0: + values = [ + img.pixels.planes[0].position_y, + img.pixels.planes[0].position_x, + ] + physical_size_y_unit = img.pixels.planes[0].position_y_unit.value + physical_size_x_unit = img.pixels.planes[0].position_x_unit.value + if values is None and "multiscales" in image.metadata.attributes: + metadata = image.metadata.attributes["multiscales"][0]["metadata"] values = [metadata["position_y"], metadata["position_x"]] physical_size_y_unit = metadata["position_y_unit"] physical_size_x_unit = metadata["position_x_unit"] @@ -420,8 +419,8 @@ def _pixel_size_from_image(image: bioio.BioImage) -> np.array: ] physical_size_y_unit = ome_metadata.images[0].pixels.physical_size_y_unit.value physical_size_x_unit = ome_metadata.images[0].pixels.physical_size_x_unit.value - elif "multiscales" in image.metadata: - metadata = image.metadata["multiscales"][0]["metadata"] + elif "multiscales" in image.metadata.attributes: + metadata = image.metadata.attributes["multiscales"][0]["metadata"] values = [metadata["physical_size_y"], metadata["physical_size_x"]] physical_size_y_unit = metadata["physical_size_y_unit"] physical_size_x_unit = metadata["physical_size_x_unit"] diff --git a/scallops/tests/test_features.py b/scallops/tests/test_features.py index cc61e832..22115398 100644 --- a/scallops/tests/test_features.py +++ b/scallops/tests/test_features.py @@ -9,6 +9,7 @@ import pandas as pd import pytest import shapely +import xarray as xr import zarr from skimage.measure import regionprops @@ -23,7 +24,7 @@ _other_features_single_channel, ) from scallops.features.find_objects import find_objects -from scallops.features.generate import _create_funcs, label_features +from scallops.features.generate import _create_funcs, _rewrite_channels, label_features from scallops.features.spots import spot_count from scallops.features.texture import pftas from scallops.io import read_image, to_label_crops @@ -112,12 +113,14 @@ def test_to_label_crops(tmp_path, array_A1_102_cells, array_A1_102_alnpheno): assert len(result_df) == 1 and result_df.index.values[0] == 2603 group = zarr.group() - intensity_image_zarr = group.create_dataset( - name="image", shape=intensity_image.shape + intensity_image_zarr = group.create_array( + name="image", shape=intensity_image.shape, dtype=intensity_image.dtype ) intensity_image_zarr[:] = intensity_image.compute() - label_image_zarr = group.create_dataset(name="label", shape=label_image.shape) + label_image_zarr = group.create_array( + name="label", shape=label_image.shape, dtype=label_image.dtype + ) label_image_zarr[:] = label_image.compute() to_label_crops( @@ -406,19 +409,60 @@ def test_features_dask(array_A1_102_cells, array_A1_102_pheno): ) +@pytest.mark.features +def test_rewrite_channels(): + funcs, all_required_channels = _create_funcs(["colocalization_1_5"], 10) + assert funcs[0].keywords["c"] == [(1, 5)] + channel_names = _rewrite_channels( + funcs, + all_required_channels=all_required_channels, + channel_names=np.arange(10).astype(str), + ) + assert channel_names == ["1", "5"] + assert funcs[0].keywords["c"] == [(0, 1)] + + funcs, all_required_channels = _create_funcs(["intensity_5,6"], 10) + assert funcs[0].keywords["c"] == (5, 6) + channel_names = _rewrite_channels( + funcs, + all_required_channels=all_required_channels, + channel_names=np.arange(10).astype(str), + ) + assert channel_names == ["5", "6"] + assert funcs[0].keywords["c"] == (0, 1) + + funcs, all_required_channels = _create_funcs(["correlationpearsonbox_5_6,8"], 10) + assert len(funcs) == 2 + assert funcs[0].keywords["c1"] == 5 + assert funcs[0].keywords["c2"] == 6 + assert funcs[1].keywords["c1"] == 5 + assert funcs[1].keywords["c2"] == 8 + channel_names = _rewrite_channels( + funcs, + all_required_channels=all_required_channels, + channel_names=np.arange(10).astype(str), + ) + + assert channel_names == ["5", "6", "8"] + assert funcs[0].keywords["c1"] == 0 + assert funcs[0].keywords["c2"] == 1 + assert funcs[1].keywords["c1"] == 0 + assert funcs[1].keywords["c2"] == 2 + + @pytest.mark.features def test_create_funcs(): - funcs, requires_intensity = _create_funcs(["colocalization_*_*"], 3) - assert requires_intensity + funcs, all_required_channels = _create_funcs(["colocalization_*_*"], 3) + assert len(all_required_channels) > 0 assert len(funcs) == 1 assert funcs[0].keywords["c"] == [(0, 1), (0, 2), (1, 2)] - funcs, requires_intensity = _create_funcs(["haralick_*_3", "haralick_*_5"], 3) - assert requires_intensity + funcs, all_required_channels = _create_funcs(["haralick_*_3", "haralick_*_5"], 3) + assert len(all_required_channels) > 0 assert len(funcs) == 2 - funcs, requires_intensity = _create_funcs(["intensitydistribution_*_4"], 3) + funcs, all_required_channels = _create_funcs(["intensitydistribution_*_4"], 3) assert funcs[0].keywords == {"c": (0, 1, 2), "bin_count": 4} - assert requires_intensity + assert len(all_required_channels) > 0 assert len(funcs) == 1 funcs, _ = _create_funcs(["colocalization_0_0"], 3) @@ -504,19 +548,16 @@ def test_features_cli_multi_images(tmp_path, array_A1_102_cells, array_A1_102_al @pytest.mark.features -def test_features_cli(tmp_path, array_A1_102_cells, array_A1_102_alnpheno): - tmp_path.mkdir(parents=True, exist_ok=True) - # test that all features run, can be saved to disk, and diff with known good output - image = ( - array_A1_102_alnpheno.transpose(*("z", "c", "t", "y", "x")).rename( - {"z": "t", "t": "z"} - ) - ).isel(t=0, z=0) # ops swaps z and t in saved tif - +def test_features_cli(tmp_path, array_A1_102_cells): labels = array_A1_102_cells.squeeze().copy() labels.values[labels.values != 17] = 0 + rng = np.random.default_rng(1) + image = xr.DataArray( + rng.integers(low=10, high=50, size=(4,) + labels.shape), dims=["c", "y", "x"] + ) zarr_path = str(tmp_path / "test.zarr") features_output_path = str(tmp_path / "features-out") + features2_output_path = str(tmp_path / "features2-out") objects_output_path = str(tmp_path / "objects-out") exp = Experiment() exp.images["test"] = image @@ -553,6 +594,24 @@ def test_features_cli(tmp_path, array_A1_102_cells, array_A1_102_alnpheno): ] check_call(cmd) + result_df = pd.read_parquet(features_output_path) + cmd = [ + "scallops", + "features", + "--images", + zarr_path, + "--labels", + zarr_path, + "--output", + features2_output_path, + "--features-cell", + "intensity_2,0", + "--objects", + objects_output_path, + ] + result2_df = pd.read_parquet(features_output_path) + check_call(cmd) + pd.testing.assert_frame_equal(result_df[result2_df.columns], result2_df) @pytest.mark.features diff --git a/scallops/tests/test_illumination_correction.py b/scallops/tests/test_illumination_correction.py index d9164e65..9c265675 100644 --- a/scallops/tests/test_illumination_correction.py +++ b/scallops/tests/test_illumination_correction.py @@ -4,7 +4,7 @@ import numpy as np import pytest import zarr -from zarr import ZipStore +from zarr.storage import ZipStore from scallops.io import read_image diff --git a/scallops/tests/test_io.py b/scallops/tests/test_io.py index 72ec4a1e..b3c3a860 100644 --- a/scallops/tests/test_io.py +++ b/scallops/tests/test_io.py @@ -186,24 +186,75 @@ def test_read_tif(use_dask): assert len(data.shape) == 5 +@pytest.mark.parametrize("zarr_format", ["ome_zarr", "zarr"]) +@pytest.mark.parametrize("use_dask", [True, False]) @pytest.mark.io -def test_write_ome_zarr_image_dask(tmp_path): - data = read_image( - "scallops/tests/data/tif/10X_c0-DAPI-p65ab_A1_Tile-7.phenotype.tif", dask=True +def test_write_zarr_attrs_coords_spacing(tmp_path, zarr_format, use_dask): + data = xr.DataArray( + np.arange(120, dtype=np.uint8).reshape((2, 3, 5, 4)), + dims=["t", "c", "y", "x"], + coords={"t": [3, 5], "c": ["foo1", "foo2", "foo3"]}, + attrs={ + "foo": "bar", + "physical_pixel_sizes": (1.3, 1.2), + "physical_pixel_units": ("angstrom", "mm"), + }, ) - data.attrs.clear() + if use_dask: + data.data = da.from_array(data.data) zarr_path = str(tmp_path / "test.zarr") - _write_zarr_image("foo", open_ome_zarr(zarr_path), data) - data2 = read_image( - "scallops/tests/data/tif/10X_c0-DAPI-p65ab_A1_Tile-7.phenotype.tif", dask=False - ) - np.testing.assert_array_equal( - data2.values, data.values, err_msg="Dask and non dask images are not equal" - ) - data3 = read_image(f"{zarr_path}/images/foo", dask=False) - np.testing.assert_array_equal( - data2.values, data3.values, err_msg="Saved image not equal" - ) + _write_zarr_image("foo", open_ome_zarr(zarr_path), data, zarr_format=zarr_format) + img = read_image(f"{zarr_path}/images/foo", dask=False) + assert img.attrs.pop("t") == [3, 5] + xr.testing.assert_identical(data, img) + assert get_image_spacing(img.attrs) == (1.3, 1.2) + assert img.attrs["foo"] == "bar" + assert img.attrs["physical_pixel_units"] == ("angstrom", "mm") + assert img.coords["t"].values.tolist() == [3, 5] + assert img.coords["c"].values.tolist() == ["foo1", "foo2", "foo3"] + g = zarr.open(f"{zarr_path}/images/foo", mode="r") + assert list(g.keys()) == ["s0"] + assert { + "channels": [ + { + "label": "foo1", + "color": "00FFFF", + "window": {"min": 0, "max": 255, "start": 0, "end": 255}, + }, + { + "label": "foo2", + "color": "FFFF00", + "window": {"min": 0, "max": 255, "start": 0, "end": 255}, + }, + { + "label": "foo3", + "color": "FF00FF", + "window": {"min": 0, "max": 255, "start": 0, "end": 255}, + }, + ] + } == g.attrs["omero"] + assert [ + { + "datasets": [ + { + "path": "s0", + "coordinateTransformations": [ + {"type": "scale", "scale": [1.0, 1.0, 1.3, 1.2]}, + {"type": "translation", "translation": [0.0, 0.0, 0.0, 0.0]}, + ], + } + ], + "name": "/images/foo", + "metadata": {"foo": "bar", "t": [3, 5]}, + "axes": [ + {"name": "t", "type": "time"}, + {"name": "c", "type": "channel"}, + {"name": "y", "type": "space", "unit": "angstrom"}, + {"name": "x", "type": "space", "unit": "mm"}, + ], + "version": "0.4", + } + ] == g.attrs["multiscales"] @pytest.mark.io @@ -225,13 +276,16 @@ def test_write_non_ome_zarr_image(tmp_path, use_dask): image.attrs = {"test": "1"} image.attrs["physical_pixel_sizes"] = (1, 1, 1) image.attrs["physical_pixel_units"] = ("mm", "mm", "mm") - zarr_path = str(tmp_path / "test.zarr") - _write_zarr_image("foo", open_ome_zarr(zarr_path), image, zarr_format="zarr") - _write_zarr_image("foo2", open_ome_zarr(zarr_path), image) + image.name = None + zarr_path = str(tmp_path / "test1.zarr") + _write_zarr_image("image1", open_ome_zarr(zarr_path), image, zarr_format="zarr") + data_zarr = read_image(f"{zarr_path}/images/image1", dask=False) + xr.testing.assert_identical(data_zarr, image) - data_zarr = read_image(f"{zarr_path}/images/foo", dask=False) - data_ome_zarr = read_image(f"{zarr_path}/images/foo2", dask=False) - xr.testing.assert_equal(data_zarr, data_ome_zarr) + ome_zarr_path = str(tmp_path / "test2.zarr") + _write_zarr_image("image1", open_ome_zarr(ome_zarr_path), image) + data_ome_zarr = read_image(f"{ome_zarr_path}/images/image1", dask=False) + xr.testing.assert_identical(data_ome_zarr, image) @pytest.mark.io @@ -374,19 +428,30 @@ def test_read_write_experiment(experiment_c, tmp_path): np.testing.assert_equal( exp.images["A1-102"].coords["c"].data, ["Channel:0:0", "Channel:0:1", "Channel:0:2", "Channel:0:3", "Channel:0:4"], + err_msg="c not equal", ) np.testing.assert_equal( - exp.images["A1-102"].coords["t"].data, [1, 2, 3, 4, 5, 7, 8, 9, 10] + exp.images["A1-102"].coords["t"].data, + [1, 2, 3, 4, 5, 7, 8, 9, 10], + err_msg="t not equal", ) np.testing.assert_equal( - exp.images["A1-102"].values, experiment_c.images["A1-102"].values + exp.images["A1-102"].values, + experiment_c.images["A1-102"].values, + err_msg="102 not equal", ) np.testing.assert_equal( - exp.images["A1-103"].values, experiment_c.images["A1-103"].values + exp.images["A1-103"].values, + experiment_c.images["A1-103"].values, + err_msg="103 not equal", + ) + assert exp.images["A1-102"].attrs["common_src"] == "10X_c-SBS_*_A1_Tile-102.sbs", ( + "common_src not equal" + ) + assert exp.images["A1-102"].attrs["group"] == dict(tile="102", well="A1"), ( + "group not equal" ) - assert exp.images["A1-102"].attrs["common_src"] == "10X_c-SBS_*_A1_Tile-102.sbs" - assert exp.images["A1-102"].attrs["group"] == dict(tile="102", well="A1") @pytest.mark.io diff --git a/scallops/tests/test_register_cli.py b/scallops/tests/test_register_cli.py index d2cf5b8d..17408e16 100644 --- a/scallops/tests/test_register_cli.py +++ b/scallops/tests/test_register_cli.py @@ -59,8 +59,8 @@ def create_itk_param_file(tmp_path): (WriteIterationInfo "false") (WriteResultImage "false") """ - path = os.path.join(tmp_path, "translation.txt") - with open(path, "wt") as out: + path = str(tmp_path / "translation.txt") + with open(path, mode="w") as out: out.write(p) return path @@ -311,7 +311,7 @@ def test_register_transform_labels_moving_only(tmp_path): img = read_image( "scallops/tests/data/experimentC/10X_c0-DAPI-p65ab/10X_c0-DAPI-p65ab_A1_Tile-102.phenotype.tif" ) - img.attrs["physical_pixel_sizes"] = (1, 1) + img.attrs["physical_pixel_sizes"] = (5, 4, 3) rng = np.random.default_rng(0) @@ -356,6 +356,8 @@ def test_register_transform_labels_moving_only(tmp_path): assert transformed_labels.max() > 0 transformed_image = read_image(output_zarr / "images" / "plateA-A1") assert transformed_image.shape[0] == 2 + transformed_image.attrs["physical_pixel_sizes"] = (4.0, 3.0) + transformed_image.attrs["physical_pixel_units"] = ("micrometer", "micrometer") @pytest.mark.registration diff --git a/scallops/tests/test_stitch.py b/scallops/tests/test_stitch.py index 38be2a06..53927a7b 100644 --- a/scallops/tests/test_stitch.py +++ b/scallops/tests/test_stitch.py @@ -301,7 +301,7 @@ def test_stitch_cli(tmp_path): z_index=0, group=fuse_result, ) - no_blend_array = fuse_result["0"][...].squeeze() + no_blend_array = fuse_result["s0"][...].squeeze() np.testing.assert_array_equal( no_blend_array, expected_value, @@ -350,7 +350,7 @@ def test_stitch_cli(tmp_path): np.testing.assert_array_equal( blending_img_cli.squeeze(), - blend_fuse_result["0"][...].squeeze(), + blend_fuse_result["s0"][...].squeeze(), err_msg="Blending images differ", ) diff --git a/scallops/utils.py b/scallops/utils.py index 9dc1f857..78e6227b 100644 --- a/scallops/utils.py +++ b/scallops/utils.py @@ -11,8 +11,6 @@ import json import logging import os -import uuid -import warnings from bisect import bisect_right from collections import Counter from collections.abc import Callable, Sequence @@ -27,17 +25,7 @@ import numpy as np import pandas as pd import skimage -from dask import is_dask_collection -from dask.array.core import ( - getter, - getter_nofancy, - graph_from_arraylike, - normalize_chunks, - slices_from_chunks, -) from dask.system import CPU_COUNT -from dask.tokenize import tokenize -from dask.utils import SerializableLock from decorator import decorator from kneed import KneeLocator from skimage import restoration @@ -546,99 +534,6 @@ def dask_chunk_stats( ).squeeze() -def _dask_from_array_no_copy( - x, - chunks="auto", - name=None, - lock=False, - asarray=False, - fancy=True, - getitem=None, - meta=None, - inline_array=False, -): - """Create dask array from something that looks like an array without copying.""" - - if isinstance(x, da.Array): - raise ValueError( - "Array is already a dask array. Use 'asarray' or 'rechunk' instead." - ) - - elif is_dask_collection(x): - warnings.warn( - "Passing an object to dask.array.from_array which is already a " - "Dask collection. This can lead to unexpected behavior." - ) - - if isinstance(x, (list, tuple, memoryview) + np.ScalarType): - x = np.array(x) - - # if is_arraylike(x) and hasattr(x, "copy"): - # x = x.copy() - - if asarray is None: - asarray = not hasattr(x, "__array_function__") - - previous_chunks = getattr(x, "chunks", None) - - chunks = normalize_chunks( - chunks, x.shape, dtype=x.dtype, previous_chunks=previous_chunks - ) - - if name in (None, True): - token = tokenize(x, chunks, lock, asarray, fancy, getitem, inline_array) - name = name or "array-" + token - elif name is False: - name = "array-" + str(uuid.uuid1()) - - if lock is True: - lock = SerializableLock() - - is_ndarray = type(x) in (np.ndarray, np.ma.core.MaskedArray) - is_single_block = all(len(c) == 1 for c in chunks) - # Always use the getter for h5py etc. Not using isinstance(x, np.ndarray) - # because np.matrix is a subclass of np.ndarray. - if is_ndarray and not is_single_block and not lock: - # eagerly slice numpy arrays to prevent memory blowup - # GH5367, GH5601 - slices = slices_from_chunks(chunks) - keys = product([name], *(range(len(bds)) for bds in chunks)) - values = [x[slc] for slc in slices] - dsk = dict(zip(keys, values)) - - elif is_ndarray and is_single_block: - # No slicing needed - dsk = {(name,) + (0,) * x.ndim: x} - else: - if getitem is None: - if fancy: - getitem = getter - else: - getitem = getter_nofancy - - dsk = graph_from_arraylike( - x, - chunks, - x.shape, - name, - getitem=getitem, - lock=lock, - asarray=asarray, - dtype=x.dtype, - inline_array=inline_array, - ) - - # Workaround for TileDB, its indexing is 1-based, - # and doesn't seems to support 0-length slicing - if x.__class__.__module__.split(".")[0] == "tiledb" and hasattr(x, "_ctx_"): - return da.Array(dsk, name, chunks, dtype=x.dtype) - - if meta is None: - meta = x - - return da.Array(dsk, name, chunks, meta=meta, dtype=getattr(x, "dtype", None)) - - def _write_img_size(file_list: list[str]): from scallops.io import _images2fov, _localize_path diff --git a/scallops/zarr_io.py b/scallops/zarr_io.py index 9f040101..6b35aa71 100644 --- a/scallops/zarr_io.py +++ b/scallops/zarr_io.py @@ -23,11 +23,12 @@ from dask.array import from_zarr from dask.delayed import Delayed from dask.graph_manipulation import bind +from ome_zarr import USE_DASK_ARRAY_KWARGS from ome_zarr.axes import KNOWN_AXES from ome_zarr.format import Format, FormatV04 from ome_zarr.io import parse_url from ome_zarr.types import JSONDict -from ome_zarr.writer import write_image +from ome_zarr.writer import write_image, write_multiscale from xarray.core.coordinates import DataArrayCoordinates from zarr.storage import StoreLike @@ -41,19 +42,28 @@ def _current_format() -> Format: return FormatV04() +def _create_array_kwargs(fmt: Format | None = None) -> dict[str, Any]: + if fmt is None: + fmt = _current_format() + zarr_array_kwargs = dict() + + if fmt.zarr_format == 2: + zarr_array_kwargs["chunk_key_encoding"] = {"name": "v2", "separator": "/"} + return zarr_array_kwargs + + def _da_to_zarr_kwargs(fmt: Format | None = None) -> dict[str, Any]: if fmt is None: fmt = _current_format() zarr_array_kwargs = dict() - if fmt.version == "0.4": + + if USE_DASK_ARRAY_KWARGS: + if fmt.zarr_format == 2: + zarr_array_kwargs["chunk_key_encoding"] = {"name": "v2", "separator": "/"} + elif fmt.zarr_format == 2: zarr_array_kwargs["dimension_separator"] = "/" - # if USE_DASK_ARRAY_KWARGS: - # if fmt.zarr_format == 2: - # zarr_array_kwargs["chunk_key_encoding"] = _chunk_key_encoding - # if fmt.zarr_format == 2: - # zarr_array_kwargs["dimension_separator"] = "/" - # if fmt.zarr_format == 2: - # zarr_array_kwargs["zarr_format"] = 2 + if fmt.zarr_format == 2: + zarr_array_kwargs["zarr_format"] = fmt.zarr_format return zarr_array_kwargs @@ -71,7 +81,7 @@ def is_anndata_zarr(store: StoreLike) -> bool: :param store: Zarr store """ try: - return isinstance(zarr.open(store, mode="r", path="uns"), zarr.Group) + return isinstance(zarr.open(store, mode="r", path="layers"), zarr.Group) except: # noqa: E722 return False @@ -118,14 +128,25 @@ def _get_sep(group: zarr.Group) -> str: return "/" -def _omero_channels(channel_names) -> dict: +def _omero_channels(channel_names, image_dtype: np.dtype) -> dict: # Napari does not like '#' at start of hex color string # Hex colors match Napari defaults colors = ["00FFFF", "FFFF00", "FF00FF", "FF0000", "008000", "0000FF"] + + max_value = ( + int(np.iinfo(image_dtype).max) + if np.issubdtype(image_dtype, np.integer) + else float(np.finfo(image_dtype).max) + ) + # Napari requires that colors are specified if channel names are specified channels = ( [ - dict(label=str(channel_names[i]), color=colors[i % len(colors)]) + dict( + label=str(channel_names[i]), + color=colors[i % len(colors)], + window=dict(min=0, max=max_value, start=0, end=max_value), + ) for i in range(len(channel_names)) ] if not np.isscalar(channel_names) @@ -135,7 +156,7 @@ def _omero_channels(channel_names) -> dict: def _create_omero_metadata( - coords: DataArrayCoordinates, dims: tuple[Hashable, ...] + coords: DataArrayCoordinates, dims: tuple[Hashable, ...], image_dtype: np.dtype ) -> dict | None: """Create OMERO metadata for a DataArray. @@ -162,18 +183,19 @@ def _create_omero_metadata( array = xr.DataArray(data, dims=dims, coords=coords) # Create OMERO metadata - omero_metadata = _create_omero_metadata(array.coords, array.dims) + omero_metadata = _create_omero_metadata(array.coords, array.dims, data.dtype) print(omero_metadata) # Output: {'channels': [{'label': 'DAPI', 'color': '00FFFF'}, {'label': 'FITC', 'color': 'FFFF00'}, ...]} """ if dims is None: return None + channel_names = coords["c"] if "c" in dims and "c" in coords else None if channel_names is not None: if isinstance(channel_names, xr.DataArray): channel_names = channel_names.values - return _omero_channels(channel_names) + return _omero_channels(channel_names, image_dtype) return None @@ -227,9 +249,12 @@ def _fix_attrs(d: dict) -> None: value[i] = str(value) -def _attrs_axes_coordinates( - image_attrs: dict, coords: DataArrayCoordinates, dims: tuple[Hashable, ...] -) -> tuple[dict, list[dict] | None, list[dict] | None]: +def _attrs_axes_scales( + image_attrs: dict, + coords: DataArrayCoordinates, + dims: tuple[Hashable, ...], + image_dtype: np.dtype, +) -> tuple[dict, list[dict] | None, dict[str, float] | None]: """Prepare attributes, axes, and coordinate transformations for Zarr storage. Processes the attributes, coordinates, and dimensions of a DataArray to generate @@ -242,10 +267,13 @@ def _attrs_axes_coordinates( :return: A tuple containing: - Updated image attributes dictionary. - List of axes dictionaries or None. - - List of coordinate transformations dictionaries or None. + - Scale dict or None. """ - omero = _create_omero_metadata(coords, dims) + omero = _create_omero_metadata(coords, dims, image_dtype) + if image_attrs is None: + image_attrs = dict() + image_attrs = image_attrs.copy() if omero is not None: image_attrs["omero"] = omero @@ -260,29 +288,37 @@ def _attrs_axes_coordinates( vals = vals.tolist() image_attrs[key] = vals - coordinate_transformations = None - physical_pixel_units = None - if "physical_pixel_sizes" in image_attrs and "physical_pixel_units" in image_attrs: - physical_pixel_sizes = image_attrs.pop("physical_pixel_sizes") - physical_pixel_units = image_attrs.pop("physical_pixel_units") - non_space_dims = [d for d in dims if d not in ("z", "y", "x")] - scale = list((1.0,) * len(non_space_dims)) + list(physical_pixel_sizes) - coordinate_transformations = [{"scale": scale, "type": "scale"}] + physical_pixel_sizes = image_attrs.pop("physical_pixel_sizes", None) + physical_pixel_units = image_attrs.pop("physical_pixel_units", None) - space_index = 0 axes = None + scale_dict = None if dims is not None: axes = [] + scale_dict = dict() + space_index = 0 + if "z" in dims: + if physical_pixel_sizes is not None and len(physical_pixel_sizes) == 2: + physical_pixel_sizes = [1.0] + list(physical_pixel_sizes) + if physical_pixel_units is not None and len(physical_pixel_units) == 2: + physical_pixel_units = ["micrometer"] + list(physical_pixel_units) + for d in dims: axis = {"name": d, "type": KNOWN_AXES.get(d)} - if physical_pixel_units is not None and axis["type"] == "space": - axis["unit"] = physical_pixel_units[space_index] + if physical_pixel_sizes is not None and axis["type"] == "space": + scale_dict[d] = physical_pixel_sizes[space_index] + unit = ( + physical_pixel_units[space_index] + if physical_pixel_units is not None + else "micrometer" + ) + axis["unit"] = unit space_index = space_index + 1 axes.append(axis) _fix_attrs(image_attrs) image_attrs = _fix_json(image_attrs) - return image_attrs, axes, coordinate_transformations + return image_attrs, axes, scale_dict def _write_zarr_image( @@ -293,6 +329,7 @@ def _write_zarr_image( group: str | None = "images", zarr_format: Literal["ome_zarr", "zarr"] = "ome_zarr", compute: bool = True, + storage_options: JSONDict | None = None, ) -> list[Delayed]: """Write image in zarr format. @@ -305,6 +342,7 @@ def _write_zarr_image( compliant images with dimensions other than (t,c,z,y,x) :param compute: If true compute immediately otherwise a list of :class:`dask.delayed.Delayed` is returned. + :param storage_options: Options to be passed on to the storage backend. :return: Empty list if the compute flag is True, otherwise it returns a list of :class:`dask.delayed.Delayed` representing the value to be computed by dask. """ @@ -325,18 +363,54 @@ def _write_zarr_image( dims = image.dims else: data = image + if image.ndim == 2: + dims = ["y", "x"] return write_zarr( grp=dest_grp, data=data, image_attrs=image_attrs, coords=coords, dims=dims, + storage_options=storage_options, metadata=metadata, zarr_format=zarr_format, compute=compute, ) +def _create_zarr_attrs(fmt, grp, dims, image_attrs, axes, scale_dict): + datasets = [{"path": "s0"}] + if scale_dict is None: + scale_dict = dict() + scale = [] + translation = [] + for d in dims: + scale.append(scale_dict.get(d, 1.0)) + translation.append(0.0) + coordinate_transformations = [ + {"type": "scale", "scale": scale}, + {"type": "translation", "translation": translation}, + ] + datasets[0]["coordinateTransformations"] = coordinate_transformations + + multiscales = [dict(version=fmt.version, datasets=datasets, name=grp.name)] + zarr_attrs = {"multiscales": multiscales} + if axes is not None: + multiscales[0]["axes"] = axes + + if fmt.version in fmt.version in ("0.5"): + omero = zarr_attrs["ome"].get("omero", {}) + omero.update(image_attrs.pop("omero", {})) + zarr_attrs["ome"]["omero"] = omero + zarr_attrs = {"ome": zarr_attrs} + else: + omero = zarr_attrs.get("omero", {}) + omero.update(image_attrs.pop("omero", {})) + zarr_attrs["omero"] = omero + multiscales[0]["metadata"] = image_attrs + return zarr_attrs + + def write_zarr( grp: zarr.Group, data: np.ndarray | da.Array | xr.DataArray, @@ -346,6 +420,7 @@ def write_zarr( metadata: dict[str, Any] | None = None, zarr_format: Literal["ome_zarr", "zarr"] = "ome_zarr", compute: bool = True, + storage_options: JSONDict | None = None, ) -> list[Delayed]: """Write data to a Zarr group with optional metadata and scaling. @@ -366,6 +441,7 @@ def write_zarr( :param compute: If True, compute immediately. Otherwise, return a list of dask.delayed. Delayed objects representing the value to be computed by dask. Default is True. + :param storage_options: Options to be passed on to the storage backend. :return: Empty list if the compute flag is True, otherwise a list of dask.delayed.Delayed objects. @@ -401,8 +477,8 @@ def write_zarr( if metadata is not None: image_attrs.update(metadata) - image_attrs, axes, coordinate_transformations = _attrs_axes_coordinates( - image_attrs, coords, dims + image_attrs, axes, scale_dict = _attrs_axes_scales( + image_attrs, coords, dims, data.dtype ) dask_delayed = [] @@ -412,36 +488,20 @@ def write_zarr( d = da.to_zarr( arr=data, url=grp.store, - component=str(Path(grp.path, "0")), + component=str(Path(grp.path, "s0")), compute=compute, **_da_to_zarr_kwargs(fmt), ) if not compute: dask_delayed.append(d) elif not isinstance(data, zarr.Array): - grp.create_dataset( - "0", data=data, overwrite=True, **_da_to_zarr_kwargs(fmt) - ) - - datasets = [{"path": "0"}] - if coordinate_transformations is not None: - datasets[0]["coordinateTransformations"] = coordinate_transformations + create_kwds = _create_array_kwargs(fmt) - multiscales = [dict(version=fmt.version, datasets=datasets, name=grp.name)] - zarr_attrs = {"multiscales": multiscales} - if axes is not None: - multiscales[0]["axes"] = axes + if storage_options.get("chunks") is not None: + create_kwds["chunks"] = storage_options["chunks"] + grp.create_array("s0", data=data, overwrite=True, **create_kwds) - if fmt.version in fmt.version in ("0.5"): - omero = zarr_attrs["ome"].get("omero", {}) - omero.update(image_attrs.pop("omero", {})) - zarr_attrs["ome"]["omero"] = omero - zarr_attrs = {"ome": zarr_attrs} - else: - omero = zarr_attrs.get("omero", {}) - omero.update(image_attrs.pop("omero", {})) - zarr_attrs["omero"] = omero - multiscales[0]["metadata"] = image_attrs + zarr_attrs = _create_zarr_attrs(fmt, grp, dims, image_attrs, axes, scale_dict) if len(dask_delayed) > 0: @@ -456,20 +516,18 @@ def _write_metadata_delayed(grp, d): grp.attrs.update(zarr_attrs) return dask_delayed else: - return write_image( - fmt=fmt, - image=data, + image_attrs = {"metadata": image_attrs} + + return write_multiscale( + pyramid=[data], group=grp, - # scale_factors=[], - scaler=None, + fmt=fmt, + storage_options=storage_options, + scale_factors=[1.0], axes=axes, compute=compute, - metadata=image_attrs, - coordinate_transformations=( - [coordinate_transformations] - if coordinate_transformations is not None - else None - ), + scale=scale_dict, + **image_attrs, ) @@ -562,8 +620,7 @@ def _write_zarr_labels( return write_image( labels, grp, - scaler=None, - # scale_factors=[], + scale_factors=[], axes=label_axes, metadata=metadata, compute=compute, @@ -596,6 +653,7 @@ def _read_zarr_attrs(attrs) -> tuple[dict, dict, list[str]]: attrs = attrs["ome"] multiscales = attrs["multiscales"] + if len(multiscales) > 0: multiscale0 = multiscales[0] else: @@ -722,8 +780,7 @@ def open_ome_zarr(url: Path | str, mode: str = "a") -> zarr.Group | None: loc = parse_url(url, mode=mode, fmt=fmt) if loc is None: return None - return zarr.open(loc.store, mode=mode) - # return zarr.open(loc.store, mode=mode, zarr_format=fmt.zarr_format) + return zarr.open(loc.store, mode=mode, zarr_format=fmt.zarr_format) except Exception as e: logger.error(f"Failed to open OME-Zarr store: {url}") raise e @@ -812,7 +869,7 @@ class _LazyLoadZarrData(_LazyLoadData): root = open_ome_zarr(store=store) group = root.create_group("test_group") - group.create_dataset("0", data=[1, 2, 3, 4, 5]) + group.create_array("0", data=[1, 2, 3, 4, 5]) # Create a _LazyLoadZarrData instance lazy_data = _LazyLoadZarrData(group, dask=True)