zea 0.0.7__py3-none-any.whl → 0.0.9__py3-none-any.whl

zea/data/file_operations.py

@@ -0,0 +1,496 @@
+"""
+This module provides some utilities to edit zea data files.
+
+Available operations
+--------------------
+
+- `sum`: Sum multiple raw data files into one.
+
+- `compound_frames`: Compound frames in a raw data file to increase SNR.
+
+- `compound_transmits`: Compound transmits in a raw data file to increase SNR.
+
+- `resave`: Resave a zea data file. This can be used to change the file format version.
+
+- `extract`: extract frames and transmits in a raw data file.
+"""
+
+import argparse
+from pathlib import Path
+
+import numpy as np
+
+from zea import Probe, Scan
+from zea.data.data_format import generate_zea_dataset, load_additional_elements, load_description
+from zea.data.file import load_file_all_data_types
+from zea.internal.checks import _IMAGE_DATA_TYPES, _NON_IMAGE_DATA_TYPES
+from zea.internal.core import DataTypes
+from zea.log import logger
+
+ALL_DATA_TYPES_EXCEPT_RAW = set(_IMAGE_DATA_TYPES + _NON_IMAGE_DATA_TYPES) - {"raw_data"}
+
+OPERATION_NAMES = [
+    "sum",
+    "compound_frames",
+    "compound_transmits",
+    "resave",
+    "extract",
+]
+
+
+def save_file(
+    path,
+    scan: Scan,
+    probe: Probe,
+    raw_data: np.ndarray = None,
+    aligned_data: np.ndarray = None,
+    beamformed_data: np.ndarray = None,
+    envelope_data: np.ndarray = None,
+    image: np.ndarray = None,
+    image_sc: np.ndarray = None,
+    additional_elements=None,
+    description="",
+    **kwargs,
+):
+    """Saves data to a zea data file (h5py file).
+
+    Args:
+        path (str, pathlike): The path to the hdf5 file.
+        raw_data (np.ndarray): The data to save.
+        scan (Scan): The scan object containing the parameters of the acquisition.
+        probe (Probe): The probe object containing the parameters of the probe.
+        additional_elements (list of DatasetElement, optional): Additional elements to save in the
+            file. Defaults to None.
+    """
+
+    generate_zea_dataset(
+        path=path,
+        raw_data=raw_data,
+        aligned_data=aligned_data,
+        beamformed_data=beamformed_data,
+        image=image,
+        image_sc=image_sc,
+        envelope_data=envelope_data,
+        probe_name="generic",
+        probe_geometry=probe.probe_geometry,
+        sampling_frequency=scan.sampling_frequency,
+        center_frequency=scan.center_frequency,
+        initial_times=scan.initial_times,
+        t0_delays=scan.t0_delays,
+        sound_speed=scan.sound_speed,
+        focus_distances=scan.focus_distances,
+        polar_angles=scan.polar_angles,
+        azimuth_angles=scan.azimuth_angles,
+        tx_apodizations=scan.tx_apodizations,
+        bandwidth_percent=scan.bandwidth_percent,
+        time_to_next_transmit=scan.time_to_next_transmit,
+        tgc_gain_curve=scan.tgc_gain_curve,
+        element_width=scan.element_width,
+        tx_waveform_indices=scan.tx_waveform_indices,
+        waveforms_one_way=scan.waveforms_one_way,
+        waveforms_two_way=scan.waveforms_two_way,
+        description=description,
+        additional_elements=additional_elements,
+    )
+
+
+def sum_data(input_paths: list[Path], output_path: Path, overwrite=False):
+    """
+    Sums multiple raw data files and saves the result to a new file.
+
+    Args:
+        input_paths (list[Path]): List of paths to the input raw data files.
+        output_path (Path): Path to the output file where the summed data will be saved.
+        overwrite (bool, optional): Whether to overwrite the output file if it exists. Defaults to
+            False.
+    """
+
+    data_dict, scan, probe = load_file_all_data_types(input_paths[0])
+    description = load_description(input_paths[0])
+    additional_elements = load_additional_elements(input_paths[0])
+
+    for file in input_paths[1:]:
+        new_data, new_scan, new_probe = load_file_all_data_types(file)
+
+        if data_dict["raw_data"] is not None:
+            _assert_shapes_equal(data_dict["raw_data"], new_data["raw_data"], "raw_data")
+            data_dict["raw_data"] += new_data["raw_data"]
+
+        if data_dict["aligned_data"] is not None:
+            _assert_shapes_equal(
+                data_dict["aligned_data"], new_data["aligned_data"], "aligned_data"
+            )
+            data_dict["aligned_data"] += new_data["aligned_data"]
+
+        if data_dict["beamformed_data"] is not None:
+            _assert_shapes_equal(
+                data_dict["beamformed_data"], new_data["beamformed_data"], "beamformed_data"
+            )
+            data_dict["beamformed_data"] += new_data["beamformed_data"]
+
+        if data_dict["envelope_data"] is not None:
+            _assert_shapes_equal(
+                data_dict["envelope_data"], new_data["envelope_data"], "envelope_data"
+            )
+            data_dict["envelope_data"] += new_data["envelope_data"]
+
+        if data_dict["image"] is not None:
+            _assert_shapes_equal(data_dict["image"], new_data["image"], "image")
+            data_dict["image"] = np.log(np.exp(new_data["image"]) + np.exp(data_dict["image"]))
+
+        if data_dict["image_sc"] is not None:
+            _assert_shapes_equal(data_dict["image_sc"], new_data["image_sc"], "image_sc")
+            data_dict["image_sc"] = np.log(
+                np.exp(new_data["image_sc"]) + np.exp(data_dict["image_sc"])
+            )
+        assert scan == new_scan, "Scan parameters do not match."
+        assert probe == new_probe, "Probe parameters do not match."
+
+    if overwrite:
+        _delete_file_if_exists(output_path)
+
+    save_file(
+        path=output_path,
+        scan=scan,
+        probe=probe,
+        additional_elements=additional_elements,
+        description=description,
+        **data_dict,
+    )
+
+
+def _assert_shapes_equal(array0, array1, name="array"):
+    shape0, shape1 = array0.shape, array1.shape
+    assert shape0 == shape1, f"{name} shapes do not match. Got {shape0} and {shape1}."
+
+
+def compound_frames(input_path: Path, output_path: Path, overwrite=False):
+    """
+    Compounds frames in a raw data file by averaging them.
+
+    Args:
+        input_path (Path): Path to the input raw data file.
+        output_path (Path): Path to the output file where the compounded data will be saved.
+        overwrite (bool, optional): Whether to overwrite the output file if it exists. Defaults to
+            False.
+    """
+
+    data_dict, scan, probe = load_file_all_data_types(input_path)
+    additional_elements = load_additional_elements(input_path)
+    description = load_description(input_path)
+
+    # Assuming the first dimension is the frame dimension
+
+    compounded_data = {}
+    for data_type in DataTypes:
+        key = data_type.value
+        if data_dict[key] is None:
+            compounded_data[key] = None
+            continue
+        if key == "image" or key == "image_sc":
+            compounded_data[key] = np.log(np.mean(np.exp(data_dict[key]), axis=0, keepdims=True))
+        else:
+            compounded_data[key] = np.mean(data_dict[key], axis=0, keepdims=True)
+
+    scan = _scan_reduce_frames(scan, [0])
+
+    if overwrite:
+        _delete_file_if_exists(output_path)
+
+    save_file(
+        path=output_path,
+        scan=scan,
+        probe=probe,
+        additional_elements=additional_elements,
+        description=description,
+        **compounded_data,
+    )
+
+
+def compound_transmits(input_path: Path, output_path: Path, overwrite=False):
+    """
+    Compounds transmits in a raw data file by averaging them.
+
+    Note
+    ----
+    This function assumes that all transmits are identical. If this is not the case the function
+    will result in incorrect scan parameters.
+
+    Args:
+        input_path (Path): Path to the input raw data file.
+        output_path (Path): Path to the output file where the compounded data will be saved.
+        overwrite (bool, optional): Whether to overwrite the output file if it exists. Defaults to
+        False.
+    """
+
+    data_dict, scan, probe = load_file_all_data_types(input_path)
+    additional_elements = load_additional_elements(input_path)
+    description = load_description(input_path)
+
+    if not _all_tx_are_identical(scan):
+        logger.warning(
+            "Not all transmits are identical. Compounding transmits may lead to unexpected results."
+        )
+
+    # Assuming the second dimension is the transmit dimension
+    for key in ["raw_data", "aligned_data"]:
+        if data_dict[key] is None:
+            continue
+        data_dict[key] = np.mean(data_dict[key], axis=1, keepdims=True)
+
+    scan.set_transmits([0])
+
+    if overwrite:
+        _delete_file_if_exists(output_path)
+
+    save_file(
+        path=output_path,
+        scan=scan,
+        probe=probe,
+        additional_elements=additional_elements,
+        description=description,
+        **data_dict,
+    )
+
+
+def _all_tx_are_identical(scan: Scan):
+    """Checks if all transmits in a Scan object are identical."""
+    attributes_to_check = [
+        scan.polar_angles,
+        scan.azimuth_angles,
+        scan.t0_delays,
+        scan.tx_apodizations,
+        scan.focus_distances,
+        scan.initial_times,
+    ]
+
+    for attr in attributes_to_check:
+        if attr is not None and not _check_all_identical(attr, axis=0):
+            return False
+    return True
+
+
+def _check_all_identical(array, axis=0):
+    """Checks if all elements along a given axis are identical."""
+    first = array.take(0, axis=axis)
+    return np.all(np.equal(array, first), axis=axis).all()
+
+
+def resave(input_path: Path, output_path: Path, overwrite=False):
+    """
+    Resaves a zea data file to a new location.
+
+    Args:
+        input_path (Path): Path to the input zea data file.
+        output_path (Path): Path to the output file where the data will be saved.
+        overwrite (bool, optional): Whether to overwrite the output file if it exists. Defaults to
+            False.
+    """
+
+    data_dict, scan, probe = load_file_all_data_types(input_path)
+    additional_elements = load_additional_elements(input_path)
+    description = load_description(input_path)
+    scan.set_transmits("all")
+
+    if overwrite:
+        _delete_file_if_exists(output_path)
+    save_file(
+        path=output_path,
+        **data_dict,
+        scan=scan,
+        probe=probe,
+        additional_elements=additional_elements,
+        description=description,
+    )
+
+
+def extract_frames_transmits(
+    input_path: Path,
+    output_path: Path,
+    frame_indices=slice(None),
+    transmit_indices=slice(None),
+    overwrite=False,
+):
+    """
+    extracts frames and transmits in a raw data file.
+
+    Note that the frame indices cannot both be lists. At least one of them must be a slice.
+    Please refer to the documentation of :func:`zea.data.file.load_file_all_data_types` for more
+    information on the supported index types.
+
+    Args:
+        input_path (Path): Path to the input raw data file.
+        output_path (Path): Path to the output file where the extracted data will be saved.
+        frame_indices (list, array-like, or slice): Indices of the frames to keep.
+        transmit_indices (list, array-like, or slice): Indices of the transmits to keep.
+        overwrite (bool, optional): Whether to overwrite the output file if it exists. Defaults to
+            False.
+    """
+    indices = (frame_indices, transmit_indices)
+    data_dict, scan, probe = load_file_all_data_types(input_path, indices=indices)
+
+    additional_elements = load_additional_elements(input_path)
+    description = load_description(input_path)
+
+    scan = _scan_reduce_frames(scan, frame_indices)
+
+    if overwrite:
+        _delete_file_if_exists(output_path)
+
+    save_file(
+        path=output_path,
+        **data_dict,
+        scan=scan,
+        probe=probe,
+        additional_elements=additional_elements,
+        description=description,
+    )
+
+
+def _delete_file_if_exists(path: Path):
+    """Deletes a file if it exists."""
+    if path.exists():
+        path.unlink()
+
+
+def _interpret_index(input_str):
+    if "-" in input_str:
+        start, end = map(int, input_str.split("-"))
+        return list(range(start, end + 1))
+    else:
+        return [int(x) for x in input_str.split(" ")]
+
+
+def _interpret_indices(input_str_list):
+    if isinstance(input_str_list, str) and input_str_list == "all":
+        return slice(None)
+
+    if len(input_str_list) == 1 and "-" in input_str_list[0]:
+        start, end = map(int, input_str_list[0].split("-"))
+        return slice(start, end + 1)
+
+    indices = []
+    for part in input_str_list:
+        indices.extend(_interpret_index(part))
+    return indices
+
+
+def _scan_reduce_frames(scan, frame_indices):
+    transmit_indices = scan.selected_transmits
+    scan.set_transmits("all")
+    if scan.time_to_next_transmit is not None:
+        scan.time_to_next_transmit = scan.time_to_next_transmit[frame_indices]
+    scan.set_transmits(transmit_indices)
+    return scan
+
+
+def get_parser():
+    """Command line argument parser with subcommands"""
+
+    parser = argparse.ArgumentParser(
+        description="Manipulate zea data files.",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    subparsers = parser.add_subparsers(dest="operation", required=True)
+    _add_parser_sum(subparsers)
+    _add_parser_compound_frames(subparsers)
+    _add_parser_compound_transmits(subparsers)
+    _add_parser_resave(subparsers)
+    _add_parser_extract(subparsers)
+
+    return parser
+
+
+def _add_parser_sum(subparsers):
+    sum_parser = subparsers.add_parser("sum", help="Sum the raw data of multiple files.")
+    sum_parser.add_argument("input_paths", type=Path, nargs="+", help="Paths to the input files.")
+    sum_parser.add_argument("output_path", type=Path, help="Output HDF5 file.")
+    sum_parser.add_argument(
+        "--overwrite", action="store_true", default=False, help="Overwrite existing output file."
+    )
+
+
+def _add_parser_compound_frames(subparsers):
+    cf_parser = subparsers.add_parser("compound_frames", help="Compound frames to increase SNR.")
+    cf_parser.add_argument("input_path", type=Path, help="Input HDF5 file.")
+    cf_parser.add_argument("output_path", type=Path, help="Output HDF5 file.")
+    cf_parser.add_argument(
+        "--overwrite", action="store_true", default=False, help="Overwrite existing output file."
+    )
+
+
+def _add_parser_compound_transmits(subparsers):
+    ct_parser = subparsers.add_parser(
+        "compound_transmits", help="Compound transmits to increase SNR."
+    )
+    ct_parser.add_argument("input_path", type=Path, help="Input HDF5 file.")
+    ct_parser.add_argument("output_path", type=Path, help="Output HDF5 file.")
+    ct_parser.add_argument(
+        "--overwrite", action="store_true", default=False, help="Overwrite existing output file."
+    )
+
+
+def _add_parser_resave(subparsers):
+    resave_parser = subparsers.add_parser("resave", help="Resave a file to change format version.")
+    resave_parser.add_argument("input_path", type=Path, help="Input HDF5 file.")
+    resave_parser.add_argument("output_path", type=Path, help="Output HDF5 file.")
+    resave_parser.add_argument(
+        "--overwrite", action="store_true", default=False, help="Overwrite existing output file."
+    )
+
+
+def _add_parser_extract(subparsers):
+    extract_parser = subparsers.add_parser("extract", help="Extract subset of frames or transmits.")
+    extract_parser.add_argument("input_path", type=Path, help="Input HDF5 file.")
+    extract_parser.add_argument("output_path", type=Path, help="Output HDF5 file.")
+    extract_parser.add_argument(
+        "--transmits",
+        type=str,
+        nargs="+",
+        default="all",
+        help="Target transmits. Can be a list of integers or ranges (e.g. 0-3 7).",
+    )
+    extract_parser.add_argument(
+        "--frames",
+        type=str,
+        nargs="+",
+        default="all",
+        help="Target frames. Can be a list of integers or ranges (e.g. 0-3 7).",
+    )
+    extract_parser.add_argument(
+        "--overwrite", action="store_true", default=False, help="Overwrite existing output file."
+    )
+
+
+if __name__ == "__main__":
+    parser = get_parser()
+    args = parser.parse_args()
+
+    if args.output_path.exists() and not args.overwrite:
+        logger.error(
+            f"Output file {args.output_path} already exists. Use --overwrite to overwrite it."
+        )
+        exit(1)
+
+    if args.operation == "compound_frames":
+        compound_frames(
+            input_path=args.input_path, output_path=args.output_path, overwrite=args.overwrite
+        )
+    elif args.operation == "compound_transmits":
+        compound_transmits(
+            input_path=args.input_path, output_path=args.output_path, overwrite=args.overwrite
+        )
+    elif args.operation == "resave":
+        resave(input_path=args.input_path, output_path=args.output_path, overwrite=args.overwrite)
+    elif args.operation == "extract":
+        extract_frames_transmits(
+            input_path=args.input_path,
+            output_path=args.output_path,
+            frame_indices=_interpret_indices(args.frames),
+            transmit_indices=_interpret_indices(args.transmits),
+            overwrite=args.overwrite,
+        )
+    else:
+        sum_data(
+            input_paths=args.input_paths, output_path=args.output_path, overwrite=args.overwrite
+        )

zea/data/preset_utils.py

@@ -1,6 +1,6 @@
 """Preset utils for zea datasets hosted on Hugging Face.
 
-See https://huggingface.co/zea/
+See https://huggingface.co/zeahub/
 """
 
 from pathlib import Path
@@ -86,7 +86,7 @@ def _download_files_in_path(
     return downloaded_files
 
 
-def _hf_resolve_path(hf_path: str, cache_dir=HF_DATASETS_DIR):
+def _hf_resolve_path(hf_path: str, cache_dir=HF_DATASETS_DIR) -> Path:
     """Resolve a Hugging Face path to a local cache directory path.
 
     Downloads files from a HuggingFace dataset repository and returns

zea/display.py

@@ -8,8 +8,7 @@ import scipy
 from keras import ops
 from PIL import Image
 
-from zea import log
-from zea.tensor_ops import translate
+from zea.func.tensor import translate
 from zea.tools.fit_scan_cone import fit_and_crop_around_scan_cone
 
 
@@ -441,11 +440,7 @@ def cartesian_to_polar_matrix(
     Returns:
         polar_matrix (Array): The image re-sampled in polar coordinates with shape `polar_shape`.
     """
-    if ops.dtype(cartesian_matrix) != "float32":
-        log.info(
-            f"Cartesian matrix with dtype {ops.dtype(cartesian_matrix)} has been cast to float32."
-        )
-        cartesian_matrix = ops.cast(cartesian_matrix, "float32")
+    assert "float" in ops.dtype(cartesian_matrix), "Input image must be float type"
 
     # Assume that polar grid is same shape as cartesian grid unless specified
     cartesian_rows, cartesian_cols = ops.shape(cartesian_matrix)
@@ -503,6 +498,7 @@ def inverse_scan_convert_2d(
     output_size=None,
     interpolation_order=1,
     find_scan_cone=True,
+    image_range: tuple | None = None,
 ):
     """
     Convert a Cartesian-format ultrasound image to a polar representation.
@@ -512,7 +508,7 @@ def inverse_scan_convert_2d(
     Optionally, it can detect and crop around the scan cone before conversion.
 
     Args:
-        cartesian_image (tensor): 2D image array in Cartesian coordinates.
+        cartesian_image (tensor): 2D image array in Cartesian coordinates of type float.
         fill_value (float): Value used to fill regions outside the original image
             during interpolation.
         angle (float): Angular field of view (in radians) used for the polar transformation.
@@ -525,12 +521,15 @@ def inverse_scan_convert_2d(
             in the Cartesian image before polar conversion, ensuring that the scan cone is
             centered without padding. Can be set to False if the image is already cropped
             and centered.
+        image_range (tuple, optional): Tuple (vmin, vmax) for display scaling
+            when detecting the scan cone.
 
     Returns:
         polar_image (Array): 2D image in polar coordinates (sector-shaped scan).
     """
     if find_scan_cone:
-        cartesian_image = fit_and_crop_around_scan_cone(cartesian_image)
+        assert image_range is not None, "image_range must be provided when find_scan_cone is True"
+        cartesian_image = fit_and_crop_around_scan_cone(cartesian_image, image_range)
     polar_image = cartesian_to_polar_matrix(
         cartesian_image,
         fill_value=fill_value,

zea/doppler.py

@@ -3,7 +3,7 @@
 import numpy as np
 from keras import ops
 
-from zea import tensor_ops
+from zea.func import tensor
 
 
 def color_doppler(
@@ -61,11 +61,11 @@ def color_doppler(
     if hamming_size[0] != 1 and hamming_size[1] != 1:
         h_row = np.hamming(hamming_size[0])
         h_col = np.hamming(hamming_size[1])
-        autocorr = tensor_ops.apply_along_axis(
-            lambda x: tensor_ops.correlate(x, h_row, mode="same"), 0, autocorr
+        autocorr = tensor.apply_along_axis(
+            lambda x: tensor.correlate(x, h_row, mode="same"), 0, autocorr
         )
-        autocorr = tensor_ops.apply_along_axis(
-            lambda x: tensor_ops.correlate(x, h_col, mode="same"), 1, autocorr
+        autocorr = tensor.apply_along_axis(
+            lambda x: tensor.correlate(x, h_col, mode="same"), 1, autocorr
         )
 
     # Doppler velocity

zea/func/__init__.py

@@ -0,0 +1,109 @@
+"""Functional API of zea.
+
+This module provides a collection of functions for various operations on tensors
+and ultrasound data. These functions can be used standalone, in contrast to the :mod:`zea.ops` module which provides operation classes for building processing pipelines.
+"""
+
+from .tensor import (
+    L1,
+    L2,
+    add_salt_and_pepper_noise,
+    apply_along_axis,
+    batch_cov,
+    boolean_mask,
+    check_patches_fit,
+    compute_required_patch_overlap,
+    compute_required_patch_shape,
+    correlate,
+    extend_n_dims,
+    find_contour,
+    flatten,
+    fori_loop,
+    func_with_one_batch_dim,
+    gaussian_filter,
+    images_to_patches,
+    interpolate_data,
+    is_jax_prng_key,
+    is_monotonic,
+    linear_sum_assignment,
+    map_indices_for_interpolation,
+    matrix_power,
+    normalize,
+    pad_array_to_divisible,
+    patches_to_images,
+    resample,
+    reshape_axis,
+    simple_map,
+    sinc,
+    split_seed,
+    split_volume_data_from_axis,
+    stack_volume_data_along_axis,
+    translate,
+    vmap,
+)
+from .ultrasound import (
+    channels_to_complex,
+    complex_to_channels,
+    compute_time_to_peak,
+    compute_time_to_peak_stack,
+    demodulate,
+    demodulate_not_jitable,
+    envelope_detect,
+    get_band_pass_filter,
+    get_low_pass_iq_filter,
+    hilbert,
+    log_compress,
+    upmix,
+)
+
+__all__ = [
+    # Tensor functions
+    "L1",
+    "L2",
+    "add_salt_and_pepper_noise",
+    "apply_along_axis",
+    "batch_cov",
+    "boolean_mask",
+    "check_patches_fit",
+    "compute_required_patch_overlap",
+    "compute_required_patch_shape",
+    "correlate",
+    "extend_n_dims",
+    "find_contour",
+    "flatten",
+    "fori_loop",
+    "func_with_one_batch_dim",
+    "gaussian_filter",
+    "images_to_patches",
+    "interpolate_data",
+    "is_jax_prng_key",
+    "is_monotonic",
+    "linear_sum_assignment",
+    "map_indices_for_interpolation",
+    "matrix_power",
+    "normalize",
+    "pad_array_to_divisible",
+    "patches_to_images",
+    "resample",
+    "reshape_axis",
+    "simple_map",
+    "sinc",
+    "split_seed",
+    "split_volume_data_from_axis",
+    "stack_volume_data_along_axis",
+    "translate",
+    "vmap",
+    # Ultrasound functions
+    "channels_to_complex",
+    "complex_to_channels",
+    "compute_time_to_peak",
+    "compute_time_to_peak_stack",
+    "demodulate",
+    "demodulate_not_jitable",
+    "envelope_detect",
+    "get_band_pass_filter",
+    "get_low_pass_iq_filter",
+    "hilbert",
+    "upmix",
+    "log_compress",
+]