zea 0.0.1__py3-none-any.whl → 0.0.2__py3-none-any.whl

zea/__init__.py

@@ -7,7 +7,7 @@ from . import log
 
 # dynamically add __version__ attribute (see pyproject.toml)
 # __version__ = __import__("importlib.metadata").metadata.version(__package__)
-__version__ = "0.0.1"
+__version__ = "0.0.2"
 
 
 def setup():

zea/beamform/beamformer.py

@@ -7,6 +7,45 @@ from zea.beamform.lens_correction import calculate_lens_corrected_delays
 from zea.tensor_ops import safe_vectorize
 
 
+def fnum_window_fn_rect(normalized_angle):
+    """Rectangular window function for f-number masking."""
+    return ops.where(normalized_angle <= 1.0, 1.0, 0.0)
+
+
+def fnum_window_fn_hann(normalized_angle):
+    """Hann window function for f-number masking."""
+    # Use a Hann window function to smoothly transition the mask
+    return ops.where(
+        normalized_angle <= 1.0,
+        0.5 * (1 + ops.cos(np.pi * normalized_angle)),
+        0.0,
+    )
+
+
+def fnum_window_fn_tukey(normalized_angle, alpha=0.5):
+    """Tukey window function for f-number masking.
+
+    Args:
+        normalized_angle (ops.Tensor): Normalized angle values in the range [0, 1].
+        alpha (float, optional): The alpha parameter for the Tukey window. 0.0 corresponds to a
+            rectangular window, 1.0 corresponds to a Hann window. Defaults to 0.5.
+    """
+    # Use a Tukey window function to smoothly transition the mask
+    normalized_angle = ops.clip(ops.abs(normalized_angle), 0.0, 1.0)
+
+    beta = 1.0 - alpha
+
+    return ops.where(
+        normalized_angle < beta,
+        1.0,
+        ops.where(
+            normalized_angle < 1.0,
+            0.5 * (1 + ops.cos(np.pi * (normalized_angle - beta) / (ops.abs(alpha) + 1e-6))),
+            0.0,
+        ),
+    )
+
+
 def tof_correction(
     data,
     flatgrid,
@@ -24,6 +63,7 @@ def tof_correction(
     apply_lens_correction=False,
     lens_thickness=1e-3,
     lens_sound_speed=1000,
+    fnum_window_fn=fnum_window_fn_rect,
 ):
     """Time-of-flight correction for a flat grid.
 
@@ -54,6 +94,9 @@ def tof_correction(
             lens correction. Defaults to 1e-3.
         lens_sound_speed (float, optional): Speed of sound in the lens in m/s. Used
             for lens correction Defaults to 1000.
+        fnum_window_fn (callable, optional): F-number function to define the transition from
+            straight in front of the element (fn(0.0)) to the largest angle within the f-number cone
+            (fn(1.0)). The function should be zero for fn(x>1.0).
 
     Returns:
         (ops.Tensor): time-of-flight corrected data
@@ -100,7 +143,7 @@ def tof_correction(
     mask = ops.cond(
         fnum == 0,
         lambda: ops.ones((n_pix, n_el, 1)),
-        lambda: apod_mask(flatgrid, probe_geometry, fnum),
+        lambda: fnumber_mask(flatgrid, probe_geometry, fnum, fnum_window_fn=fnum_window_fn),
     )
 
     def _apply_delays(data_tx, txdel):
@@ -408,64 +451,48 @@ def distance_Tx_generic(
     return dist
 
 
-def apod_mask(grid, probe_geometry, f_number):
+def fnumber_mask(flatgrid, probe_geometry, f_number, fnum_window_fn):
     """Apodization mask for the receive beamformer.
 
-    Computes a binary mask to disregard pixels outside of the vision cone of a
+    Computes a mask to disregard pixels outside of the vision cone of a
     transducer element. Transducer elements can only accurately measure
     signals within some range of incidence angles. Waves coming in from the
     side do not register correctly leading to a worse image.
 
     Args:
-        grid (ops.Tensor): The flattened image grid `(n_pix, 3)`.
+        flatgrid (ops.Tensor): The flattened image grid `(n_pix, 3)`.
         probe_geometry (ops.Tensor): The transducer element positions of shape
             `(n_el, 3)`.
         f_number (int): The receive f-number. Set to zero to not use masking and
             return 1. (The f-number is the  ratio between distance from the transducer
             and the size of the aperture below which transducer elements contribute to
             the signal for a pixel.).
+        fnum_window_fn (callable): F-number function to define the transition from
+            straight in front of the element (fn(0.0)) to the largest angle within the f-number cone
+            (fn(1.0)). The function should be zero for fn(x>1.0).
+
 
     Returns:
         Tensor: Mask of shape `(n_pix, n_el, 1)`
     """
-    n_pix = ops.shape(grid)[0]
-    n_el = ops.shape(probe_geometry)[0]
-
-    # Get the depth of every pixel
-    z_pixel = grid[:, 2]
-    # Get the lateral location of each pixel
-    x_pixel = grid[:, 0]
-    # Get the lateral location of each element
-    x_element = ops.cast(probe_geometry[:, 0], dtype="float32")
-
-    # Compute the aperture size for every pixel
-    # The f-number is by definition f=z/aperture
-    aperture = z_pixel / f_number
-
-    # Use matrix multiplication to expand aperture tensor, x_pixel tensor, and
-    # x_element tensor to shape (n_pix, n_el)
-    ones_aperture = ops.ones(
-        (1, n_el), dtype=ops.dtype(aperture)
-    )  # getting error here? pip install -U keras ;)
-    ones_x_pixel = ops.ones((1, n_el), dtype=ops.dtype(x_pixel))
-    ones_x_element = ops.ones((n_pix, 1), dtype=ops.dtype(x_element))
-
-    aperture = ops.matmul(aperture[..., None], ones_aperture)
-    expanded_x_pixel = ops.matmul(x_pixel[..., None], ones_x_pixel)
-    expanded_x_element = ops.matmul(ones_x_element, x_element[None])
-
-    # Compute the lateral distance between elements and pixels
-    # Of shape (n_pix, n_el)
-    distance = ops.abs(expanded_x_pixel - expanded_x_element)
-
-    # Compute binary mask for which the lateral pixel distance is less than
-    # half
-    # the aperture i.e. where the pixel lies within the vision cone of the
-    # element
-    mask = distance <= aperture / 2
-    mask = ops.cast(mask, "float32")
-
-    # Add dummy dimension for RF/IQ channel channel
+
+    grid_relative_to_probe = flatgrid[:, None] - probe_geometry[None]
+
+    grid_relative_to_probe_norm = ops.linalg.norm(grid_relative_to_probe, axis=-1)
+
+    grid_relative_to_probe_z = grid_relative_to_probe[..., 2] / (grid_relative_to_probe_norm + 1e-6)
+
+    alpha = ops.arccos(grid_relative_to_probe_z)
+
+    # The f-number is fnum = z/aperture = 1/(2 * tan(alpha))
+    # Rearranging gives us alpha = arctan(1/(2 * fnum))
+    # We can use this to compute the maximum angle alpha that is allowed
+    max_alpha = ops.arctan(1 / (2 * f_number))
+
+    normalized_angle = alpha / max_alpha
+    mask = fnum_window_fn(normalized_angle)
+
+    # Add dummy channel dimension
     mask = mask[..., None]
 
     return mask

zea/beamform/phantoms.py

@@ -0,0 +1,43 @@
+import numpy as np
+
+
+def fish():
+    """Returns a scatterer phantom for ultrasound simulation tests.
+
+    Returns:
+        ndarray: The scatterer positions of shape (n_scat, 3).
+    """
+    # The size is the height of the fish
+    size = 11e-3
+    z_offset = 2.0 * size
+
+    # See https://en.wikipedia.org/wiki/Fish_curve
+    def fish_curve(t, size=1):
+        x = size * (np.cos(t) - np.sin(t) ** 2 / np.sqrt(2))
+        y = size * np.cos(t) * np.sin(t)
+        return x, y
+
+    scat_x, scat_z = fish_curve(np.linspace(0, 2 * np.pi, 100), size=size)
+
+    scat_x = np.concatenate(
+        [
+            scat_x,
+            np.array([size * 0.7]),
+            np.array([size * 1.1]),
+            np.array([size * 1.4]),
+            np.array([size * 1.2]),
+        ]
+    )
+    scat_y = np.zeros_like(scat_x)
+    scat_z = np.concatenate(
+        [
+            scat_z,
+            np.array([-size * 0.1]),
+            np.array([-size * 0.25]),
+            np.array([-size * 0.6]),
+            np.array([-size * 1.0]),
+        ]
+    )
+
+    scat_z += z_offset
+    return np.stack([scat_x, scat_y, scat_z], axis=1)

zea/beamform/pixelgrid.py

@@ -7,67 +7,27 @@ from zea import log
 eps = 1e-10
 
 
-def get_grid(
-    xlims,
-    zlims,
-    Nx: int,
-    Nz: int,
-    sound_speed,
-    center_frequency,
-    pixels_per_wavelength,
-    verbose=False,
-):
-    """Creates a pixelgrid based on scan class parameters."""
-
-    if Nx and Nz:
-        grid = cartesian_pixel_grid(xlims, zlims, Nx=int(Nx), Nz=int(Nz))
-    else:
-        wvln = sound_speed / center_frequency
-        dx = wvln / pixels_per_wavelength
-        dz = dx
-        grid = cartesian_pixel_grid(xlims, zlims, dx=dx, dz=dz)
-        if verbose:
-            print(
-                f"Pixelgrid was set automatically to Nx: {grid.shape[1]}, Nz: {grid.shape[0]}, "
-                f"using {pixels_per_wavelength} pixels per wavelength."
-            )
-    return grid
-
-
 def check_for_aliasing(scan):
     """Checks if the scan class parameters will cause spatial aliasing due to a too low pixel
     density. If so, a warning is printed with a suggestion to increase the pixel density by either
     increasing the number of pixels, or decreasing the pixel spacing, depending on which parameter
     was set by the user."""
-    wvln = scan.sound_speed / scan.center_frequency
-    dx = wvln / scan.pixels_per_wavelength
-    dz = dx
-
     width = scan.xlims[1] - scan.xlims[0]
     depth = scan.zlims[1] - scan.zlims[0]
-
-    if scan.Nx and scan.Nz:
-        if width / scan.Nx > wvln / 2:
-            log.warning(
-                f"width/Nx = {width / scan.Nx:.7f} < wvln/2 = {wvln / 2}. "
-                f"Consider increasing scan.Nx to {int(np.ceil(width / (wvln / 2)))} or more."
-            )
-        if depth / scan.Nz > wvln / 2:
-            log.warning(
-                f"depth/Nz = {depth / scan.Nz:.7f} < wvln/2 = {wvln / 2:.7f}. "
-                f"Consider increasing scan.Nz to {int(np.ceil(depth / (wvln / 2)))} or more."
-            )
-    else:
-        if dx > wvln / 2:
-            log.warning(
-                f"dx = {dx:.7f} > wvln/2 = {wvln / 2:.7f}. "
-                f"Consider increasing scan.pixels_per_wavelength to 2 or more"
-            )
-        if dz > wvln / 2:
-            log.warning(
-                f"dz = {dz:.7f} > wvln/2 = {wvln / 2:.7f}. "
-                f"Consider increasing scan.pixels_per_wavelength to 2 or more"
-            )
+    wvln = scan.wavelength
+
+    if width / scan.Nx > wvln / 2:
+        log.warning(
+            f"width/Nx = {width / scan.Nx:.7f} < wavelength/2 = {wvln / 2}. "
+            f"Consider either increasing scan.Nx to {int(np.ceil(width / (wvln / 2)))} or more, or "
+            "increasing scan.pixels_per_wavelength to 2 or more."
+        )
+    if depth / scan.Nz > wvln / 2:
+        log.warning(
+            f"depth/Nz = {depth / scan.Nz:.7f} < wavelength/2 = {wvln / 2:.7f}. "
+            f"Consider either increasing scan.Nz to {int(np.ceil(depth / (wvln / 2)))} or more, or "
+            "increasing scan.pixels_per_wavelength to 2 or more."
+        )
 
 
 def cartesian_pixel_grid(xlims, zlims, Nx=None, Nz=None, dx=None, dz=None):
@@ -130,7 +90,7 @@ def radial_pixel_grid(rlims, dr, oris, dirs):
             Cartesian coordinates (x, y, z), with nr being the number of radial pixels.
     """
     # Get focusing positions in rho-theta coordinates
-    r = np.arange(rlims[0], rlims[1] + eps, dr)  # Depth rho
+    r = np.arange(rlims[0], rlims[1], dr)  # Depth rho
     t = dirs[:, 0]  # Use azimuthal angle theta (ignore elevation angle)
     tt, rr = np.meshgrid(t, r, indexing="ij")
 
@@ -140,3 +100,31 @@ def radial_pixel_grid(rlims, dr, oris, dirs):
     yy = 0 * xx
     grid = np.stack((xx, yy, zz), axis=-1)
     return grid
+
+
+def polar_pixel_grid(polar_limits, zlims, Nz: int, Nr: int):
+    """Generate a polar grid.
+
+    Uses radial_pixel_grid but based on parameters that are present in the scan class.
+
+    Args:
+        polar_limits (tuple): Azimuthal limits of pixel grid ([azimuth_min, azimuth_max])
+        zlims (tuple): Depth limits of pixel grid ([zmin, zmax])
+        Nz (int, optional): Number of depth pixels.
+        Nr (int, optional): Number of azimuthal pixels.
+
+    Returns:
+        grid (np.ndarray): Pixel grid of size (Nz, Nr, 3) in Cartesian coordinates (x, y, z)
+    """
+    assert len(polar_limits) == 2, "polar_limits must be a tuple of length 2."
+    assert len(zlims) == 2, "zlims must be a tuple of length 2."
+
+    dr = (zlims[1] - zlims[0]) / Nz
+
+    oris = np.array([0, 0, 0])
+    oris = np.tile(oris, (Nr, 1))
+    dirs_az = np.linspace(*polar_limits, Nr)
+
+    dirs_el = np.zeros(Nr)
+    dirs = np.vstack((dirs_az, dirs_el)).T
+    return radial_pixel_grid(zlims, dr, oris, dirs).transpose(1, 0, 2)

zea/data/__main__.py

@@ -0,0 +1,31 @@
+"""Command-line interface for copying a zea.Folder to a new location.
+
+Usage:
+    python -m zea.data <source_folder> <destination_folder> <key>
+"""
+
+import argparse
+
+from zea import Folder
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Copy a zea.Folder to a new location.")
+    parser.add_argument("src", help="Source folder path")
+    parser.add_argument("dst", help="Destination folder path")
+    parser.add_argument("key", help="Key to access in the hdf5 files")
+    parser.add_argument(
+        "--mode",
+        default="a",
+        choices=["a", "w", "r+", "x"],
+        help="Mode in which to open the destination files (default: 'a')",
+    )
+
+    args = parser.parse_args()
+
+    src_folder = Folder(args.src, args.key, validate=False)
+    src_folder.copy(args.dst, args.key, mode=args.mode)
+
+
+if __name__ == "__main__":
+    main()

zea/data/data_format.py

@@ -84,10 +84,12 @@ def generate_example_dataset(
         aligned_data = np.ones((n_frames, n_tx, n_ax, n_el, n_ch))
         envelope_data = np.ones((n_frames, n_z, n_x))
         beamformed_data = np.ones((n_frames, n_z, n_x, n_ch))
+        image_sc = np.ones_like(image)
     else:
         aligned_data = None
         envelope_data = None
         beamformed_data = None
+        image_sc = None
 
     generate_zea_dataset(
         path,
@@ -96,6 +98,7 @@ def generate_example_dataset(
         envelope_data=envelope_data,
         beamformed_data=beamformed_data,
         image=image,
+        image_sc=image_sc,
         probe_geometry=probe_geometry,
         sampling_frequency=sampling_frequency,
         center_frequency=center_frequency,

zea/data/datasets.py

@@ -113,14 +113,14 @@ class H5FileHandleCache:
             self._file_handle_cache = OrderedDict()
 
 
-def find_h5_files(paths: str | list, key: str, search_file_tree_kwargs: dict | None = None):
+def find_h5_files(paths: str | list, key: str = None, search_file_tree_kwargs: dict | None = None):
     """
-    Find HDF5 files from a directory or list of directories and retrieve their shapes.
+    Find HDF5 files from a directory or list of directories and optionally retrieve their shapes.
 
     Args:
         paths (str or list): A single directory path, a list of directory paths,
             or a single HDF5 file path.
-        key (str): The key to access the HDF5 dataset.
+        key (str, optional): The key to get the file shapes for.
         search_file_tree_kwargs (dict, optional): Additional keyword arguments for the
             search_file_tree function. Defaults to None.
 
@@ -147,7 +147,8 @@ def find_h5_files(paths: str | list, key: str, search_file_tree_kwargs: dict | N
         if Path(path).is_file():
             path = Path(path)
             # If the path is a file, get its shape directly
-            file_shapes.append(File.get_shape(path, key))
+            if key is not None:
+                file_shapes.append(File.get_shape(path, key))
             file_paths.append(str(path))
             continue
 
@@ -171,7 +172,7 @@ class Folder:
     def __init__(
         self,
         folder_path: list[str] | list[Path],
-        key: str,
+        key: str = None,
         search_file_tree_kwargs: dict | None = None,
         validate: bool = True,
         hf_cache_dir: str = HF_DATASETS_DIR,
@@ -192,7 +193,7 @@ class Folder:
 
         super().__init__(**kwargs)
 
-        self.folder_path = folder_path
+        self.folder_path = Path(folder_path)
         self.key = key
         self.search_file_tree_kwargs = search_file_tree_kwargs
         self.validate = validate
@@ -346,6 +347,64 @@ class Folder:
     def __str__(self):
         return f"Folder with {self.n_files} files in '{self.folder_path}' (key='{self.key}')"
 
+    def copy(self, to_path: str | Path, key: str = None, mode: str | None = None):
+        """Copy the data for all or a specific key to a new location.
+
+        Has the option to copy all keys or only a specific key. By default, it only copies if the
+        destination file does not already contain the key. You can change the mode to 'w' to
+        overwrite the destination file. Will always copy metadata such as dataset attributes and
+        scan object.
+
+        Args:
+            to_path (str or Path): The destination path where files will be copied.
+            key (str, optional): The key to copy from the source files.
+                If 'all' or '*', all keys will be copied. Defaults to None, which
+                uses the key set in the Folder instance.
+            mode (str): The mode in which to open the destination files.
+                Defaults to 'a' (append mode), and 'w' (write mode) if key is 'all' or '*'.
+                See: https://docs.h5py.org/en/stable/high/file.html#opening-creating-files
+        """
+        if key is None and self.key is None:
+            raise ValueError(
+                "No key specified. Please provide a key to copy the data for, or set the "
+                "key attribute of the Folder instance."
+            )
+        elif key is None:
+            key = self.key
+
+        all_keys = key == "all" or key == "*"
+
+        if mode is None:
+            mode = "a" if not all_keys else "w"
+
+        if all_keys:
+            key_msg = "Including all keys."
+            assert mode in ["w", "x"], (
+                "If you want to copy all keys, the destination file must be opened "
+                "in 'w' or 'x' mode, which means it will be overwritten or created."
+            )
+        else:
+            key_msg = f"Only copying key '{key}'."
+            assert mode in ["a", "w", "r+", "x"], (
+                f"Invalid mode '{mode}'. Must be one of 'a', 'w', 'r+', or 'x'."
+            )
+
+        to_path = Path(to_path)
+        to_path.mkdir(parents=True, exist_ok=True)
+
+        for file_path in tqdm.tqdm(
+            self.file_paths,
+            total=self.n_files,
+            desc=f"Copying dataset from {self.folder_path} to {to_path}. {key_msg}",
+        ):
+            dst_path = Path(file_path).relative_to(self.folder_path)
+            with File(file_path) as src, File(to_path / dst_path, mode) as dst:
+                if all_keys:
+                    for obj in src.keys():
+                        src.copy(obj, dst)
+                else:
+                    src.copy_key(key, dst)
+
 
 class Dataset(H5FileHandleCache):
     """Iterate over File(s) and Folder(s)."""
@@ -582,3 +641,10 @@ def count_samples_per_directory(file_names, directories):
     )
 
     return counts
+
+
+if __name__ == "__main__":
+    src_folder = Folder(
+        "/mnt/z/Ultrasound-BMd/data/USBMD_datasets/CAMUS/val/patient0450", "image", validate=False
+    )
+    src_folder.copy("./CAMIUS", key="all")

zea/data/file.py

@@ -175,6 +175,8 @@ class File(h5py.File):
         if isinstance(key, enum.Enum):
             key = key.value
 
+        assert isinstance(key, str), f"Key must be a string, got {type(key)}. "
+
         # Return the key if it is in the file
         if key in self.keys():
             return key
@@ -448,6 +450,35 @@ class File(h5py.File):
     def __str__(self):
         return f"zea HDF5 File: '{self.path.name}' (mode={self.mode})"
 
+    def copy_key(self, key: str, dst: "File"):
+        """Copy a specific key to another file.
+
+        Will always copy the attributes and the scan data if it exists. Will warn if the key is
+        not in this file or if the key already exists in the destination file.
+
+        Args:
+            key (str): The key to copy.
+            dst (File): The destination file to copy the key to.
+        """
+        key = self.format_key(key)
+
+        # Copy the key if it does not already exist in the destination file
+        if key in dst:
+            log.warning(f"Skipping key '{key}' because it already exists in dst file {dst.path}.")
+        elif key in self:
+            self.copy(key, dst, name=key)
+        else:
+            log.warning(f"Key '{key}' not found in src file {self.path}. Skipping copy.")
+
+        # Copy attributes from src to dst
+        for attr_key, attr_value in self.attrs.items():
+            dst[key].attrs[attr_key] = attr_value
+
+        # Copy scan data if requested
+        if "scan" in self and "scan" not in dst:
+            # Copy the scan data if it exists
+            self.copy("scan", dst)
+
     def summary(self):
         """Print the contents of the file."""
         _print_hdf5_attrs(self)

zea/display.py

@@ -3,6 +3,7 @@
 from functools import partial
 from typing import Tuple, Union
 
+import keras
 import numpy as np
 import scipy
 from keras import ops
@@ -357,6 +358,11 @@ def map_coordinates(inputs, coordinates, order, fill_mode="constant", fill_value
 
 def _interpolate_batch(images, coordinates, fill_value=0.0, order=1, vectorize=True):
     """Interpolate a batch of images."""
+
+    # TODO: figure out why tensorflow map_coordinates is broken
+    if keras.backend.backend() == "tensorflow":
+        assert order > 1, "Some bug in tensorflow in map_coordinates, set order > 1 to use scipy."
+
     image_shape = images.shape
     num_image_dims = coordinates.shape[0]
 

zea/internal/cache.py

@@ -80,7 +80,7 @@ _CACHE_DIR = ZEA_CACHE_DIR / "cached_funcs"
 _CACHE_DIR.mkdir(parents=True, exist_ok=True)
 
 
-def serialize_elements(key_elements: list):
+def serialize_elements(key_elements: list, shorten: bool = False) -> str:
     """Serialize elements of a list to generate a cache key.
 
     In general uses the string representation of the elements unless
@@ -90,18 +90,18 @@ def serialize_elements(key_elements: list):
     Args:
         key_elements (list): List of elements to serialize. Can be nested lists
             or tuples. In this case the elements are serialized recursively.
+        shorten (bool): If True, the serialized string is hashed to a shorter
+            representation using MD5. Defaults to False.
 
     Returns:
-        list[str]: List of serialized elements. In cases where the elements were
-            lists of tuples those are combined into a single string.
+        str: A serialized string representation of the elements, joined by underscores.
 
     """
     serialized_elements = []
     for element in key_elements:
         if isinstance(element, (list, tuple)):
             # If element is a list or tuple, serialize its elements recursively
-            element = serialize_elements(element)
-            serialized_elements.append("_".join(element))
+            serialized_elements.append(serialize_elements(element))
         elif hasattr(element, "serialized"):
             # Use the serialized attribute if it exists (e.g. for zea.core.Object)
             serialized_elements.append(str(element.serialized))
@@ -120,7 +120,10 @@ def serialize_elements(key_elements: list):
             element = hashlib.md5(element).hexdigest()
             serialized_elements.append(element)
 
-    return serialized_elements
+    serialized = "_".join(serialized_elements)
+    if shorten:
+        return hashlib.md5(serialized.encode()).hexdigest()
+    return serialized
 
 
 def get_function_source(func):
@@ -182,8 +185,10 @@ def generate_cache_key(func, args, kwargs, arg_names):
             if name in bound_args.arguments:
                 key_elements.append(bound_args.arguments[name])
 
-    key = "_".join(serialize_elements(key_elements))
-    return f"{func.__qualname__}_" + hashlib.md5(key.encode()).hexdigest()
+    # Add keras backend
+    key_elements.append(keras.backend.backend())
+
+    return f"{func.__qualname__}_" + serialize_elements(key_elements, shorten=True)
 
 
 def cache_output(*arg_names, verbose=False):