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

zea/interface.py

@@ -3,15 +3,15 @@
 Example usage
 ^^^^^^^^^^^^^^
 
-.. code-block:: python
+.. doctest::
 
-    import zea
-    from zea.internal.setup_zea import setup_config
+    >>> import zea
+    >>> from zea.internal.setup_zea import setup_config
 
-    config = setup_config("hf://zeahub/configs/config_camus.yaml")
+    >>> config = setup_config("hf://zeahub/configs/config_camus.yaml")
 
-    interface = zea.Interface(config)
-    interface.run(plot=True)
+    >>> interface = zea.Interface(config)
+    >>> interface.run(plot=True)  # doctest: +SKIP
 
 """
 
@@ -31,15 +31,15 @@ from zea.data.file import File
 from zea.datapaths import format_data_path
 from zea.display import to_8bit
 from zea.internal.core import DataTypes
+from zea.internal.utils import keep_trying
 from zea.internal.viewer import (
     ImageViewerMatplotlib,
     ImageViewerOpenCV,
     filename_from_window_dialog,
     running_in_notebook,
 )
-from zea.io_lib import matplotlib_figure_to_numpy
+from zea.io_lib import matplotlib_figure_to_numpy, save_video
 from zea.ops import Pipeline
-from zea.utils import keep_trying, save_to_gif, save_to_mp4
 
 
 class Interface:
@@ -266,10 +266,11 @@ class Interface:
         save = self.config.plot.save
 
         if self.frame_no == "all":
-            if not asyncio.get_event_loop().is_running():
-                asyncio.run(self.run_movie(save))
-            else:
-                asyncio.create_task(self.run_movie(save))
+            try:
+                loop = asyncio.get_running_loop()
+                loop.create_task(self.run_movie(save))  # already running loop
+            except RuntimeError:
+                asyncio.run(self.run_movie(save))  # no loop yet
 
         else:
             if plot:
@@ -520,10 +521,7 @@ class Interface:
 
         fps = self.config.plot.fps
 
-        if self.config.plot.video_extension == "gif":
-            save_to_gif(images, path, fps=fps)
-        elif self.config.plot.video_extension == "mp4":
-            save_to_mp4(images, path, fps=fps)
+        save_video(images, path, fps=fps)
 
         if self.verbose:
             log.info(f"Video saved to {log.yellow(path)}")

zea/internal/_generate_keras_ops.py

@@ -3,11 +3,10 @@ and :mod:`keras.ops.image` functions.
 
 They can be used in zea pipelines like any other :class:`zea.Operation`, for example:
 
-.. code-block:: python
+.. doctest::
 
-    from zea.keras_ops import Squeeze
-
-    op = Squeeze(axis=1)
+    >>> from zea.keras_ops import Squeeze
+    >>> op = Squeeze(axis=1)
 """
 
 import inspect
@@ -77,11 +76,11 @@ and :mod:`keras.ops.image` functions.
 
 They can be used in zea pipelines like any other :class:`zea.Operation`, for example:
 
-.. code-block:: python
+.. doctest::
 
-    from zea.keras_ops import Squeeze
+    >>> from zea.keras_ops import Squeeze
 
-    op = Squeeze(axis=1)
+    >>> op = Squeeze(axis=1)
 
 This file is generated automatically. Do not edit manually.
 Generated with Keras {keras.__version__}

zea/internal/cache.py

@@ -21,10 +21,8 @@
 
 import ast
 import atexit
-import hashlib
 import inspect
 import os
-import pickle
 import tempfile
 import textwrap
 from pathlib import Path
@@ -33,6 +31,7 @@ import joblib
 import keras
 
 from zea import log
+from zea.internal.core import hash_elements
 
 _DEFAULT_ZEA_CACHE_DIR = Path.home() / ".cache" / "zea"
 
@@ -80,52 +79,6 @@ _CACHE_DIR = ZEA_CACHE_DIR / "cached_funcs"
 _CACHE_DIR.mkdir(parents=True, exist_ok=True)
 
 
-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
-    the element has a `serialized` attribute, in which case it uses that.
-    For instance this is useful for custom classes that inherit from `zea.core.Object`.
-
-    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:
-        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
-            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))
-        elif isinstance(element, str):
-            # If element is a string, use it as is
-            serialized_elements.append(element)
-        elif isinstance(element, keras.random.SeedGenerator):
-            # If element is a SeedGenerator, use the state
-            element = keras.ops.convert_to_numpy(element.state.value)
-            element = pickle.dumps(element)
-            element = hashlib.md5(element).hexdigest()
-            serialized_elements.append(element)
-        else:
-            # Otherwise, serialize the element using pickle and hash it
-            element = pickle.dumps(element)
-            element = hashlib.md5(element).hexdigest()
-            serialized_elements.append(element)
-
-    serialized = "_".join(serialized_elements)
-    if shorten:
-        return hashlib.md5(serialized.encode()).hexdigest()
-    return serialized
-
-
 def get_function_source(func):
     """Recursively get the source code of a function and its nested functions."""
     try:
@@ -188,7 +141,7 @@ def generate_cache_key(func, args, kwargs, arg_names):
     # Add keras backend
     key_elements.append(keras.backend.backend())
 
-    return f"{func.__qualname__}_" + serialize_elements(key_elements, shorten=True)
+    return f"{func.__qualname__}_" + hash_elements(key_elements)
 
 
 def cache_output(*arg_names, verbose=False):

zea/internal/checks.py

@@ -64,8 +64,7 @@ def _check_raw_data(data=None, shape=None, with_batch_dim=None):
         shape (tuple, optional): shape of the data. Defaults to None.
             either data or shape must be provided.
         with_batch_dim (bool, optional): whether data has frame dimension at the start.
-            Setting this to True requires the data to have 5 dimensions. Defaults to
-            False.
+            Setting this to True requires the data to have 5 dimensions. Defaults to None.
 
     Raises:
         AssertionError: if data does not have expected shape
@@ -105,8 +104,7 @@ def _check_aligned_data(data=None, shape=None, with_batch_dim=None):
         shape (tuple, optional): shape of the data. Defaults to None.
             either data or shape must be provided.
         with_batch_dim (bool, optional): whether data has frame dimension at the start.
-            Setting this to True requires the data to have 5 dimensions. Defaults to
-            False.
+            Setting this to True requires the data to have 5 dimensions. Defaults to None.
 
     Raises:
         AssertionError: if data does not have expected shape
@@ -147,8 +145,7 @@ def _check_beamformed_data(data=None, shape=None, with_batch_dim=None):
         shape (tuple, optional): shape of the data. Defaults to None.
             either data or shape must be provided.
         with_batch_dim (bool, optional): whether data has frame dimension at the start.
-            Setting this to True requires the data to have 4 dimensions. Defaults to
-            False.
+            Setting this to True requires the data to have 4 dimensions. Defaults to None.
 
     Raises:
         AssertionError: if data does not have expected shape
@@ -190,8 +187,7 @@ def _check_envelope_data(data=None, shape=None, with_batch_dim=None):
         shape (tuple, optional): shape of the data. Defaults to None.
             either data or shape must be provided.
         with_batch_dim (bool, optional): whether data has frame dimension at the start.
-            Setting this to True requires the data to have 4 dimensions. Defaults to
-            False.
+            Setting this to True requires the data to have 3 dimensions. Defaults to None.
 
     Raises:
         AssertionError: if data does not have expected shape
@@ -227,8 +223,7 @@ def _check_image(data=None, shape=None, with_batch_dim=None):
         shape (tuple, optional): shape of the data. Defaults to None.
             either data or shape must be provided.
         with_batch_dim (bool, optional): whether data has frame dimension at the start.
-            Setting this to True requires the data to have 4 dimensions. Defaults to
-            False.
+            Setting this to True requires the data to have 3 dimensions. Defaults to None.
 
     Raises:
         AssertionError: if data does not have expected shape.
@@ -264,8 +259,7 @@ def _check_image_sc(data=None, shape=None, with_batch_dim=None):
         shape (tuple, optional): shape of the data. Defaults to None.
             either data or shape must be provided.
         with_batch_dim (bool, optional): whether data has frame dimension at the start.
-            Setting this to True requires the data to have 4 dimensions. Defaults to
-            False.
+            Setting this to True requires the data to have 3 dimensions. Defaults to None.
 
     Raises:
         AssertionError: if data does not have expected shape.

zea/internal/config/validation.py

@@ -15,9 +15,8 @@ from pathlib import Path
 
 from schema import And, Optional, Or, Regex, Schema
 
-import zea.metrics  # noqa: F401
 from zea.internal.checks import _DATA_TYPES
-from zea.internal.registry import metrics_registry
+from zea.metrics import metrics_registry
 
 # predefined checks, later used in schema to check validity of parameter
 any_number = Or(

zea/internal/core.py

@@ -1,6 +1,7 @@
 """Base classes for the toolbox"""
 
 import enum
+import hashlib
 import json
 import pickle
 from copy import deepcopy
@@ -8,7 +9,8 @@ from copy import deepcopy
 import keras
 import numpy as np
 
-from zea.utils import reduce_to_signature, update_dictionary
+from zea.internal.utils import reduce_to_signature
+from zea.utils import update_dictionary
 
 CONVERT_TO_KERAS_TYPES = (np.ndarray, int, float, list, tuple, bool)
 BASE_FLOAT_PRECISION = "float32"
@@ -76,7 +78,7 @@ class Object:
             attributes.pop(
                 "_serialized", None
             )  # Remove the cached serialized attribute to avoid recursion
-            self._serialized = pickle.dumps(attributes)
+            self._serialized = serialize_elements([attributes])
         return self._serialized
 
     def __setattr__(self, name: str, value):
@@ -167,9 +169,7 @@ def _skip_to_tensor(value):
     # Skip str (because JIT does not support it)
     # Skip methods and functions
     # Skip byte strings
-    if isinstance(value, str) or callable(value) or isinstance(value, bytes):
-        return True
-    return False
+    return isinstance(value, str) or callable(value) or isinstance(value, bytes)
 
 
 def dict_to_tensor(dictionary, keep_as_is=None):
@@ -184,8 +184,9 @@ def dict_to_tensor(dictionary, keep_as_is=None):
         # Get the value from the dictionary
         value = dictionary[key]
 
-        if isinstance(value, Object):
+        if isinstance(value, Object) and hasattr(value, "to_tensor"):
             snapshot[key] = value.to_tensor(keep_as_is=keep_as_is)
+            continue
 
         # Skip certain types
         if _skip_to_tensor(value):
@@ -288,3 +289,65 @@ class ZEADecoderJSON(json.JSONDecoder):
                 obj[key] = self._MOD_TYPES_MAP[value] if value is not None else None
 
         return obj
+
+
+def serialize_elements(key_elements: list) -> str:
+    """Serialize elements of a list to a string.
+
+    Generally, uses the pickle representation of the elements.
+
+    Args:
+        key_elements (list): List of elements to serialize. Can be nested lists
+            or tuples. In this case the elements are serialized recursively.
+
+    Returns:
+        str: A serialized string representation of the elements, joined by underscores.
+    """
+
+    def _serialize(element) -> str:
+        return pickle.dumps(element).hex()
+
+    def _serialize_element(element) -> str:
+        if isinstance(element, (list, tuple)):
+            # If element is a list or tuple, serialize its elements recursively
+            element = serialize_elements(element)
+        elif isinstance(element, Object) and hasattr(element, "serialized"):
+            # Use the serialized attribute if it exists
+            element = str(element.serialized)
+        elif isinstance(element, keras.random.SeedGenerator):
+            # If element is a SeedGenerator, use the state
+            element = keras.ops.convert_to_numpy(element.state.value)
+            element = _serialize(element)
+        elif isinstance(element, dict):
+            # If element is a dictionary, sort its keys and serialize its values recursively.
+            # This is needed to ensure the internal state and ordering of the dictionary does
+            # not affect the serialization.
+            keys = list(sorted(element.keys()))
+            values = [element[k] for k in keys]
+            keys = serialize_elements(keys)
+            values = serialize_elements(values)
+            element = f"k_{keys}_v_{values}"
+        else:
+            # Otherwise, serialize the element directly
+            element = _serialize(element)
+
+        return element
+
+    serialized_elements = []
+    for element in key_elements:
+        serialized_elements.append(_serialize_element(element))
+
+    return "_".join(serialized_elements)
+
+
+def hash_elements(key_elements: list) -> str:
+    """Generate an MD5 hash of the elements.
+
+    Args:
+        key_elements (list): List of elements to serialize and hash.
+
+    Returns:
+        str: An MD5 hash of the serialized elements.
+    """
+    serialized = serialize_elements(key_elements)
+    return hashlib.md5(serialized.encode()).hexdigest()

zea/internal/device.py

@@ -377,7 +377,11 @@ def init_device(
     allow_preallocate: bool = True,
     verbose: bool = True,
 ):
-    """Selects a GPU or CPU device based on the config.
+    """Automatically selects a GPU or CPU device.
+
+    Useful to call at the start of a script to set the device for
+    tensorflow, jax or pytorch. The function will select a GPU based
+    on available memory, or fall back to CPU if no GPU is available.
 
     Args:
         backend (str): String indicating which backend to use. Can be
@@ -412,7 +416,7 @@ def init_device(
     elif backend in ["numpy", "cpu"]:
         device = "cpu"
     else:
-        raise ValueError(f"Unknown backend ({backend}) in config.")
+        raise ValueError(f"Unknown backend ({backend}).")
 
     # Early exit if device is CPU
     if device == "cpu":