|
| 1 | +# Copyright 2024 (C) BioVisionCenter, University of Zurich |
| 2 | +# |
| 3 | +# Original authors: |
| 4 | + |
| 5 | +# |
| 6 | +# This file is derived from a Fractal task core task developed by |
| 7 | +# Tommaso Comparin & Marco Franzon |
| 8 | +"""Task to remove singelton T dimension from an OME-Zarr.""" |
| 9 | +import logging |
| 10 | +from pathlib import Path |
| 11 | +from typing import Any, Sequence |
| 12 | + |
| 13 | +import dask.array as da |
| 14 | +import zarr |
| 15 | +from fractal_tasks_core.ngff import load_NgffImageMeta |
| 16 | +from fractal_tasks_core.pyramids import build_pyramid |
| 17 | +from pydantic.decorator import validate_arguments |
| 18 | + |
| 19 | +logger = logging.getLogger(__name__) |
| 20 | + |
| 21 | + |
| 22 | +def get_attrs_without_t(zarr_url: str): |
| 23 | + """ |
| 24 | + Generate zattrs without the t dimension. |
| 25 | +
|
| 26 | + Args: |
| 27 | + zarr_url: Path to the zarr image |
| 28 | + """ |
| 29 | + image_group = zarr.open_group(zarr_url) |
| 30 | + zattrs = image_group.attrs.asdict() |
| 31 | + # print(zattrs) |
| 32 | + for multiscale in zattrs["multiscales"]: |
| 33 | + # Update axes |
| 34 | + multiscale["axes"] = multiscale["axes"][1:] |
| 35 | + # Update coordinate Transforms |
| 36 | + for dataset in multiscale["datasets"]: |
| 37 | + for transform in dataset["coordinateTransformations"]: |
| 38 | + if transform["type"] == "scale": |
| 39 | + transform["scale"] = transform["scale"][1:] |
| 40 | + return zattrs |
| 41 | + |
| 42 | + |
| 43 | +@validate_arguments |
| 44 | +def drop_t_dimension( |
| 45 | + *, |
| 46 | + input_paths: Sequence[str], |
| 47 | + output_path: str, |
| 48 | + component: str, |
| 49 | + metadata: dict[str, Any], |
| 50 | + suffix: str = "no_T", |
| 51 | + # overwrite_input: bool = False, |
| 52 | +) -> dict[str, Any]: |
| 53 | + """ |
| 54 | + Drops singleton t dimension. |
| 55 | +
|
| 56 | + Args: |
| 57 | + input_paths: This parameter is not used by this task. |
| 58 | + This task only supports a single input path. |
| 59 | + (standard argument for Fractal tasks, managed by Fractal server). |
| 60 | + output_path: Path were the output of this task is stored. |
| 61 | + Example: `"/some/path/"` => puts the new OME-Zarr file in that |
| 62 | + folder. |
| 63 | + (standard argument for Fractal tasks, managed by Fractal server). |
| 64 | + component: Path to the OME-Zarr image in the OME-Zarr plate that |
| 65 | + is processed. Component is typically changed by the `copy_ome_zarr` |
| 66 | + task before, to point to a new mip Zarr file. |
| 67 | + Example: `"some_plate_mip.zarr/B/03/0"`. |
| 68 | + (standard argument for Fractal tasks, managed by Fractal server). |
| 69 | + metadata: Dictionary containing metadata about the OME-Zarr. |
| 70 | + This task requires the key `copy_ome_zarr` to be present in the |
| 71 | + metadata (as defined in `copy_ome_zarr` task). |
| 72 | + (standard argument for Fractal tasks, managed by Fractal server). |
| 73 | + suffix: Suffix to be used for the new Zarr image. If overwrite_input |
| 74 | + is True, this file is only temporary. |
| 75 | + """ |
| 76 | + # Preliminary checks |
| 77 | + if len(input_paths) > 1: |
| 78 | + raise NotImplementedError |
| 79 | + |
| 80 | + zarrurl_old = (Path(input_paths[0]).resolve() / component).as_posix() |
| 81 | + new_component = component + "_" + suffix |
| 82 | + zarrurl_new = (Path(input_paths[0]).resolve() / new_component).as_posix() |
| 83 | + logger.info(f"{zarrurl_old=}") |
| 84 | + logger.info(f"{zarrurl_new=}") |
| 85 | + |
| 86 | + # Read some parameters from metadata |
| 87 | + ngff_image = load_NgffImageMeta(zarrurl_old) |
| 88 | + |
| 89 | + # Check that T axis is the first axis: |
| 90 | + if not ngff_image.multiscale.axes[0].name == "t": |
| 91 | + logger.warning( |
| 92 | + f"The Zarr image {zarrurl_old} did not contain a T axis as its " |
| 93 | + f"first axis. The axes were: {ngff_image.multiscale.axes} \n" |
| 94 | + "The Drop T axis task is skipped" |
| 95 | + ) |
| 96 | + return {} |
| 97 | + |
| 98 | + # TODO: Generate attrs without the T dimension |
| 99 | + new_attrs = get_attrs_without_t(zarrurl_old) |
| 100 | + new_image_group = zarr.group(zarrurl_new) |
| 101 | + new_image_group.attrs.put(new_attrs) |
| 102 | + |
| 103 | + # TODO: Check if image contains labels & raise error (or even copy them) |
| 104 | + # FIXME: Check if image contains ROI tables & copy them |
| 105 | + |
| 106 | + # Load 0-th level |
| 107 | + data_tczyx = da.from_zarr(zarrurl_old + "/0") |
| 108 | + new_data = data_tczyx[0, ...] |
| 109 | + # Write to disk (triggering execution) |
| 110 | + logger.debug(f"Writing Zarr without T dimension to {zarrurl_new}") |
| 111 | + new_data.to_zarr( |
| 112 | + f"{zarrurl_new}/0", |
| 113 | + overwrite=True, |
| 114 | + dimension_separator="/", |
| 115 | + write_empty_chunks=False, |
| 116 | + ) |
| 117 | + logger.debug(f"Finished writing Zarr without T dimension to {zarrurl_new}") |
| 118 | + build_pyramid( |
| 119 | + zarrurl=zarrurl_new, |
| 120 | + overwrite=True, |
| 121 | + num_levels=ngff_image.num_levels, |
| 122 | + coarsening_xy=ngff_image.coarsening_xy, |
| 123 | + chunksize=new_data.chunksize, |
| 124 | + ) |
| 125 | + return {} |
| 126 | + |
| 127 | + |
| 128 | +if __name__ == "__main__": |
| 129 | + from fractal_tasks_core.tasks._utils import run_fractal_task |
| 130 | + |
| 131 | + run_fractal_task( |
| 132 | + task_function=drop_t_dimension, |
| 133 | + logger_name=logger.name, |
| 134 | + ) |
0 commit comments