Coverage for src/bioimageio/core/sample.py: 87%
188 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-08 15:59 +0000
1from __future__ import annotations
3import collections.abc
4from dataclasses import dataclass
5from math import ceil, floor
6from types import MappingProxyType
7from typing import (
8 Any,
9 Callable,
10 Dict,
11 Generic,
12 Iterable,
13 Literal,
14 Mapping,
15 Optional,
16 Sequence,
17 Tuple,
18 TypeVar,
19 Union,
20)
22import numpy as np
23import pandas as pd
24import pydantic
25import xarray as xr
26from numpy.typing import NDArray
27from typing_extensions import Self
29from ._common_annotations import PerMemberAnno
30from .axis import AxisId, PerAxis
31from .block import Block
32from .block_meta import (
33 BlockMeta,
34 LinearAxisTransform,
35 split_multiple_shapes_into_blocks,
36)
37from .common import (
38 BlockIndex,
39 Halo,
40 HaloLike,
41 MemberId,
42 PadMode,
43 PadWidthLike,
44 PerMember,
45 SampleId,
46 SliceInfo,
47 TotalNumberOfBlocks,
48)
49from .stat_measures import Stat
50from .tensor import Tensor
52# TODO: allow for lazy samples to read/write to disk
55@dataclass
56class Sample:
57 """A dataset sample.
59 A `Sample` has `members`, which allows to combine multiple tensors into a single
60 sample.
61 For example a `Sample` from a dataset with masked images may contain a
62 `MemberId("raw")` and `MemberId("mask")` image.
63 """
65 members: Dict[MemberId, Tensor]
66 """The sample's tensors"""
68 stat: Stat
69 """Sample and dataset statistics"""
71 id: SampleId
72 """Identifies the `Sample` within the dataset -- typically a number or a string."""
74 def __getitem__(
75 self,
76 key: PerMember[
77 Union[
78 SliceInfo,
79 slice,
80 int,
81 PerAxis[Union[SliceInfo, slice, int]],
82 Tensor,
83 xr.DataArray,
84 ]
85 ],
86 ) -> Self:
87 return self.__class__(
88 members={m: t[key[m]] for m, t in self.members.items() if m in key},
89 stat=self.stat,
90 id=self.id,
91 )
93 @property
94 def batch_multi_index(self) -> Optional["pd.MultiIndex"]:
95 """Return the batch multi-index of the sample, if it has one.
97 Returns:
98 The batch multi-index of the sample, or `None` if the sample does not have a batch dimension.
99 """
100 if not self.members:
101 return None
103 for tensor in self.members.values():
104 idx = tensor.data.indexes.get(AxisId("batch")) # pyright: ignore[reportUnknownVariableType]
105 if isinstance(idx, pd.MultiIndex):
106 return idx
108 return None
110 def set_block(self, block: SampleBlock) -> None:
111 """Set values of `block`.
113 Note:
114 - Updates only existing sample members (extra block members are ignored)
115 - Ignores missing block members (i.e. members in the sample but not in the block are not modified)
117 Raises:
118 ValueError if block and sample members do not overlap at all.
119 """
120 no_overlap = True
121 for m in self.members:
122 if m not in block.blocks:
123 continue
124 b = block.blocks[m]
125 self.members[m][b.inner_slice] = b.inner_data
126 no_overlap = False
128 if no_overlap:
129 raise ValueError(
130 f"block with members {list(block.blocks)} does not overlap with sample members {list(self.members)}"
131 )
133 @property
134 def shape(self) -> PerMember[PerAxis[int]]:
135 return {tid: t.sizes for tid, t in self.members.items()}
137 def as_arrays(self) -> Dict[MemberId, NDArray[Any]]:
138 """Return sample as dictionary of arrays."""
139 return {m: t.to_numpy() for m, t in self.members.items()}
141 def split_into_blocks(
142 self,
143 block_shapes: PerMember[PerAxis[int]],
144 halo: PerMember[PerAxis[HaloLike]],
145 pad_mode: Union[PadMode, PerMember[PadMode]],
146 broadcast: bool = False,
147 ) -> Tuple[TotalNumberOfBlocks, Iterable[SampleBlockWithOrigin]]:
148 assert not (missing := [m for m in block_shapes if m not in self.members]), (
149 f"`block_shapes` specified for unknown members: {missing}"
150 )
151 assert not (missing := [m for m in halo if m not in block_shapes]), (
152 f"`halo` specified for members without `block_shape`: {missing}"
153 )
155 n_blocks, blocks = split_multiple_shapes_into_blocks(
156 shapes=self.shape,
157 block_shapes=block_shapes,
158 halo=halo,
159 broadcast=broadcast,
160 )
161 return n_blocks, sample_block_generator(blocks, origin=self, pad_mode=pad_mode)
163 def as_single_block(self, halo: Optional[PerMember[PerAxis[Halo]]] = None):
164 if halo is None:
165 halo = {}
166 return SampleBlockWithOrigin(
167 sample_shape=self.shape,
168 sample_id=self.id,
169 blocks={
170 m: Block(
171 sample_shape=self.shape[m],
172 data=data,
173 inner_slice={
174 a: SliceInfo(0, s) for a, s in data.tagged_shape.items()
175 },
176 halo=halo.get(m, {}),
177 block_index=0,
178 blocks_in_sample=1,
179 )
180 for m, data in self.members.items()
181 },
182 stat=self.stat,
183 origin=self,
184 block_index=0,
185 blocks_in_sample=1,
186 )
188 @classmethod
189 def from_blocks(
190 cls,
191 sample_blocks: Iterable[SampleBlock],
192 *,
193 fill_value: float = float("nan"),
194 ) -> Self:
195 """Create a `Sample` from an iterable of `SampleBlock`s.
197 Note:
198 All sample blocks must have the same `sample_id`.
200 Args:
201 sample_blocks: The blocks to create the sample from.
202 fill_value: The value to fill missing values with (default: `nan`).
203 """
204 output = None
205 for output in cls.from_blocks_yield_intermediates(
206 sample_blocks, fill_value=fill_value
207 ):
208 pass
210 if output is None:
211 raise ValueError("no sample blocks provided")
213 return output
215 @classmethod
216 def from_blocks_yield_intermediates(
217 cls,
218 sample_blocks: Iterable[SampleBlock],
219 *,
220 fill_value: float = float("nan"),
221 ):
222 """Create a `Sample` from an iterable of `SampleBlock`s, yielding the intermediate sample after each block.
224 Args:
225 sample_blocks: The blocks to create the sample from.
226 fill_value: The value to fill missing values with (default: `nan`).
227 """
228 output = cls(members={}, stat={}, id=None)
229 for sample_block in sample_blocks:
230 if output.id is None:
231 output.id = sample_block.sample_id
232 else:
233 assert output.id == sample_block.sample_id, (
234 "sample id changed between sample blocks"
235 )
237 output.stat = sample_block.stat
239 for m, block in sample_block.blocks.items():
240 if m not in output.members:
241 if -1 in block.sample_shape.values():
242 raise NotImplementedError(
243 "merging blocks with data dependent axis not yet implemented"
244 )
246 output.members[m] = Tensor(
247 np.full(
248 tuple(block.sample_shape[a] for a in block.data.dims),
249 fill_value,
250 dtype=block.data.dtype,
251 ),
252 dims=block.data.dims,
253 )
255 output.members[m][block.inner_slice] = block.inner_data
256 yield output
258 yield output
260 def pad(
261 self,
262 pad_width: PerMember[PerAxis[Union[int, PadWidthLike]]],
263 mode: Union[PerMember[PadMode], PadMode],
264 ) -> Self:
265 """Convenience method to pad sample members."""
266 default_mode = "symmetric"
267 if isinstance(mode, collections.abc.Mapping):
268 mode_per_member = mode
269 else:
270 mode_per_member: Mapping[MemberId, PadMode] = {}
271 default_mode = mode
273 return self.__class__(
274 members={
275 m: t.pad(
276 pad_width=pad_width.get(m, {}),
277 mode=mode_per_member.get(m, default_mode),
278 )
279 for m, t in self.members.items()
280 },
281 stat=self.stat,
282 id=self.id,
283 )
285 def transpose(
286 self,
287 axes: PerMember[Sequence[AxisId]],
288 *,
289 extra_dims: Literal["raise", "squeeze", "stack", "squeeze_or_stack"] = "raise",
290 missing_dims: Literal[
291 "raise", "expand", "unstack", "unstack_or_expand"
292 ] = "raise",
293 ) -> Self:
294 """Return a new sample with transposed sample members.
296 Raises:
297 ValueError: If not all batch dimensions have the same length after transposition (and possibly stacking/unstacking extra dimensions).
299 """
300 if any((unknown := [m not in self.members for m in axes])):
301 raise ValueError(f"Axes specified for unknown members: {unknown}")
303 members = {
304 m: t
305 if m not in axes
306 else t.transpose(
307 axes=axes[m],
308 extra_dims=extra_dims,
309 missing_dims=missing_dims,
310 )
311 for m, t in self.members.items()
312 }
314 if (
315 len(
316 (
317 batch_lengths := {
318 t.sizes[AxisId("batch")]
319 for t in members.values()
320 if AxisId("batch") in t.dims
321 }
322 )
323 )
324 > 1
325 ):
326 raise ValueError(
327 f"Transposed sample members have incompatible batch lengths: {batch_lengths}."
328 )
330 return self.__class__(members=members, stat=dict(self.stat), id=self.id)
332 def assign_batch_multi_index(self, multi_index: "pd.MultiIndex") -> Self:
333 """Return a new sample with the batch multi-index assigned to all sample members.
335 Raises:
336 ValueError: If not all sample members have a batch dimension.
337 """
338 if len(
339 no_batch := [
340 m for m, t in self.members.items() if AxisId("batch") not in t.dims
341 ]
342 ) == len(self.members):
343 raise ValueError(f"No member has a batch dimension: {no_batch}")
345 return self.__class__(
346 members={
347 m: t
348 if AxisId("batch") not in t.dims
349 else t.assign_batch_multi_index(multi_index)
350 for m, t in self.members.items()
351 },
352 stat=dict(self.stat),
353 id=self.id,
354 )
356 def unstack_batch_multi_index(
357 self, *, errors: Literal["raise", "ignore"] = "raise"
358 ) -> Self:
359 """Unstack the batch multi-index of all sample members.
361 Args:
362 errors: Whether to raise an error if a member does not have a batch multi-index. Default is "raise".
364 Returns:
365 A new `Sample` with unstacked batch multi-index for all members.
366 """
367 if (
368 len(
369 no_batch := [
370 m for m, t in self.members.items() if AxisId("batch") not in t.dims
371 ]
372 )
373 == len(self.members)
374 and errors == "raise"
375 ):
376 raise ValueError(f"No member has a batch dimension: {no_batch}")
378 members = {
379 m: t
380 if AxisId("batch") not in t.dims
381 else t.unstack_batch_multi_index(errors=errors)
382 for m, t in self.members.items()
383 }
384 if (
385 len(
386 batch_lengths := {
387 t.sizes.get(AxisId("batch"))
388 for t in members.values()
389 if AxisId("batch") in t.dims
390 }
391 )
392 > 1
393 ):
394 raise ValueError(
395 f"Different batch lengths after unstacking: {batch_lengths}"
396 )
398 stat: Stat = {
399 k: v.unstack_batch_multi_index(errors="ignore")
400 if isinstance(v, Tensor)
401 else float(v)
402 for k, v in self.stat.items()
403 }
404 return self.__class__(
405 members=members,
406 stat=stat,
407 id=self.id,
408 )
411BlockT = TypeVar("BlockT", bound=BlockMeta)
414@pydantic.dataclasses.dataclass(frozen=True)
415class SampleBlockBase(Generic[BlockT]):
416 """base class for `SampleBlockMeta` and `SampleBlock`"""
418 sample_shape: PerMemberAnno[PerAxis[int]]
419 """the sample shape this block represents a part of"""
421 sample_id: SampleId
422 """identifier for the sample within its dataset"""
424 blocks: PerMemberAnno[BlockT]
425 """Individual tensor blocks comprising this sample block"""
427 block_index: BlockIndex
428 """the n-th block of the sample"""
430 blocks_in_sample: TotalNumberOfBlocks
431 """total number of blocks in the sample"""
433 @property
434 def shape(self) -> PerMember[PerAxis[int]]:
435 return MappingProxyType({mid: b.shape for mid, b in self.blocks.items()})
437 @property
438 def inner_shape(self) -> PerMember[PerAxis[int]]:
439 return MappingProxyType({mid: b.inner_shape for mid, b in self.blocks.items()})
442@dataclass
443class LinearSampleAxisTransform(LinearAxisTransform):
444 member: MemberId
447@pydantic.dataclasses.dataclass(frozen=True)
448class SampleBlockMeta(SampleBlockBase[BlockMeta]):
449 """Meta data of a dataset sample block"""
451 def get_transformed(
452 self, new_axes: PerMember[PerAxis[Union[LinearSampleAxisTransform, int]]]
453 ) -> Self:
454 sample_shape = {
455 m: {
456 a: (
457 trf
458 if isinstance(trf, int)
459 else trf.compute(self.sample_shape[trf.member][trf.axis])
460 )
461 for a, trf in new_axes[m].items()
462 }
463 for m in new_axes
464 }
466 def get_member_halo(m: MemberId, round: Callable[[float], int]):
467 return {
468 a: (
469 Halo(0, 0)
470 if isinstance(trf, int)
471 or trf.axis not in self.blocks[trf.member].halo
472 else Halo(
473 round(self.blocks[trf.member].halo[trf.axis].left * trf.scale),
474 round(self.blocks[trf.member].halo[trf.axis].right * trf.scale),
475 )
476 )
477 for a, trf in new_axes[m].items()
478 }
480 halo: Dict[MemberId, Dict[AxisId, Halo]] = {}
481 for m in new_axes:
482 halo[m] = get_member_halo(m, floor)
483 if halo[m] != get_member_halo(m, ceil):
484 raise ValueError(
485 f"failed to unambiguously scale halo {halo[m]} with {new_axes[m]}"
486 + f" for {m}."
487 )
489 inner_slice = {
490 m: {
491 a: (
492 SliceInfo(0, trf)
493 if isinstance(trf, int)
494 else SliceInfo(
495 trf.compute(
496 self.blocks[trf.member].inner_slice[trf.axis].start
497 ),
498 trf.compute(self.blocks[trf.member].inner_slice[trf.axis].stop),
499 )
500 )
501 for a, trf in new_axes[m].items()
502 }
503 for m in new_axes
504 }
505 return self.__class__(
506 blocks={
507 m: BlockMeta(
508 sample_shape=sample_shape[m],
509 inner_slice=inner_slice[m],
510 halo=halo[m],
511 block_index=self.block_index,
512 blocks_in_sample=self.blocks_in_sample,
513 )
514 for m in new_axes
515 },
516 sample_shape=sample_shape,
517 sample_id=self.sample_id,
518 block_index=self.block_index,
519 blocks_in_sample=self.blocks_in_sample,
520 )
522 def with_data(self, data: PerMember[Tensor], *, stat: Stat) -> SampleBlock:
523 return SampleBlock(
524 sample_shape={
525 m: {
526 a: data[m].tagged_shape[a] if s == -1 else s
527 for a, s in member_shape.items()
528 }
529 for m, member_shape in self.sample_shape.items()
530 },
531 sample_id=self.sample_id,
532 blocks={
533 m: Block.from_meta(b, data=data[m]) for m, b in self.blocks.items()
534 },
535 stat=stat,
536 block_index=self.block_index,
537 blocks_in_sample=self.blocks_in_sample,
538 )
541@dataclass(frozen=True)
542class SampleBlock(SampleBlockBase[Block]):
543 """A block of a dataset sample"""
545 blocks: Dict[MemberId, Block]
546 """Individual tensor blocks comprising this sample block"""
548 stat: Stat
549 """computed statistics"""
551 @property
552 def members(self) -> PerMember[Tensor]:
553 """the sample block's tensors"""
554 return {m: b.data for m, b in self.blocks.items()}
556 def get_transformed_meta(
557 self, new_axes: PerMember[PerAxis[Union[LinearSampleAxisTransform, int]]]
558 ) -> SampleBlockMeta:
559 return SampleBlockMeta(
560 sample_id=self.sample_id,
561 blocks=dict(self.blocks),
562 sample_shape=self.sample_shape,
563 block_index=self.block_index,
564 blocks_in_sample=self.blocks_in_sample,
565 ).get_transformed(new_axes)
567 @classmethod
568 def from_meta(
569 cls, meta: SampleBlockMeta, data: PerMember[Tensor], stat: Stat
570 ) -> Self:
571 return cls(
572 sample_shape=meta.sample_shape,
573 sample_id=meta.sample_id,
574 blocks={
575 m: Block.from_meta(b, data=data[m]) for m, b in meta.blocks.items()
576 },
577 stat=stat,
578 block_index=meta.block_index,
579 blocks_in_sample=meta.blocks_in_sample,
580 )
582 def get_meta(self) -> SampleBlockMeta:
583 return SampleBlockMeta(
584 sample_id=self.sample_id,
585 blocks={m: b.get_meta() for m, b in self.blocks.items()},
586 sample_shape=self.sample_shape,
587 block_index=self.block_index,
588 blocks_in_sample=self.blocks_in_sample,
589 )
591 def as_sample(self) -> Sample:
592 """Convert this sample block to a `Sample` with the shape of this block.
594 Note:
595 If you want to convert one or more sample block to a sample with the shape of the original, whole sample,
596 use `Sample.from_blocks()` instead.
597 """
598 return Sample(
599 members=dict(self.members),
600 stat=dict(self.stat),
601 id=self.sample_id,
602 )
605@dataclass(frozen=True)
606class SampleBlockWithOrigin(SampleBlock):
607 """A `SampleBlock` with a reference (`origin`) to the whole `Sample`"""
609 origin: Sample
610 """the sample this sample block was taken from"""
613class _ConsolidatedMemberBlocks:
614 def __init__(self, blocks: PerMember[BlockMeta]):
615 super().__init__()
616 block_indices = {b.block_index for b in blocks.values()}
617 assert len(block_indices) == 1
618 self.block_index = block_indices.pop()
619 blocks_in_samples = {b.blocks_in_sample for b in blocks.values()}
620 assert len(blocks_in_samples) == 1
621 self.blocks_in_sample = blocks_in_samples.pop()
624def sample_block_meta_generator(
625 blocks: Iterable[PerMember[BlockMeta]],
626 *,
627 sample_shape: PerMember[PerAxis[int]],
628 sample_id: SampleId,
629):
630 for member_blocks in blocks:
631 cons = _ConsolidatedMemberBlocks(member_blocks)
632 yield SampleBlockMeta(
633 blocks=dict(member_blocks),
634 sample_shape=sample_shape,
635 sample_id=sample_id,
636 block_index=cons.block_index,
637 blocks_in_sample=cons.blocks_in_sample,
638 )
641def sample_block_generator(
642 blocks: Iterable[PerMember[BlockMeta]],
643 *,
644 origin: Sample,
645 pad_mode: Union[PadMode, PerMember[PadMode]],
646) -> Iterable[SampleBlockWithOrigin]:
647 for member_blocks in blocks:
648 cons = _ConsolidatedMemberBlocks(member_blocks)
649 yield SampleBlockWithOrigin(
650 blocks={
651 m: Block.from_sample_member(
652 origin.members[m],
653 block=member_blocks[m],
654 pad_mode=pad_mode.get(m, "symmetric")
655 if isinstance(pad_mode, collections.abc.Mapping)
656 else pad_mode,
657 )
658 for m in origin.members
659 },
660 sample_shape=origin.shape,
661 origin=origin,
662 stat=origin.stat,
663 sample_id=origin.id,
664 block_index=cons.block_index,
665 blocks_in_sample=cons.blocks_in_sample,
666 )