Skip to content

zarr_source

anomaly_match.datasets.zarr_source

Zarr data source for training from pixel-array containers.

ZarrSource(cfg, *, lazy=False)

Bases: TrainingDataSource

Load training data from one or more Zarr arrays with optional metadata.

Supports two layouts:

  • Single store -- cfg.data_dir points directly to a .zarr directory containing an images array.
  • Multi-store directory -- cfg.data_dir points to a parent directory containing multiple .zarr stores, each with an images array. All stores are indexed transparently.

When lazy=True, the expensive O(N) filename index and shuffled-index permutation are skipped. Unlabeled sampling uses random (store, index) generation instead. Labeled images must come from a LabeledDataCache, not from this source.

Parameters:

Name Type Description Default
cfg DotMap

Configuration; cfg.data_dir must point to a .zarr directory or a parent directory containing .zarr stores.

 required
lazy bool

If True, skip O(N) indexing for billion-scale sources.

 False
Source code in anomaly_match/datasets/zarr_source.py 
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
def __init__(self, cfg: DotMap, *, lazy: bool = False) -> None:
    super().__init__(cfg)
    self._zarr_path = Path(cfg.data_dir)
    self._lazy = lazy

    # Discover stores: single .zarr or parent with multiple .zarr children
    self._stores: list[_ZarrStoreInfo] = self._discover_stores()
    if not self._stores:
        raise ValueError(f"No Zarr stores with 'images' array found in {self._zarr_path}")

    self._total = sum(s.count for s in self._stores)

    if lazy:
        # Lightweight: only store sizes, no filename or shuffle index
        self._name_to_loc = None
        self._filenames = None
        self._shuffled_indices = None
        store_desc = ", ".join(f"{s.path.name}({s.count})" for s in self._stores)
        logger.debug(f"ZarrSource(lazy): {self._total} images from [{store_desc}]")
    else:
        # Full mode: build O(N) filename index + shuffled indices
        self._name_to_loc: dict[str, tuple[int, int]] = {}
        self._filenames: list[str] = []
        for si, store in enumerate(self._stores):
            for li, fn in enumerate(store.filenames):
                self._name_to_loc[fn] = (si, li)
                self._filenames.append(fn)

        rng = np.random.RandomState(cfg.seed)
        self._shuffled_indices = rng.permutation(self._total)

        store_desc = ", ".join(f"{s.path.name}({s.count})" for s in self._stores)
        logger.debug(f"ZarrSource: indexed {self._total} images from [{store_desc}]")

get_labeled_images(label_df)

Load images for filenames present in label_df.

Not available in lazy mode -- use a LabeledDataCache instead.

Returns:

Type Description
list[tuple[str, ndarray]]

List of (filename, image_array) tuples.

Raises:

Type Description
RuntimeError

If called in lazy mode.
Source code in anomaly_match/datasets/zarr_source.py 
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
def get_labeled_images(self, label_df: pd.DataFrame) -> list[tuple[str, np.ndarray]]:
    """Load images for filenames present in *label_df*.

    Not available in lazy mode -- use a ``LabeledDataCache`` instead.

    Returns:
        List of (filename, image_array) tuples.

    Raises:
        RuntimeError: If called in lazy mode.
    """
    if self._lazy:
        raise RuntimeError(
            "get_labeled_images() not available in lazy mode. Use a LabeledDataCache instead."
        )

    results: list[tuple[str, np.ndarray]] = []
    for fn in label_df["id"]:
        if fn in self._name_to_loc:
            si, li = self._name_to_loc[fn]
            img = decode_zarr_image(self._stores[si].images[li], self._cfg)
            results.append((fn, img))

    logger.debug(f"ZarrSource: loaded {len(results)} labeled images")
    return results

get_unlabeled_batch(n, exclude=None, exclude_positions=None)

Load up to n unlabeled images.

In normal mode, skips filenames in exclude. In lazy mode, skips (store_idx, local_idx) pairs in exclude_positions using random sampling without a pre-computed permutation.

Parameters:

Name Type Description Default
n int

Maximum number of images to return.

 required
exclude set[str] | None

Filenames to skip (normal mode).

 None
exclude_positions set[tuple[int, int]] | None

(store_idx, local_idx) pairs to skip (lazy mode).

 None

Returns:

Type Description
list[tuple[str, ndarray]]

List of (filename, image_array) tuples.
Source code in anomaly_match/datasets/zarr_source.py 
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
def get_unlabeled_batch(
    self,
    n: int,
    exclude: set[str] | None = None,
    exclude_positions: set[tuple[int, int]] | None = None,
) -> list[tuple[str, np.ndarray]]:
    """Load up to *n* unlabeled images.

    In normal mode, skips filenames in *exclude*. In lazy mode, skips
    ``(store_idx, local_idx)`` pairs in *exclude_positions* using
    random sampling without a pre-computed permutation.

    Args:
        n: Maximum number of images to return.
        exclude: Filenames to skip (normal mode).
        exclude_positions: ``(store_idx, local_idx)`` pairs to skip
            (lazy mode).

    Returns:
        List of (filename, image_array) tuples.
    """
    if self._lazy:
        return self._get_unlabeled_batch_lazy(n, exclude_positions or set())

    results: list[tuple[str, np.ndarray]] = []
    exclude = exclude or set()
    for idx in self._shuffled_indices:
        if len(results) >= n:
            break
        fn = self._filenames[idx]
        if fn not in exclude:
            si, li = self._name_to_loc[fn]
            img = decode_zarr_image(self._stores[si].images[li], self._cfg)
            results.append((fn, img))

    logger.debug(f"ZarrSource: loaded {len(results)} unlabeled images")
    return results

get_total_count()

Return total number of images across all Zarr stores.

Returns:

Type Description
int

Image count.
Source code in anomaly_match/datasets/zarr_source.py 
297
298
299
300
301
302
303
def get_total_count(self) -> int:
    """Return total number of images across all Zarr stores.

    Returns:
        Image count.
    """
    return self._total

detect_num_channels()

Infer channel count from the first store's array shape.

Returns:

Type Description
int | None

Number of channels, or None if the array is empty.
Source code in anomaly_match/datasets/zarr_source.py 
305
306
307
308
309
310
311
312
313
314
315
316
317
318
def detect_num_channels(self) -> int | None:
    """Infer channel count from the first store's array shape.

    Returns:
        Number of channels, or ``None`` if the array is empty.
    """
    if self._total == 0:
        return None
    shape = self._stores[0].images.shape
    # (N, H, W, C) or (N, C, H, W)
    if len(shape) == 4:
        # Smaller of last vs second dim is likely channels
        return min(shape[1], shape[3])
    return None