Source code for ewokstomo.tasks.convert_volume

import logging
import shutil
from pathlib import Path

import numpy as np
from ewokscore import Task
from ewokscore.model import BaseInputModel, BaseOutputModel
from nabu.io.cast_volume import cast_volume
from nabu.io.utils import get_first_hdf5_entry
from pydantic import Field
from tomoscan.esrf.volume import EDFVolume, HDF5Volume, JP2KVolume, MultiTIFFVolume
from typing import Literal

volume_formats = {
    "hdf5": HDF5Volume,
    "h5": HDF5Volume,
    "tiff": MultiTIFFVolume,
    "edf": EDFVolume,
    "jp2": JP2KVolume,
}

logger = logging.getLogger(__name__)




class ConvertVolumeTo16BitInputsModel(BaseInputModel):
    reconstructed_volume_path: str = Field(
        ...,
        description="The file path to the reconstructed volume.",
    )
    nabu_dict: dict = Field(
        ...,
        description=(
            "The Nabu configuration dictionary used for reconstruction. "
            "Used to determine the output file format for the converted volume."
        ),
    )
    output_format: Literal["hdf5", "tiff", "edf", "jp2"] = Field(
        "tiff",
        description=(
            "The desired output file format for the converted volume. "
            "Supported formats include 'hdf5', 'tiff', 'edf', 'jp2'. "
        ),
    )





class ConvertVolumeTo16BitOutputsModel(BaseOutputModel):
    converted_volume_path: str = Field(
        ..., description="The file path of the converted 16-bit volume."
    )





class ConvertVolumeTo16Bit(  # type: ignore[call-arg]
    Task,
    input_model=ConvertVolumeTo16BitInputsModel,
    output_model=ConvertVolumeTo16BitOutputsModel,
):
    """
    Task to convert volume to 16-bit format and save it to disk.
    The output file format is determined by the Nabu configuration dictionary provided as input.
    The task also removes the original 32-bit volume after conversion to save disk space.
    """



    def run(self):
        logger.info("Starting volume conversion to 16-bit.")
        input_volume = self.get_input_value("reconstructed_volume_path")
        logger.info(f"Input volume path: {input_volume}")
        nabu_dict = self.get_input_value("nabu_dict")

        # Determine input and output file formats
        input_format_str = nabu_dict.get("output", {}).get("file_format", "hdf5")
        output_format_str = self.get_input_value("output_format")

        # Construct output file path
        input_path = Path(input_volume)
        output_filename = input_path.stem
        if output_format_str == "hdf5":
            output_filename = output_filename + ".hdf5"
        elif output_format_str == "tiff":
            output_filename = output_filename + ".tiff"

        input_parent_str = str(input_path.parent)
        if "32Bit" in input_parent_str:
            output_volume = input_parent_str.replace("32Bit", "16Bit").replace(
                input_format_str, output_format_str
            )
        else:
            logger.warning(
                f"'32Bit' not found in input path '{input_parent_str}'. "
                "Falling back to a '16Bit' subdirectory next to the input."
            )
            output_volume = (
                input_parent_str.replace(input_format_str, output_format_str) + "_16Bit"
            )

        output_volume = Path(output_volume) / output_filename
        output_volume.parent.mkdir(parents=True, exist_ok=True)

        self.outputs.converted_volume_path = str(output_volume)

        # Create volume objects for input and output formats
        input_volume_class = volume_formats.get(input_format_str)
        input_vol_obj = self._make_input_volume(
            file_format=input_volume_class, file_path=input_path
        )

        output_volume_class = volume_formats.get(output_format_str)
        output_vol_obj = self._make_output_volume(
            file_format=output_volume_class, file_path=output_volume
        )

        cast_volume(
            input_volume=input_vol_obj,
            output_volume=output_vol_obj,
            output_data_type=np.uint16,
            remove_input_volume=True,
        )
        if input_format_str == "hdf5" and Path(output_vol_obj.url.file_path()).exists():
            if "32Bit" in str(input_path.parent) and not input_path.exists():
                shutil.rmtree(input_path.parent)
            else:
                logger.warning(
                    f"Skipping removal of '{input_path.parent}' because it does not contain '32Bit' in its path."
                )
        logger.info("Volume conversion completed.")


    def _make_input_volume(self, file_format, file_path):
        if file_format == HDF5Volume:
            entry = get_first_hdf5_entry(str(file_path))
            entry_path = f"{entry}/reconstruction"
            return HDF5Volume(file_path=file_path, data_path=entry_path)
        elif file_format == MultiTIFFVolume:
            return MultiTIFFVolume(file_path=file_path)
        else:
            constructor = file_format
            volume = constructor(
                folder=str(file_path.parent), volume_basename=file_path.name
            )
            return volume

    def _make_output_volume(self, file_format, file_path):
        if file_format == HDF5Volume:
            return HDF5Volume(file_path=file_path, data_path="/volume")
        elif file_format == MultiTIFFVolume:
            return MultiTIFFVolume(file_path=file_path)
        else:
            constructor = file_format
            volume = constructor(
                folder=str(file_path.parent), volume_basename=file_path.name
            )
            return volume