Rhapso 0.1.92__py3-none-any.whl

Rhapso/fusion/multiscale/aind_z1_radial_correction/utils/utils.py

@@ -0,0 +1,266 @@
+"""
+Utility functions for the radial correction step
+"""
+
+import json
+import os
+from pathlib import Path
+from typing import List
+from urllib.parse import urlparse
+
+# Commented out - only needed for generate_processing function which is not used by worker
+# from aind_data_schema.core.processing import (
+#     DataProcess,
+#     PipelineProcess,
+#     Processing,
+# )
+# from packaging import version
+
+
+# def generate_processing(
+#     data_processes: List[DataProcess],
+#     dest_processing: str,
+#     prefix: str,
+#     processor_full_name: str,
+#     pipeline_version: str,
+# ):
+#     """
+#     Generates data description for the output folder.
+#
+#     Parameters
+#     ------------------------
+#
+#     data_processes: List[dict]
+#         List with the processes aplied in the pipeline.
+#
+#     dest_processing: PathLike
+#         Path where the processing file will be placed.
+#
+#     processor_full_name: str
+#         Person in charged of running the pipeline
+#         for this data asset
+#
+#     pipeline_version: str
+#         Terastitcher pipeline version
+#
+#     """
+#     # flake8: noqa: E501
+#     processing_pipeline = PipelineProcess(
+#         data_processes=data_processes,
+#         processor_full_name=processor_full_name,
+#         pipeline_version=pipeline_version,
+#         pipeline_url="",
+#         note="Metadata for radial correction",
+#     )
+#
+#     processing = Processing(
+#         processing_pipeline=processing_pipeline,
+#         notes="This processing only contains metadata about radial correction \
+#             and needs to be compiled with other steps at the end",
+#     )
+#
+#     processing.write_standard_file(
+#         output_directory=dest_processing, prefix=prefix
+#     )
+
+
+def read_json_as_dict(filepath: str) -> dict:
+    """
+    Reads a json as dictionary.
+
+    Parameters
+    ------------------------
+
+    filepath: PathLike
+        Path where the json is located.
+
+    Returns
+    ------------------------
+
+    dict:
+        Dictionary with the data the json has.
+
+    """
+
+    dictionary = {}
+
+    if os.path.exists(filepath):
+        with open(filepath) as json_file:
+            dictionary = json.load(json_file)
+
+    return dictionary
+
+
+def get_voxel_resolution(acquisition_path: Path) -> List[float]:
+    """
+    Get the voxel resolution from an acquisition.json file.
+
+    Parameters
+    ----------
+    acquisition_path: Path
+        Path to the acquisition.json file.
+    Returns
+    -------
+    List[float]
+        Voxel resolution in the format [z, y, x].
+    """
+    if not Path(acquisition_path).is_file():
+        raise FileNotFoundError(
+            f"acquisition.json file not found at: {acquisition_path}"
+        )
+
+    acquisition_config = read_json_as_dict(str(acquisition_path))
+
+    schema_version = acquisition_config.get("schema_version")
+    print(f"Schema version: {schema_version}")
+
+    if version.parse(schema_version) >= version.parse("2.0.0"):
+        return _get_voxel_resolution_v2(acquisition_config)
+    else:
+        return _get_voxel_resolution_v1(acquisition_config)
+
+
+def _get_voxel_resolution_v1(acquisition_config: dict) -> List[float]:
+    """
+    Get the voxel resolution from an acquisition.json file.
+
+    Parameters
+    ----------
+    acquisition_config: Dict
+        Dictionary with the acquisition.json data.
+    Returns
+    -------
+    List[float]
+        Voxel resolution in the format [z, y, x].
+    """
+
+    if not acquisition_config:
+        raise ValueError("acquisition.json file is empty or invalid.")
+
+    # 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_voxel_resolution_v2(acquisition_config: dict) -> List[float]:
+    """
+    Get the voxel resolution from an acquisition.json in
+    aind-data-schema v2 format.
+
+    Parameters
+    ----------
+    acquisition_config: Dict
+        Dictionary with the acquisition.json data.
+
+    Returns
+    -------
+    List[float]
+        Voxel resolution in the format [z, y, x].
+    """
+    try:
+        data_stream = acquisition_config.get("data_streams", [])[0]
+        configuration = data_stream.get("configurations", [])[0]
+        image = configuration.get("images", [])[0]
+        image_to_acquisition_transform = image[
+            "image_to_acquisition_transform"
+        ]
+    except (IndexError, AttributeError, KeyError) as e:
+        raise ValueError(
+            "acquisition_config structure is invalid or missing "
+            "required fields"
+        ) from e
+
+    scale_transform = [
+        x["scale"]
+        for x in image_to_acquisition_transform
+        if x["object_type"] == "Scale"
+    ][0]
+
+    x = float(scale_transform[0])
+    y = float(scale_transform[1])
+    z = float(scale_transform[2])
+
+    return [z, y, x]
+
+
+def validate_capsule_inputs(input_elements: List[str]) -> List[str]:
+    """
+    Validates input elemts for a capsule in
+    Code Ocean.
+
+    Parameters
+    -----------
+    input_elements: List[str]
+        Input elements for the capsule. This
+        could be sets of files or folders.
+
+    Returns
+    -----------
+    List[str]
+        List of missing files
+    """
+
+    missing_inputs = []
+    for required_input_element in input_elements:
+        required_input_element = Path(required_input_element)
+
+        if not required_input_element.exists():
+            missing_inputs.append(str(required_input_element))
+
+    return missing_inputs
+
+
+def is_s3_path(path: str) -> bool:
+    """
+    Checks if a path is an s3 path
+
+    Parameters
+    ----------
+    path: str
+        Provided path
+
+    Returns
+    -------
+    bool
+        True if it is a S3 path,
+        False if not.
+    """
+    parsed = urlparse(str(path))
+    return parsed.scheme == "s3"
+
+
+def get_parent_path(path: str) -> str:
+    """
+    Gets parent path
+
+    Parameters
+    ----------
+    path: str
+        Provided path
+
+    Returns
+    -------
+    str
+        Parent path
+    """
+    parsed = urlparse(path)
+    if parsed.scheme == "s3":
+        # Remove the last part of the S3 key
+        parts = parsed.path.strip("/").split("/")
+        parent_key = "/".join(parts[:-1])
+        return f"s3://{parsed.netloc}/{parent_key}"
+    else:
+        # Local path fallback
+        return str(Path(path).parent)

Rhapso/fusion/multiscale/aind_z1_radial_correction/worker.py

@@ -0,0 +1,89 @@
+"""
+Worker script to run multiscale conversion on a zarr dataset
+"""
+
+import os
+import sys
+from pathlib import Path
+import dask.array as da
+
+from Rhapso.fusion.multiscale.aind_z1_radial_correction.array_to_zarr import convert_array_to_zarr
+from Rhapso.fusion.multiscale.aind_z1_radial_correction.utils import utils
+
+
+def run():
+    """
+    Main run function for multiscale conversion
+    """
+    
+    # Input and output paths
+    input_zarr_path = "s3://martin-test-bucket/output/channel_488.zarr"
+    output_zarr_path = "s3://martin-test-bucket/output/multiscale_channel_488.zarr"
+    
+    print(f"Loading data from: {input_zarr_path}")
+    
+    # Load the zarr dataset
+    # Assuming the data is in the root or scale "0" of the zarr
+    try:
+        # Try loading from scale "0" first (common for OME-Zarr)
+        dataset = da.from_zarr(f"{input_zarr_path}/0")
+        print(f"Loaded data from {input_zarr_path}/0")
+    except Exception as e:
+        # If scale "0" doesn't exist, try loading from root
+        print(f"Could not load from scale 0: {e}")
+        try:
+            dataset = da.from_zarr(input_zarr_path)
+            print(f"Loaded data from {input_zarr_path}")
+        except Exception as e2:
+            raise ValueError(
+                f"Could not load data from {input_zarr_path} or {input_zarr_path}/0. Error: {e2}"
+            )
+    
+    print(f"Dataset shape: {dataset.shape}")
+    print(f"Dataset dtype: {dataset.dtype}")
+    print(f"Dataset chunks: {dataset.chunks}")
+    
+    # Compute the array (load into memory)
+    # For large datasets, you might want to process this differently
+    print("Computing array...")
+    array = dataset.compute()
+    print("Array computed successfully")
+    
+    # Set parameters for multiscale conversion
+    # Adjust these parameters based on your data characteristics
+    chunk_size = [128, 128, 128]  # Chunk size for the output zarr
+    voxel_size = [1.0, 1.0, 1.0]  # Voxel size in micrometers (adjust if known)
+    n_lvls = 6  # Number of pyramid levels
+    scale_factor = [2, 2, 2]  # Downsampling factor per level
+    
+    compressor_kwargs = {
+        "cname": "zstd",
+        "clevel": 3,
+        "shuffle": 2,  # Blosc.SHUFFLE
+    }
+    
+    print(f"Converting to multiscale format and writing to: {output_zarr_path}")
+    print(f"Chunk size: {chunk_size}")
+    print(f"Voxel size: {voxel_size}")
+    print(f"Number of levels: {n_lvls}")
+    print(f"Scale factor: {scale_factor}")
+    
+    # Convert to multiscale zarr
+    convert_array_to_zarr(
+        array=array,
+        chunk_size=chunk_size,
+        output_path=output_zarr_path,
+        voxel_size=voxel_size,
+        n_lvls=n_lvls,
+        scale_factor=scale_factor,
+        compressor_kwargs=compressor_kwargs,
+        target_size_mb=24000,
+    )
+    
+    print(f"Multiscale conversion completed successfully!")
+    print(f"Output written to: {output_zarr_path}")
+
+
+if __name__ == "__main__":
+    run()
+

Rhapso/fusion/multiscale_worker.py

@@ -0,0 +1,113 @@
+"""
+Worker script to run multiscale conversion on a zarr dataset
+"""
+
+import os
+import sys
+from pathlib import Path
+import dask.array as da
+import logging
+
+from Rhapso.fusion.multiscale.aind_z1_radial_correction.array_to_zarr import convert_array_to_zarr
+
+# Setup logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+
+def run():
+    """
+    Main run function for multiscale conversion
+    """
+    
+    # Input and output paths
+    input_zarr_path = "s3://martin-test-bucket/output7/channel_488.zarr"
+    output_zarr_path = "s3://martin-test-bucket/output7/multiscale_channel_488.zarr"
+    
+    # Set parameters for multiscale conversion
+    # Adjust these parameters based on your data characteristics
+    chunk_size = [128, 128, 128]  # Chunk size for the output zarr
+    voxel_size = [1.0, 1.0, 1.0]  # Voxel size in micrometers (adjust if known)
+    n_lvls = 6  # Number of pyramid levels
+    scale_factor = [2, 2, 2]  # Downsampling factor per level
+
+    logger.info(f"Starting multiscale conversion")
+    logger.info(f"Input: {input_zarr_path}")
+    logger.info(f"Output: {output_zarr_path}")
+    
+    # Load the zarr dataset
+    # Assuming the data is in the root or scale "0" of the zarr
+    try:
+        # Try loading from scale "0" first (common for OME-Zarr)
+        logger.info(f"Attempting to load from {input_zarr_path}/0...")
+        sys.stdout.flush()
+        dataset = da.from_zarr(f"{input_zarr_path}/0")
+        logger.info(f"Successfully loaded data from {input_zarr_path}/0")
+    except Exception as e:
+        # If scale "0" doesn't exist, try loading from root
+        logger.warning(f"Could not load from scale 0: {e}")
+        try:
+            logger.info(f"Attempting to load from {input_zarr_path}...")
+            sys.stdout.flush()
+            dataset = da.from_zarr(input_zarr_path)
+            logger.info(f"Successfully loaded data from {input_zarr_path}")
+        except Exception as e2:
+            logger.error(f"Failed to load data: {e2}")
+            raise ValueError(
+                f"Could not load data from {input_zarr_path} or {input_zarr_path}/0. Error: {e2}"
+            )
+    
+    logger.info(f"Dataset shape: {dataset.shape}")
+    logger.info(f"Dataset dtype: {dataset.dtype}")
+    logger.info(f"Dataset chunks: {dataset.chunks}")
+    
+    # Calculate dataset size
+    import numpy as np
+    dtype_bytes = np.dtype(dataset.dtype).itemsize
+    total_size_gb = np.prod(dataset.shape) * dtype_bytes / (1024**3)
+    logger.info(f"Dataset size: {total_size_gb:.2f} GB")
+    
+    # Use dask array directly instead of computing (don't load into memory)
+    logger.info("Using Dask array for lazy/chunked processing (not loading into memory)")
+    array = dataset
+    
+    
+    compressor_kwargs = {
+        "cname": "zstd",
+        "clevel": 3,
+        "shuffle": 2,  # Blosc.SHUFFLE
+    }
+    
+    logger.info("=" * 60)
+    logger.info("Starting multiscale conversion with parameters:")
+    logger.info(f"  Output path: {output_zarr_path}")
+    logger.info(f"  Chunk size: {chunk_size}")
+    logger.info(f"  Voxel size: {voxel_size}")
+    logger.info(f"  Number of levels: {n_lvls}")
+    logger.info(f"  Scale factor: {scale_factor}")
+    logger.info("=" * 60)
+    sys.stdout.flush()
+    
+    # Convert to multiscale zarr
+    convert_array_to_zarr(
+        array=array,
+        chunk_size=chunk_size,
+        output_path=output_zarr_path,
+        voxel_size=voxel_size,
+        n_lvls=n_lvls,
+        scale_factor=scale_factor,
+        compressor_kwargs=compressor_kwargs,
+        target_size_mb=24000,
+    )
+    
+    logger.info("=" * 60)
+    logger.info("MULTISCALE CONVERSION COMPLETED SUCCESSFULLY!")
+    logger.info(f"Output written to: {output_zarr_path}")
+    logger.info("=" * 60)
+
+
+if __name__ == "__main__":
+    run()

Rhapso/fusion/neuroglancer_link_gen/__init__.py

@@ -0,0 +1,8 @@
+"""
+Package for the generation of neuroglancer links
+"""
+__version__ = "1.0.17"
+
+# flake8: noqa: F401
+from .ng_layer import NgLayer
+from .ng_state import NgState

Rhapso/fusion/neuroglancer_link_gen/dispim_link.py

@@ -0,0 +1,235 @@
+"""
+Library for generating dispim link.
+"""
+import pathlib
+
+from . import link_utils
+import numpy as np
+from ng_state import NgState
+from parsers import XmlParser
+from utils import transfer
+from typing import List, Dict, Tuple
+
+
+def apply_deskewing(matrix_3x4: np.ndarray, theta: float = 45) -> np.ndarray:
+    """
+    Compounds deskewing transform to input 3x4 matrix:
+        Deskewing @ Input_3x4_Matrix.
+
+    Parameters
+    ------------------------
+    matrix_3x4: np.ndarray
+        3x4 numpy array representing transformation matrix applied to tile.
+
+    theta: float
+        Angle of lens during acquisition.
+
+    Returns
+    ------------------------
+    np.ndarray:
+        3x4 numpy array composite transform.
+
+    """
+
+    # Deskewing
+    # X vector => XZ direction
+    deskew_factor = np.tan(np.deg2rad(theta))
+    deskew = np.array([[1, 0, 0], [0, 1, 0], [deskew_factor, 0, 1]])
+    matrix_3x4 = deskew @ matrix_3x4
+
+    return matrix_3x4
+
+
+def generate_dispim_link(
+    base_channel_xml_path: str,
+    # cross_channel_xml_path: str,
+    s3_path: str,
+    max_dr: int = 800,
+    opacity: float = 0.5,
+    blend: str = "additive",
+    deskew_angle: int = 45,
+    output_json_path: str = ".",
+    spim_foldername="SPIM.ome.zarr",
+) -> None:
+    """
+    Creates an neuroglancer link to visualize
+    registration transforms on dispim dataset pre-fusion.
+
+    Parameters
+    ------------------------
+    base_channel_xml_path: str
+        Path to xml file acquired from tile-to-tile
+        registration within the base channel.
+        These registrations are reused for
+        registering tiles in all other channels.
+
+    cross_channel_xml_path: str
+        Path to xml file acquired from channel-to-channel registration.
+        These registrations are prepended to each tile registration.
+
+    s3_path: str
+        Path of s3 bucket where dipim dataset is located.
+
+    output_json_path: str
+        Local path to write process_output.json file that nueroglancer reads.
+
+    Returns
+    ------------------------
+    None
+    """
+
+    # Gather base channel xml info
+    vox_sizes: Tuple[float, float, float] = XmlParser.extract_tile_vox_size(
+        base_channel_xml_path
+    )
+    tile_paths: Dict[int, str] = XmlParser.extract_tile_paths(
+        base_channel_xml_path
+    )
+    tile_transforms: Dict[int, List[Dict]] = XmlParser.extract_tile_transforms(
+        base_channel_xml_path
+    )
+    intertile_transforms: Dict[
+        int, np.ndarray
+    ] = link_utils.calculate_net_transforms(tile_transforms)
+    base_channel: int = link_utils.extract_channel_from_tile_path(
+        tile_paths[0]
+    )
+
+    channels: List[int] = link_utils.get_unique_channels_for_dataset(
+        s3_path + spim_foldername
+    )
+
+    # Generate input config
+    layers = []  # Represent Neuroglancer Tabs
+    input_config = {
+        "dimensions": {
+            "x": {"voxel_size": vox_sizes[0], "unit": "microns"},
+            "y": {"voxel_size": vox_sizes[1], "unit": "microns"},
+            # reverse the order from bigstitcher again
+            "z": {"voxel_size": vox_sizes[2], "unit": "microns"},
+            "c'": {"voxel_size": 1, "unit": ""},
+            "t": {"voxel_size": 0.001, "unit": "seconds"},
+        },
+        "layers": layers,
+        "showScaleBar": False,
+        "showAxisLines": False,
+    }
+
+    for channel in channels:
+        # Determine color of this layer
+        hex_val: int = link_utils.wavelength_to_hex(channel)
+        hex_str = f"#{str(hex(hex_val))[2:]}"
+
+        # Init new list of sources for each channel
+        sources = []  # Represent Tiles w/in Tabs
+        layers.append(
+            {
+                "type": "image",  # Optional
+                "source": sources,
+                "channel": 0,  # Optional
+                "shaderControls": {
+                    "normalized": {"range": [90, max_dr]}
+                },  # Optional
+                "shader": {"color": hex_str, "emitter": "RGB", "vec": "vec3",},
+                "visible": True,  # Optional
+                "opacity": opacity,
+                "name": f"CH_{channel}",
+                "blend": blend,
+            }
+        )
+
+        for tile_id in range(len(intertile_transforms)):
+            # Get base tile path, modify path across channels
+            base_t_path = tile_paths[tile_id]
+            t_path = base_t_path.replace(f"{base_channel}", f"{channel}")
+
+            # Get net transform
+            intertile_tf = intertile_transforms[tile_id]
+            i_matrix_3x3 = intertile_tf[:, 0:3]
+            i_translation = intertile_tf[:, 3]
+
+            net_matrix_3x3 = i_matrix_3x3  # NOTE: Right-multiply
+            net_translation = i_translation
+            net_tf = np.hstack((net_matrix_3x3, net_translation.reshape(3, 1)))
+
+            # net_tf = apply_deskewing(net_tf, deskew_angle)
+
+            # Add (path, transform) source entry
+            if s3_path.endswith("/"):
+                url = f"{s3_path}{spim_foldername}/{t_path}"
+            else:
+                url = f"{s3_path}/{spim_foldername}/{t_path}"
+
+            final_transform = link_utils.convert_matrix_3x4_to_5x6(net_tf)
+            sources.append(
+                {"url": url, "transform_matrix": final_transform.tolist()}
+            )
+
+    bucket_name, prefix = s3_path.replace("s3://", "").split("/", 1)
+    prefix = prefix[:-1]  # remove trailing '/'
+    # Generate the link
+    neuroglancer_link = NgState(
+        input_config=input_config,
+        mount_service="s3",
+        bucket_path=f"{bucket_name}",
+        output_dir=output_json_path,
+        base_url="https://aind-neuroglancer-sauujisjxq-uw.a.run.app/",
+    )
+    neuroglancer_link.save_state_as_json()
+    print(neuroglancer_link.get_url_link())
+    return neuroglancer_link.get_url_link()
+
+
+def ingest_xml_and_write_ng_link(
+    xml_path: str, s3_bucket: str = "aind-open-data"
+):
+    """A wrapper function that autogenerates the s3_path
+     for dispim_link.generate_dispim_link
+
+    Automatically saves process_output.json, which can be
+     manually uploaded to S3 bucket/dataset.
+
+    Parameters:
+    ----------
+    xml_path: str
+        Relative path to xml file (bigstitcher format) that
+        contains tile position information
+
+    s3_bucket:str
+        name of s3 bucket where the dataset lives
+
+    Return:
+    -------
+    link: str
+    Neuroglancer link for xml dataset.
+
+
+    """
+    # read_xml and get dataset prefix for S3
+    dataset_path = XmlParser.extract_dataset_path(xml_path)
+    dataset_name = dataset_path.split("/")[2]
+
+    # print(f"dataset_path {dataset_path}")
+    # print(f"dataset_name {dataset_name}")
+
+    s3_path = f"s3://{s3_bucket}/{dataset_name}/"
+
+    output_folder = f"/results/{dataset_name}/"
+
+    if not pathlib.Path(output_folder).exists():
+        pathlib.Path(output_folder).mkdir(parents=True, exist_ok=True)
+
+    link = generate_dispim_link(
+        xml_path,
+        s3_path,
+        max_dr=400,
+        opacity=1.0,
+        blend="additive",
+        output_json_path=output_folder,
+    )
+
+    # copy output json to s3 bucket dataset
+
+    transfer.copy_to_s3(output_folder + "process_output.json", s3_path)
+
+    return link