nabu 2024.1.10__py3-none-any.whl → 2024.2.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (152) hide show
  1. nabu/__init__.py +1 -1
  2. nabu/app/bootstrap.py +2 -3
  3. nabu/app/cast_volume.py +4 -2
  4. nabu/app/cli_configs.py +5 -0
  5. nabu/app/composite_cor.py +1 -1
  6. nabu/app/create_distortion_map_from_poly.py +5 -6
  7. nabu/app/diag_to_pix.py +7 -19
  8. nabu/app/diag_to_rot.py +14 -29
  9. nabu/app/double_flatfield.py +32 -44
  10. nabu/app/parse_reconstruction_log.py +3 -0
  11. nabu/app/reconstruct.py +53 -15
  12. nabu/app/reconstruct_helical.py +2 -2
  13. nabu/app/stitching.py +27 -13
  14. nabu/app/tests/__init__.py +0 -0
  15. nabu/app/tests/test_reduce_dark_flat.py +4 -1
  16. nabu/cuda/kernel.py +11 -2
  17. nabu/cuda/processing.py +2 -2
  18. nabu/cuda/src/cone.cu +77 -0
  19. nabu/cuda/src/hierarchical_backproj.cu +271 -0
  20. nabu/cuda/utils.py +0 -6
  21. nabu/estimation/alignment.py +5 -19
  22. nabu/estimation/cor.py +173 -599
  23. nabu/estimation/cor_sino.py +356 -26
  24. nabu/estimation/focus.py +63 -11
  25. nabu/estimation/tests/test_cor.py +124 -58
  26. nabu/estimation/tests/test_focus.py +6 -6
  27. nabu/estimation/tilt.py +2 -1
  28. nabu/estimation/utils.py +5 -33
  29. nabu/io/__init__.py +1 -1
  30. nabu/io/cast_volume.py +1 -1
  31. nabu/io/reader.py +416 -21
  32. nabu/io/tests/test_readers.py +422 -0
  33. nabu/io/tests/test_writers.py +1 -102
  34. nabu/io/writer.py +4 -433
  35. nabu/opencl/kernel.py +14 -3
  36. nabu/opencl/processing.py +8 -0
  37. nabu/pipeline/config_validators.py +5 -2
  38. nabu/pipeline/datadump.py +12 -5
  39. nabu/pipeline/estimators.py +162 -188
  40. nabu/pipeline/fullfield/chunked.py +168 -92
  41. nabu/pipeline/fullfield/chunked_cuda.py +7 -3
  42. nabu/pipeline/fullfield/computations.py +2 -7
  43. nabu/pipeline/fullfield/dataset_validator.py +0 -4
  44. nabu/pipeline/fullfield/nabu_config.py +37 -13
  45. nabu/pipeline/fullfield/processconfig.py +22 -13
  46. nabu/pipeline/fullfield/reconstruction.py +13 -9
  47. nabu/pipeline/helical/helical_chunked_regridded.py +1 -1
  48. nabu/pipeline/helical/helical_chunked_regridded_cuda.py +1 -0
  49. nabu/pipeline/helical/helical_reconstruction.py +1 -1
  50. nabu/pipeline/params.py +21 -1
  51. nabu/pipeline/processconfig.py +1 -12
  52. nabu/pipeline/reader.py +146 -0
  53. nabu/pipeline/tests/test_estimators.py +44 -72
  54. nabu/pipeline/utils.py +4 -2
  55. nabu/pipeline/writer.py +10 -2
  56. nabu/preproc/ccd_cuda.py +1 -1
  57. nabu/preproc/ctf.py +14 -7
  58. nabu/preproc/ctf_cuda.py +2 -3
  59. nabu/preproc/double_flatfield.py +5 -12
  60. nabu/preproc/double_flatfield_cuda.py +2 -2
  61. nabu/preproc/flatfield.py +5 -1
  62. nabu/preproc/flatfield_cuda.py +5 -1
  63. nabu/preproc/phase.py +24 -73
  64. nabu/preproc/phase_cuda.py +5 -8
  65. nabu/preproc/tests/test_ctf.py +11 -7
  66. nabu/preproc/tests/test_flatfield.py +67 -122
  67. nabu/preproc/tests/test_paganin.py +54 -30
  68. nabu/processing/azim.py +206 -0
  69. nabu/processing/convolution_cuda.py +1 -1
  70. nabu/processing/fft_cuda.py +15 -17
  71. nabu/processing/histogram.py +2 -0
  72. nabu/processing/histogram_cuda.py +2 -1
  73. nabu/processing/kernel_base.py +3 -0
  74. nabu/processing/muladd_cuda.py +1 -0
  75. nabu/processing/padding_opencl.py +1 -1
  76. nabu/processing/roll_opencl.py +1 -0
  77. nabu/processing/rotation_cuda.py +2 -2
  78. nabu/processing/tests/test_fft.py +17 -10
  79. nabu/processing/unsharp_cuda.py +1 -1
  80. nabu/reconstruction/cone.py +104 -40
  81. nabu/reconstruction/fbp.py +3 -0
  82. nabu/reconstruction/fbp_base.py +7 -2
  83. nabu/reconstruction/filtering.py +20 -7
  84. nabu/reconstruction/filtering_cuda.py +7 -1
  85. nabu/reconstruction/hbp.py +424 -0
  86. nabu/reconstruction/mlem.py +99 -0
  87. nabu/reconstruction/reconstructor.py +2 -0
  88. nabu/reconstruction/rings_cuda.py +19 -19
  89. nabu/reconstruction/sinogram_cuda.py +1 -0
  90. nabu/reconstruction/sinogram_opencl.py +3 -1
  91. nabu/reconstruction/tests/test_cone.py +10 -5
  92. nabu/reconstruction/tests/test_deringer.py +7 -6
  93. nabu/reconstruction/tests/test_fbp.py +124 -10
  94. nabu/reconstruction/tests/test_filtering.py +13 -11
  95. nabu/reconstruction/tests/test_halftomo.py +30 -4
  96. nabu/reconstruction/tests/test_mlem.py +91 -0
  97. nabu/reconstruction/tests/test_reconstructor.py +8 -3
  98. nabu/resources/dataset_analyzer.py +142 -92
  99. nabu/resources/gpu.py +1 -0
  100. nabu/resources/nxflatfield.py +134 -125
  101. nabu/resources/templates/id16a_fluo.conf +42 -0
  102. nabu/resources/tests/test_extract.py +10 -0
  103. nabu/resources/tests/test_nxflatfield.py +2 -2
  104. nabu/stitching/alignment.py +80 -24
  105. nabu/stitching/config.py +105 -68
  106. nabu/stitching/definitions.py +1 -0
  107. nabu/stitching/frame_composition.py +68 -60
  108. nabu/stitching/overlap.py +91 -51
  109. nabu/stitching/single_axis_stitching.py +32 -0
  110. nabu/stitching/slurm_utils.py +6 -6
  111. nabu/stitching/stitcher/__init__.py +0 -0
  112. nabu/stitching/stitcher/base.py +124 -0
  113. nabu/stitching/stitcher/dumper/__init__.py +3 -0
  114. nabu/stitching/stitcher/dumper/base.py +94 -0
  115. nabu/stitching/stitcher/dumper/postprocessing.py +356 -0
  116. nabu/stitching/stitcher/dumper/preprocessing.py +60 -0
  117. nabu/stitching/stitcher/post_processing.py +555 -0
  118. nabu/stitching/stitcher/pre_processing.py +1068 -0
  119. nabu/stitching/stitcher/single_axis.py +484 -0
  120. nabu/stitching/stitcher/stitcher.py +0 -0
  121. nabu/stitching/stitcher/y_stitcher.py +13 -0
  122. nabu/stitching/stitcher/z_stitcher.py +45 -0
  123. nabu/stitching/stitcher_2D.py +278 -0
  124. nabu/stitching/tests/test_config.py +12 -37
  125. nabu/stitching/tests/test_frame_composition.py +33 -59
  126. nabu/stitching/tests/test_overlap.py +149 -7
  127. nabu/stitching/tests/test_utils.py +1 -1
  128. nabu/stitching/tests/test_y_preprocessing_stitching.py +132 -0
  129. nabu/stitching/tests/{test_z_stitching.py → test_z_postprocessing_stitching.py} +167 -561
  130. nabu/stitching/tests/test_z_preprocessing_stitching.py +431 -0
  131. nabu/stitching/utils/__init__.py +1 -0
  132. nabu/stitching/utils/post_processing.py +281 -0
  133. nabu/stitching/utils/tests/test_post-processing.py +21 -0
  134. nabu/stitching/{utils.py → utils/utils.py} +79 -52
  135. nabu/stitching/y_stitching.py +27 -0
  136. nabu/stitching/z_stitching.py +32 -2281
  137. nabu/testutils.py +1 -152
  138. nabu/thirdparty/tomocupy_remove_stripe.py +43 -9
  139. nabu/utils.py +158 -61
  140. {nabu-2024.1.10.dist-info → nabu-2024.2.0.dist-info}/METADATA +24 -17
  141. {nabu-2024.1.10.dist-info → nabu-2024.2.0.dist-info}/RECORD +145 -121
  142. {nabu-2024.1.10.dist-info → nabu-2024.2.0.dist-info}/WHEEL +1 -1
  143. nabu/io/tiffwriter_zmm.py +0 -99
  144. nabu/pipeline/fallback_utils.py +0 -149
  145. nabu/pipeline/helical/tests/test_accumulator.py +0 -158
  146. nabu/pipeline/helical/tests/test_pipeline_elements_full.py +0 -355
  147. nabu/pipeline/helical/tests/test_strategy.py +0 -61
  148. nabu/pipeline/helical/utils.py +0 -51
  149. nabu/pipeline/tests/test_chunk_reader.py +0 -74
  150. {nabu-2024.1.10.dist-info → nabu-2024.2.0.dist-info}/LICENSE +0 -0
  151. {nabu-2024.1.10.dist-info → nabu-2024.2.0.dist-info}/entry_points.txt +0 -0
  152. {nabu-2024.1.10.dist-info → nabu-2024.2.0.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,94 @@
1
+ import h5py
2
+ import numpy
3
+ from typing import Union, Optional
4
+ from tomoscan.identifier import BaseIdentifier
5
+ from nabu.stitching.config import StitchingConfiguration
6
+ from tomoscan.volumebase import VolumeBase
7
+ from contextlib import AbstractContextManager
8
+
9
+
10
+ class DumperBase:
11
+ """
12
+ Base class to define all the functions that can be used to save a stitching
13
+ """
14
+
15
+ def __init__(self, configuration) -> None:
16
+ assert isinstance(configuration, StitchingConfiguration)
17
+ self._configuration = configuration
18
+
19
+ @property
20
+ def configuration(self):
21
+ return self._configuration
22
+
23
+ @property
24
+ def output_identifier(self) -> BaseIdentifier:
25
+ raise NotImplementedError("Base class")
26
+
27
+ def save_stitched_frame(
28
+ self,
29
+ stitched_frame: numpy.ndarray,
30
+ i_frame: int,
31
+ axis: int,
32
+ **kwargs,
33
+ ):
34
+ self.save_frame_to_disk(
35
+ output_dataset=self.output_dataset,
36
+ index=i_frame,
37
+ stitched_frame=stitched_frame,
38
+ axis=axis,
39
+ region_start=0,
40
+ region_end=None,
41
+ )
42
+
43
+ @property
44
+ def output_dataset(self) -> Optional[Union[h5py.VirtualLayout, h5py.Dataset, VolumeBase]]:
45
+ return self._output_dataset
46
+
47
+ @output_dataset.setter
48
+ def output_dataset(self, dataset: Optional[Union[h5py.VirtualLayout, h5py.Dataset, VolumeBase]]):
49
+ self._output_dataset = dataset
50
+
51
+ @staticmethod
52
+ def save_frame_to_disk(
53
+ output_dataset: Union[h5py.Dataset, h5py.VirtualLayout],
54
+ index: int,
55
+ stitched_frame: Union[numpy.ndarray, h5py.VirtualSource],
56
+ axis: int,
57
+ region_start: int,
58
+ region_end: int,
59
+ ):
60
+ if not isinstance(output_dataset, (h5py.VirtualLayout, h5py.Dataset, numpy.ndarray)):
61
+ raise TypeError(
62
+ f"'output_dataset' should be a 'h5py.Dataset' or a 'h5py.VirtualLayout'. Get {type(output_dataset)}"
63
+ )
64
+ if not isinstance(stitched_frame, (h5py.VirtualSource, numpy.ndarray)):
65
+ raise TypeError(
66
+ f"'stitched_frame' should be a 'numpy.ndarray' or a 'h5py.VirtualSource'. Get {type(stitched_frame)}"
67
+ )
68
+ if isinstance(output_dataset, h5py.VirtualLayout) and not isinstance(stitched_frame, h5py.VirtualSource):
69
+ raise TypeError(
70
+ "output_dataset is an instance of h5py.VirtualLayout and stitched_frame not an instance of h5py.VirtualSource"
71
+ )
72
+ if axis == 0:
73
+ if region_end is not None:
74
+ output_dataset[index, region_start:region_end] = stitched_frame
75
+ else:
76
+ output_dataset[index, region_start:] = stitched_frame
77
+ elif axis == 1:
78
+ if region_end is not None:
79
+ output_dataset[region_start:region_end, index, :] = stitched_frame
80
+ else:
81
+ output_dataset[region_start:, index, :] = stitched_frame
82
+ elif axis == 2:
83
+ if region_end is not None:
84
+ output_dataset[region_start:region_end, :, index] = stitched_frame
85
+ else:
86
+ output_dataset[region_start:, :, index] = stitched_frame
87
+ else:
88
+ raise ValueError(f"provided axis ({axis}) is invalid")
89
+
90
+ def create_output_dataset(self):
91
+ """
92
+ function called at the beginning of the stitching to prepare output dataset
93
+ """
94
+ raise NotImplementedError
@@ -0,0 +1,356 @@
1
+ import h5py
2
+ import numpy
3
+ import logging
4
+ from typing import Optional
5
+ from .base import DumperBase
6
+ from nabu.stitching.config import PostProcessedSingleAxisStitchingConfiguration
7
+ from nabu import version as nabu_version
8
+ from nabu.io.writer import get_datetime
9
+ from tomoscan.identifier import VolumeIdentifier
10
+ from tomoscan.volumebase import VolumeBase
11
+ from tomoscan.esrf.volume import HDF5Volume
12
+ from tomoscan.io import HDF5File
13
+ from contextlib import AbstractContextManager
14
+
15
+
16
+ _logger = logging.getLogger(__name__)
17
+
18
+
19
+ class OutputVolumeContext(AbstractContextManager):
20
+ """
21
+ Utils class to Manage the data volume creation and save it (data only !). target: used for volume stitching
22
+ In the case of HDF5 we want to save this directly in the file to avoid
23
+ keeping the full volume in memory.
24
+ Insure also contain processing will be common between the different processing
25
+
26
+ If stitching_sources_arr_shapes is provided this mean that we want to create stitching region and then create a VDS to avoid data duplication
27
+ """
28
+
29
+ def __init__(
30
+ self,
31
+ volume: VolumeBase,
32
+ volume_shape: tuple,
33
+ dtype: numpy.dtype,
34
+ dumper,
35
+ ) -> None:
36
+ super().__init__()
37
+ if not isinstance(volume, VolumeBase):
38
+ raise TypeError(f"Volume is expected to be an instance of {VolumeBase}. {type(volume)} provided instead")
39
+
40
+ self._volume = volume
41
+ self._volume_shape = volume_shape
42
+ self.__file_handler = None
43
+ self._dtype = dtype
44
+ self._dumper = dumper
45
+
46
+ @property
47
+ def _file_handler(self):
48
+ return self.__file_handler
49
+
50
+ def _build_hdf5_output(self):
51
+ return self._file_handler.create_dataset(
52
+ self._volume.data_url.data_path(),
53
+ shape=self._volume_shape,
54
+ dtype=self._dtype,
55
+ )
56
+
57
+ def _create_stitched_volume_dataset(self):
58
+ # handle the specific case of HDF5. Goal: avoid getting the full stitched volume in memory
59
+ if isinstance(self._volume, HDF5Volume):
60
+ self.__file_handler = HDF5File(self._volume.data_url.file_path(), mode="a")
61
+ # if need to delete an existing dataset
62
+ if self._volume.overwrite and self._volume.data_path in self._file_handler:
63
+ try:
64
+ del self._file_handler[self._volume.data_path]
65
+ except Exception as e:
66
+ _logger.error(f"Fail to overwrite data. Reason is {e}")
67
+ data = None
68
+ self._file_handler.close()
69
+ self._duplicate_data = True
70
+ # avoid creating a dataset for stitched volume as creation of the stitched_volume failed
71
+ return data
72
+
73
+ # create dataset
74
+ try:
75
+ data = self._build_hdf5_output()
76
+ except Exception as e2:
77
+ _logger.error(f"Fail to create final dataset. Reason is {e2}")
78
+ data = None
79
+ self._file_handler.close()
80
+ self._duplicate_data = True
81
+ # avoid creating a dataset for stitched volume as creation of the stitched_volume failed
82
+ else:
83
+ raise TypeError("only HDF5 output is handled")
84
+ # else:
85
+ # # for other file format: create the full dataset in memory before dumping it
86
+ # data = numpy.empty(self._volume_shape, dtype=self._dtype)
87
+ # self._volume.data = data
88
+ return data
89
+
90
+ def __enter__(self):
91
+ assert self._dumper.output_dataset is None
92
+ self._dumper.output_dataset = self._create_stitched_volume_dataset()
93
+ return self._dumper.output_dataset
94
+
95
+ def __exit__(self, exc_type, exc_value, traceback):
96
+ if self._file_handler is not None:
97
+ return self._file_handler.close()
98
+ else:
99
+ self._volume.save_data()
100
+
101
+
102
+ class OutputVolumeNoDDContext(OutputVolumeContext):
103
+ """
104
+ Dedicated output volume context for saving a volume without Data Duplication (DD)
105
+ """
106
+
107
+ def __init__(
108
+ self,
109
+ volume: VolumeBase,
110
+ volume_shape: tuple,
111
+ dtype: numpy.dtype,
112
+ dumper,
113
+ stitching_sources_arr_shapes: Optional[tuple],
114
+ ) -> None:
115
+ if not isinstance(dumper, PostProcessingStitchingDumperNoDD):
116
+ raise TypeError
117
+ # TODO: compute volume_shape from here
118
+ self._stitching_sources_arr_shapes = stitching_sources_arr_shapes
119
+
120
+ super().__init__(volume, volume_shape, dtype, dumper)
121
+
122
+ def __enter__(self):
123
+ dataset = super().__enter__()
124
+ assert isinstance(self._dumper, PostProcessingStitchingDumperNoDD)
125
+ self._dumper.stitching_regions_hdf5_dataset = self._create_stitched_sub_region_datasets()
126
+ return dataset
127
+
128
+ def _build_hdf5_output(self):
129
+ return h5py.VirtualLayout(
130
+ shape=self._volume_shape,
131
+ dtype=self._dtype,
132
+ )
133
+
134
+ def __exit__(self, exc_type, exc_value, traceback):
135
+ # in the case of no data duplication we need to create the virtual dataset at the end
136
+ if not isinstance(self._dumper.output_dataset, h5py.VirtualLayout):
137
+ raise TypeError("dumper output_dataset should be a virtual layout")
138
+ self._file_handler.create_virtual_dataset(self._volume.data_url.data_path(), layout=self._dumper.output_dataset)
139
+ super().__exit__(exc_type=exc_type, exc_value=exc_value, traceback=traceback)
140
+
141
+ def _create_stitched_sub_region_datasets(self):
142
+ # create datasets to store overlaps if needed
143
+ if not isinstance(self._volume, HDF5Volume):
144
+ raise TypeError("Avoid Data Duplication is only available for HDF5 output volume")
145
+
146
+ stitching_regions_hdf5_dataset = []
147
+ for i_region, overlap_shape in enumerate(self._stitching_sources_arr_shapes):
148
+ data_path = f"{self._volume.data_path}/stitching_regions/region_{i_region}"
149
+ if self._volume.overwrite and data_path in self._file_handler:
150
+ del self._file_handler[data_path]
151
+ stitching_regions_hdf5_dataset.append(
152
+ self._file_handler.create_dataset(
153
+ name=data_path,
154
+ shape=overlap_shape,
155
+ dtype=self._dtype,
156
+ )
157
+ )
158
+ self._dumper.stitching_regions_hdf5_dataset = stitching_regions_hdf5_dataset
159
+ return stitching_regions_hdf5_dataset
160
+
161
+
162
+ class PostProcessingStitchingDumper(DumperBase):
163
+ """
164
+ dumper to be used when save data durint post-processing stitching (on recosntructed volume). Output is expected to be an NXtomo
165
+ """
166
+
167
+ OutputDatasetContext = OutputVolumeContext
168
+
169
+ def __init__(self, configuration) -> None:
170
+ if not isinstance(configuration, PostProcessedSingleAxisStitchingConfiguration):
171
+ raise TypeError(
172
+ f"configuration is expected to be an instance of {PostProcessedSingleAxisStitchingConfiguration}. Get {type(configuration)} instead"
173
+ )
174
+ super().__init__(configuration)
175
+ self._output_dataset = None
176
+ self._input_volumes = configuration.input_volumes
177
+
178
+ def save_configuration(self):
179
+ voxel_size = self._input_volumes[0].voxel_size
180
+
181
+ def get_position():
182
+ # the z-serie is z-ordered from higher to lower. We can reuse this with pixel size and shape to
183
+ # compute the position of the stitched volume
184
+ if voxel_size is None:
185
+ return None
186
+ return numpy.array(self._input_volumes[0].position) + voxel_size * (
187
+ numpy.array(self._input_volumes[0].get_volume_shape()) / 2.0
188
+ - numpy.array(self.configuration.output_volume.get_volume_shape()) / 2.0
189
+ )
190
+
191
+ self.configuration.output_volume.voxel_size = voxel_size or ""
192
+ try:
193
+ self.configuration.output_volume.position = get_position()
194
+ except Exception:
195
+ self.configuration.output_volume.position = numpy.array([0, 0, 0])
196
+
197
+ self.configuration.output_volume.metadata.update(
198
+ {
199
+ "about": {
200
+ "program": "nabu-stitching",
201
+ "version": nabu_version,
202
+ "date": get_datetime(),
203
+ },
204
+ "configuration": self.configuration.to_dict(),
205
+ }
206
+ )
207
+ self.configuration.output_volume.save_metadata()
208
+
209
+ @property
210
+ def output_identifier(self) -> VolumeIdentifier:
211
+ return self.configuration.output_volume.get_identifier()
212
+
213
+ def create_output_dataset(self):
214
+ """
215
+ function called at the beginning of the stitching to prepare output dataset
216
+ """
217
+ self._dataset = h5py.VirtualLayout(
218
+ shape=self._volume_shape,
219
+ dtype=self._dtype,
220
+ )
221
+
222
+
223
+ class PostProcessingStitchingDumperNoDD(PostProcessingStitchingDumper):
224
+ """
225
+ same as PostProcessingStitchingDumper but prevent to do data duplication.
226
+ In this case we need to work on HDF5 file only
227
+ """
228
+
229
+ OutputDatasetContext = OutputVolumeNoDDContext
230
+
231
+ def __init__(self, configuration) -> None:
232
+ if not isinstance(configuration, PostProcessedSingleAxisStitchingConfiguration):
233
+ raise TypeError(
234
+ f"configuration is expected to be an instance of {PostProcessedSingleAxisStitchingConfiguration}. Get {type(configuration)} instead"
235
+ )
236
+ super().__init__(configuration)
237
+ self._stitching_regions_hdf5_dataset = None
238
+ self._raw_regions_hdf5_dataset = None
239
+
240
+ def create_output_dataset(self):
241
+ """
242
+ function called at the beginning of the stitching to prepare output dataset
243
+ """
244
+ self._dataset = h5py.VirtualLayout(
245
+ shape=self._volume_shape,
246
+ dtype=self._dtype,
247
+ )
248
+
249
+ @staticmethod
250
+ def create_subset_selection(dataset: h5py.Dataset, slices: tuple) -> h5py.VirtualSource:
251
+ assert isinstance(dataset, h5py.Dataset), f"dataset is expected to be a h5py.Dataset. Get {type(dataset)}"
252
+ assert isinstance(slices, tuple), f"slices is expected to be a tuple of slices. Get {type(slices)} instead"
253
+ import h5py._hl.selections as selection
254
+
255
+ virtual_source = h5py.VirtualSource(dataset)
256
+ sel = selection.select(dataset.shape, slices, dataset=dataset)
257
+ virtual_source.sel = sel
258
+ return virtual_source
259
+
260
+ @PostProcessingStitchingDumper.output_dataset.setter
261
+ def output_dataset(self, dataset: Optional[h5py.VirtualLayout]):
262
+ if dataset is not None and not isinstance(dataset, h5py.VirtualLayout):
263
+ raise TypeError("in the case we want to avoid data duplication 'output_dataset' must be a VirtualLayout")
264
+ self._output_dataset = dataset
265
+
266
+ @property
267
+ def stitching_regions_hdf5_dataset(self) -> Optional[tuple]:
268
+ """hdf5 dataset storing the stitched regions"""
269
+ return self._stitching_regions_hdf5_dataset
270
+
271
+ @stitching_regions_hdf5_dataset.setter
272
+ def stitching_regions_hdf5_dataset(self, datasets: tuple):
273
+ self._stitching_regions_hdf5_dataset = datasets
274
+
275
+ @property
276
+ def raw_regions_hdf5_dataset(self) -> Optional[tuple]:
277
+ """hdf5 raw dataset"""
278
+ return self._raw_regions_hdf5_dataset
279
+
280
+ @raw_regions_hdf5_dataset.setter
281
+ def raw_regions_hdf5_dataset(self, datasets: tuple):
282
+ self._raw_regions_hdf5_dataset = datasets
283
+
284
+ def save_stitched_frame(
285
+ self,
286
+ stitched_frame: numpy.ndarray,
287
+ composition_cls: dict,
288
+ i_frame: int,
289
+ axis: int,
290
+ ):
291
+ """
292
+ Save the full stitched frame to disk
293
+ """
294
+ output_dataset = self.output_dataset
295
+ if output_dataset is None:
296
+ raise ValueError("output_dataset must be set before calling any frame stitching")
297
+ stitching_regions_hdf5_dataset = self.stitching_regions_hdf5_dataset
298
+ if stitching_regions_hdf5_dataset is None:
299
+ raise ValueError("stitching_region_hdf5_dataset must be set before calling any frame stitching")
300
+ raw_regions_hdf5_dataset = self.raw_regions_hdf5_dataset
301
+
302
+ # save stitched region
303
+ stitching_regions = composition_cls["overlap_composition"]
304
+ for (_, _, region_start, region_end), stitching_region_hdf5_dataset in zip(
305
+ stitching_regions.browse(), stitching_regions_hdf5_dataset
306
+ ):
307
+ assert isinstance(output_dataset, h5py.VirtualLayout)
308
+ assert isinstance(stitching_region_hdf5_dataset, h5py.Dataset)
309
+ stitching_region_array = stitched_frame[region_start:region_end]
310
+ self.save_frame_to_disk(
311
+ output_dataset=stitching_region_hdf5_dataset,
312
+ index=i_frame,
313
+ stitched_frame=stitching_region_array,
314
+ axis=1,
315
+ region_start=0,
316
+ region_end=None,
317
+ )
318
+ vs = self.create_subset_selection(
319
+ dataset=stitching_region_hdf5_dataset,
320
+ slices=(
321
+ slice(0, stitching_region_hdf5_dataset.shape[0]),
322
+ slice(i_frame, i_frame + 1),
323
+ slice(0, stitching_region_hdf5_dataset.shape[2]),
324
+ ),
325
+ )
326
+
327
+ self.save_frame_to_disk(
328
+ output_dataset=output_dataset,
329
+ index=i_frame,
330
+ axis=axis,
331
+ region_start=region_start,
332
+ region_end=region_end,
333
+ stitched_frame=vs,
334
+ )
335
+
336
+ # create virtual source of the raw data
337
+ raw_regions = composition_cls["raw_composition"]
338
+ for (frame_start, frame_end, region_start, region_end), raw_region_hdf5_dataset in zip(
339
+ raw_regions.browse(), raw_regions_hdf5_dataset
340
+ ):
341
+ vs = self.create_subset_selection(
342
+ dataset=raw_region_hdf5_dataset,
343
+ slices=(
344
+ slice(frame_start, frame_end),
345
+ slice(i_frame, i_frame + 1),
346
+ slice(0, raw_region_hdf5_dataset.shape[2]),
347
+ ),
348
+ )
349
+ self.save_frame_to_disk(
350
+ output_dataset=output_dataset,
351
+ index=i_frame,
352
+ axis=1,
353
+ region_start=region_start,
354
+ region_end=region_end,
355
+ stitched_frame=vs,
356
+ )
@@ -0,0 +1,60 @@
1
+ import h5py
2
+ import numpy
3
+ import logging
4
+ from .base import DumperBase
5
+ from nabu.stitching.config import PreProcessedSingleAxisStitchingConfiguration
6
+ from nabu import version as nabu_version
7
+ from nabu.io.writer import get_datetime
8
+ from silx.io.dictdump import dicttonx
9
+ from tomoscan.identifier import ScanIdentifier
10
+
11
+
12
+ _logger = logging.getLogger(__name__)
13
+
14
+
15
+ class PreProcessingStitchingDumper(DumperBase):
16
+ """
17
+ dumper to be used when save data durint pre-processing stitching (on projections). Output is expected to be an NXtomo
18
+ """
19
+
20
+ def __init__(self, configuration) -> None:
21
+ if not isinstance(configuration, PreProcessedSingleAxisStitchingConfiguration):
22
+ raise TypeError(
23
+ f"configuration is expected to be an instance of {PreProcessedSingleAxisStitchingConfiguration}. Get {type(configuration)} instead"
24
+ )
25
+ super().__init__(configuration)
26
+
27
+ def save_frame_to_disk(self, output_dataset: h5py.Dataset, index: int, stitched_frame: numpy.ndarray, **kwargs):
28
+ output_dataset[index] = stitched_frame
29
+
30
+ def save_configuration(self):
31
+ """dump configuration used for stitching at the NXtomo entry"""
32
+ process_name = "stitching_configuration"
33
+ config_dict = self.configuration.to_dict()
34
+ # adding nabu specific information
35
+ nabu_process_info = {
36
+ "@NX_class": "NXentry",
37
+ f"{process_name}@NX_class": "NXprocess",
38
+ f"{process_name}/program": "nabu-stitching",
39
+ f"{process_name}/version": nabu_version,
40
+ f"{process_name}/date": get_datetime(),
41
+ f"{process_name}/configuration": config_dict,
42
+ }
43
+
44
+ dicttonx(
45
+ nabu_process_info,
46
+ h5file=self.configuration.output_file_path,
47
+ h5path=self.configuration.output_data_path,
48
+ update_mode="replace",
49
+ mode="a",
50
+ )
51
+
52
+ @property
53
+ def output_identifier(self) -> ScanIdentifier:
54
+ return self.configuration.get_output_object().get_identifier()
55
+
56
+ def create_output_dataset(self):
57
+ """
58
+ function called at the beginning of the stitching to prepare output dataset
59
+ """
60
+ raise NotImplementedError