From 136ddb54a1be3dc7477317062092fee9073d6862 Mon Sep 17 00:00:00 2001 From: Matt McCormick Date: Mon, 4 Mar 2024 23:15:33 -0500 Subject: [PATCH] WIP: ENH: Use cuCIM in CLI conversion --- ngff_zarr/cli.py | 27 +++++++ ngff_zarr/cli_input_to_ngff_image.py | 18 +---- ngff_zarr/cucim_image_to_multiscales.py | 99 +++++++++++++++++++++++++ ngff_zarr/cucim_image_to_ngff_image.py | 22 ++++++ ngff_zarr/detect_cli_io_backend.py | 2 + 5 files changed, 152 insertions(+), 16 deletions(-) create mode 100644 ngff_zarr/cucim_image_to_multiscales.py create mode 100644 ngff_zarr/cucim_image_to_ngff_image.py diff --git a/ngff_zarr/cli.py b/ngff_zarr/cli.py index 259c45d..044cb8f 100755 --- a/ngff_zarr/cli.py +++ b/ngff_zarr/cli.py @@ -28,6 +28,8 @@ from .cli_input_to_ngff_image import cli_input_to_ngff_image from .config import config +from .cucim_image_to_multiscales import cucim_image_to_multiscales +from .cucim_image_to_ngff_image import cucim_image_to_ngff_image from .detect_cli_io_backend import ( ConversionBackend, conversion_backends_values, @@ -336,6 +338,31 @@ def shutdown_client(sig_id, frame): # noqa: ARG001 _multiscales_to_ngff_zarr( live, args, output_store, rich_dask_progress, multiscales ) + elif input_backend is ConversionBackend.CUCIM: + try: + import cucim + + cuimage = cucim.CuImage(str(args.input[0])) + if args.chunks is None: + # Present the existing chunks and resolution levels + multiscales = cucim_image_to_multiscales(cuimage) + else: + ngff_image = cucim_image_to_ngff_image(cuimage) + multiscales = _ngff_image_to_multiscales( + live, + ngff_image, + args, + progress, + rich_dask_progress, + subtitle, + method, + ) + _multiscales_to_ngff_zarr( + live, args, output_store, rich_dask_progress, multiscales + ) + except ImportError: + sys.stdout.write("[red]Please install the [i]cucim[/i] package.\n") + sys.exit(1) elif input_backend is ConversionBackend.TIFFFILE: try: import tifffile diff --git a/ngff_zarr/cli_input_to_ngff_image.py b/ngff_zarr/cli_input_to_ngff_image.py index a85904f..ba82c9c 100644 --- a/ngff_zarr/cli_input_to_ngff_image.py +++ b/ngff_zarr/cli_input_to_ngff_image.py @@ -1,14 +1,13 @@ import sys -import numpy as np import zarr from dask.array.image import imread as daimread from rich import print +from .cucim_image_to_ngff_image import cucim_image_to_ngff_image from .detect_cli_io_backend import ConversionBackend from .from_ngff_zarr import from_ngff_zarr from .itk_image_to_ngff_image import itk_image_to_ngff_image -from .methods._support import _spatial_dims from .ngff_image import NgffImage from .to_ngff_image import to_ngff_image @@ -55,20 +54,7 @@ def imread(filename): import cucim cuimage = cucim.CuImage(str(input[0])) - data = np.array(cuimage) - dims = tuple(d for d in cuimage.dims.lower()) - spatial_dims = set(dims).intersection(_spatial_dims) - spatial_dims = [d for d in dims if d in spatial_dims] - spatial_dims_str = "".join(spatial_dims).upper() - translation = {d: 0.0 for d in spatial_dims} - for idx, dim in enumerate(spatial_dims): - # cucim: Should origin have a dim_order argument like spacing? - translation[dim] = cuimage.origin[idx] - spacing = cuimage.spacing(spatial_dims_str) - scale = {d: 1.0 for d in spatial_dims} - for idx, dim in enumerate(spatial_dims): - scale[dim] = spacing[idx] - return to_ngff_image(data, dims=dims, translation=translation, scale=scale) + return cucim_image_to_ngff_image(cuimage) except ImportError: print("[red]Please install the [i]cucim[/i] package.") sys.exit(1) diff --git a/ngff_zarr/cucim_image_to_multiscales.py b/ngff_zarr/cucim_image_to_multiscales.py new file mode 100644 index 0000000..4743c64 --- /dev/null +++ b/ngff_zarr/cucim_image_to_multiscales.py @@ -0,0 +1,99 @@ +import dask +import dask.array as da +import numpy as np + +from .methods._support import _spatial_dims +from .multiscales import Multiscales +from .to_ngff_image import to_ngff_image +from .zarr_metadata import Axis, Dataset, Metadata, Scale, Translation + + +def _read_region_data(cuimage, location, size, level): + return np.array(cuimage.read_region(location, size, level)) + + +def cucim_image_to_multiscales(cuimage) -> Multiscales: + dims = tuple(d for d in cuimage.dims.lower()) + spatial_dims = set(dims).intersection(_spatial_dims) + spatial_dims = [d for d in dims if d in spatial_dims] + spatial_dims_str = "".join(spatial_dims).upper() + + images = [] + + axes = [] + for dim in dims: + unit = None + if dim in {"x", "y", "z"}: + axis = Axis(name=dim, type="space", unit=unit) + elif dim == "c": + axis = Axis(name=dim, type="channel", unit=unit) + elif dim == "t": + axis = Axis(name=dim, type="time", unit=unit) + else: + msg = f"Dimension identifier is not valid: {dim}" + raise KeyError(msg) + axes.append(axis) + + for level in range(cuimage.resolutions["level_count"]): + scale_dimensions = cuimage.resolutions["level_dimensions"][level] + scale_downsamples = cuimage.resolutions["level_downsamples"][level] + scale_tile_size = cuimage.resolutions["level_tile_sizes"][level] + # hard coded for 2d + blocks = [] + for ii in range(scale_dimensions[0] // scale_tile_size[0]): + block_row = [] + for jj in range(scale_dimensions[1] // scale_tile_size[1]): + location = (ii * scale_tile_size[0], jj * scale_tile_size[1]) + size = scale_tile_size + block_row.append( + dask.delayed(_read_region_data)(cuimage, location, size, level) + ) + blocks.append(block_row) + data = da.block(blocks) + + spacing = cuimage.spacing(spatial_dims_str) + scale = {d: 1.0 for d in spatial_dims} + for idx, dim in enumerate(spatial_dims): + scale[dim] = spacing[idx] * scale_downsamples + + translation = {d: 0.0 for d in spatial_dims} + for idx, dim in enumerate(spatial_dims): + # cucim: Should origin have a dim_order argument like spacing? + translation[dim] = ( + cuimage.origin[idx] + spacing[idx] * scale_downsamples / 2 + ) + + image = to_ngff_image(data, dims=dims, translation=translation, scale=scale) + images.append(image) + + datasets = [] + for index, image in enumerate(images): + path = f"scale{index}/{image.name}" + scale = [] + for dim in image.dims: + if dim in image.scale: + scale.append(image.scale[dim]) + else: + scale.append(1.0) + translation = [] + for dim in image.dims: + if dim in image.translation: + translation.append(image.translation[dim]) + else: + translation.append(1.0) + coordinateTransformations = [Scale(scale), Translation(translation)] + dataset = Dataset( + path=path, coordinateTransformations=coordinateTransformations + ) + datasets.append(dataset) + metadata = Metadata( + axes=axes, + datasets=datasets, + name=image.name, + coordinateTransformations=None, + ) + return Multiscales( + images=images, + metadata=metadata, + scale_factors=cuimage.resolutions["level_downsamples"], + ) diff --git a/ngff_zarr/cucim_image_to_ngff_image.py b/ngff_zarr/cucim_image_to_ngff_image.py new file mode 100644 index 0000000..5aca358 --- /dev/null +++ b/ngff_zarr/cucim_image_to_ngff_image.py @@ -0,0 +1,22 @@ +import numpy as np + +from .methods._support import _spatial_dims +from .ngff_image import NgffImage +from .to_ngff_image import to_ngff_image + + +def cucim_image_to_ngff_image(cuimage) -> NgffImage: + data = np.array(cuimage) + dims = tuple(d for d in cuimage.dims.lower()) + spatial_dims = set(dims).intersection(_spatial_dims) + spatial_dims = [d for d in dims if d in spatial_dims] + spatial_dims_str = "".join(spatial_dims).upper() + translation = {d: 0.0 for d in spatial_dims} + for idx, dim in enumerate(spatial_dims): + # cucim: Should origin have a dim_order argument like spacing? + translation[dim] = cuimage.origin[idx] + spacing = cuimage.spacing(spatial_dims_str) + scale = {d: 1.0 for d in spatial_dims} + for idx, dim in enumerate(spatial_dims): + scale[dim] = spacing[idx] + return to_ngff_image(data, dims=dims, translation=translation, scale=scale) diff --git a/ngff_zarr/detect_cli_io_backend.py b/ngff_zarr/detect_cli_io_backend.py index ab293b7..5f821bd 100644 --- a/ngff_zarr/detect_cli_io_backend.py +++ b/ngff_zarr/detect_cli_io_backend.py @@ -105,6 +105,8 @@ def detect_cli_io_backend(input: List[str]) -> ConversionBackend: if extension in itk_supported_extensions: return ConversionBackend.ITK + extension = Path(input[0]).suffixes[-1].lower() + if importlib.util.find_spec("cucim") is not None: cucim_supported_extensions = (".svs", ".tif", ".tiff")