JOSS review @tomelse

docs/source/notebooks/data/zea_data_example.ipynb

@@ -9,7 +9,7 @@
     "\n",
     "1. Loading data from single file with `zea.File`\n",
     "2. Loading data from a group of files with `zea.Dataset`\n",
-    "3. Loading data in batches with dataloading utilities with `zea.backend.tensorflow.make_dataloader`"
+    "3. Loading data in batches with dataloading utilities via `zea.backend.tensorflow.make_dataloader`. This requires a working TensorFlow installation!"
    ]
   },
   {
@@ -86,12 +86,12 @@
     }
    ],
    "source": [
+    "import keras\n",
     "import matplotlib.pyplot as plt\n",
     "\n",
     "import zea\n",
     "from zea import init_device, load_file\n",
-    "from zea.visualize import set_mpl_style\n",
-    "from zea.backend.tensorflow import make_dataloader"
+    "from zea.visualize import set_mpl_style"
    ]
   },
   {
@@ -380,7 +380,11 @@
    "source": [
     "## Loading data with `make_dataloader`\n",
     "\n",
-    "In machine and deep learning workflows, we often want more features like batching, shuffling, and parallel data loading. The `zea.backend.tensorflow.make_dataloader` function provides a convenient way to create a TensorFlow data loader from a zea  dataset. This does require a working TensorFlow installation, but does work in combination with any other backend as well. This dataloader is particularly useful for training models. It is important that there is some consistency in the dataset, which is not the case for [PICMUS](https://www.creatis.insa-lyon.fr/Challenge/IEEE_IUS_2016/home). Therefore in this example we will use a small part of the [CAMUS](https://www.creatis.insa-lyon.fr/Challenge/camus/) dataset."
+    "In machine and deep learning workflows, we often want more features like batching, shuffling, and parallel data loading. The `zea.backend.tensorflow.make_dataloader` function provides a convenient way to create a TensorFlow data loader from a zea  dataset. \n",
+    "\n",
+    "🚨 Note! This does require a working TensorFlow installation, but does work in combination with any other backend as well. We are [working on migrating](https://github.com/tue-bmd/zea/pull/256) to [Grain](https://github.com/google/grain), which will provide a backend-agnostic dataloader in the near future.\n",
+    "\n",
+    "This dataloader is particularly useful for training models. It is important that there is some consistency in the dataset, which is not the case for [PICMUS](https://www.creatis.insa-lyon.fr/Challenge/IEEE_IUS_2016/home). Therefore in this example we will use a small part of the [CAMUS](https://www.creatis.insa-lyon.fr/Challenge/camus/) dataset."
    ]
   },
   {
@@ -459,6 +463,8 @@
     }
    ],
    "source": [
+    "from zea.backend.tensorflow import make_dataloader  # needs TensorFlow to be installed\n",
+    "\n",
     "dataset_path = \"hf://zeahub/camus-sample/val\"\n",
     "dataloader = make_dataloader(\n",
     "    dataset_path,\n",
@@ -556,7 +562,8 @@
     "pipeline = zea.Pipeline.from_config(config)\n",
     "parameters = pipeline.prepare_parameters(probe=probe, scan=scan)\n",
     "\n",
-    "images = pipeline(data=data, **parameters)[\"data\"]"
+    "images = pipeline(data=data, **parameters)[\"data\"]\n",
+    "images = keras.ops.convert_to_numpy(images)"
    ]
   },
   {

zea/data/file.py

@@ -172,7 +172,7 @@ def load_data(
         self,
         data_type,
         indices: Tuple[Union[list, slice, int], ...] | List[int] | int | None = None,
-    ):
+    ) -> np.ndarray:
         """Load data from the file.
 
         .. include:: ../common/file_indexing.rst
@@ -537,7 +537,7 @@ def load_file(
     data_type="raw_data",
     indices: Tuple[Union[list, slice, int], ...] | List[int] | int | None = None,
     scan_kwargs: dict = None,
-):
+) -> Tuple[np.ndarray, Scan, Probe]:
     """Loads a zea data files (h5py file).
 
     Returns the data together with a scan object containing the parameters

zea/ops/pipeline.py

@@ -1,5 +1,5 @@
 import json
-from typing import Dict, List, Union
+from typing import Any, Dict, List, Union
 
 import keras
 import numpy as np
@@ -303,7 +303,7 @@ def _get_timed_operations(self):
         self.timer = FunctionTimer()
         return [self.timer(op, name=op.__class__.__name__) for op in self._pipeline_layers]
 
-    def call(self, **inputs):
+    def call(self, **inputs) -> Dict[str, Any]:
         """Process input data through the pipeline."""
         for operation in self._callable_layers:
             try:
@@ -325,7 +325,7 @@ def call(self, **inputs):
             inputs = outputs
         return outputs
 
-    def __call__(self, return_numpy=False, **inputs):
+    def __call__(self, return_numpy=False, **inputs) -> Dict[str, Any]:
         """Process input data through the pipeline."""
 
         if any(key in inputs for key in ["probe", "scan", "config"]) or any(