Commit 475ad90f authored by Pierre Paleo's avatar Pierre Paleo
Browse files

Compute DFF on host for CudaFullFieldPipelineLimitedMemory

parent 55bbe459
Pipeline #27042 passed with stages
in 3 minutes and 5 seconds
......@@ -184,6 +184,7 @@ class FullFieldPipeline:
self.radios = self.chunk_reader.files_data
self.radios_shape = self.radios.shape
self._sinobuilder_radios_shape = self.radios_shape
self._dff_radios_shape = self.radios_shape
def _process_finalize(self):
......@@ -244,7 +245,7 @@ class FullFieldPipeline:
def _init_double_flatfield(self):
options = self.processing_options["double_flatfield"]
self.double_flatfield = self.DoubleFlatFieldClass(
self.radios_shape,
self._dff_radios_shape,
result_url=None,
input_is_mlog=False,
output_is_mlog=False,
......@@ -417,8 +418,10 @@ class FullFieldPipeline:
self.flatfield.normalize_radios(self.radios)
@pipeline_step("double_flatfield", "Applying double flat-field")
def _double_flatfield(self):
self.double_flatfield.apply_double_flatfield(self.radios)
def _double_flatfield(self, radios=None):
if radios is None:
radios = self.radios
self.double_flatfield.apply_double_flatfield(radios)
@pipeline_step("phase_retrieval", "Performing phase retrieval")
def _retrieve_phase(self):
......
from math import ceil
import numpy as np
from ..preproc.ccd_cuda import CudaFlatField, CudaLog
from ..preproc.double_flatfield import DoubleFlatField
from ..preproc.double_flatfield_cuda import CudaDoubleFlatField
from ..preproc.phase_cuda import CudaPaganinPhaseRetrieval
from ..preproc.sinogram_cuda import CudaSinoProcessing
......@@ -104,6 +105,8 @@ class CudaFullFieldPipelineLimitedMemory(CudaFullFieldPipeline):
# In this case, the "build sinogram" step is best done on host to avoid
# extraneous CPU<->GPU copies, and save some GPU memory.
SinoProcessingClass = SinoProcessing
# Same for DoubleFlatField
DoubleFlatFieldClass = DoubleFlatField
def __init__(self, process_config, sub_region, chunk_size, logger=None, extra_options=None, cuda_options=None):
"""
......@@ -160,6 +163,7 @@ class CudaFullFieldPipelineLimitedMemory(CudaFullFieldPipeline):
assert (self.delta_z % self.chunk_size) == 0, "delta_z must be a multiple of chunk_size"
self._allocate_radios()
self._h_recs = None
self._old_flatfield = None
def _init_reader_finalize(self):
......@@ -178,7 +182,10 @@ class CudaFullFieldPipelineLimitedMemory(CudaFullFieldPipeline):
self.radios = self.chunk_reader.files_data
# (group_size, delta_z, n_x) - fits in GPU mem
self.radios_group_shape = (self.radios_group_size, ) + self.radios.shape[1:]
self.radios_shape = self.radios_group_shape # passed to CCD processing classes
# passed to CCD processing classes
self.radios_shape = self.radios_group_shape
# DoubleFF is done on host. Here self.radios is still the host array (n_angles, delta_z, width)
self._dff_radios_shape = self.radios.shape
if "flatfield" in self.processing_steps:
# We need to update the projections indices passed to FlatField
# for each group of radios
......@@ -212,6 +219,51 @@ class CudaFullFieldPipelineLimitedMemory(CudaFullFieldPipeline):
# re-allocate _d_radios for processing a new chunk
self.radios = self._h_radios
self._allocate_radios()
self.flatfield = self._old_flatfield
def _reinit_flatfield(self, start_idx, end_idx, transfer_size):
if "flatfield" in self.processing_steps:
# We need to update the projections indices passed to FlatField
# for each group of radios
ff_opts = self.processing_options["flatfield"]
ff_opts["projs_indices"] = self._ff_proj_indices[start_idx:end_idx]
self._init_flatfield(shape=(transfer_size, ) + self.radios_shape[1:])
def _flatfield_radios_group(self, start_idx, end_idx, transfer_size):
self._reinit_flatfield(start_idx, end_idx, transfer_size)
self._flatfield()
def _apply_flatfield_and_dff(self, n_groups, group_size, n_images):
"""
If double flat-field is activated, apply flat-field + double flat-field.
Otherwise, do nothing and leave the flat-field for later "group processing"
"""
if "double_flatfield" not in self.processing_steps:
return
for i in range(n_groups):
self.logger.debug("processing group %d/%d" % (i+1, n_groups))
start_idx = i * group_size
end_idx = min((i + 1) * group_size, n_images)
transfer_size = end_idx - start_idx
# Copy H2D
self._d_radios[:transfer_size, :, :] = self._h_radios[start_idx:end_idx, :, :]
# Process a group of radios (radios_group_size, delta_z, width)
self._old_radios = self.radios
self.radios = self.radios[:transfer_size]
self._flatfield_radios_group(start_idx, end_idx, transfer_size)
self.radios = self._old_radios
# Copy D2H
self._d_radios[:transfer_size, :, :].get(ary=self._h_radios[start_idx:end_idx])
# Here flat-field has been applied on all radios (n_angles, delta_z, n_x).
# Now apply double FF on host.
self._double_flatfield(radios=self._h_radios)
if "flatfield" in self.processing_steps:
self.processing_steps.remove("flatfield")
self._old_flatfield = self.flatfield
self.flatfield = None
def _process_chunk_ccd(self):
......@@ -221,6 +273,9 @@ class CudaFullFieldPipelineLimitedMemory(CudaFullFieldPipeline):
n_groups = self._n_radios_groups
group_size = self.radios_group_size
n_images = self._h_radios.shape[0]
self._apply_flatfield_and_dff(n_groups, group_size, n_images)
for i in range(n_groups):
self.logger.debug("processing group %d/%d" % (i+1, n_groups))
start_idx = i * group_size
......@@ -229,16 +284,9 @@ class CudaFullFieldPipelineLimitedMemory(CudaFullFieldPipeline):
# Copy H2D
self._d_radios[:transfer_size, :, :] = self._h_radios[start_idx:end_idx, :, :]
# Process a group of radios (radios_group_size, delta_z, width)
if "flatfield" in self.processing_steps:
# We need to update the projections indices passed to FlatField
# for each group of radios
ff_opts = self.processing_options["flatfield"]
ff_opts["projs_indices"] = self._ff_proj_indices[start_idx:end_idx]
self._init_flatfield(shape=(transfer_size, ) + self.radios_shape[1:])
self._old_radios = self.radios
self.radios = self.radios[:transfer_size]
self._flatfield()
self._double_flatfield()
self._flatfield_radios_group(start_idx, end_idx, transfer_size)
self._retrieve_phase()
self._apply_unsharp()
self._take_log()
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment