Move dataloader to Grain

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 with `zea.data.dataloader.Dataloader`"
    ]
   },
   {
@@ -91,7 +91,7 @@
     "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.data.dataloader import Dataloader"
    ]
   },
   {
@@ -378,9 +378,9 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Loading data with `make_dataloader`\n",
+    "## Loading data with `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.data.dataloader.Dataloader` class provides a convenient way to create a high-performance data loader from a zea dataset. It is built on Grain and does not require TensorFlow. This dataloader is particularly useful for training models. Consistency of shape is preferred, 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."
    ]
   },
   {
@@ -460,7 +460,7 @@
    ],
    "source": [
     "dataset_path = \"hf://zeahub/camus-sample/val\"\n",
-    "dataloader = make_dataloader(\n",
+    "dataloader = Dataloader(\n",
     "    dataset_path,\n",
     "    key=\"data/image_sc\",\n",
     "    batch_size=4,\n",

docs/source/notebooks/metrics.rst

@@ -1,5 +1,5 @@
 Metrics
-=======
+========
 
 .. toctree::
     :maxdepth: 1

docs/source/notebooks/metrics/lpips_example.ipynb

@@ -78,7 +78,7 @@
     "from keras import ops\n",
     "\n",
     "from zea import init_device\n",
-    "from zea.backend.tensorflow.dataloader import make_dataloader\n",
+    "from zea.data.dataloader import Dataloader\n",
     "from zea.models.lpips import LPIPS\n",
     "from zea.visualize import plot_image_grid, set_mpl_style"
    ]
@@ -165,7 +165,7 @@
    ],
    "source": [
     "n_imgs = 9\n",
-    "val_dataset = make_dataloader(\n",
+    "val_dataset = Dataloader(\n",
     "    \"hf://zeahub/camus-sample/val\",\n",
     "    key=\"data/image\",\n",
     "    batch_size=n_imgs,\n",

docs/source/notebooks/metrics/myocardial_quality_example.ipynb

@@ -62,7 +62,7 @@
     "import numpy as np\n",
     "\n",
     "from zea import init_device\n",
-    "from zea.backend.tensorflow.dataloader import make_dataloader\n",
+    "from zea.data.dataloader import Dataloader\n",
     "from zea.visualize import set_mpl_style\n",
     "from zea.io_lib import matplotlib_figure_to_numpy, save_video\n",
     "\n",
@@ -142,7 +142,7 @@
     "# Load a batch and run both models.\n",
     "n_imgs = 1\n",
     "INFERENCE_SIZE = 256\n",
-    "val_dataset = make_dataloader(\n",
+    "val_dataset = Dataloader(\n",
     "    \"hf://zeahub/camus-sample/val\",\n",
     "    key=\"data/image_sc\",\n",
     "    batch_size=n_imgs,\n",

docs/source/notebooks/models/hvae_model_example.ipynb

@@ -79,7 +79,7 @@
     "from zea.display import scan_convert_2d\n",
     "from zea.agent.selection import UniformRandomLines\n",
     "from zea.visualize import set_mpl_style, plot_image_grid\n",
-    "from zea.backend.tensorflow.dataloader import make_dataloader\n",
+    "from zea.data.dataloader import Dataloader\n",
     "\n",
     "init_device(verbose=False)\n",
     "set_mpl_style()"
@@ -135,7 +135,7 @@
     }
    ],
    "source": [
-    "val_dataset = make_dataloader(\n",
+    "val_dataset = Dataloader(\n",
     "    \"hf://zeahub/camus-sample/val\",\n",
     "    key=\"data/image\",\n",
     "    batch_size=batch_size,\n",

docs/source/notebooks/models/left_ventricle_segmentation_example.ipynb

@@ -70,7 +70,7 @@
     "from zea import init_device\n",
     "import matplotlib.pyplot as plt\n",
     "from keras import ops\n",
-    "from zea.backend.tensorflow.dataloader import make_dataloader\n",
+    "from zea.data.dataloader import Dataloader\n",
     "from zea.visualize import plot_shape_from_mask\n",
     "from zea.func import translate\n",
     "from zea.visualize import plot_image_grid, set_mpl_style\n",
@@ -116,7 +116,7 @@
    "source": [
     "n_imgs = 16\n",
     "INFERENCE_SIZE = 256  # Used for both models\n",
-    "val_dataset = make_dataloader(\n",
+    "val_dataset = Dataloader(\n",
     "    \"hf://zeahub/camus-sample/val\",\n",
     "    key=\"data/image_sc\",\n",
     "    batch_size=n_imgs,\n",

docs/source/notebooks/models/taesd_autoencoder_example.ipynb

@@ -81,7 +81,7 @@
     "\n",
     "import zea\n",
     "from zea import init_device\n",
-    "from zea.backend.tensorflow.dataloader import make_dataloader\n",
+    "from zea.data.dataloader import Dataloader\n",
     "from zea.models.taesd import TinyAutoencoder\n",
     "from zea.visualize import plot_image_grid, set_mpl_style"
    ]
@@ -143,7 +143,7 @@
    ],
    "source": [
     "n_imgs = 4\n",
-    "val_dataset = make_dataloader(\n",
+    "val_dataset = Dataloader(\n",
     "    \"hf://zeahub/camus-sample/val\",\n",
     "    key=\"data/image\",\n",
     "    batch_size=n_imgs,\n",

docs/source/notebooks/models/unet_example.ipynb

@@ -77,7 +77,7 @@
     "from keras import ops\n",
     "\n",
     "from zea import init_device, log\n",
-    "from zea.backend.tensorflow.dataloader import make_dataloader\n",
+    "from zea.data.dataloader import Dataloader\n",
     "from zea.models.unet import UNet\n",
     "from zea.models.lpips import LPIPS\n",
     "from zea.agent.masks import random_uniform_lines\n",
@@ -142,7 +142,7 @@
    "source": [
     "n_imgs = 8\n",
     "\n",
-    "val_dataset = make_dataloader(\n",
+    "val_dataset = Dataloader(\n",
     "    \"hf://zeahub/camus-sample/val\",\n",
     "    key=\"data/image\",\n",
     "    batch_size=n_imgs,\n",