"""
Online reconstruction tasks for real-time tomography processing.
This module provides tasks and utilities for performing tomographic
reconstruction in real-time as data is acquired from the beamline.
"""
import logging
import numpy as np
import time
from pathlib import Path
import h5py
from ewokscore import Task
from nabu.cuda.utils import __has_cupy__
from nabu.estimation.cor import CenterOfRotation
from nabu.preproc.flatfield import FlatField
from nabu.resources.dataset_inmemory import InMemoryDataset
from blissdata.redis_engine.store import DataStore
from blissdata.beacon.data import BeaconData
from blissdata.redis_engine.scan import ScanState
from ewokstomo.tasks.utils import wait_for_scan_state
from ewokstomo.tasks.utils import _get_technique_from_nabu_config
from ewokstomo.tasks.online.nabu_pipeline import get_nabu_pipeline
from ewokstomo.tasks.online.utils import get_image_stream_name, load_reduced_frames
from ewokscore.model import BaseInputModel
from ewokscore.model import BaseOutputModel
from pydantic import Field
logger = logging.getLogger(__name__)
if __has_cupy__:
logger.info("Using CUDA for reconstruction")
else:
logger.info("Using OpenCL for reconstruction")
def _projection_batches(total_nb_projection: int, batch_size: int) -> list[list[int]]:
if total_nb_projection < 2:
raise ValueError("FBP reconstruction requires at least 2 projections")
batch_size = max(2, batch_size)
batches = []
proj_idx = 0
while proj_idx < total_nb_projection:
end_idx = min(proj_idx + batch_size, total_nb_projection)
if total_nb_projection - end_idx == 1:
end_idx = total_nb_projection
batches.append(list(range(proj_idx, end_idx)))
proj_idx = end_idx
return batches
[docs]
class OnlineReconstructSliceOutputModel(BaseOutputModel):
reconstructed_slices_directory: str = Field(
..., description="Path to the saved reconstructed slice files"
)
[docs]
class OnlineReconstructSlice( # type: ignore[call-arg]
Task,
input_model=OnlineReconstructSliceInputModel,
output_model=OnlineReconstructSliceOutputModel,
):
"""
(ESRF-only) Perform real-time tomography reconstruction on streaming data.
This task connects to a live scan stream and performs incremental
reconstruction by processing projections in batches. It includes
flat field correction, phase retrieval, and FBP reconstruction.
The reconstruction is performed on a single slice (middle by default) to minimize
computational cost while providing real-time feedback.
"""
STREAM_TIMEOUT = 5 # seconds to wait for new data after scan stops
WAITING_INTERVAL = 1 # seconds between checks for new data
[docs]
def run(self):
"""Execute the online reconstruction pipeline."""
# Get required inputs
scan_key = self.get_input_value("scan_key")
output_path = self.get_input_value("output_path")
rotation_motor = self.get_input_value("rotation_motor")
batch_size = int(self.get_input_value("batch_size", 100))
reduced_dark_path = self.get_input_value("reduced_dark_path")
reduced_flat_path = self.get_input_value("reduced_flat_path")
pixel_size_m = float(self.get_input_value("pixel_size_m"))
distance_m = float(self.get_input_value("distance_m"))
energy_keV = float(self.get_input_value("energy_keV"))
conf_dict = self.get_input_value("config_dict", {})
conf_dict.setdefault("reconstruction", {})
# Get optional parameters
extra_options = self.get_input_value("extra_options", {"centered_axis": True})
slice_index = self.get_input_value("slice_index", "middle")
# Connect to beacon and load scan
logger.info(f"Connecting to scan: {scan_key}")
beacon_client = BeaconData()
redis_url = beacon_client.get_redis_data_db()
data_store = DataStore(redis_url)
scan = data_store.load_scan(scan_key)
# Get number of projections
total_nb_projection = int(scan.info["technique"]["proj"]["proj_n"])
# Wait for scan to start
wait_for_scan_state(scan, ScanState.STARTED)
# Prepare output directory
self.outputs.reconstructed_slices_directory = output_path
Path(output_path).mkdir(parents=True, exist_ok=True)
# Populate dataset_info structure
dataset_info = InMemoryDataset("")
dataset_info.energy = energy_keV
dataset_info.distance = distance_m
dataset_info.pixel_size = pixel_size_m * 1e6
dataset_info.is_halftomo = conf_dict["reconstruction"].get(
"enable_halftomo", False
)
dataset_info.radio_dims = tuple(
scan.streams[get_image_stream_name(scan)].info["shape"][::-1]
)
dataset_info.current_projections_indices = np.arange(
batch_size, dtype=np.uint64
)
dataset_info.n_angles = batch_size
logger.info("Loading flat field correction data")
dataset_info.darks = load_reduced_frames(reduced_dark_path, "darks")
dataset_info.flats = load_reduced_frames(reduced_flat_path, "flats")
dataset_info.n_frames = (
total_nb_projection + len(dataset_info.darks) + len(dataset_info.flats)
)
conf_dict["reconstruction"]["method"] = "FBP"
conf_dict["reconstruction"]["start_z"] = slice_index
conf_dict["reconstruction"]["end_z"] = slice_index
conf_dict["reconstruction"]["centered_axis"] = extra_options.get(
"centered_axis", True
)
conf_dict["output"]["location"] = output_path
conf_dict["output"]["file_prefix"] = "rec"
technique = _get_technique_from_nabu_config(conf_dict)
sample_name = scan.info["sample"]["name"]
# Estimate rotation axis position if not furnished
rotation_axis_position = conf_dict["reconstruction"].get(
"rotation_axis_position", None
)
if rotation_axis_position is None:
idx_180 = self._wait_for_180_degrees_projection(scan, rotation_motor)
stream_name = self._get_image_stream_name(scan)
detector_stream = scan.streams[stream_name]
self._wait_for_stream_length(
scan, detector_stream, idx_180 + 1, stream_name
)
projections = np.array(
[detector_stream[0], detector_stream[idx_180]], dtype=np.float32
)
projections = self._flat_field_correct(
projections, dataset_info, radios_indices=[0, idx_180]
)
rotation_axis_position = self._estimate_cor(projections[0], projections[1])
conf_dict["reconstruction"]["rotation_axis_position"] = (
rotation_axis_position
)
reco_pipeline = get_nabu_pipeline(
conf_dict, dataset_info, batch_size, pipeline_backend="cuda"
)
logger.info("Starting online reconstruction...")
# Process batches
projection_batches = _projection_batches(total_nb_projection, batch_size)
for batch_index, projections_indices in enumerate(projection_batches, start=1):
logger.info(f"Processing batch {batch_index}/{len(projection_batches)}")
try:
proj_idx = projections_indices[0]
end_idx = projections_indices[-1] + 1
projections = self._get_projections_from_stream(
scan, projections_indices, output_dtype=np.float32
)
angles = self._get_angles_from_stream(
scan, rotation_motor, projections_indices, output_dtype=np.float32
)
# Pad last batch if it's smaller than batch_size to avoid issues
actual_batch_size = end_idx - proj_idx
if actual_batch_size < batch_size:
pad_width = batch_size - actual_batch_size
projections = np.pad(
projections, ((0, pad_width), (0, 0), (0, 0)), mode="edge"
)
angles = np.pad(angles, (0, pad_width), mode="edge")
projections_indices = (
list(projections_indices)
+ [projections_indices[-1]] * pad_width
)
dataset_info.data = projections
dataset_info.rotation_angles = angles
dataset_info.n_angles = len(projections_indices)
dataset_info.current_projections_indices = np.array(projections_indices)
reco_pipeline.process_config.processing_options["save"]["location"] = (
str(output_path)
)
reco_pipeline.process_config.processing_options["save"][
"file_prefix"
] = f"{sample_name}_{technique}_xy_online_{proj_idx}_{end_idx - 1}"
reco_pipeline.process_chunk(reco_pipeline.target_sub_region)
reco_pipeline.save_reconstruction()
except RuntimeError as e:
logger.error(f"Stopping reconstruction due to error: {e}")
break
# Final flush: persist the last accumulated reconstruction once all
# batches have been processed (the in-loop call saves per batch).
reco_pipeline.save_reconstruction()
logger.info("Online reconstruction completed")
[docs]
def save_reconstructed_slice(
self, output_path: str, reconstructed_slice: np.ndarray
):
"""Save the reconstructed slice to an HDF5 file.
Parameters
----------
output_path : str
Path to save the reconstructed slice
reconstructed_slice : np.ndarray
The reconstructed slice data
"""
try:
with h5py.File(output_path, "w") as h5f:
h5f.create_dataset("reconstructed_slice", data=reconstructed_slice)
logger.info(f"Reconstructed slice saved to {output_path}")
except Exception as e:
logger.error(f"Error saving reconstructed slice: {e}")
raise
def _wait_for_180_degrees_projection(self, scan, rotation_motor: str) -> int:
"""Wait until the rotation motor reaches a difference of 180 degrees from start."""
logger.info("Waiting to reach 180 degrees for COR estimation...")
motor_stream = scan.streams[self._get_motor_stream_name(scan, rotation_motor)]
waited_time = 0
while True:
if len(motor_stream) > 0:
angles = np.array(motor_stream[:], dtype=np.float32)
diffs = np.abs(angles - angles[0])
if np.any(diffs >= 180.0):
idx_180 = int(np.argmax(diffs >= 180.0))
logger.info(
f"Reached 180 degrees at index {idx_180} (angle: {angles[idx_180]})"
)
return idx_180
if scan.state >= ScanState.STOPPED:
waited_time += 1
if waited_time > self.STREAM_TIMEOUT:
if len(motor_stream) > 0:
idx_180 = int(np.argmax(diffs))
logger.warning(
f"Scan stopped early. Using maximum difference index {idx_180} (angle: {angles[idx_180]})"
)
return idx_180
raise RuntimeError(
"Scan stopped before any motor positions were acquired."
)
time.sleep(self.WAITING_INTERVAL)
scan.update(block=False)
def _flat_field_correct(self, projections, dataset_info, radios_indices):
"""Apply Nabu flat-field correction in place and return the projections."""
flat_field = FlatField(
projections.shape,
flats=dataset_info.flats,
darks=dataset_info.darks,
radios_indices=radios_indices,
)
return flat_field.normalize_radios(projections)
def _estimate_cor(self, proj_0, proj_180) -> float:
"""Estimate the center of rotation with Nabu."""
logger.info("Estimating center of rotation using Nabu...")
try:
cor = CenterOfRotation().find_shift(
proj_0,
np.fliplr(proj_180),
return_relative_to_middle=False,
)
logger.info(f"Estimated center of rotation: {cor}")
return float(cor)
except Exception as e:
logger.error(f"Failed to estimate center of rotation: {e}.")
raise RuntimeError("Center of rotation estimation failed.")
def _get_image_stream_name(self, scan) -> str:
"""Get the image stream name from the scan."""
return get_image_stream_name(scan)
def _get_motor_stream_name(self, scan, motor_name: str) -> str:
"""Get the motor stream name from the scan."""
for name in scan.streams:
if motor_name in name and "axis:" in name:
return name
raise ValueError(f"No motor stream found for {motor_name} in scan")
def _wait_for_stream_length(self, scan, stream, idx_end, stream_name):
"""Block until ``stream`` holds at least ``idx_end`` items.
Polls the scan for new data and raises ``TimeoutError`` if the scan
stops before enough items arrive, so a missing or short stream can
never hang the task indefinitely.
"""
waited_time = 0
while len(stream) < idx_end:
logger.debug(
f"Waiting for data in stream {stream_name} ({len(stream)}/{idx_end})..."
)
time.sleep(self.WAITING_INTERVAL)
scan.update(block=False)
# Only start the timeout counter after the scan has stopped.
if scan.state >= ScanState.STOPPED:
waited_time += 1
if waited_time > self.STREAM_TIMEOUT:
raise TimeoutError(
f"Not enough data in stream {stream_name} before scan "
f"stopped. Got {len(stream)}, needed {idx_end}"
)
def _get_projections_from_stream(
self, scan, projections_indices, output_dtype=np.float32
) -> np.ndarray:
"""Get projections from the image stream."""
stream_name = get_image_stream_name(scan)
detector_stream = scan.streams[stream_name]
idx_start, idx_end = projections_indices[0], projections_indices[-1] + 1
self._wait_for_stream_length(scan, detector_stream, idx_end, stream_name)
return np.array(detector_stream[idx_start:idx_end], dtype=output_dtype)
def _get_angles_from_stream(
self, scan, rotation_motor: str, projections_indices, output_dtype=np.float32
) -> np.ndarray:
"""Get angles from the motor stream."""
motor_stream_name = self._get_motor_stream_name(scan, rotation_motor)
motor_stream = scan.streams[motor_stream_name]
idx_start, idx_end = projections_indices[0], projections_indices[-1] + 1
self._wait_for_stream_length(scan, motor_stream, idx_end, motor_stream_name)
angles_deg = np.array(motor_stream[idx_start:idx_end], dtype=output_dtype)
angles_rad = np.deg2rad(angles_deg)
return angles_rad
