nabu 2022.3.0a1__py3-none-any.whl → 2023.1.0a2__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 (96) hide show
  1. nabu/__init__.py +1 -1
  2. nabu/app/bootstrap.py +7 -1
  3. nabu/app/cast_volume.py +8 -2
  4. nabu/app/cli_configs.py +69 -0
  5. nabu/app/composite_cor.py +97 -0
  6. nabu/app/create_distortion_map_from_poly.py +118 -0
  7. nabu/app/nx_z_splitter.py +1 -1
  8. nabu/app/prepare_weights_double.py +21 -16
  9. nabu/app/reconstruct_helical.py +0 -1
  10. nabu/app/utils.py +10 -5
  11. nabu/cuda/processing.py +1 -0
  12. nabu/cuda/tests/test_padding.py +1 -0
  13. nabu/cuda/utils.py +1 -0
  14. nabu/distributed/__init__.py +0 -0
  15. nabu/distributed/utils.py +57 -0
  16. nabu/distributed/worker.py +543 -0
  17. nabu/estimation/cor.py +3 -7
  18. nabu/estimation/cor_sino.py +2 -1
  19. nabu/estimation/distortion.py +6 -4
  20. nabu/io/cast_volume.py +10 -1
  21. nabu/io/detector_distortion.py +305 -0
  22. nabu/io/reader.py +37 -7
  23. nabu/io/reader_helical.py +0 -3
  24. nabu/io/tests/test_cast_volume.py +16 -4
  25. nabu/io/tests/test_detector_distortion.py +178 -0
  26. nabu/io/tests/test_writers.py +2 -2
  27. nabu/io/tiffwriter_zmm.py +2 -3
  28. nabu/io/writer.py +84 -1
  29. nabu/io/writer_BACKUP_193259.py +556 -0
  30. nabu/io/writer_BACKUP_193381.py +556 -0
  31. nabu/io/writer_BASE_193259.py +548 -0
  32. nabu/io/writer_BASE_193381.py +548 -0
  33. nabu/io/writer_LOCAL_193259.py +550 -0
  34. nabu/io/writer_LOCAL_193381.py +550 -0
  35. nabu/io/writer_REMOTE_193259.py +557 -0
  36. nabu/io/writer_REMOTE_193381.py +557 -0
  37. nabu/misc/fourier_filters.py +2 -0
  38. nabu/misc/rotation.py +0 -1
  39. nabu/misc/tests/test_rotation.py +1 -0
  40. nabu/pipeline/config_validators.py +10 -0
  41. nabu/pipeline/datadump.py +1 -1
  42. nabu/pipeline/dataset_validator.py +0 -1
  43. nabu/pipeline/detector_distortion_provider.py +20 -0
  44. nabu/pipeline/estimators.py +35 -21
  45. nabu/pipeline/fallback_utils.py +1 -1
  46. nabu/pipeline/fullfield/chunked.py +30 -15
  47. nabu/pipeline/fullfield/chunked_black.py +881 -0
  48. nabu/pipeline/fullfield/chunked_cuda.py +34 -4
  49. nabu/pipeline/fullfield/chunked_fb.py +966 -0
  50. nabu/pipeline/fullfield/chunked_google.py +921 -0
  51. nabu/pipeline/fullfield/chunked_pep8.py +920 -0
  52. nabu/pipeline/fullfield/computations.py +7 -6
  53. nabu/pipeline/fullfield/dataset_validator.py +1 -1
  54. nabu/pipeline/fullfield/grouped_cuda.py +6 -0
  55. nabu/pipeline/fullfield/nabu_config.py +15 -3
  56. nabu/pipeline/fullfield/processconfig.py +5 -0
  57. nabu/pipeline/fullfield/reconstruction.py +1 -2
  58. nabu/pipeline/helical/gridded_accumulator.py +1 -8
  59. nabu/pipeline/helical/helical_chunked_regridded.py +48 -33
  60. nabu/pipeline/helical/helical_reconstruction.py +1 -9
  61. nabu/pipeline/helical/nabu_config.py +11 -14
  62. nabu/pipeline/helical/span_strategy.py +11 -4
  63. nabu/pipeline/helical/tests/test_accumulator.py +0 -3
  64. nabu/pipeline/helical/tests/test_pipeline_elements_full.py +0 -6
  65. nabu/pipeline/helical/tests/test_strategy.py +0 -1
  66. nabu/pipeline/helical/weight_balancer.py +0 -1
  67. nabu/pipeline/params.py +4 -0
  68. nabu/pipeline/processconfig.py +6 -2
  69. nabu/pipeline/writer.py +9 -4
  70. nabu/preproc/distortion.py +4 -3
  71. nabu/preproc/double_flatfield.py +16 -4
  72. nabu/preproc/double_flatfield_cuda.py +3 -2
  73. nabu/preproc/double_flatfield_variable_region.py +13 -4
  74. nabu/preproc/flatfield.py +29 -7
  75. nabu/preproc/flatfield_cuda.py +0 -1
  76. nabu/preproc/flatfield_variable_region.py +5 -2
  77. nabu/preproc/phase.py +0 -1
  78. nabu/preproc/phase_cuda.py +0 -1
  79. nabu/preproc/tests/test_ctf.py +4 -3
  80. nabu/preproc/tests/test_flatfield.py +6 -7
  81. nabu/reconstruction/fbp_opencl.py +1 -1
  82. nabu/reconstruction/filtering.py +0 -1
  83. nabu/reconstruction/tests/test_fbp.py +1 -0
  84. nabu/resources/dataset_analyzer.py +0 -1
  85. nabu/resources/templates/bm05_pag.conf +34 -0
  86. nabu/resources/templates/id16_ctf.conf +2 -1
  87. nabu/resources/tests/test_nxflatfield.py +0 -1
  88. nabu/resources/tests/test_units.py +0 -1
  89. nabu/stitching/frame_composition.py +7 -1
  90. {nabu-2022.3.0a1.dist-info → nabu-2023.1.0a2.dist-info}/METADATA +2 -7
  91. {nabu-2022.3.0a1.dist-info → nabu-2023.1.0a2.dist-info}/RECORD +96 -75
  92. {nabu-2022.3.0a1.dist-info → nabu-2023.1.0a2.dist-info}/WHEEL +1 -1
  93. {nabu-2022.3.0a1.dist-info → nabu-2023.1.0a2.dist-info}/entry_points.txt +2 -1
  94. {nabu-2022.3.0a1.dist-info → nabu-2023.1.0a2.dist-info}/LICENSE +0 -0
  95. {nabu-2022.3.0a1.dist-info → nabu-2023.1.0a2.dist-info}/top_level.txt +0 -0
  96. {nabu-2022.3.0a1.dist-info → nabu-2023.1.0a2.dist-info}/zip-safe +0 -0
@@ -0,0 +1,881 @@
1
+ from os import path
2
+ from time import time
3
+ import numpy as np
4
+ from silx.io.url import DataUrl
5
+ from ...resources.logger import LoggerOrPrint
6
+ from ...resources.utils import is_hdf5_extension, extract_parameters
7
+ from ...io.reader import ChunkReader, HDF5Loader, get_hdf5_dataset_shape
8
+ from ...preproc.ccd import Log, CCDFilter
9
+ from ...preproc.flatfield import FlatFieldDataUrls
10
+ from ...preproc.distortion import DistortionCorrection
11
+ from ...preproc.shift import VerticalShift
12
+ from ...preproc.double_flatfield import DoubleFlatField
13
+ from ...preproc.phase import PaganinPhaseRetrieval
14
+ from ...reconstruction.sinogram import SinoBuilder, SinoNormalization
15
+ from ...misc.rotation import Rotation
16
+ from ...preproc.rings import MunchDeringer
17
+ from ...misc.unsharp import UnsharpMask
18
+ from ...misc.histogram import PartialHistogram, hist_as_2Darray
19
+ from ..utils import use_options, pipeline_step, WriterConfigurator
20
+
21
+ # For now we don't have a plain python/numpy backend for reconstruction
22
+ try:
23
+ from ...reconstruction.fbp_opencl import Backprojector
24
+ except:
25
+ Backprojector = None
26
+
27
+
28
+ class ChunkedPipeline:
29
+ """
30
+ Pipeline for "regular" full-field tomography.
31
+ Data is processed by chunks. A chunk consists in K contiguous lines of all the radios.
32
+ In parallel geometry, a chunk of K radios lines gives K sinograms,
33
+ and equivalently K reconstructed slices.
34
+ """
35
+
36
+ backend = "numpy"
37
+ FlatFieldClass = FlatFieldDataUrls
38
+ DoubleFlatFieldClass = DoubleFlatField
39
+ CCDCorrectionClass = CCDFilter
40
+ PaganinPhaseRetrievalClass = PaganinPhaseRetrieval
41
+ UnsharpMaskClass = UnsharpMask
42
+ ImageRotationClass = Rotation
43
+ VerticalShiftClass = VerticalShift
44
+ SinoBuilderClass = SinoBuilder
45
+ SinoDeringerClass = MunchDeringer
46
+ MLogClass = Log
47
+ SinoNormalizationClass = SinoNormalization
48
+ FBPClass = Backprojector
49
+ HistogramClass = PartialHistogram
50
+
51
+ def __init__(self, process_config, sub_region, logger=None, extra_options=None, phase_margin=None):
52
+ """
53
+ Initialize a "Chunked" pipeline.
54
+
55
+ Parameters
56
+ ----------
57
+ processing_config: `ProcessConfig`
58
+ Process configuration.
59
+ sub_region: tuple
60
+ Sub-region to process in the volume for this worker, in the format
61
+ `(start_x, end_x, start_z, end_z)`.
62
+ logger: `nabu.app.logger.Logger`, optional
63
+ Logger class
64
+ extra_options: dict, optional
65
+ Advanced extra options.
66
+ phase_margin: tuple, optional
67
+ Margin to use when performing phase retrieval, in the form ((up, down), (left, right)).
68
+ See also the documentation of PaganinPhaseRetrieval.
69
+ If not provided, no margin is applied.
70
+
71
+
72
+ Notes
73
+ ------
74
+ Using a `phase_margin` results in a lesser number of reconstructed slices.
75
+ More specifically, if `phase_margin = (V, H)`, then there will be `delta_z - 2*V`
76
+ reconstructed slices (if the sub-region is in the middle of the volume)
77
+ or `delta_z - V` reconstructed slices (if the sub-region is on top or bottom
78
+ of the volume).
79
+ """
80
+ self.logger = LoggerOrPrint(logger)
81
+ self._set_params(process_config, sub_region, extra_options, phase_margin)
82
+ self.set_subregion(sub_region)
83
+ self._init_pipeline()
84
+
85
+ def _set_params(self, process_config, sub_region, extra_options, phase_margin):
86
+ self.process_config = process_config
87
+ self.dataset_info = self.process_config.dataset_info
88
+ self.dataset_infos = self.process_config.dataset_info # shorthand - deprecated
89
+ self.processing_steps = self.process_config.processing_steps.copy()
90
+ self.processing_options = self.process_config.processing_options
91
+ self.sub_region = self._check_subregion(sub_region)
92
+ self.delta_z = sub_region[-1] - sub_region[-2]
93
+ self.chunk_size = self.delta_z
94
+ self._set_phase_margin(phase_margin)
95
+ self._set_extra_options(extra_options)
96
+ self._callbacks = {}
97
+ self._steps_name2component = {}
98
+ self._steps_component2name = {}
99
+ self._data_dump = {}
100
+ self._resume_from_step = None
101
+
102
+ @staticmethod
103
+ def _check_subregion(sub_region):
104
+ if len(sub_region) < 4:
105
+ assert len(sub_region) == 2
106
+ sub_region = (None, None) + sub_region
107
+ if None in sub_region[-2:]:
108
+ raise ValueError("Cannot set z_min or z_max to None")
109
+ return sub_region
110
+
111
+ def _set_extra_options(self, extra_options):
112
+ if extra_options is None:
113
+ extra_options = {}
114
+ advanced_options = {}
115
+ advanced_options.update(extra_options)
116
+ self.extra_options = advanced_options
117
+
118
+ def _set_phase_margin(self, phase_margin):
119
+ if phase_margin is None:
120
+ phase_margin = ((0, 0), (0, 0))
121
+ self._phase_margin_up = phase_margin[0][0]
122
+ self._phase_margin_down = phase_margin[0][1]
123
+ self._phase_margin_left = phase_margin[1][0]
124
+ self._phase_margin_right = phase_margin[1][1]
125
+
126
+ def set_subregion(self, sub_region):
127
+ """
128
+ Set a sub-region to process.
129
+
130
+ Parameters
131
+ ----------
132
+ sub_region: tuple
133
+ Sub-region to process in the volume, in the format
134
+ `(start_x, end_x, start_z, end_z)` or `(start_z, end_z)`.
135
+ """
136
+ sub_region = self._check_subregion(sub_region)
137
+ dz = sub_region[-1] - sub_region[-2]
138
+ if dz != self.delta_z:
139
+ raise ValueError(
140
+ "Class was initialized for delta_z = %d but provided sub_region has delta_z = %d" % (self.delta_z, dz)
141
+ )
142
+ self.sub_region = sub_region
143
+ self.z_min = sub_region[-2]
144
+ self.z_max = sub_region[-1]
145
+
146
+ def _compute_phase_kernel_margin(self):
147
+ """
148
+ Get the "margin" to pass to classes like PaganinPhaseRetrieval.
149
+ In order to have a good accuracy for filter-based phase retrieval methods,
150
+ we need to load extra data around the edges of each image. Otherwise,
151
+ a default padding type is applied.
152
+ """
153
+ if not (self.use_radio_processing_margin):
154
+ self._phase_margin = None
155
+ return
156
+ up_margin = self._phase_margin_up
157
+ down_margin = self._phase_margin_down
158
+ # Horizontal margin is not implemented
159
+ left_margin, right_margin = (0, 0)
160
+ self._phase_margin = ((up_margin, down_margin), (left_margin, right_margin))
161
+
162
+ @property
163
+ def use_radio_processing_margin(self):
164
+ return ("phase" in self.processing_steps) or ("unsharp_mask" in self.processing_steps)
165
+
166
+ def _get_phase_margin(self):
167
+ if not (self.use_radio_processing_margin):
168
+ return ((0, 0), (0, 0))
169
+ return self._phase_margin
170
+
171
+ def _get_cropped_radios(self):
172
+ ((up_margin, down_margin), (left_margin, right_margin)) = self._phase_margin
173
+ zslice = slice(up_margin or None, -down_margin or None)
174
+ xslice = slice(left_margin or None, -right_margin or None)
175
+ self._radios_cropped = self.radios[:, zslice, xslice]
176
+ return self._radios_cropped
177
+
178
+ @property
179
+ def phase_margin(self):
180
+ """
181
+ Return the margin for phase retrieval in the form ((up, down), (left, right))
182
+ """
183
+ return self._get_phase_margin()
184
+
185
+ @property
186
+ def n_recs(self):
187
+ """
188
+ Return the final number of reconstructed slices.
189
+ """
190
+ n_recs = self.delta_z
191
+ n_recs -= sum(self._get_phase_margin()[0])
192
+ return n_recs
193
+
194
+ def _get_process_name(self, kind="reconstruction"):
195
+ # In the future, might be something like "reconstruction-<ID>"
196
+ if kind == "reconstruction":
197
+ return "reconstruction"
198
+ elif kind == "histogram":
199
+ return "histogram"
200
+ return kind
201
+
202
+ def _configure_dump(self, step_name):
203
+ if step_name not in self.processing_steps:
204
+ if step_name == "sinogram" and self.process_config._dump_sinogram:
205
+ fname_full = self.process_config._dump_sinogram_file
206
+ else:
207
+ return
208
+ else:
209
+ if not self.processing_options[step_name].get("save", False):
210
+ return
211
+ fname_full = self.processing_options[step_name]["save_steps_file"]
212
+
213
+ fname, ext = path.splitext(fname_full)
214
+ dirname, file_prefix = path.split(fname)
215
+ output_dir = path.join(dirname, file_prefix)
216
+ file_prefix += str("_%04d" % self._get_image_start_index())
217
+
218
+ self._data_dump[step_name] = WriterConfigurator(
219
+ output_dir,
220
+ file_prefix,
221
+ file_format="hdf5",
222
+ overwrite=True,
223
+ logger=self.logger,
224
+ nx_info={
225
+ "process_name": step_name,
226
+ "processing_index": 0, # TODO
227
+ "config": {
228
+ "processing_options": self.processing_options,
229
+ "nabu_config": self.process_config.nabu_config,
230
+ },
231
+ "entry": getattr(self.dataset_info.dataset_scanner, "entry", None),
232
+ },
233
+ )
234
+
235
+ def _configure_data_dumps(self):
236
+ for step_name in self.processing_steps:
237
+ self._configure_dump(step_name)
238
+ # sinogram is a special keyword: not in processing_steps, but guaranteed to be before sinogram generation
239
+ if self.process_config._dump_sinogram:
240
+ self._configure_dump("sinogram")
241
+
242
+ #
243
+ # Callbacks
244
+ #
245
+
246
+ def register_callback(self, step_name, callback):
247
+ """
248
+ Register a callback for a pipeline processing step.
249
+
250
+ Parameters
251
+ ----------
252
+ step_name: str
253
+ processing step name
254
+ callback: callable
255
+ A function. It will be executed once the processing step `step_name`
256
+ is finished. The function takes only one argument: the class instance.
257
+ """
258
+ if step_name not in self.processing_steps:
259
+ raise ValueError("'%s' is not in processing steps %s" % (step_name, self.processing_steps))
260
+ if step_name in self._callbacks:
261
+ self._callbacks[step_name].append(callback)
262
+ else:
263
+ self._callbacks[step_name] = [callback]
264
+
265
+ def _reshape_radios_after_phase(self):
266
+ """
267
+ Callback executed after phase retrieval, if margin != (0, 0).
268
+ It modifies self.radios so that further processing will be done
269
+ on the "inner part".
270
+ """
271
+ if sum(self._get_phase_margin()[0]) <= 0:
272
+ return
273
+ self._orig_radios = self.radios
274
+ self.logger.debug("Reshaping radios from %s to %s" % (str(self.radios.shape), str(self._radios_cropped.shape)))
275
+ self.radios = self._radios_cropped
276
+
277
+ #
278
+ # Overwritten in inheriting classes
279
+ #
280
+
281
+ def _get_shape(self, step_name):
282
+ """
283
+ Get the shape to provide to the class corresponding to step_name.
284
+ """
285
+ if step_name == "flatfield":
286
+ shape = self.radios.shape
287
+ elif step_name == "double_flatfield":
288
+ shape = self.radios.shape
289
+ elif step_name == "rotate_projections":
290
+ shape = self.radios.shape[1:]
291
+ elif step_name == "phase":
292
+ shape = self.radios.shape[1:]
293
+ elif step_name == "ccd_correction":
294
+ shape = self.radios.shape[1:]
295
+ elif step_name == "unsharp_mask":
296
+ shape = self.radios.shape[1:]
297
+ elif step_name == "take_log":
298
+ shape = self._radios_cropped_shape
299
+ elif step_name == "radios_movements":
300
+ shape = self._radios_cropped_shape
301
+ elif step_name == "sino_normalization":
302
+ shape = self._radios_cropped_shape
303
+ elif step_name == "build_sino":
304
+ shape = self._radios_cropped_shape
305
+ elif step_name == "sino_rings_correction":
306
+ shape = self.sino_builder.output_shape
307
+ elif step_name == "reconstruction":
308
+ shape = self.sino_builder.output_shape[1:]
309
+ else:
310
+ raise ValueError("Unknown processing step %s" % step_name)
311
+ self.logger.debug("Data shape for %s is %s" % (step_name, str(shape)))
312
+ return shape
313
+
314
+ def _get_phase_output_shape(self):
315
+ if not (self.use_radio_processing_margin):
316
+ self._radios_cropped_shape = self.radios.shape
317
+ return
318
+ ((up_margin, down_margin), (left_margin, right_margin)) = self._phase_margin
319
+ self._radios_cropped_shape = (
320
+ self.radios.shape[0],
321
+ self.radios.shape[1] - (up_margin + down_margin),
322
+ self.radios.shape[2] - (left_margin + right_margin),
323
+ )
324
+
325
+ def _allocate_array(self, shape, dtype, name=None):
326
+ return np.zeros(shape, dtype=dtype)
327
+
328
+ def _allocate_sinobuilder_output(self):
329
+ return self._allocate_array(self.sino_builder.output_shape, "f", name="sinos")
330
+
331
+ def _allocate_recs(self, ny, nx):
332
+ self.n_slices = self.radios.shape[1] # TODO modify with vertical shifts
333
+ if self.use_radio_processing_margin:
334
+ self.n_slices -= sum(self.phase_margin[0])
335
+ self.recs = self._allocate_array((self.n_slices, ny, nx), "f", name="recs")
336
+
337
+ def _reset_memory(self):
338
+ pass
339
+
340
+ def _get_read_dump_subregion(self):
341
+ read_opts = self.processing_options["read_chunk"]
342
+ if read_opts.get("process_file", None) is None:
343
+ return None
344
+ dump_start_z, dump_end_z = read_opts["dump_start_z"], read_opts["dump_end_z"]
345
+ relative_start_z = self.z_min - dump_start_z
346
+ relative_end_z = relative_start_z + self.delta_z
347
+ # (n_angles, n_z, n_x)
348
+ subregion = (None, None, relative_start_z, relative_end_z, None, None)
349
+ return subregion
350
+
351
+ def _check_resume_from_step(self):
352
+ if self._resume_from_step is None:
353
+ return
354
+ read_opts = self.processing_options["read_chunk"]
355
+ expected_radios_shape = get_hdf5_dataset_shape(
356
+ read_opts["process_file"],
357
+ read_opts["process_h5_path"],
358
+ sub_region=self._get_read_dump_subregion(),
359
+ )
360
+ # TODO check
361
+
362
+ def _init_reader_finalize(self):
363
+ """
364
+ Method called after _init_reader
365
+ """
366
+ self._check_resume_from_step()
367
+ self.radios = self.chunk_reader.data
368
+ self._compute_phase_kernel_margin()
369
+ self._get_phase_output_shape()
370
+
371
+ def _process_finalize(self):
372
+ """
373
+ Method called once the pipeline has been executed
374
+ """
375
+ if sum(self._get_phase_margin()[0]) > 0:
376
+ self.radios = self._orig_radios
377
+
378
+ def _get_slice_start_index(self):
379
+ return self.z_min + self._phase_margin_up
380
+
381
+ _get_image_start_index = _get_slice_start_index
382
+
383
+ #
384
+ # Pipeline initialization
385
+ #
386
+
387
+ def _init_pipeline(self):
388
+ self._init_reader()
389
+ self._init_flatfield()
390
+ self._init_double_flatfield()
391
+ self._init_ccd_corrections()
392
+ self._init_radios_rotation()
393
+ self._init_phase()
394
+ self._init_unsharp()
395
+ self._init_radios_movements()
396
+ self._init_mlog()
397
+ self._init_sino_normalization()
398
+ self._init_sino_builder()
399
+ self._init_sino_rings_correction()
400
+ self._prepare_reconstruction()
401
+ self._init_reconstruction()
402
+ self._init_histogram()
403
+ self._init_writer()
404
+ self._configure_data_dumps()
405
+
406
+ @use_options("read_chunk", "chunk_reader")
407
+ def _init_reader(self):
408
+ if "read_chunk" not in self.processing_steps:
409
+ raise ValueError("Cannot proceed without reading data")
410
+ options = self.processing_options["read_chunk"]
411
+ process_file = options.get("process_file", None)
412
+ if process_file is None:
413
+ # Standard case - start pipeline from raw data
414
+ # ChunkReader always take a non-subsampled dictionary "files"
415
+ self.chunk_reader = ChunkReader(
416
+ options["files"],
417
+ sub_region=self.sub_region,
418
+ convert_float=True,
419
+ binning=options["binning"],
420
+ dataset_subsampling=options["dataset_subsampling"],
421
+ )
422
+ else:
423
+ # Resume pipeline from dumped intermediate step
424
+ self.chunk_reader = HDF5Loader(
425
+ process_file, options["process_h5_path"], sub_region=self._get_read_dump_subregion()
426
+ )
427
+ self._resume_from_step = options["step_name"]
428
+ self.logger.debug(
429
+ "Load subregion %s from file %s" % (str(self.chunk_reader.sub_region), self.chunk_reader.fname)
430
+ )
431
+ self._init_reader_finalize()
432
+
433
+ @use_options("flatfield", "flatfield")
434
+ def _init_flatfield(self, shape=None):
435
+ if shape is None:
436
+ shape = self._get_shape("flatfield")
437
+ options = self.processing_options["flatfield"]
438
+
439
+ distortion_correction = None
440
+ if options["do_flat_distortion"]:
441
+ self.logger.info("Flats distortion correction will be applied")
442
+ estimation_kwargs = {}
443
+ estimation_kwargs.update(options["flat_distortion_params"])
444
+ estimation_kwargs["logger"] = self.logger
445
+ distortion_correction = DistortionCorrection(
446
+ estimation_method="fft-correlation", estimation_kwargs=estimation_kwargs, correction_method="interpn"
447
+ )
448
+
449
+ # FlatField parameter "radios_indices" must account for subsampling
450
+ self.flatfield = self.FlatFieldClass(
451
+ shape,
452
+ flats=self.dataset_info.flats,
453
+ darks=self.dataset_info.darks,
454
+ radios_indices=options["projs_indices"],
455
+ interpolation="linear",
456
+ distortion_correction=distortion_correction,
457
+ sub_region=self.sub_region,
458
+ binning=options["binning"],
459
+ convert_float=True,
460
+ )
461
+
462
+ @use_options("double_flatfield", "double_flatfield")
463
+ def _init_double_flatfield(self):
464
+ options = self.processing_options["double_flatfield"]
465
+ avg_is_on_log = options["sigma"] is not None
466
+ result_url = None
467
+ if options["processes_file"] not in (None, ""):
468
+ result_url = DataUrl(
469
+ file_path=options["processes_file"],
470
+ data_path=(self.dataset_info.hdf5_entry or "entry") + "/double_flatfield/results/data",
471
+ )
472
+ self.logger.info("Loading double flatfield from %s" % result_url.file_path())
473
+ self.double_flatfield = self.DoubleFlatFieldClass(
474
+ self._get_shape("double_flatfield"),
475
+ result_url=result_url,
476
+ sub_region=self.sub_region,
477
+ input_is_mlog=False,
478
+ output_is_mlog=False,
479
+ average_is_on_log=avg_is_on_log,
480
+ sigma_filter=options["sigma"],
481
+ )
482
+
483
+ @use_options("ccd_correction", "ccd_correction")
484
+ def _init_ccd_corrections(self):
485
+ options = self.processing_options["ccd_correction"]
486
+ self.ccd_correction = self.CCDCorrectionClass(
487
+ self._get_shape("ccd_correction"), median_clip_thresh=options["median_clip_thresh"]
488
+ )
489
+
490
+ @use_options("phase", "phase_retrieval")
491
+ def _init_phase(self):
492
+ options = self.processing_options["phase"]
493
+ # If unsharp mask follows phase retrieval, then it should be done
494
+ # before cropping to the "inner part".
495
+ # Otherwise, crop the data just after phase retrieval.
496
+ if "unsharp_mask" in self.processing_steps:
497
+ margin = None
498
+ else:
499
+ margin = self._phase_margin
500
+ self.phase_retrieval = self.PaganinPhaseRetrievalClass(
501
+ self._get_shape("phase"),
502
+ distance=options["distance_m"],
503
+ energy=options["energy_kev"],
504
+ delta_beta=options["delta_beta"],
505
+ pixel_size=options["pixel_size_m"],
506
+ padding=options["padding_type"],
507
+ margin=margin,
508
+ fftw_num_threads=0, # TODO tune in advanced params of nabu config file
509
+ )
510
+ if self.phase_retrieval.use_fftw:
511
+ self.logger.debug(
512
+ "PaganinPhaseRetrieval using FFTW with %d threads" % self.phase_retrieval.fftw.num_threads
513
+ )
514
+ if "unsharp_mask" not in self.processing_steps:
515
+ self.register_callback("phase", ChunkedPipeline._reshape_radios_after_phase)
516
+
517
+ @use_options("unsharp_mask", "unsharp_mask")
518
+ def _init_unsharp(self):
519
+ options = self.processing_options["unsharp_mask"]
520
+ self.unsharp_mask = self.UnsharpMaskClass(
521
+ self._get_shape("unsharp_mask"),
522
+ options["unsharp_sigma"],
523
+ options["unsharp_coeff"],
524
+ mode="reflect",
525
+ method=options["unsharp_method"],
526
+ )
527
+ self.register_callback("unsharp_mask", ChunkedPipeline._reshape_radios_after_phase)
528
+
529
+ @use_options("take_log", "mlog")
530
+ def _init_mlog(self):
531
+ options = self.processing_options["take_log"]
532
+ self.mlog = self.MLogClass(
533
+ self._get_shape("take_log"), clip_min=options["log_min_clip"], clip_max=options["log_max_clip"]
534
+ )
535
+
536
+ @use_options("rotate_projections", "projs_rot")
537
+ def _init_radios_rotation(self):
538
+ options = self.processing_options["rotate_projections"]
539
+ center = options["center"]
540
+ if center is None:
541
+ nx, ny = self.dataset_info.radio_dims
542
+ center = (nx / 2 - 0.5, ny / 2 - 0.5)
543
+ center = (center[0], center[1] - self.z_min)
544
+ self.projs_rot = self.ImageRotationClass(
545
+ self._get_shape("rotate_projections"), options["angle"], center=center, mode="edge", reshape=False
546
+ )
547
+ self._tmp_rotated_radio = self._allocate_array(
548
+ self._get_shape("rotate_projections"), "f", name="tmp_rotated_radio"
549
+ )
550
+
551
+ @use_options("radios_movements", "radios_movements")
552
+ def _init_radios_movements(self):
553
+ options = self.processing_options["radios_movements"]
554
+ self._vertical_shifts = options["translation_movements"][:, 1]
555
+ self.radios_movements = self.VerticalShiftClass(self._get_shape("radios_movements"), self._vertical_shifts)
556
+
557
+ @use_options("sino_normalization", "sino_normalization")
558
+ def _init_sino_normalization(self):
559
+ options = self.processing_options["sino_normalization"]
560
+ self.sino_normalization = self.SinoNormalizationClass(
561
+ kind=options["method"],
562
+ radios_shape=self._get_shape("sino_normalization"),
563
+ normalization_array=options["normalization_array"],
564
+ )
565
+
566
+ @use_options("build_sino", "sino_builder")
567
+ def _init_sino_builder(self):
568
+ options = self.processing_options["build_sino"]
569
+ self.sino_builder = self.SinoBuilderClass(
570
+ radios_shape=self._get_shape("build_sino"),
571
+ rot_center=options["rotation_axis_position"],
572
+ halftomo=options["enable_halftomo"],
573
+ )
574
+ if not (options["enable_halftomo"]):
575
+ self._sinobuilder_copy = False
576
+ self._sinobuilder_output = None
577
+ self.sinos = None
578
+ else:
579
+ self._sinobuilder_copy = True
580
+ self.sinos = self._allocate_sinobuilder_output()
581
+ self._sinobuilder_output = self.sinos
582
+
583
+ @use_options("sino_rings_correction", "sino_deringer")
584
+ def _init_sino_rings_correction(self):
585
+ options = self.processing_options["sino_rings_correction"]
586
+ fw_params = extract_parameters(options["user_options"])
587
+ fw_sigma = fw_params.pop("sigma", 1.0)
588
+ self.sino_deringer = self.SinoDeringerClass(
589
+ fw_sigma, sinos_shape=self._get_shape("sino_rings_correction"), **fw_params
590
+ )
591
+
592
+ # this should be renamed, as it could be confused with _init_reconstruction. What about _get_reconstruction_array ?
593
+ @use_options("reconstruction", "reconstruction")
594
+ def _prepare_reconstruction(self):
595
+ options = self.processing_options["reconstruction"]
596
+ x_s, x_e = options["start_x"], options["end_x"]
597
+ y_s, y_e = options["start_y"], options["end_y"]
598
+ self._rec_roi = (x_s, x_e + 1, y_s, y_e + 1)
599
+ self._allocate_recs(y_e - y_s + 1, x_e - x_s + 1)
600
+
601
+ @use_options("reconstruction", "reconstruction")
602
+ def _init_reconstruction(self):
603
+ options = self.processing_options["reconstruction"]
604
+ if self.sino_builder is None:
605
+ raise ValueError("Reconstruction cannot be done without build_sino")
606
+ if self.FBPClass is None:
607
+ raise ValueError("No usable FBP module was found")
608
+
609
+ if options["enable_halftomo"]:
610
+ rot_center = options["rotation_axis_position_halftomo"]
611
+ else:
612
+ rot_center = options["rotation_axis_position"]
613
+ if self.sino_builder._halftomo_flip:
614
+ rot_center = self.sino_builder.rot_center
615
+
616
+ self.reconstruction = self.FBPClass(
617
+ self._get_shape("reconstruction"),
618
+ angles=options["angles"],
619
+ rot_center=rot_center,
620
+ filter_name=options["fbp_filter_type"],
621
+ slice_roi=self._rec_roi,
622
+ padding_mode=options["padding_type"],
623
+ extra_options={
624
+ "scale_factor": 1.0 / options["pixel_size_cm"],
625
+ "axis_correction": options["axis_correction"],
626
+ "centered_axis": options["centered_axis"],
627
+ "clip_outer_circle": options["clip_outer_circle"],
628
+ "filter_cutoff": options["fbp_filter_cutoff"],
629
+ },
630
+ )
631
+ if options["fbp_filter_type"] is None:
632
+ self.reconstruction.fbp = self.reconstruction.backproj
633
+
634
+ @use_options("histogram", "histogram")
635
+ def _init_histogram(self):
636
+ options = self.processing_options["histogram"]
637
+ self.histogram = self.HistogramClass(method="fixed_bins_number", num_bins=options["histogram_bins"])
638
+
639
+ @use_options("save", "writer")
640
+ def _init_writer(self):
641
+ # TODO: henri: have a look to simplify
642
+ options = self.processing_options["save"]
643
+ file_prefix = options["file_prefix"]
644
+ output_dir = path.join(options["location"], file_prefix)
645
+ nx_info = None
646
+ self._hdf5_output = is_hdf5_extension(options["file_format"], errors=False)
647
+ if self._hdf5_output:
648
+ fname_start_index = None
649
+ file_prefix += str("_%04d" % self._get_slice_start_index())
650
+ entry = getattr(self.dataset_info.dataset_scanner, "entry", None)
651
+ nx_info = {
652
+ "process_name": self._get_process_name(),
653
+ "processing_index": 0,
654
+ "config": {
655
+ # "processing_options": self.processing_options, # Takes too much time to write, not useful for partial files
656
+ "nabu_config": self.process_config.nabu_config,
657
+ },
658
+ "entry": entry,
659
+ }
660
+ self._histogram_processing_index = nx_info["processing_index"] + 1
661
+ else:
662
+ fname_start_index = self._get_slice_start_index()
663
+ self._histogram_processing_index = 1
664
+ writer_options = {}
665
+ if options["tiff_single_file"]:
666
+ writer_options = {
667
+ "tiff_single_file": options["tiff_single_file"],
668
+ "single_tiff_initialized": getattr(self.process_config, "single_tiff_initialized", False),
669
+ }
670
+ self.process_config.single_tiff_initialized = True
671
+ self._writer_configurator = WriterConfigurator(
672
+ output_dir,
673
+ file_prefix,
674
+ file_format=options["file_format"],
675
+ overwrite=options["overwrite"],
676
+ start_index=fname_start_index,
677
+ logger=self.logger,
678
+ nx_info=nx_info,
679
+ write_histogram=("histogram" in self.processing_steps),
680
+ histogram_entry=getattr(self.dataset_info.dataset_scanner, "entry", "entry"),
681
+ writer_options=writer_options,
682
+ extra_options={
683
+ "jpeg2000_compression_ratio": options["jpeg2000_compression_ratio"],
684
+ "float_clip_values": options["float_clip_values"],
685
+ },
686
+ )
687
+ self.writer = self._writer_configurator.writer
688
+ self._writer_exec_args = self._writer_configurator._writer_exec_args
689
+ self._writer_exec_kwargs = self._writer_configurator._writer_exec_kwargs
690
+ self.histogram_writer = self._writer_configurator.get_histogram_writer()
691
+
692
+ #
693
+ # Pipeline re-initialization
694
+ #
695
+
696
+ def _reset_sub_region(self, sub_region):
697
+ self.set_subregion(sub_region)
698
+ self._reset_reader_subregion()
699
+ self._reset_flatfield()
700
+
701
+ def _reset_flatfield(self):
702
+ self._init_flatfield()
703
+
704
+ #
705
+ # Pipeline execution
706
+ #
707
+
708
+ @pipeline_step("chunk_reader", "Reading data")
709
+ def _read_data(self):
710
+ self.logger.debug("Region = %s" % str(self.sub_region))
711
+ t0 = time()
712
+ self.chunk_reader.load_data()
713
+ el = time() - t0
714
+
715
+ shp = self.chunk_reader.data.shape
716
+ itemsize = self.chunk_reader.dtype.itemsize if hasattr(self.chunk_reader, "dtype") else 4
717
+ GB = np.prod(shp) * itemsize / 1e9
718
+ self.logger.info("Read subvolume %s in %.2f s" % (str(shp), el))
719
+
720
+ def _reset_reader_subregion(self):
721
+ if self._resume_from_step is None:
722
+ self.chunk_reader._set_subregion(self.sub_region)
723
+ self.chunk_reader._init_reader()
724
+ self.chunk_reader._loaded = False
725
+ else:
726
+ self.chunk_reader._set_subregion(self._get_read_dump_subregion())
727
+ self.chunk_reader._loaded = False
728
+
729
+ @pipeline_step("flatfield", "Applying flat-field")
730
+ def _flatfield(self):
731
+ self.flatfield.normalize_radios(self.radios)
732
+
733
+ @pipeline_step("double_flatfield", "Applying double flat-field")
734
+ def _double_flatfield(self, radios=None):
735
+ if radios is None:
736
+ radios = self.radios
737
+ self.double_flatfield.apply_double_flatfield(radios)
738
+
739
+ @pipeline_step("ccd_correction", "Applying CCD corrections")
740
+ def _ccd_corrections(self, radios=None):
741
+ if radios is None:
742
+ radios = self.radios
743
+ _tmp_radio = self._allocate_array(radios.shape[1:], "f", name="tmp_ccdcorr_radio")
744
+ for i in range(radios.shape[0]):
745
+ self.ccd_correction.median_clip_correction(radios[i], output=_tmp_radio)
746
+ radios[i][:] = _tmp_radio[:]
747
+
748
+ @pipeline_step("projs_rot", "Rotating projections")
749
+ def _rotate_projections(self, radios=None):
750
+ if radios is None:
751
+ radios = self.radios
752
+ tmp_radio = self._tmp_rotated_radio
753
+ for i in range(radios.shape[0]):
754
+ self.projs_rot.rotate(radios[i], output=tmp_radio)
755
+ radios[i][:] = tmp_radio[:]
756
+
757
+ @pipeline_step("phase_retrieval", "Performing phase retrieval")
758
+ def _retrieve_phase(self):
759
+ if "unsharp_mask" in self.processing_steps:
760
+ output = self.radios
761
+ else:
762
+ self._get_cropped_radios()
763
+ output = self._radios_cropped
764
+ for i in range(self.radios.shape[0]):
765
+ self.phase_retrieval.apply_filter(self.radios[i], output=output[i])
766
+
767
+ @pipeline_step("unsharp_mask", "Performing unsharp mask")
768
+ def _apply_unsharp(self):
769
+ for i in range(self.radios.shape[0]):
770
+ self.radios[i] = self.unsharp_mask.unsharp(self.radios[i])
771
+ self._get_cropped_radios()
772
+
773
+ @pipeline_step("mlog", "Taking logarithm")
774
+ def _take_log(self):
775
+ self.mlog.take_logarithm(self.radios)
776
+
777
+ @pipeline_step("radios_movements", "Applying radios movements")
778
+ def _radios_movements(self, radios=None):
779
+ if radios is None:
780
+ radios = self.radios
781
+ self.radios_movements.apply_vertical_shifts(radios, list(range(radios.shape[0])))
782
+
783
+ @pipeline_step("sino_normalization", "Normalizing sinograms")
784
+ def _normalize_sinos(self, radios=None):
785
+ if radios is None:
786
+ radios = self.radios
787
+ sinos = radios.transpose((1, 0, 2))
788
+ self.sino_normalization.normalize(sinos)
789
+
790
+ def _dump_sinogram(self, radios=None):
791
+ if radios is None:
792
+ radios = self.radios
793
+ self._dump_data_to_file("sinogram", data=radios)
794
+
795
+ @pipeline_step("sino_builder", "Building sinograms")
796
+ def _build_sino(self, radios=None):
797
+ if radios is None:
798
+ radios = self.radios
799
+ # Either a new array (previously allocated in "_sinobuilder_output"),
800
+ # or a view of "radios"
801
+ self.sinos = self.sino_builder.radios_to_sinos(
802
+ radios, output=self._sinobuilder_output, copy=self._sinobuilder_copy
803
+ )
804
+
805
+ @pipeline_step("sino_deringer", "Removing rings on sinograms")
806
+ def _destripe_sinos(self, sinos=None):
807
+ if sinos is None:
808
+ sinos = self.sinos
809
+ self.sino_deringer.remove_rings(sinos)
810
+
811
+ @pipeline_step("reconstruction", "Reconstruction")
812
+ def _reconstruct(self, sinos=None):
813
+ if sinos is None:
814
+ sinos = self.sinos
815
+ for i in range(sinos.shape[0]):
816
+ self.reconstruction.fbp(sinos[i], output=self.recs[i])
817
+
818
+ @pipeline_step("histogram", "Computing histogram")
819
+ def _compute_histogram(self, data=None):
820
+ if data is None:
821
+ data = self.recs
822
+ self.recs_histogram = self.histogram.compute_histogram(data)
823
+
824
+ @pipeline_step("writer", "Saving data")
825
+ def _write_data(self, data=None):
826
+ if data is None:
827
+ data = self.recs
828
+ self.writer.write(data, *self._writer_exec_args, **self._writer_exec_kwargs)
829
+ self.logger.info("Wrote %s" % self.writer.get_filename())
830
+ self._write_histogram()
831
+
832
+ def _write_histogram(self):
833
+ if "histogram" not in self.processing_steps:
834
+ return
835
+ self.logger.info("Saving histogram")
836
+ self.histogram_writer.write(
837
+ hist_as_2Darray(self.recs_histogram),
838
+ self._get_process_name(kind="histogram"),
839
+ processing_index=self._histogram_processing_index,
840
+ config={
841
+ "file": path.basename(self.writer.get_filename()),
842
+ "bins": self.processing_options["histogram"]["histogram_bins"],
843
+ },
844
+ )
845
+
846
+ def _dump_data_to_file(self, step_name, data=None):
847
+ if step_name not in self._data_dump:
848
+ return
849
+ if data is None:
850
+ data = self.radios
851
+ writer = self._data_dump[step_name]
852
+ self.logger.info("Dumping data to %s" % writer.fname)
853
+ writer.write_data(data)
854
+
855
+ def _process_chunk(self):
856
+ self._flatfield()
857
+ self._double_flatfield()
858
+ self._ccd_corrections()
859
+ self._rotate_projections()
860
+ self._retrieve_phase()
861
+ self._apply_unsharp()
862
+ self._take_log()
863
+ self._radios_movements()
864
+ self._normalize_sinos()
865
+ self._dump_sinogram()
866
+ self._build_sino()
867
+ self._destripe_sinos()
868
+ self._reconstruct()
869
+ self._compute_histogram()
870
+ self._write_data()
871
+ self._process_finalize()
872
+
873
+ def process_chunk(self, sub_region=None):
874
+ if sub_region is not None:
875
+ self._reset_sub_region(sub_region)
876
+ self._reset_memory()
877
+ self._init_writer()
878
+ self._init_double_flatfield()
879
+ self._configure_data_dumps()
880
+ self._read_data()
881
+ self._process_chunk()