From 53840b84b30d09bf6f4c181baf273064f7caf8bb Mon Sep 17 00:00:00 2001 From: lorenzo Date: Thu, 26 Sep 2024 23:48:09 +0200 Subject: [PATCH] multiple bug fixes --- src/ngio/core/image_like_handler.py | 54 +++++++++ src/ngio/core/label_handler.py | 141 ++++++++++++++++++++--- src/ngio/core/ngff_image.py | 85 ++++++++++++-- src/ngio/core/utils.py | 102 +++++++++++++++- src/ngio/ngff_meta/fractal_image_meta.py | 40 +++++++ src/ngio/ngff_meta/v04/zarr_utils.py | 10 +- src/ngio/tables/tables_group.py | 7 +- src/ngio/tables/v1/feature_tables.py | 1 - 8 files changed, 408 insertions(+), 32 deletions(-) diff --git a/src/ngio/core/image_like_handler.py b/src/ngio/core/image_like_handler.py index 5528574..f4ee53d 100644 --- a/src/ngio/core/image_like_handler.py +++ b/src/ngio/core/image_like_handler.py @@ -1,5 +1,6 @@ """Generic class to handle Image-like data in a OME-NGFF file.""" +from functools import partial from pathlib import Path from typing import Literal from warnings import warn @@ -9,6 +10,7 @@ import zarr from dask.delayed import Delayed from dask.distributed import Lock +from scipy.ndimage import zoom from ngio._common_types import ArrayLike from ngio.core.dimensions import Dimensions @@ -382,3 +384,55 @@ def set_data( x=x, y=y, z=z, t=t, c=c, preserve_dimensions=preserve_dimensions ) self._set_pipe(data_pipe=data_pipe, patch=patch) + + def consolidate(self, order: int = 1) -> None: + """Consolidate the Zarr array.""" + processed_paths = [self] + + todo_image = [ + ImageLike(store=self.group, path=_path) + for _path in self.metadata.levels_paths + if _path != self.path + ] + + while todo_image: + dist_matrix = np.zeros((len(processed_paths), len(todo_image))) + for i, image in enumerate(todo_image): + for j, processed_image in enumerate(processed_paths): + dist_matrix[j, i] = np.sqrt( + np.sum( + [ + (s1 - s2) ** 2 + for s1, s2 in zip( + image.shape, processed_image.shape, strict=False + ) + ] + ) + ) + + source, target = np.unravel_index(dist_matrix.argmin(), dist_matrix.shape) + + source_image = processed_paths[source] + target_image = todo_image.pop(target) + + zoom_factors = tuple( + [ + s1 / s2 + for s1, s2 in zip( + target_image.shape, source_image.shape, strict=False + ) + ] + ) + + partial_zoom = partial(zoom, zoom=zoom_factors, order=order) + out_image = da.map_blocks( + partial_zoom, + source_image.dask_array, + # dtype = source_image.dask_array, + chunks=target_image.array.chunks, + ) + out_image = out_image.astype(target_image.array.dtype) + target_image.array[...] = out_image.compute() + # compute the transformation + + processed_paths.append(target_image) diff --git a/src/ngio/core/label_handler.py b/src/ngio/core/label_handler.py index ae85f9e..ce725db 100644 --- a/src/ngio/core/label_handler.py +++ b/src/ngio/core/label_handler.py @@ -1,6 +1,10 @@ """A module to handle OME-NGFF images stored in Zarr format.""" +import zarr + +from ngio.core.image_handler import Image from ngio.core.image_like_handler import ImageLike +from ngio.core.utils import create_empty_ome_zarr_label from ngio.io import StoreLike, StoreOrGroup from ngio.ngff_meta.fractal_image_meta import LabelMeta, PixelSize @@ -13,6 +17,7 @@ class Label(ImageLike): """ def __init__( + self, store: StoreOrGroup, *, path: str | None = None, @@ -47,6 +52,7 @@ def __init__( cache=cache, ) + @property def metadata(self) -> LabelMeta: """Return the metadata of the image.""" return super().metadata @@ -55,23 +61,132 @@ def metadata(self) -> LabelMeta: class LabelGroup: """A class to handle the /labels group in an OME-NGFF file.""" - def __init__(self, group: StoreLike) -> None: + def __init__( + self, + group: StoreLike | zarr.Group, + image_ref: Image | None = None, + ) -> None: """Initialize the LabelGroupHandler.""" - self._group = group + if not isinstance(group, zarr.Group): + group = zarr.open_group(group, mode="a") - @property - def group_name(self) -> str: - """Return the name of the group.""" - return "labels" + if "labels" not in group: + self._group = group.create_group("labels") + self._group.attrs["labels"] = [] # initialize the labels attribute + else: + self._group: zarr.Group = group["labels"] + + self._imgage_ref = image_ref def list(self) -> list[str]: """List all labels in the group.""" - return list(self._group.array_keys()) + return self._group.attrs.get("labels", []) + + def get( + self, + name: str, + path: str | None = None, + pixel_size: PixelSize | None = None, + highest_resolution: bool = True, + ) -> Label: + """Geta a Label from the group. + + Args: + name (str): The name of the label. + path (str | None, optional): The path to the level. + pixel_size (tuple[float, ...] | list[float] | None, optional): The pixel + size of the level. + highest_resolution (bool, optional): Whether to get the highest + resolution level + """ + if name not in self.list(): + raise ValueError(f"Label {name} not found in the group.") - def get(self, name: str) -> Label: - """Get a label from the group.""" - raise NotImplementedError("Not yet implemented.") + if path is not None or pixel_size is not None: + highest_resolution = False + + return Label( + store=self._group[name], + path=path, + pixel_size=pixel_size, + highest_resolution=highest_resolution, + ) + + def derive( + self, + name: str, + overwrite: bool = False, + **kwargs, + ) -> Label: + """Derive a new label from an existing label. + + Args: + name (str): The name of the new label. + overwrite (bool): If True, the label will be overwritten if it exists. + Default is False. + **kwargs: Additional keyword arguments to pass to the new label. + """ + list_of_labels = self.list() + + if overwrite and name in list_of_labels: + self._group.attrs["label"] = [ + label for label in list_of_labels if label != name + ] + elif not overwrite and name in list_of_labels: + raise ValueError(f"Label {name} already exists in the group.") + + # create the new label + new_label_group = self._group.create_group(name, overwrite=overwrite) + + if self._imgage_ref is None: + label_0 = self.get(list_of_labels[0]) + metadata = label_0.metadata + on_disk_shape = label_0.on_disk_shape + chunks = label_0.array.chunks + dataset = label_0.dataset + else: + label_0 = self._imgage_ref + metadata = label_0.metadata + channel_index = metadata.index_mapping.get("c", None) + if channel_index is not None: + on_disk_shape = ( + label_0.on_disk_shape[:channel_index] + + label_0.on_disk_shape[channel_index + 1 :] + ) + chunks = ( + label_0.array.chunks[:channel_index] + + label_0.array.chunks[channel_index + 1 :] + ) + else: + on_disk_shape = label_0.on_disk_shape + chunks = label_0.array.chunks + + metadata = metadata.remove_axis("c") + dataset = metadata.get_highest_resolution_dataset() + + default_kwargs = { + "store": new_label_group, + "shape": on_disk_shape, + "chunks": chunks, + "dtype": label_0.array.dtype, + "on_disk_axis": dataset.on_disk_axes_names, + "pixel_sizes": dataset.pixel_size, + "xy_scaling_factor": metadata.xy_scaling_factor, + "z_scaling_factor": metadata.z_scaling_factor, + "time_spacing": dataset.time_spacing, + "time_units": dataset.time_axis_unit, + "num_levels": metadata.num_levels, + "name": name, + "overwrite": overwrite, + "version": metadata.version, + } + + default_kwargs.update(kwargs) + + create_empty_ome_zarr_label( + **default_kwargs, + ) - def write(self, name: str, data: Label) -> None: - """Create a label in the group.""" - raise NotImplementedError("Not yet implemented.") + if name not in self.list(): + self._group.attrs["labels"] = [*list_of_labels, name] + return self.get(name) diff --git a/src/ngio/core/ngff_image.py b/src/ngio/core/ngff_image.py index 8f7f61e..5f33eae 100644 --- a/src/ngio/core/ngff_image.py +++ b/src/ngio/core/ngff_image.py @@ -1,10 +1,16 @@ """Abstract class for handling OME-NGFF images.""" -from typing import Protocol, TypeVar +from typing import Any, Protocol, TypeVar + +import numpy as np from ngio.core.image_handler import Image +from ngio.core.label_handler import LabelGroup +from ngio.core.utils import create_empty_ome_zarr_image from ngio.io import StoreLike, open_group_wrapper -from ngio.ngff_meta import FractalImageLabelMeta, get_ngff_image_meta_handler +from ngio.ngff_meta import get_ngff_image_meta_handler +from ngio.ngff_meta.fractal_image_meta import ImageMeta, PixelSize +from ngio.tables.tables_group import TableGroup T = TypeVar("T") @@ -70,8 +76,11 @@ def __init__(self, store: StoreLike) -> None: self.group, meta_mode="image", cache=False ) + self.table = TableGroup(self.group) + self.label = LabelGroup(self.group, image_ref=self.get_image()) + @property - def image_meta(self) -> FractalImageLabelMeta: + def image_meta(self) -> ImageMeta: """Get the image metadata.""" return self._image_meta.load_meta() @@ -88,14 +97,14 @@ def levels_paths(self) -> list[str]: def get_image( self, *, - level_path: str | int | None = None, - pixel_size: tuple[float, ...] | list[float] | None = None, + path: str | None = None, + pixel_size: PixelSize | None = None, highest_resolution: bool = True, ) -> Image: """Get an image handler for the given level. Args: - level_path (str | int | None, optional): The path to the level. + path (str | None, optional): The path to the level. pixel_size (tuple[float, ...] | list[float] | None, optional): The pixel size of the level. highest_resolution (bool, optional): Whether to get the highest @@ -104,23 +113,83 @@ def get_image( Returns: ImageHandler: The image handler. """ + if path is not None or pixel_size is not None: + highest_resolution = False + return Image( store=self.group, - level_path=level_path, + path=path, pixel_size=pixel_size, highest_resolution=highest_resolution, ) + def _update_omero_window(self) -> None: + """Update the OMERO window.""" + meta = self.image_meta + image = self.get_image(highest_resolution=True) + max_dtype = np.iinfo(image.array.dtype).max + start, end = image.dask_array.min().compute(), image.dask_array.max().compute() + + channel_list = meta.omero.channels + + new_channel_list = [] + for channel in channel_list: + channel.extra_fields["window"] = { + "start": start, + "end": end, + "min": 0, + "max": max_dtype, + } + new_channel_list.append(channel) + + meta.omero.channels = new_channel_list + self._image_meta.write_meta(meta) + def derive_new_image( self, store: StoreLike, + name: str, + overwrite: bool = True, + **kwargs, ) -> "NgffImage": """Derive a new image from the current image. Args: store (StoreLike): The store to create the new image in. + name (str): The name of the new image. + overwrite (bool): Whether to overwrite the image if it exists + **kwargs: Additional keyword arguments. + Follow the same signature as `create_empty_ome_zarr_image`. Returns: NgffImage: The new image. """ - raise NotImplementedError("Deriving new images is not yet implemented.") + image_0 = self.get_image(highest_resolution=True) + + default_kwargs = { + "store": store, + "shape": image_0.on_disk_shape, + "chunks": image_0.array.chunks, + "dtype": image_0.array.dtype, + "on_disk_axis": image_0.dataset.on_disk_axes_names, + "pixel_sizes": image_0.pixel_size, + "xy_scaling_factor": self.image_meta.xy_scaling_factor, + "z_scaling_factor": self.image_meta.z_scaling_factor, + "time_spacing": image_0.dataset.time_spacing, + "time_units": image_0.dataset.time_axis_unit, + "num_levels": self.num_levels, + "name": name, + "channel_labels": image_0.channel_labels, + "channel_wavelengths": None, + "channel_kwargs": None, + "omero_kwargs": None, + "overwrite": overwrite, + "version": self.image_meta.version, + } + + default_kwargs.update(kwargs) + + create_empty_ome_zarr_image( + **default_kwargs, + ) + return NgffImage(store=store) diff --git a/src/ngio/core/utils.py b/src/ngio/core/utils.py index 25c473d..9fd5acd 100644 --- a/src/ngio/core/utils.py +++ b/src/ngio/core/utils.py @@ -4,7 +4,11 @@ from ngio.core.image_handler import Image from ngio.io import StoreLike -from ngio.ngff_meta import create_image_metadata, get_ngff_image_meta_handler +from ngio.ngff_meta import ( + create_image_metadata, + get_ngff_image_meta_handler, + create_label_metadata, +) from ngio.ngff_meta.fractal_image_meta import ( PixelSize, TimeUnits, @@ -30,7 +34,7 @@ def create_empty_ome_zarr_image( omero_kwargs: dict[str, Any] | None = None, overwrite: bool = True, version: str = "0.4", -) -> Image: +) -> None: """Create an empty OME-Zarr image with the given shape and metadata.""" if len(shape) != len(on_disk_axis): raise ValueError( @@ -88,9 +92,99 @@ def create_empty_ome_zarr_image( for dataset in image_meta.datasets: path = dataset.path - group.create_array( - name=path, fill_value=0, shape=shape, dtype=dtype, chunks=chunks + # V3 + # group.create_array( + # name=path, fill_value=0, shape=shape, dtype=dtype, chunks=chunks, + # ) + + group.zeros( + name=path, + shape=shape, + dtype=dtype, + chunks=chunks, + dimension_separator="/", ) # Todo redo this with when a proper build of pyramid id implemente shape = [int(s / sc) for s, sc in zip(shape, scaling_factor, strict=True)] + + if chunks is not None: + chunks = [int(c / sc) for c, sc in zip(chunks, scaling_factor, strict=True)] + + +def create_empty_ome_zarr_label( + store: StoreLike, + shape: list[int], + chunks: list[int] | None = None, + dtype: str = "uint16", + on_disk_axis: list[str] = ("t", "z", "y", "x"), + pixel_sizes: PixelSize | None = None, + xy_scaling_factor: float = 2.0, + z_scaling_factor: float = 1.0, + time_spacing: float = 1.0, + time_units: TimeUnits | str = TimeUnits.s, + num_levels: int = 5, + name: str | None = None, + overwrite: bool = True, + version: str = "0.4", +) -> None: + """Create an empty OME-Zarr image with the given shape and metadata.""" + if len(shape) != len(on_disk_axis): + raise ValueError( + "The number of dimensions in the shape must match the number of " + "axes in the on-disk axis." + ) + + image_meta = create_label_metadata( + on_disk_axis=on_disk_axis, + pixel_sizes=pixel_sizes, + xy_scaling_factor=xy_scaling_factor, + z_scaling_factor=z_scaling_factor, + time_spacing=time_spacing, + time_units=time_units, + num_levels=num_levels, + name=name, + version=version, + ) + + # Open the store (if it is not empty, fail) + mode = "w" if overwrite else "w-" + meta_handler = get_ngff_image_meta_handler( + store=store, version=version, meta_mode="label", mode=mode + ) + meta_handler.write_meta(image_meta) + group = meta_handler.group + group.attrs["image-label"] = {"version": version, "source": {"image": "../../"}} + # Create the empty image at each level in the pyramid + + # Return the an Image object + scaling_factor = [] + for ax in on_disk_axis: + if ax in ["x", "y"]: + scaling_factor.append(xy_scaling_factor) + elif ax == "z": + scaling_factor.append(z_scaling_factor) + else: + scaling_factor.append(1.0) + + for dataset in image_meta.datasets: + path = dataset.path + + # V3 + # group.create_array( + # name=path, fill_value=0, shape=shape, dtype=dtype, chunks=chunks, + # ) + + group.zeros( + name=path, + shape=shape, + dtype=dtype, + chunks=chunks, + dimension_separator="/", + ) + + # Todo redo this with when a proper build of pyramid id implemente + shape = [int(s / sc) for s, sc in zip(shape, scaling_factor, strict=True)] + + if chunks is not None: + chunks = [int(c / sc) for c, sc in zip(chunks, scaling_factor, strict=True)] diff --git a/src/ngio/ngff_meta/fractal_image_meta.py b/src/ngio/ngff_meta/fractal_image_meta.py index 04de513..3737edc 100644 --- a/src/ngio/ngff_meta/fractal_image_meta.py +++ b/src/ngio/ngff_meta/fractal_image_meta.py @@ -401,6 +401,15 @@ def scale(self) -> list[float]: """Get the scale transformation of the dataset in the canonical order.""" return [self._scale[i] for i in self._index_mapping.values() if i is not None] + @property + def time_spacing(self) -> float: + """Get the time spacing of the dataset.""" + if "t" not in self.axes_names: + return 1.0 + + scale_t = self.scale[self.index_mapping.get("t")] + return scale_t + @property def on_disk_scale(self) -> list[float]: """Get the scale transformation of the dataset in the on-disk order.""" @@ -717,6 +726,37 @@ def scale(self, path: str | None = None, idx: int | None = None) -> list[float]: """ return self.get_dataset(path=path, idx=idx).scale + def scaling_factors(self) -> list[float]: + scaling_factors = [] + for d1, d2 in zip(self.datasets[1:], self.datasets[:-1], strict=True): + scaling_factors.append([d1 / d2 for d1, d2 in zip(d1.scale, d2.scale)]) + + for sf in scaling_factors: + assert ( + sf == scaling_factors[0] + ), "Inconsistent scaling factors not well supported." + return scaling_factors[0] + + @property + def xy_scaling_factor(self) -> float: + """Get the xy scaling factor of the dataset.""" + scaling_factors = self.scaling_factors() + x_scaling_f = scaling_factors[self.index_mapping.get("x")] + y_scaling_f = scaling_factors[self.index_mapping.get("y")] + + if not np.allclose(x_scaling_f, y_scaling_f): + raise ValueError("Inconsistent xy scaling factor.") + return x_scaling_f + + @property + def z_scaling_factor(self) -> float: + """Get the z scaling factor of the dataset.""" + scaling_factors = self.scaling_factors() + if "z" not in self.axes_names: + return 1.0 + z_scaling_f = scaling_factors[self.index_mapping.get("z")] + return z_scaling_f + def translation( self, path: str | None = None, idx: int | None = None ) -> list[float] | None: diff --git a/src/ngio/ngff_meta/v04/zarr_utils.py b/src/ngio/ngff_meta/v04/zarr_utils.py index 9ddd95c..702c778 100644 --- a/src/ngio/ngff_meta/v04/zarr_utils.py +++ b/src/ngio/ngff_meta/v04/zarr_utils.py @@ -33,6 +33,7 @@ def check_ngff_image_meta_v04(store: StoreOrGroup) -> bool: """Check if a Zarr Group contains the OME-NGFF v0.4.""" store = open_group_wrapper(store=store, mode="r", zarr_format=2) attrs = dict(store.attrs) + print(attrs) multiscales = attrs.get("multiscales", None) if multiscales is None: return False @@ -158,6 +159,10 @@ def fractal_ngff_image_meta_to_vanilla_v04( datasets=dataset04, version="0.4", ) + + if isinstance(meta, LabelMeta): + return NgffImageMeta04(multiscales=[multiscale04]) + omero04 = None if meta.omero is None else Omero04(**meta.omero.model_dump()) return NgffImageMeta04( multiscales=[multiscale04], @@ -183,10 +188,11 @@ def write_ngff_image_meta_v04(group: Group, meta: ImageLabelMeta) -> None: raise ValueError( "The Zarr store does not contain the correct OME-Zarr version." ) - if meta.omero is not None: + + if isinstance(meta, ImageMeta) and meta.omero is not None: for c in meta.omero.channels: if "color" not in c.extra_fields: - c.extra_fields["color"] = "0x000000" + c.extra_fields["color"] = "00FFFF" meta04 = fractal_ngff_image_meta_to_vanilla_v04(meta=meta) group.attrs.update(meta04.model_dump(exclude=None)) diff --git a/src/ngio/tables/tables_group.py b/src/ngio/tables/tables_group.py index cfaff00..4c6097e 100644 --- a/src/ngio/tables/tables_group.py +++ b/src/ngio/tables/tables_group.py @@ -80,9 +80,9 @@ def __init__(self, group: StoreLike | zarr.Group) -> None: group = zarr.open_group(group, mode="a") if "tables" not in group: - raise ValueError("The NGFF image contains no 'tables' Group.") - - self._group: zarr.Group = group["tables"] + self._group = group.create_group("tables") + else: + self._group: zarr.Group = group["tables"] def _validate_list_of_tables(self, list_of_tables: list[str]) -> None: """Validate the list of tables.""" @@ -122,7 +122,6 @@ def list( ) list_of_typed_tables = [] for table_name in list_of_tables: - print(table_name) table = self._group[table_name] common_meta = CommonMeta(**table.attrs) if common_meta.type == type: diff --git a/src/ngio/tables/v1/feature_tables.py b/src/ngio/tables/v1/feature_tables.py index f08afe6..5c80904 100644 --- a/src/ngio/tables/v1/feature_tables.py +++ b/src/ngio/tables/v1/feature_tables.py @@ -134,7 +134,6 @@ def label_image_name(self, get_full_path: bool = False) -> str: def write(self) -> None: """Write the crrent state of the table to the Zarr file.""" - self._validate_feature_table(table=self.data_frame) self._write(group=self.table_group, table=self.data_frame, meta=self.meta) @staticmethod