zea 0.0.6__py3-none-any.whl → 0.0.8__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/layers.py

@@ -2,7 +2,7 @@
 
 import keras
 import numpy as np
-from keras.src.layers.preprocessing.tf_data_layer import TFDataLayer
+from keras.src.layers.preprocessing.data_layer import DataLayer
 
 from zea.ops import Pad as PadOp
 from zea.utils import map_negative_indices
@@ -11,7 +11,7 @@ from zea.utils import map_negative_indices
 class Pad(PadOp):
     """Pad layer for padding tensors to a specified shape which can be used in tf.data pipelines."""
 
-    __call__ = TFDataLayer.__call__
+    __call__ = DataLayer.__call__
 
     def call(self, inputs):
         """
@@ -20,7 +20,7 @@ class Pad(PadOp):
         return super().call(data=inputs)["data"]
 
 
-class Resizer(TFDataLayer):
+class Resizer(DataLayer):
     """
     Resize layer for resizing images. Can deal with N-dimensional images.
     Can do resize, center_crop, random_crop and crop_or_pad.

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

zea/datapaths.py

@@ -11,12 +11,24 @@ to set up your local data paths.
 Example usage
 ^^^^^^^^^^^^^
 
-.. code-block:: python
+.. doctest::
 
-    from zea.datapaths import set_data_paths
+    >>> import yaml
+    >>> from zea.datapaths import set_data_paths
 
-    user = set_data_paths("users.yaml")
-    print(user.data_root)
+    >>> user_config = {"data_root": "/path/to/data", "output": "/path/to/output"}
+    >>> with open("users.yaml", "w", encoding="utf-8") as file:
+    ...     yaml.dump(user_config, file)
+
+    >>> user = set_data_paths("users.yaml")
+    >>> print(user.data_root)
+    /path/to/data
+
+.. testcleanup::
+
+    import os
+
+    os.remove("users.yaml")
 
 """
 

zea/display.py

@@ -8,9 +8,8 @@ import scipy
 from keras import ops
 from PIL import Image
 
-from zea import log
+from zea.tensor_ops import translate
 from zea.tools.fit_scan_cone import fit_and_crop_around_scan_cone
-from zea.utils import translate
 
 
 def to_8bit(image, dynamic_range: Union[None, tuple] = None, pillow: bool = True):
@@ -340,12 +339,14 @@ def scan_convert(
 def map_coordinates(inputs, coordinates, order, fill_mode="constant", fill_value=0):
     """map_coordinates using keras.ops or scipy.ndimage when order > 1."""
     if order > 1:
-        inputs = ops.convert_to_numpy(inputs)
-        coordinates = ops.convert_to_numpy(coordinates)
+        # Preserve original dtype before conversion
+        original_dtype = ops.dtype(inputs)
+        inputs_np = ops.convert_to_numpy(inputs).astype(np.float32)
+        coordinates_np = ops.convert_to_numpy(coordinates).astype(np.float32)
         out = scipy.ndimage.map_coordinates(
-            inputs, coordinates, order=order, mode=fill_mode, cval=fill_value
+            inputs_np, coordinates_np, order=order, mode=fill_mode, cval=fill_value
         )
-        return ops.convert_to_tensor(out)
+        return ops.convert_to_tensor(out.astype(original_dtype))
     else:
         return ops.image.map_coordinates(
             inputs,
@@ -439,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)
@@ -501,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.
@@ -510,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.
@@ -523,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,