diff --git a/agml/data/__init__.py b/agml/data/__init__.py index 4c239a85..2a7ac751 100644 --- a/agml/data/__init__.py +++ b/agml/data/__init__.py @@ -19,3 +19,4 @@ from .public import download_public_dataset, public_data_sources, source from .tools import coco_to_bboxes, convert_bbox_format from .hf_loader import HuggingFaceDataLoader +from .multispectral_hf_loader import MultispectralDataLoader diff --git a/agml/data/multispectral_hf_loader.py b/agml/data/multispectral_hf_loader.py new file mode 100644 index 00000000..fff58ffc --- /dev/null +++ b/agml/data/multispectral_hf_loader.py @@ -0,0 +1,97 @@ +import numpy as np + +try: + from datasets import DatasetDict +except ImportError: + raise ImportError( + "The `datasets` library is required to use the MultispectralDataLoader. " + "Please install it using `pip install datasets`." + ) + +from agml.data.hf_loader import HuggingFaceDataLoader + + +class MultispectralDataLoader(HuggingFaceDataLoader): + """A `HuggingFaceDataLoader` variant for multispectral datasets. + + In addition to the base loader's functionality, this loader decodes a + `bands` column stored on the Hub as raw binary buffers back into `numpy` + arrays. The dataset is expected to expose the following columns: + + - `bands`: `binary`, the raw bytes of the array's buffer. + - `bands_shape`: `list[int64]`, the array shape as `[channels, H, W]`. + - `bands_dtype`: `string`, the `numpy` dtype name to reinterpret the buffer as. + + Decoding is applied lazily (mirroring how `Image` columns are decoded on + access), so the raw bytes are only materialized into arrays when an + example is actually retrieved. + + Parameters + ---------- + dataset_name : str + The name of a dataset on the Hugging Face Hub. + config : str, optional + The dataset configuration (subset) name for datasets that expose multiple + configs (e.g. an augmented variant). + cache_dir : str, optional + The local directory to cache the dataset in. + """ + + def _cast_features(self): + super()._cast_features() + if isinstance(self._hf_dataset, DatasetDict): + for split_name, ds in self._hf_dataset.items(): + self._hf_dataset[split_name] = self._decode_bands(ds) + else: + self._hf_dataset = self._decode_bands(self._hf_dataset) + + def _decode_bands(self, ds): + """Attaches a lazy transform that decodes the `bands` column into `numpy` arrays.""" + if "bands" not in ds.features: + return ds + return ds.with_transform(self._bands_transform) + + @staticmethod + def _bands_transform(batch): + """Decodes raw `bands` byte buffers into `numpy` arrays using `bands_shape`/`bands_dtype`.""" + raw_bands = batch.get("bands") + if raw_bands is None: + return batch + + shapes = batch.get("bands_shape") + dtypes = batch.get("bands_dtype") + + decoded = [] + for i, raw in enumerate(raw_bands): + dtype = dtypes[i] if dtypes is not None else "float32" + # `frombuffer` returns a read-only view over `raw`; copy so consumers + # (e.g. in-place normalization/augmentation) get a writable array. + arr = np.frombuffer(raw, dtype=np.dtype(dtype)).copy() + shape = shapes[i] if shapes is not None else None + if shape is not None: + arr = arr.reshape(shape) + decoded.append(arr) + + batch["bands"] = decoded + return batch + + def split(self, *args, **kwargs) -> DatasetDict: + """Splits the dataset, re-attaching the `bands` decoding transform to each split.""" + # `HuggingFaceDataLoader.split` internally calls `.map()`/`.unique()` on the + # dataset to build/inspect the stratification key. Those operations read + # examples through the current format, so the lazy `bands` decode transform + # must be cleared beforehand — otherwise raw bytes get replaced by decoded + # arrays mid-split, corrupting the underlying binary column. + self._hf_dataset = self._reset_bands_format(self._hf_dataset) + dataset_dict = super().split(*args, **kwargs) + for split_name, ds in dataset_dict.items(): + dataset_dict[split_name] = self._decode_bands(ds) + self._hf_dataset = dataset_dict + return dataset_dict + + @staticmethod + def _reset_bands_format(ds): + """Clears any `bands` decoding transform, restoring the raw binary column.""" + if isinstance(ds, DatasetDict): + return DatasetDict({name: split.with_format(None) for name, split in ds.items()}) + return ds.with_format(None) diff --git a/tests/data/test_multispectral_hf_loader.py b/tests/data/test_multispectral_hf_loader.py new file mode 100644 index 00000000..0db62992 --- /dev/null +++ b/tests/data/test_multispectral_hf_loader.py @@ -0,0 +1,161 @@ +import numpy as np +import pytest +from unittest.mock import patch +from datasets import Dataset, DatasetDict + +from agml.data.multispectral_hf_loader import MultispectralDataLoader + + +PATCH_TARGET = "agml.data.hf_loader.load_dataset" + + +# ── Helpers ────────────────────────────────────────────────────────────────── + +def _band_arrays(n, shape=(4, 8, 8), dtype="float32", seed=0): + rng = np.random.default_rng(seed) + return [rng.standard_normal(shape).astype(dtype) for _ in range(n)] + + +def _multispectral_dataset(n=100, shape=(4, 8, 8), dtype="float32", n_classes=5): + """In-memory Dataset with a raw-bytes `bands` column — no network required.""" + arrays = _band_arrays(n, shape=shape, dtype=dtype) + labels = (list(range(n_classes)) * (n // n_classes + 1))[:n] + return Dataset.from_dict( + { + "bands": [a.tobytes() for a in arrays], + "bands_shape": [list(shape)] * n, + "bands_dtype": [dtype] * n, + "label": labels, + } + ), arrays + + +def _presplit_multispectral_dataset_dict(n=100, shape=(4, 8, 8), dtype="float32", n_classes=5): + ds, arrays = _multispectral_dataset(n, shape, dtype, n_classes) + split = ds.train_test_split(test_size=0.2, seed=0) + return DatasetDict({"train": split["train"], "test": split["test"]}), arrays + + +def _loader(dataset=None): + if dataset is None: + dataset, _ = _multispectral_dataset() + with patch(PATCH_TARGET, return_value=dataset): + return MultispectralDataLoader("org/dataset") + + +# ── Band decoding ──────────────────────────────────────────────────────────── + +class TestBandDecoding: + def test_bands_decoded_to_numpy_array(self): + loader = _loader() + example = loader.dataset[0] + assert isinstance(example["bands"], np.ndarray) + + def test_bands_decoded_shape(self): + shape = (6, 16, 16) + ds, _ = _multispectral_dataset(shape=shape) + loader = _loader(ds) + example = loader.dataset[0] + assert example["bands"].shape == shape + + def test_bands_decoded_dtype(self): + ds, _ = _multispectral_dataset(dtype="float64") + loader = _loader(ds) + example = loader.dataset[0] + assert example["bands"].dtype == np.dtype("float64") + + def test_bands_values_round_trip(self): + ds, arrays = _multispectral_dataset(n=10) + loader = _loader(ds) + for i in range(10): + np.testing.assert_array_equal(loader.dataset[i]["bands"], arrays[i]) + + def test_bands_array_is_writable(self): + loader = _loader() + example = loader.dataset[0] + example["bands"][0, 0, 0] = 123.0 + assert example["bands"][0, 0, 0] == 123.0 + + def test_bands_decoded_on_dataset_dict(self): + fake, arrays = _presplit_multispectral_dataset_dict() + with patch(PATCH_TARGET, return_value=fake): + loader = MultispectralDataLoader("org/dataset") + example = loader.dataset["train"][0] + assert isinstance(example["bands"], np.ndarray) + + def test_batched_access_decodes_all_items(self): + ds, arrays = _multispectral_dataset(n=5) + loader = _loader(ds) + batch = loader.dataset[0:5] + assert len(batch["bands"]) == 5 + for arr in batch["bands"]: + assert isinstance(arr, np.ndarray) + + def test_other_columns_unaffected(self): + ds, _ = _multispectral_dataset(n=10) + loader = _loader(ds) + example = loader.dataset[0] + assert "label" in example + assert isinstance(example["label"], int) + + def test_no_bands_column_is_noop(self): + ds = Dataset.from_dict({"label": [0, 1, 0, 1]}) + loader = _loader(ds) + example = loader.dataset[0] + assert "bands" not in example + + +# ── Inheritance ────────────────────────────────────────────────────────────── + +class TestInheritance: + def test_is_subclass_of_hf_loader(self): + from agml.data.hf_loader import HuggingFaceDataLoader + assert issubclass(MultispectralDataLoader, HuggingFaceDataLoader) + + def test_dataset_property_returns_internal(self): + loader = _loader() + assert loader.dataset is loader._hf_dataset + + def test_raises_on_load_failure(self): + with patch(PATCH_TARGET, side_effect=Exception("not found")): + with pytest.raises(RuntimeError, match="Failed to load Hugging Face dataset"): + MultispectralDataLoader("nonexistent/dataset") + + +# ── Split ──────────────────────────────────────────────────────────────────── + +class TestSplit: + def test_split_preserves_band_decoding(self): + ds, _ = _multispectral_dataset(n=100) + loader = _loader(ds) + result = loader.split(val_size=0.2) + assert set(result.keys()) == {"train", "val"} + for split in result.values(): + assert isinstance(split[0]["bands"], np.ndarray) + + def test_split_updates_internal_dataset(self): + loader = _loader() + result = loader.split(val_size=0.2) + assert loader._hf_dataset is result + + def test_val_and_test_split_decode_bands(self): + ds, _ = _multispectral_dataset(n=100) + loader = _loader(ds) + result = loader.split(val_size=0.15, test_size=0.15) + for split_name in ("train", "val", "test"): + assert isinstance(result[split_name][0]["bands"], np.ndarray) + + def test_stratified_split_preserves_band_decoding(self): + ds, _ = _multispectral_dataset(n=100, n_classes=5) + loader = _loader(ds) + result = loader.split(val_size=0.2, stratify_cols="label") + assert isinstance(result["train"][0]["bands"], np.ndarray) + assert isinstance(result["val"][0]["bands"], np.ndarray) + + def test_bands_shapes_correct_after_split(self): + shape = (3, 10, 10) + ds, _ = _multispectral_dataset(n=50, shape=shape) + loader = _loader(ds) + result = loader.split(val_size=0.2) + assert result["train"][0]["bands"].shape == shape + assert result["val"][0]["bands"].shape == shape