Coverage for src/bioimageio/core/_sample_serializer.py: 78%

1from abc import ABC, abstractmethod 

2from typing import ( 

3 Generic, 

4 Iterable, 

5 Tuple, 

6 TypeVar, 

7 Union, 

8) 

9 

10from bioimageio.spec.model import v0_5 

11 

12from .axis import PerAxis 

13from .common import HaloLike, PadMode, PerMember 

14from .digest_spec import split_sample_into_blocks_for_model 

15from .sample import Sample, SampleBlock 

16 

17SerializedSampleBlockType = TypeVar("SerializedSampleBlockType") 

18 

19 

20class SampleSerializer(ABC, Generic[SerializedSampleBlockType]): 

21 @classmethod 

22 def serialize_sample( 

23 cls, 

24 sample: Sample, 

25 ) -> Tuple[SerializedSampleBlockType]: 

26 """Serialize a sample as a single block""" 

27 return (cls.serialize_sample_block(sample.as_single_block()),) 

28 

29 @classmethod 

30 def deserialize_sample( 

31 cls, 

32 serialized: Iterable[SerializedSampleBlockType], 

33 fill_value: float = float("nan"), 

34 ) -> Sample: 

35 return Sample.from_blocks( 

36 (cls.deserialize_sample_block(s) for s in serialized), fill_value=fill_value 

37 ) 

38 

39 def serialize_sample_blockwise( 

40 self, 

41 sample: Sample, 

42 *, 

43 model: v0_5.ModelDescr, 

44 blocksize_parameter: int, 

45 batch_size: int = 1, 

46 ) -> Iterable[SerializedSampleBlockType]: 

47 """Split a sample into blocks according to the model's input specifications and `blocksize_parameter` and serialize each block.""" 

48 

49 _n_blocks, blocks = split_sample_into_blocks_for_model( 

50 sample, 

51 model=model, 

52 blocksize_parameter=blocksize_parameter, 

53 batch_size=batch_size, 

54 ) 

55 for block in blocks: 

56 yield self.serialize_sample_block(block) 

57 

58 @classmethod 

59 def serialize_sample_with_fixed_blocking( 

60 cls, 

61 sample: Sample, 

62 *, 

63 block_shapes: PerMember[PerAxis[int]], 

64 halo: PerMember[PerAxis[HaloLike]], 

65 pad_mode: Union[PadMode, PerMember[PadMode]] = "symmetric", 

66 ) -> Iterable[SerializedSampleBlockType]: 

67 

68 _n_blocks, input_blocks = sample.split_into_blocks( 

69 block_shapes=block_shapes, 

70 halo=halo, 

71 pad_mode=pad_mode, 

72 ) 

73 for block in input_blocks: 

74 yield cls.serialize_sample_block(block) 

75 

76 @staticmethod 

77 @abstractmethod 

78 def serialize_sample_block( 

79 sample_block: SampleBlock, 

80 ) -> SerializedSampleBlockType: ... 

81 

82 @staticmethod 

83 @abstractmethod 

84 def deserialize_sample_block(serialized: SerializedSampleBlockType) -> SampleBlock: 

85 """Deserialize a sample block into a new sample or merge it into `output_sample` if provided."""