Source code for ewokstomo.tasks.reconstruct_volume

from pathlib import Path
import warnings

from ewokstomo.tasks.utils import (
    _get_technique_from_nabu_config,
    _get_output_format_from_nabu_config,
)
from ewokscore import Task
from ewokscore.model import BaseInputModel
from ewokscore.model import BaseOutputModel
from pydantic import Field
from nabu.pipeline.fullfield.reconstruction import FullFieldReconstructor
from nabu.pipeline.fullfield.processconfig import ProcessConfig
import gc




class ReconstructVolumeInputsModel(BaseInputModel):
    nx_path: Path = Field(
        ...,
        description="Path to the input NX file containing the tomography data.",
    )
    config_dict: dict = Field(
        ...,
        description=(
            "A dictionary containing parameters used to override Nabu's default configuration. "
            "Must include at least 'dataset' -> 'location', pointing to the input NX file. "
            "(see https://www.silx.org/pub/nabu/doc/nabu_config_items.html)"
        ),
    )
    slice_index_range: tuple[int, int] | None = Field(
        None,
        description=(
            "Optional tuple specifying the (start, end) indices of slices to reconstruct. "
            "If not provided, the entire volume will be reconstructed."
        ),
    )





class ReconstructVolumeOutputsModel(BaseOutputModel):
    reconstructed_volume_path: str = Field(
        ...,
        description="The file path to the saved reconstructed volume.",
    )
    nabu_dict: dict = Field(
        ...,
        description="The Nabu configuration dictionary used for reconstruction.",
    )
    processing_options: dict = Field(
        ...,
        description="The resolved Nabu processing options used by ProcessConfig.",
    )





class ReconstructVolume(  # type: ignore[call-arg]
    Task,
    input_model=ReconstructVolumeInputsModel,
    output_model=ReconstructVolumeOutputsModel,
):


    def run(self):
        """
        Task that reconstructs a volume from full-field tomography data using Nabu:
        """
        warnings.filterwarnings(
            "ignore",
            message="invalid value encountered in divide",
            category=RuntimeWarning,
            module="nabu\\.preproc\\.flatfield",
        )

        overwritten_config_fields = self.get_input_value("config_dict")
        nx_path = Path(self.get_input_value("nx_path"))
        slice_index_range = self.get_input_value("slice_index_range", None)

        if not nx_path.exists():
            raise FileNotFoundError(f"NX file not found: {nx_path}")

        technique = _get_technique_from_nabu_config(overwritten_config_fields)
        output_format = _get_output_format_from_nabu_config(overwritten_config_fields)

        output_dir = nx_path.parent.parent / "volumes"
        output_dir.mkdir(exist_ok=True)
        output_dir = output_dir / f"{nx_path.stem}_{technique}_32Bit_{output_format}"
        output_dir.mkdir(exist_ok=True)

        # Prepare the configuration for nabu
        overwritten_config_fields["dataset"]["location"] = str(nx_path)
        overwritten_config_fields["output"]["location"] = str(output_dir)

        if slice_index_range is not None:
            start_idx, end_idx = slice_index_range
            overwritten_config_fields["reconstruction"]["start_z"] = start_idx
            overwritten_config_fields["reconstruction"]["end_z"] = end_idx
        else:
            overwritten_config_fields["reconstruction"]["start_z"] = 0
            overwritten_config_fields["reconstruction"]["end_z"] = -1

        proc = ProcessConfig(conf_dict=overwritten_config_fields)

        file_prefix = f"{nx_path.stem}"
        proc.processing_options["save"]["file_prefix"] = file_prefix
        proc.nabu_config["output"]["file_prefix"] = file_prefix
        overwritten_config_fields["output"]["file_prefix"] = file_prefix

        reconstructor = FullFieldReconstructor(proc)
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", DeprecationWarning)
            reconstructor.reconstruct()
        reconstructor.finalize_files_saving()
        del reconstructor
        try:
            import cupy

            cupy.get_default_memory_pool().free_all_blocks()
            cupy.get_default_pinned_memory_pool().free_all_blocks()
        except Exception:
            print("Could not free GPU memory")
        gc.collect()
        save_options = proc.processing_options["save"]

        save_options = proc.processing_options["save"]
        output_location = Path(save_options["location"])
        file_prefix = save_options["file_prefix"]
        file_format = save_options["file_format"].lower()
        if file_format in ("hdf5", "h5"):
            reconstructed_path = str(output_location / f"{file_prefix}.{file_format}")
        else:
            reconstructed_path = str(output_location / f"{file_prefix}")

        self.outputs.reconstructed_volume_path = reconstructed_path
        self.outputs.nabu_dict = overwritten_config_fields
        self.outputs.processing_options = proc.processing_options