Rhapso 0.1.92__py3-none-any.whl

Rhapso/fusion/multiscale/aind_hcr_data_transformation/zeiss_job.py

@@ -0,0 +1,249 @@
+"""Module to handle zeiss data compression"""
+
+import logging
+import os
+import shutil
+import sys
+from pathlib import Path
+from time import time
+from typing import Any, List, Optional
+
+from aind_data_transformation.core import GenericEtl, JobResponse, get_parser
+from numcodecs.blosc import Blosc
+
+from .compress.czi_to_zarr import (
+    czi_stack_zarr_writer,
+)
+from .models import (
+    CompressorName,
+    ZeissJobSettings,
+)
+from .utils import utils
+
+logging.basicConfig(level=os.getenv("LOG_LEVEL", "WARNING"))
+
+
+class ZeissCompressionJob(GenericEtl[ZeissJobSettings]):
+    """Job to handle compressing and uploading Zeiss data."""
+
+    @staticmethod
+    def partition_list(
+        lst: List[Any], num_of_partitions: int
+    ) -> List[List[Any]]:
+        """Partitions a list"""
+        accumulated_list = []
+        for _ in range(num_of_partitions):
+            accumulated_list.append([])
+        for list_item_index, list_item in enumerate(lst):
+            a_index = list_item_index % num_of_partitions
+            accumulated_list[a_index].append(list_item)
+        return accumulated_list
+
+    def _get_partitioned_list_of_stack_paths(self) -> List[List[Path]]:
+        """
+        Scans through the input source and partitions a list of stack
+        paths that it finds there.
+        """
+        all_stack_paths = []
+        total_counter = 0
+        for p in (
+            Path(self.job_settings.input_source).joinpath("SPIM").glob("*.czi")
+        ):
+            if p.is_file():
+                total_counter += 1
+                all_stack_paths.append(p)
+
+        # Important to sort paths so every node computes the same list
+        all_stack_paths.sort(key=lambda x: str(x))
+        return self.partition_list(
+            all_stack_paths, self.job_settings.num_of_partitions
+        )
+
+    @staticmethod
+    def _get_voxel_resolution(acquisition_path: Path) -> List[float]:
+        """Get the voxel resolution from an acquisition.json file."""
+
+        if not acquisition_path.is_file():
+            raise FileNotFoundError(
+                f"acquisition.json file not found at: {acquisition_path}"
+            )
+
+        acquisition_config = utils.read_json_as_dict(acquisition_path)
+
+        # Grabbing a tile with metadata from acquisition - we assume all
+        # dataset was acquired with the same resolution
+        tile_coord_transforms = acquisition_config["tiles"][0][
+            "coordinate_transformations"
+        ]
+
+        scale_transform = [
+            x["scale"] for x in tile_coord_transforms if x["type"] == "scale"
+        ][0]
+
+        x = float(scale_transform[0])
+        y = float(scale_transform[1])
+        z = float(scale_transform[2])
+
+        return [z, y, x]
+
+    def _get_compressor(self) -> Optional[Blosc]:
+        """
+        Utility method to construct a compressor class.
+        Returns
+        -------
+        Blosc | None
+          An instantiated Blosc compressor. Return None if not set in configs.
+
+        """
+        if self.job_settings.compressor_name == CompressorName.BLOSC:
+            return Blosc(**self.job_settings.compressor_kwargs)
+        else:
+            return None
+
+    def _write_stacks(self, stacks_to_process: List) -> None:
+        """
+        Write a list of stacks.
+        Parameters
+        ----------
+        stacks_to_process : List
+
+        Returns
+        -------
+        None
+
+        """
+
+        if not len(stacks_to_process):
+            logging.info("No stacks to process!")
+            return
+
+        compressor = self._get_compressor()
+
+        # Acquisition path in root folder
+        acquisition_path = self.job_settings.input_source.joinpath(
+            "acquisition.json"
+        )
+
+        # Getting voxel resolution
+        voxel_size_zyx = self._get_voxel_resolution(
+            acquisition_path=acquisition_path
+        )
+
+        # Converting CZI tiles to Multiscale OMEZarr
+        for stack in stacks_to_process:
+            logging.info(f"Converting {stack}")
+            stack_name = stack.stem
+
+            output_path = Path(self.job_settings.output_directory)
+
+            msg = (
+                f"Voxel resolution ZYX {voxel_size_zyx} for {stack} "
+                f"with name {stack_name} - output: {output_path}"
+            )
+            logging.info(msg)
+
+            czi_stack_zarr_writer(
+                czi_path=str(stack),
+                output_path=output_path,
+                voxel_size=voxel_size_zyx,
+                final_chunksize=self.job_settings.chunk_size,
+                scale_factor=self.job_settings.scale_factor,
+                n_lvls=self.job_settings.downsample_levels,
+                channel_name=stack_name,
+                stack_name=f"{stack_name}.ome.zarr",
+                logger=logging,
+                writing_options=compressor,
+                target_size_mb=self.job_settings.target_size_mb,
+            )
+
+            if self.job_settings.s3_location is not None:
+                channel_zgroup_file = output_path / ".zgroup"
+                s3_channel_zgroup_file = (
+                    f"{self.job_settings.s3_location}/.zgroup"
+                )
+                logging.info(
+                    f"Uploading {channel_zgroup_file} to "
+                    f"{s3_channel_zgroup_file}"
+                )
+                utils.copy_file_to_s3(
+                    channel_zgroup_file, s3_channel_zgroup_file
+                )
+                ome_zarr_stack_name = f"{stack_name}.ome.zarr"
+                ome_zarr_stack_path = output_path.joinpath(ome_zarr_stack_name)
+                s3_stack_dir = (
+                    f"{self.job_settings.s3_location}/"
+                    f"{ome_zarr_stack_name}"
+                )
+                logging.info(
+                    f"Uploading {ome_zarr_stack_path} to {s3_stack_dir}"
+                )
+                utils.sync_dir_to_s3(ome_zarr_stack_path, s3_stack_dir)
+                logging.info(f"Removing: {ome_zarr_stack_path}")
+                # Remove stack if uploaded to s3. We can potentially do all
+                # the stacks in the partition in parallel using dask to speed
+                # this up
+                shutil.rmtree(ome_zarr_stack_path)
+
+    def _upload_derivatives_folder(self):
+        """
+        Uploads the derivatives folder inside of
+        the SPIM folder in the cloud.
+        """
+        s3_derivatives_dir = f"{self.job_settings.s3_location}/derivatives"
+        derivatives_path = Path(self.job_settings.input_source).joinpath(
+            "derivatives"
+        )
+
+        if not derivatives_path.exists():
+            raise FileNotFoundError(f"{derivatives_path} does not exist.")
+
+        if self.job_settings.s3_location is not None:
+            logging.info(
+                f"Uploading {derivatives_path} to {s3_derivatives_dir}"
+            )
+            utils.sync_dir_to_s3(derivatives_path, s3_derivatives_dir)
+            logging.info(f"{derivatives_path} uploaded to s3.")
+
+    def run_job(self):
+        """Main entrypoint to run the job."""
+        job_start_time = time()
+
+        # Reading data within the SPIM folder
+        partitioned_list = self._get_partitioned_list_of_stack_paths()
+
+        # Upload derivatives folder
+        if self.job_settings.partition_to_process == 0:
+            self._upload_derivatives_folder()
+
+        stacks_to_process = partitioned_list[
+            self.job_settings.partition_to_process
+        ]
+
+        self._write_stacks(stacks_to_process=stacks_to_process)
+        total_job_duration = time() - job_start_time
+        return JobResponse(
+            status_code=200, message=f"Job finished in {total_job_duration}"
+        )
+
+
+# TODO: Add this to core aind_data_transformation class
+def job_entrypoint(sys_args: list):
+    """Main function"""
+    parser = get_parser()
+    cli_args = parser.parse_args(sys_args)
+    if cli_args.job_settings is not None:
+        job_settings = ZeissJobSettings.model_validate_json(
+            cli_args.job_settings
+        )
+    elif cli_args.config_file is not None:
+        job_settings = ZeissJobSettings.from_config_file(cli_args.config_file)
+    else:
+        # Construct settings from env vars
+        job_settings = ZeissJobSettings()
+    job = ZeissCompressionJob(job_settings=job_settings)
+    job_response = job.run_job()
+    logging.info(job_response.model_dump_json())
+
+
+if __name__ == "__main__":
+    job_entrypoint(sys.argv[1:])

Rhapso/fusion/multiscale/aind_z1_radial_correction/__init__.py

@@ -0,0 +1,21 @@
+"""Z1 radial correction."""
+
+__version__ = "0.0.4"
+__authors__ = ["Camilo Laiton", "Carson Berry", "Tim Wang"]
+__author_emails__ = [
+    "camilo.laiton@alleninstitute.org",
+    "carson.berry@alleninstitute.org",
+    "tim.wang@alleninstitute.org",
+]
+__license__ = "MIT"
+__description__ = "Image preprocessing for radial correction"
+__url__ = "https://github.com/AllenNeuralDynamics/aind-z1-radial-correction"
+
+__maintainers__ = ["Camilo Laiton", "Carson Berry"]
+__maintainer_emails__ = [
+    "carson.berry@alleninstitute.org",
+    "camilo.laiton@alleninstitute.org",
+]
+__title__ = "aind-z1-radial-correction"
+__status__ = "Development"  # 'Development' 'Production', 'Beta'
+__pipeline_version__ = "0.0.1"

Rhapso/fusion/multiscale/aind_z1_radial_correction/array_to_zarr.py

@@ -0,0 +1,257 @@
+"""
+Writes a multiscale zarrv3 dataset from an array
+"""
+
+import logging
+from pathlib import Path
+from typing import Dict, List, Optional
+import dask.array as da
+import numpy as np
+import zarr
+from Rhapso.fusion.multiscale.aind_hcr_data_transformation.compress.czi_to_zarr import (
+    _get_pyramid_metadata,
+    compute_pyramid,
+    write_ome_ngff_metadata,
+)
+from Rhapso.fusion.multiscale.aind_hcr_data_transformation.compress.zarr_writer import (
+    BlockedArrayWriter,
+)
+from Rhapso.fusion.multiscale.aind_hcr_data_transformation.utils.utils import pad_array_n_d
+from numcodecs.blosc import Blosc
+from numpy.typing import ArrayLike
+from ome_zarr.io import parse_url
+from zarr.errors import ContainsGroupError
+from zarr.storage import FSStore
+from .utils.utils import get_parent_path, is_s3_path
+
+
+def safe_create_zarr_group(store, path: str = "", **kwargs):
+    """
+    Safe creation of the zarr group.
+
+    Parameters
+    ----------
+    store
+        Zarr store
+    path: str
+        Path to the creation of the zarr group
+        Default: ''
+
+    Returns
+    -------
+    Zarr.group
+        Zarr group pointing to where the data is written
+    """
+    try:
+        return zarr.group(store=store, path=path, overwrite=False, **kwargs)
+    except ContainsGroupError:
+        # Group already exists, which is expected with multiple workers
+        return zarr.open_group(store=store, path=path, mode="r+")
+
+
+def convert_array_to_zarr(
+    array: ArrayLike,
+    chunk_size: List[int],
+    output_path: str,
+    voxel_size: List[float],
+    n_lvls: Optional[int] = 6,
+    scale_factor: Optional[List[int]] = [2, 2, 2],
+    compressor_kwargs: Optional[Dict] = {
+        "cname": "zstd",
+        "clevel": 3,
+        "shuffle": Blosc.SHUFFLE,
+    },
+    target_size_mb: Optional[int] = 24000,
+):
+    """
+    Converts an array to zarr format
+
+    Parameters
+    ----------
+    array: ArrayLike
+        Array to convert to zarr v3
+
+    chunk_size: List[int]
+        Chunksize in each shard
+
+    output_path: str
+        Output path. It must contain the ome.zarr
+        extension attached.
+
+    voxel_size: List[float]
+        Voxel size
+
+    n_lvls: Optional[int]
+        Number of downsampled levels to write.
+        Default: 6
+
+    scale_factor: Optional[List[int]]
+        Scaling factor per axis. Default: [2, 2, 2]
+
+    compressor_kwargs: Optional[Dict]
+        Compressor parameters
+        Default: {"cname": "zstd", "clevel": 3, "shuffle": "shuffle"}
+    """
+    logger = logging.getLogger(__name__)
+    array = pad_array_n_d(array)
+    dataset_shape = tuple(i for i in array.shape if i != 1)
+    extra_axes = (1,) * (5 - len(dataset_shape))
+    dataset_shape = extra_axes + dataset_shape
+    chunk_size = ([1] * (5 - len(chunk_size))) + chunk_size
+
+    #verify that the chunksize is not larger than the dataset shape
+    for i, val in enumerate(dataset_shape):
+        if chunk_size[i] > val:
+            chunk_size[i] = val
+
+    compressor = Blosc(
+        cname=compressor_kwargs["cname"],
+        clevel=compressor_kwargs["clevel"],
+        shuffle=compressor_kwargs["shuffle"],
+        blocksize=0,
+    )
+
+    # Getting channel color
+    channel_colors = None
+    stack_name = Path(output_path).name
+    parent_path = get_parent_path(output_path)
+    # Creating Zarr dataset in s3 or local
+    if is_s3_path(output_path):
+        store = FSStore(parent_path, mode="w", dimension_separator="/")
+    else:
+        store = parse_url(path=parent_path, mode="w").store
+
+    root_group = safe_create_zarr_group(store=store)
+
+    # Using 1 thread since is in single machine.
+    # Avoiding the use of multithreaded due to GIL
+    if np.issubdtype(array.dtype, np.integer):
+        np_info_func = np.iinfo
+
+    else:
+        # Floating point
+        np_info_func = np.finfo
+
+    # Getting min max metadata for the dtype
+    channel_minmax = [
+        (
+            # int(np_info_func(array.dtype).min),
+            int(0.0),
+            # int(np_info_func(array.dtype).max),
+            int(1.0),
+        )
+        for _ in range(dataset_shape[1])
+    ]
+
+    # Setting values for CZI
+    # Ideally we would use da.percentile(image_data, (0.1, 95))
+    # However, it would take so much time and resources and it is
+    # not used that much on neuroglancer
+    channel_startend = [(int(0), int(1.0)) for _ in range(dataset_shape[1])]
+
+    # Writing OME-NGFF metadata
+    scale_factor = [int(s) for s in scale_factor]
+    voxel_size = [float(v) for v in voxel_size]
+
+    new_channel_group = root_group.create_group(
+        name=stack_name, overwrite=True
+    )
+
+    # Writing OME-NGFF metadata
+    write_ome_ngff_metadata(
+        group=new_channel_group,
+        arr_shape=dataset_shape,
+        image_name=stack_name,
+        n_lvls=n_lvls,
+        scale_factors=scale_factor,
+        voxel_size=voxel_size,
+        channel_names=None,
+        channel_colors=channel_colors,
+        channel_minmax=channel_minmax,
+        channel_startend=channel_startend,
+        metadata=_get_pyramid_metadata(),
+        final_chunksize=chunk_size,
+        origin = [0,0,0]
+    )
+
+    # Writing first multiscale by default
+    pyramid_group = new_channel_group.create_dataset(
+        name="0",
+        shape=dataset_shape,
+        chunks=chunk_size,
+        dtype=array.dtype,
+        compressor=compressor,
+        dimension_separator="/",
+        overwrite=True,
+    )
+
+    # Writing multiscales
+    # Handle both numpy arrays and dask arrays
+    if isinstance(array, da.Array):
+        # Already a dask array, rechunk if needed
+        previous_scale = da.rechunk(array, chunks=pyramid_group.chunks)
+    else:
+        # Convert numpy array to dask array
+        previous_scale = da.from_array(array, pyramid_group.chunks)
+
+    block_shape = list(
+        BlockedArrayWriter.get_block_shape(
+            arr=previous_scale,
+            target_size_mb=target_size_mb,
+            chunks=chunk_size,
+        )
+    )
+    block_shape = extra_axes + tuple(block_shape)
+
+    logger.info(f"Writing {n_lvls} pyramid levels...")
+    
+    for level in range(0, n_lvls):
+        if not level:
+            array_to_write = previous_scale
+            logger.info(f"Level {level}/{n_lvls-1}: Writing full resolution - shape {array_to_write.shape}")
+
+        else:
+            previous_scale = da.from_zarr(pyramid_group, pyramid_group.chunks)
+            new_scale_factor = (
+                [1] * (len(previous_scale.shape) - len(scale_factor))
+            ) + scale_factor
+
+            logger.info(f"Level {level}/{n_lvls-1}: Computing downsampled pyramid (scale factor: {scale_factor})...")
+            previous_scale_pyramid, _ = compute_pyramid(
+                data=previous_scale,
+                scale_axis=new_scale_factor,
+                chunks=chunk_size,
+                n_lvls=2,
+            )
+            array_to_write = previous_scale_pyramid[-1]
+
+            logger.info(f"Level {level}/{n_lvls-1}: Downsampled to shape {array_to_write.shape}")
+
+            pyramid_group = new_channel_group.create_dataset(
+                name=str(level),
+                shape=array_to_write.shape,
+                chunks=chunk_size,
+                dtype=array_to_write.dtype,
+                compressor=compressor,
+                dimension_separator="/",
+                overwrite=True,
+            )
+
+        logger.info(f"Level {level}/{n_lvls-1}: Writing to storage...")
+        BlockedArrayWriter.store(array_to_write, pyramid_group, block_shape)
+        logger.info(f"Level {level}/{n_lvls-1}: ✓ Complete ({level+1}/{n_lvls} levels done)")
+
+if __name__ == "__main__":
+    BASE_PATH = "/data"
+    tilename = "Tile_X_0000_Y_0011_Z_0000_ch_488.ome.zarr"
+    test_dataset = f"HCR_785830_2025-03-19_17-00-00/SPIM/{tilename}"
+    scale = "0"
+
+    dataset = da.from_zarr(f"{BASE_PATH}/{test_dataset}/{scale}").compute()
+    convert_array_to_zarr(
+        array=dataset,
+        voxel_size=[1.0] * 3,
+        shard_size=[512] * 3,
+        chunk_size=[128] * 3,
+        output_path="/results/test.ome.zarr",
+    )