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