Remove custom json for dataloader

Author: wesselvannierop

tests/data/test_dataloader.py

@@ -14,7 +14,6 @@
 from zea.data.datasets import Dataset
 from zea.data.file import File
 from zea.data.layers import Resizer
-from zea.data.utils import json_loads
 from zea.tools.hf import HFPath
 
 from .. import DEFAULT_TEST_SEED
@@ -239,6 +238,7 @@ def test_h5_dataset_return_filename(
     validate = directory != "dummy_hdf5"
     directory = request.getfixturevalue(directory)
 
+    N_AXIS = 3  # n_frames, height, width
     dataset = Dataloader(
         directory,
         key=key,
@@ -259,18 +259,31 @@ def test_h5_dataset_return_filename(
 
     _, file_dict = batch
 
-    assert len(file_dict) == batch_size, (
-        "The file_dict should contain the same number of elements as the batch size"
-    )
+    # Check keys
+    keys = ["filename", "fullpath", "indices"]
+    for key in keys:
+        assert key in file_dict, f"The file_dict should contain the key '{key}'"
 
-    file_dict = file_dict[0]  # get the first file_dict of the batch
-    file_dict = json_loads(file_dict)
+    # Check batch size and types
+    keys = ["filename", "fullpath"]
+    for key in keys:
+        assert len(file_dict[key]) == batch_size, (
+            f"The file_dict['{key}'] should contain the same number of elements as the batch size"
+        )
+        for path in file_dict[key]:
+            assert isinstance(path, str), f"Each path in file_dict['{key}'] should be a string"
 
-    filename = file_dict["filename"]
-    assert isinstance(filename, str), "The filename should be a string"
-    fullpath = file_dict["fullpath"]
-    assert isinstance(fullpath, str), "The fullpath should be a string"
-    assert "indices" in file_dict, "The file_dict should contain indices"
+    # indices nests one deeper, because it has one element per axis (n_frames, height, width)
+    indices = file_dict["indices"]
+    assert len(indices) == N_AXIS, (
+        f"The file_dict['indices'] should contain {N_AXIS} elements in this test"
+    )
+
+    for idx in indices:
+        assert len(idx) == batch_size, (
+            "Each axis in file_dict['indices'] should contain the same number of elements "
+            "as the batch size"
+        )
 
 
 @pytest.mark.parametrize(

zea/data/dataloader.py

@@ -34,7 +34,6 @@
 from zea.data.datasets import Dataset, H5FileHandleCache, count_samples_per_directory
 from zea.data.file import File
 from zea.data.layers import Resizer
-from zea.data.utils import json_dumps
 from zea.utils import map_negative_indices
 
 DEFAULT_NORMALIZATION_RANGE = (0, 1)
@@ -84,12 +83,12 @@ def generate_h5_indices(
                 (
                     "/folder/path_to_file.hdf5",
                     "data/image",
-                    (range(0, 1), slice(None, 256, None), slice(None, 256, None)),
+                    (slice(0, 1, 1), slice(None, 256, None), slice(None, 256, None)),
                 ),
                 (
                     "/folder/path_to_file.hdf5",
                     "data/image",
-                    (range(1, 2), slice(None, 256, None), slice(None, 256, None)),
+                    (slice(1, 2, 1), slice(None, 256, None), slice(None, 256, None)),
                 ),
                 ...,
             ]
@@ -136,7 +135,7 @@ def axis_indices_files():
             # Optionally limit frames to load from each file
             n_frames_in_file = min(n_frames_in_file, limit_n_frames)
             indices = [
-                list(range(i, i + block_size, frame_index_stride))
+                slice(i, i + block_size, frame_index_stride)
                 for i in range(0, n_frames_in_file - block_size + 1, block_step_size)
             ]
             yield [indices]
@@ -293,18 +292,16 @@ def __getitem__(self, index: int):
             return self._data_cache[index]
 
         file_name, key, indices = self.indices[index]
-        file_cache = self._get_cache()
-        file = file_cache.get_file(file_name)
+        file_handle_cache = self._get_file_handle_cache()
+        file = file_handle_cache.get_file(file_name)
         image = self._load(file, key, indices)
 
         if self.return_filename:
-            file_data = json_dumps(
-                {
-                    "fullpath": file.filename,
-                    "filename": Path(file_name).stem,
-                    "indices": indices,
-                }
-            )
+            file_data = {
+                "fullpath": file.filename,
+                "filename": Path(file_name).stem,
+                "indices": indices,
+            }
             result = (image, file_data)
         else:
             result = image
@@ -321,7 +318,7 @@ def __repr__(self) -> str:
 
     # -- internals -------------------------------------------------------------
 
-    def _get_cache(self) -> H5FileHandleCache:
+    def _get_file_handle_cache(self) -> H5FileHandleCache:
         """Return the file-handle cache for the current thread."""
         if not hasattr(self._local, "cache"):
             self._local.cache = H5FileHandleCache()
@@ -335,14 +332,10 @@ def _load(self, file: File, key: str, indices):
             images = file.load_data(key, indices)
         except (OSError, IOError):
             # Invalidate cache entry and retry once
-            cache = self._get_cache()
             fname = file.filename
-            cache._file_handle_cache.pop(fname, None)
-            try:
-                file.close()
-            except Exception:
-                pass
-            file = cache.get_file(fname)
+            file_handle_cache = self._get_file_handle_cache()
+            file_handle_cache.pop(fname)
+            file = file_handle_cache.get_file(fname)
             images = file.load_data(key, indices)
 
         if self.insert_frame_axis:

zea/data/datasets.py

@@ -104,6 +104,15 @@ def get_file(self, file_path) -> File:
 
         return self._file_handle_cache[file_path]
 
+    def pop(self, file_path):
+        """Pop a file from the cache and close it."""
+        file = self._file_handle_cache.pop(file_path, None)
+        if file is not None:
+            try:
+                file.close()
+            except Exception:
+                pass  # swallow exceptions during close
+
     def close(self):
         """Close all cached file handles."""
         cache: OrderedDict = getattr(self, "_file_handle_cache", None)