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

zea/data/data_format.py

@@ -256,8 +256,6 @@ def _write_datasets(
         n_tx = _first_not_none_shape([raw_data, aligned_data], axis=1)
     if n_ax is None:
         n_ax = _first_not_none_shape([raw_data, aligned_data, beamformed_data], axis=-3)
-    if n_ax is None:
-        n_ax = _first_not_none_shape([envelope_data, image, image_sc], axis=-2)
     if n_el is None:
         n_el = _first_not_none_shape([raw_data], axis=-2)
     if n_ch is None:

zea/data/file.py

@@ -43,13 +43,16 @@ class File(h5py.File):
             **kwargs: Additional keyword arguments to pass to h5py.File.
         """
 
+        # Resolve huggingface path
         if str(name).startswith(HF_PREFIX):
             name = _hf_resolve_path(str(name))
 
+        # Disable locking for read mode by default
         if "locking" not in kwargs and "mode" in kwargs and kwargs["mode"] == "r":
             # If the file is opened in read mode, disable locking
             kwargs["locking"] = False
 
+        # Initialize the h5py.File
         super().__init__(name, *args, **kwargs)
 
     @property
@@ -285,6 +288,31 @@ class File(h5py.File):
         else:
             log.warning("Could not find scan parameters in file.")
 
+        scan_parameters = self._check_focus_distances(scan_parameters)
+
+        return scan_parameters
+
+    def _check_focus_distances(self, scan_parameters):
+        if "focus_distances" in scan_parameters:
+            focus_distances = scan_parameters["focus_distances"]
+            # check if focus distances are in wavelengths
+            if np.any(np.logical_and(focus_distances >= 1, focus_distances != np.inf)):
+                log.warning(
+                    f"We have detected that focus distances in '{self.path}' are "
+                    "(probably) stored wavelengths. Please update your file! "
+                    "Converting to meters automatically for now."
+                )
+                assert "sound_speed" in scan_parameters, (
+                    "Cannot convert focus distances from wavelengths to meters "
+                    "because sound_speed is not defined in the scan parameters."
+                )
+                assert "center_frequency" in scan_parameters, (
+                    "Cannot convert focus distances from wavelengths to meters "
+                    "because center_frequency is not defined in the scan parameters."
+                )
+                wavelength = scan_parameters["sound_speed"] / scan_parameters["center_frequency"]
+                focus_distances = focus_distances * wavelength
+                scan_parameters["focus_distances"] = focus_distances
         return scan_parameters
 
     def get_scan_parameters(self, event=None) -> dict:

zea/data/preset_utils.py

@@ -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,7 +8,7 @@ import scipy
 from keras import ops
 from PIL import Image
 
-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
 
 

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",
+]

zea/{tensor_ops.py → func/tensor.py}

@@ -329,7 +329,7 @@ def _map(fun, in_axes=0, out_axes=0, map_fn=None, _use_torch_vmap=False):
     For jax, this uses the native vmap implementation.
     For other backends, this a wrapper that uses `ops.vectorized_map` under the hood.
 
-    Probably you want to use `zea.tensor_ops.vmap` instead, which uses this function
+    Probably you want to use `zea.func.vmap` instead, which uses this function
     with additional batching/chunking support.
 
     Args:
@@ -845,7 +845,7 @@ def stack_volume_data_along_axis(data, batch_axis: int, stack_axis: int, number:
         .. doctest::
 
             >>> import keras
-            >>> from zea.tensor_ops import stack_volume_data_along_axis
+            >>> from zea.func import stack_volume_data_along_axis
 
             >>> data = keras.random.uniform((10, 20, 30))
             >>> # stacking along 1st axis with 2 frames per block
@@ -889,7 +889,7 @@ def split_volume_data_from_axis(data, batch_axis: int, stack_axis: int, number:
         .. doctest::
 
             >>> import keras
-            >>> from zea.tensor_ops import split_volume_data_from_axis
+            >>> from zea.func import split_volume_data_from_axis
 
             >>> data = keras.random.uniform((20, 10, 30))
             >>> split_data = split_volume_data_from_axis(data, 0, 1, 2, 2)
@@ -1012,7 +1012,7 @@ def check_patches_fit(
     Example:
         .. doctest::
 
-            >>> from zea.tensor_ops import check_patches_fit
+            >>> from zea.func import check_patches_fit
             >>> image_shape = (10, 10)
             >>> patch_shape = (4, 4)
             >>> overlap = (2, 2)
@@ -1080,7 +1080,7 @@ def images_to_patches(
         .. doctest::
 
             >>> import keras
-            >>> from zea.tensor_ops import images_to_patches
+            >>> from zea.func import images_to_patches
 
             >>> images = keras.random.uniform((2, 8, 8, 3))
             >>> patches = images_to_patches(images, patch_shape=(4, 4), overlap=(2, 2))
@@ -1166,7 +1166,7 @@ def patches_to_images(
         .. doctest::
 
             >>> import keras
-            >>> from zea.tensor_ops import patches_to_images
+            >>> from zea.func import patches_to_images
 
             >>> patches = keras.random.uniform((2, 3, 3, 4, 4, 3))
             >>> images = patches_to_images(patches, image_shape=(8, 8, 3), overlap=(2, 2))
@@ -1254,7 +1254,7 @@ def reshape_axis(data, newshape: tuple, axis: int):
         .. doctest::
 
             >>> import keras
-            >>> from zea.tensor_ops import reshape_axis
+            >>> from zea.func import reshape_axis
 
             >>> data = keras.random.uniform((3, 4, 5))
             >>> newshape = (2, 2)
@@ -1641,7 +1641,7 @@ def find_contour(binary_mask):
     Example:
         .. doctest::
 
-            >>> from zea.tensor_ops import find_contour
+            >>> from zea.func import find_contour
             >>> import keras
             >>> mask = keras.ops.zeros((10, 10))
             >>> mask = keras.ops.scatter_update(
@@ -1700,3 +1700,27 @@ def translate(array, range_from=None, range_to=(0, 255)):
 
     # Convert the 0-1 range into a value in the right range.
     return right_min + (value_scaled * (right_max - right_min))
+
+
+def normalize(data, output_range, input_range=None):
+    """Normalize data to a given range.
+
+    Equivalent to `translate` with clipping.
+
+    Args:
+        data (ops.Tensor): Input data to normalize.
+        output_range (tuple): Range to which data should be mapped, e.g., (0, 1).
+        input_range (tuple, optional): Range of input data.
+            If None, the range will be computed from the data.
+            Defaults to None.
+    """
+    if input_range is None:
+        input_range = (None, None)
+    minval, maxval = input_range
+    if minval is None:
+        minval = ops.min(data)
+    if maxval is None:
+        maxval = ops.max(data)
+    data = ops.clip(data, minval, maxval)
+    normalized_data = translate(data, (minval, maxval), output_range)
+    return normalized_data