Source code for trx.io

# -*- coding: utf-8 -*-
"""Unified I/O interface for tractogram file formats."""

import logging
import os
import sys
import tempfile

try:
    import dipy  # noqa: F401



    dipy_available = True

except ImportError:
    dipy_available = False

from trx.utils import split_name_with_gz




def get_trx_tmp_dir():
    """Return a temporary directory honoring the ``TRX_TMPDIR`` setting.

    When the ``TRX_TMPDIR`` environment variable is set to ``"use_working_dir"``
    the current working directory is used. Otherwise, the value of
    ``TRX_TMPDIR`` is used directly. If the variable is not set, the system
    temporary directory is used.

    Returns
    -------
    tempfile.TemporaryDirectory
        Context-managed temporary directory placed according to the environment
        configuration.
    """
    if os.getenv("TRX_TMPDIR") is not None:
        if os.getenv("TRX_TMPDIR") == "use_working_dir":
            trx_tmp_dir = os.getcwd()
        else:
            trx_tmp_dir = os.getenv("TRX_TMPDIR")
    else:
        trx_tmp_dir = tempfile.gettempdir()

    if sys.version_info[1] >= 10:
        return tempfile.TemporaryDirectory(
            dir=trx_tmp_dir, prefix="trx_", ignore_cleanup_errors=True
        )
    else:
        return tempfile.TemporaryDirectory(dir=trx_tmp_dir, prefix="trx_")





def load_sft_with_reference(filepath, reference=None, bbox_check=True):
    """Load a tractogram as a StatefulTractogram with an explicit reference.

    Parameters
    ----------
    filepath : str
        Path to the tractogram file (.trk, .tck, .fib, .vtk, .dpy).
    reference : str or nibabel.Nifti1Image, optional
        Reference image used for formats without embedded affine information.
        Pass ``"same"`` to reuse the header embedded in .trk files.
    bbox_check : bool, optional
        If True, validate that streamlines lie within the reference bounding
        box. Defaults to True.

    Returns
    -------
    StatefulTractogram or None
        Loaded tractogram. Returns ``None`` when ``dipy`` is unavailable.

    Raises
    ------
    IOError
        If the file format is unsupported or a required reference is missing.
    """
    if not dipy_available:
        logging.error(
            "Dipy library is missing, cannot use functions related "
            "to the StatefulTractogram."
        )
        return None
    from dipy.io.streamline import load_tractogram

    # Force the usage of --reference for all file formats without an header
    _, ext = os.path.splitext(filepath)
    if ext == ".trk":
        if reference is not None and reference != "same":
            logging.warning(
                "Reference is discarded for this file format {}.".format(filepath)
            )
        sft = load_tractogram(filepath, "same", bbox_valid_check=bbox_check)
    elif ext in [".tck", ".fib", ".vtk", ".dpy"]:
        if reference is None or reference == "same":
            raise IOError(
                "--reference is required for this file format {}.".format(filepath)
            )
        else:
            sft = load_tractogram(filepath, reference, bbox_valid_check=bbox_check)

    else:
        raise IOError("{} is an unsupported file format".format(filepath))

    return sft





def load(tractogram_filename, reference):
    """Load a tractogram from disk and return a TRX or StatefulTractogram.

    Parameters
    ----------
    tractogram_filename : str
        Path to the input tractogram. TRX directories are supported.
    reference : str or nibabel.Nifti1Image
        Reference image used for formats without embedded affine information.

    Returns
    -------
    TrxFile or StatefulTractogram
        TRX file handle for ``.trx`` inputs, otherwise a StatefulTractogram.
    """
    import trx.trx_file_memmap as tmm

    in_ext = split_name_with_gz(tractogram_filename)[1]
    if in_ext != ".trx" and not os.path.isdir(tractogram_filename):
        tractogram_obj = load_sft_with_reference(
            tractogram_filename, reference, bbox_check=False
        )
    else:
        tractogram_obj = tmm.load(tractogram_filename)

    return tractogram_obj





def save(tractogram_obj, tractogram_filename, bbox_valid_check=False):
    """Save a tractogram object to disk.

    Parameters
    ----------
    tractogram_obj : TrxFile or StatefulTractogram
        Tractogram to persist. Non-TRX inputs are converted to StatefulTractogram
        before saving to non-TRX formats.
    tractogram_filename : str
        Destination file name. ``.trx`` will be saved using the TRX writer; all
        other extensions are handled by ``dipy.save_tractogram``.
    bbox_valid_check : bool, optional
        If True, validate that streamlines lie within the reference bounding
        box when saving non-TRX formats. Defaults to False.

    Returns
    -------
    None
        The function writes to disk and returns ``None``. Returns ``None``
        immediately when ``dipy`` is unavailable.
    """
    if not dipy_available:
        logging.error(
            "Dipy library is missing, cannot use functions related "
            "to the StatefulTractogram."
        )
        return None
    from dipy.io.stateful_tractogram import StatefulTractogram
    from dipy.io.streamline import save_tractogram

    import trx.trx_file_memmap as tmm

    out_ext = split_name_with_gz(tractogram_filename)[1]

    if out_ext != ".trx":
        if not isinstance(tractogram_obj, StatefulTractogram):
            tractogram_obj = tractogram_obj.to_sft()
        save_tractogram(
            tractogram_obj, tractogram_filename, bbox_valid_check=bbox_valid_check
        )
    else:
        if not isinstance(tractogram_obj, tmm.TrxFile):
            tractogram_obj = tmm.TrxFile.from_sft(tractogram_obj)
        tmm.save(tractogram_obj, tractogram_filename)
        tractogram_obj.close()