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