nabu 2024.1.0rc4__py3-none-any.whl → 2024.1.1__py3-none-any.whl

build/lib/build/lib/build/lib/nabu/io/tests/test_writers.py

@@ -0,0 +1,187 @@
+from os import path
+from tempfile import TemporaryDirectory
+import pytest
+import numpy as np
+from tifffile import TiffReader
+from silx.io.dictdump import h5todict, dicttoh5
+from nabu.misc.utils import psnr
+from nabu.io.writer import NXProcessWriter, TIFFWriter, JP2Writer, __have_jp2k__
+from nabu.io.reader import import_h5_to_dict
+from nabu.testutils import get_data
+
+if __have_jp2k__:
+    from glymur import Jp2k
+
+
+@pytest.fixture(scope="class")
+def bootstrap(request):
+    cls = request.cls
+    cls._tmpdir = TemporaryDirectory(prefix="nabu_")
+    cls.tempdir = cls._tmpdir.name
+    cls.sino_data = get_data("mri_sino500.npz")["data"].astype(np.uint16)
+    cls.data = cls.sino_data
+    yield
+    # teardown
+    cls._tmpdir.cleanup()
+
+
+@pytest.mark.usefixtures("bootstrap")
+class TestTiff:
+    def _check_tif_file(self, fname, expected_data, n_expected_images):
+        with TiffReader(fname) as tif:
+            assert len(tif.pages) == n_expected_images
+            for i in range(n_expected_images):
+                data_read = tif.pages[i].asarray()
+                if expected_data.ndim == 3:
+                    expected_data_ = expected_data[i]
+                else:
+                    expected_data_ = expected_data
+                assert np.allclose(data_read, expected_data_)
+
+    def test_2D(self):
+        # TODO use start_index=None (by default) in TIFFWriter
+        # if start_index is None:
+        # url = DataUrl(file_path=dirname(self.fname), data_path="basename(self.fname).tiff", scheme="tifffile")
+        # volume = TIFFVolume(volume_basename=basename(self.fname), data_url=url)
+        # volume.data = img
+        # volume.save()
+        pytest.skip("Writing a single 2D tiff is disabled for now")
+
+        fname = path.join(self.tempdir, "test_tiff2D.tif")
+        data = np.arange(100 * 101, dtype="f").reshape((100, 101))
+        nabu_tif = TIFFWriter(fname)
+        nabu_tif.write(data)
+        self._check_tif_file(fname, data, 1)
+
+    def test_3D_data_split_in_multiple_files(self):
+        fname = path.join(self.tempdir, "test_tiff3D_single.tif")
+        data = np.arange(11 * 100 * 101, dtype="f").reshape((11, 100, 101))
+        nabu_tif = TIFFWriter(fname, multiframe=False, start_index=500)
+        nabu_tif.write(data)
+
+        assert not (path.isfile(fname)), "found %s" % fname
+
+        prefix, ext = path.splitext(fname)
+        for i in range(data.shape[0]):
+            curr_rel_fname = prefix + str("_%06d" % (i + nabu_tif.start_index)) + ext
+            curr_fname = path.join(self.tempdir, curr_rel_fname)
+            self._check_tif_file(curr_fname, data[i], 1)
+
+    def test_3D_data_in_one_file(self):
+        fname = path.join(self.tempdir, "test_tiff3D_multi.tif")
+        data = np.arange(11 * 100 * 101, dtype="f").reshape((11, 100, 101))
+        nabu_tif = TIFFWriter(fname, multiframe=True)
+        nabu_tif.write(data)
+
+        assert path.isfile(fname)
+        self._check_tif_file(fname, data, data.shape[0])
+
+
+@pytest.mark.skipif(not (__have_jp2k__), reason="Need openjpeg2000/glymur for this test")
+@pytest.mark.usefixtures("bootstrap")
+class TestJP2:
+    def _check_jp2_file(self, fname, expected_data, expected_psnr=None):
+        data = Jp2k(fname)[:]
+        if expected_psnr is None:
+            assert np.allclose(data, expected_data)
+        else:
+            computed_psnr = psnr(data, expected_data)
+            assert np.abs(computed_psnr - expected_psnr) < 1
+
+    def test_2D_lossless(self):
+        data = get_data("mri_sino500.npz")["data"].astype(np.uint16)
+        fname = path.join(self.tempdir, "sino500.jp2")
+        nabu_jp2 = JP2Writer(fname, psnr=[0])
+        nabu_jp2.write(data)
+        self._check_jp2_file(fname, data)
+
+    def test_2D_lossy(self):
+        fname = path.join(self.tempdir, "sino500_lossy.jp2")
+        nabu_jp2 = JP2Writer(fname, psnr=[80])
+        nabu_jp2.write(self.sino_data)
+        self._check_jp2_file(fname, self.sino_data, expected_psnr=80)
+
+    def test_3D(self):
+        fname = path.join(self.tempdir, "sino500_multi.jp2")
+        n_images = 5
+        data = np.tile(self.sino_data, (n_images, 1, 1))
+        for i in range(n_images):
+            data[i] += i
+
+        nabu_jp2 = JP2Writer(fname, start_index=10)
+        nabu_jp2.write(data)
+
+        assert not (path.isfile(fname))
+
+        prefix, ext = path.splitext(fname)
+        for i in range(data.shape[0]):
+            curr_rel_fname = prefix + str("_%06d" % (i + nabu_jp2.start_index)) + ext
+            curr_fname = path.join(self.tempdir, curr_rel_fname)
+            self._check_jp2_file(curr_fname, data[i])
+
+
+@pytest.fixture(scope="class")
+def bootstrap_h5(request):
+    cls = request.cls
+    cls._tmpdir = TemporaryDirectory(prefix="nabu_")
+    cls.tempdir = cls._tmpdir.name
+    cls.data = get_data("mri_sino500.npz")["data"].astype(np.uint16)
+    cls.h5_config = {
+        "key1": "value1",
+        "some_int": 1,
+        "some_float": 1.0,
+        "some_dict": {
+            "numpy_array": np.ones((5, 6), dtype="f"),
+            "key2": "value2",
+        },
+    }
+    yield
+    # teardown
+    cls._tmpdir.cleanup()
+
+
+@pytest.mark.usefixtures("bootstrap_h5")
+class TestNXWriter:
+    def test_write_simple(self):
+        fname = path.join(self.tempdir, "sino500.h5")
+        writer = NXProcessWriter(fname, entry="entry0000")
+        writer.write(self.data, "test_write_simple")
+
+    def test_write_with_config(self):
+        fname = path.join(self.tempdir, "sino500_cfg.h5")
+        writer = NXProcessWriter(fname, entry="entry0000")
+        writer.write(self.data, "test_write_with_config", config=self.h5_config)
+
+    def test_overwrite(self):
+        fname = path.join(self.tempdir, "sino500_overwrite.h5")
+        writer = NXProcessWriter(fname, entry="entry0000", overwrite=True)
+        writer.write(self.data, "test_overwrite", config=self.h5_config)
+
+        writer2 = NXProcessWriter(fname, entry="entry0001", overwrite=True)
+        writer2.write(self.data, "test_overwrite", config=self.h5_config)
+
+        # overwrite entry0000
+        writer3 = NXProcessWriter(fname, entry="entry0000", overwrite=True)
+        new_data = self.data.copy()
+        new_data += 1
+        new_config = self.h5_config.copy()
+        new_config["key1"] = "modified value"
+        writer3.write(new_data, "test_overwrite", config=new_config)
+
+        res = import_h5_to_dict(fname, "/")
+        assert "entry0000" in res
+        assert "entry0001" in res
+        assert np.allclose(res["entry0000"]["test_overwrite"]["results"]["data"], self.data + 1)
+        rec_cfg = res["entry0000"]["test_overwrite"]["configuration"]
+        assert rec_cfg["key1"] == "modified value"
+
+    def test_no_overwrite(self):
+        fname = path.join(self.tempdir, "sino500_no_overwrite.h5")
+        writer = NXProcessWriter(fname, entry="entry0000", overwrite=False)
+        writer.write(self.data, "test_no_overwrite")
+
+        writer2 = NXProcessWriter(fname, entry="entry0000", overwrite=False)
+        with pytest.raises((RuntimeError, OSError)) as ex:
+            writer2.write(self.data, "test_no_overwrite")
+
+        message = "Error should have been raised for trying to overwrite, but got the following: %s" % str(ex.value)

build/lib/build/lib/build/lib/nabu/io/tiffwriter_zmm.py

@@ -0,0 +1,99 @@
+from . import *
+from .writer import TIFFWriter as StandardTIFFWriter
+from os import path
+from tifffile import TiffWriter
+import numpy as np
+
+
+class TIFFWriter(StandardTIFFWriter):  # pylint: disable=E0102
+    def __init__(
+        self,
+        fname,
+        multiframe=False,
+        start_index=0,
+        heights_above_stage_mm=None,
+        filemode=None,
+        append=False,
+        big_tiff=None,
+    ):
+        """
+        Tiff writer.
+
+        Parameters
+        -----------
+        fname: str
+            Path to the output file name
+        multiframe: bool, optional
+            Whether to write all data in one single file. Default is False.
+        start_index: int, optional
+            When writing a stack of images, each image is written in a dedicated file
+            (unless multiframe is set to True).
+            In this case, the output is a series of files `filename_0000.tif`,
+            `filename_0001.tif`, etc. This parameter is the starting index for
+            file names.
+            This option is ignored when multiframe is True.
+        heights_above_stage_mm: None or a list of heights
+            if this parameters is given, the file names will be indexed with the height
+        filemode: str, optional
+            DEPRECATED. Will be ignored. Please refer to 'append'
+        append: bool, optional
+            Whether to append data to the file rather than overwriting. Default is False.
+        big_tiff: bool, optional
+            Whether to write in "big tiff" format: https://www.awaresystems.be/imaging/tiff/bigtiff.html
+            Default is True when multiframe is True.
+            Note that default "standard" tiff cannot exceed 4 GB.
+
+        Notes
+        ------
+        If multiframe is False (default), then each image will be written in a
+        dedicated tiff file.
+        """
+        super().__init__(
+            fname, multiframe=multiframe, start_index=start_index, filemode=filemode, append=append, big_tiff=big_tiff
+        )
+        self.heights_above_stage_mm = heights_above_stage_mm
+
+    def _write_tiff(self, data, config=None, filename=None):
+        # TODO metadata
+        filename = filename or self.fname
+        with TiffWriter(filename, bigtiff=self.big_tiff, append=self.append) as tif:
+            tif.write(data)
+
+    def write(self, data, *args, config=None, **kwargs):
+        # Single image, or multiple image in the same file
+        if self.multiframe:
+            self._write_tiff(data, config=config)
+        # Multiple image, one file per image
+        else:
+            if len(data.shape) == 2:
+                data = np.array([data])
+            dirname, rel_filename = path.split(self.fname)
+            prefix, ext = path.splitext(rel_filename)
+            for i in range(data.shape[0]):
+                if self.heights_above_stage_mm is None:
+                    curr_rel_filename = prefix + str("_%06d" % (self.start_index + i)) + ext
+                else:
+                    value_mm = self.heights_above_stage_mm[i]
+                    if value_mm < 0:
+                        sign = "-"
+                        value_mm = -value_mm
+                    else:
+                        sign = ""
+
+                    part_mm = int(value_mm)
+                    rest_um = (value_mm - part_mm) * 1000
+                    part_um = int(rest_um)
+                    rest_nm = (rest_um - part_um) * 10
+                    part_nm = int(rest_nm)
+
+                    curr_rel_filename = prefix + "_{}{:04d}p{:03d}{:1d}".format(sign, part_mm, part_um, part_nm) + ext
+
+                fname = path.join(dirname, curr_rel_filename)
+
+                self._write_tiff(data[i], filename=fname, config=None)
+
+    def get_filename(self):
+        if self.multiframe:
+            return self.fname
+        else:
+            return path.dirname(self.fname)

build/lib/build/lib/build/lib/nabu/io/utils.py

@@ -0,0 +1,263 @@
+import os
+from typing import Optional
+import contextlib
+import h5py
+import numpy as np
+from silx.io.url import DataUrl
+from tomoscan.volumebase import VolumeBase
+from tomoscan.esrf import EDFVolume, HDF5Volume, TIFFVolume, JP2KVolume, MultiTIFFVolume
+from tomoscan.io import HDF5File
+
+
+# This function might be moved elsewhere
+def get_compacted_dataslices(urls, subsampling=None, begin=0):
+    """
+    Regroup urls to get the data more efficiently.
+    Build a structure mapping files indices to information on
+    how to load the data: `{indices_set: data_location}`
+    where `data_location` contains contiguous indices.
+
+    Parameters
+    -----------
+    urls: dict
+        Dictionary where the key is an integer and the value is a silx `DataUrl`.
+    subsampling: int, optional
+        Subsampling factor when reading the frames. If an integer `n` is provided,
+        then one frame out of `n` will be read.
+
+    Returns
+    --------
+    merged_urls: dict
+        Dictionary with the same keys as the `urls` parameter, and where the
+        values are the corresponding `silx.io.url.DataUrl` with merged data_slice.
+    """
+    subsampling = subsampling or 1
+
+    def _convert_to_slice(idx):
+        if np.isscalar(idx):
+            return slice(idx, idx + 1)
+        # otherwise, assume already slice object
+        return idx
+
+    def is_contiguous_slice(slice1, slice2, step=1):
+        if np.isscalar(slice1):
+            slice1 = slice(slice1, slice1 + step)
+        if np.isscalar(slice2):
+            slice2 = slice(slice2, slice2 + step)
+        return slice2.start == slice1.stop
+
+    def merge_slices(slice1, slice2, step=1):
+        return slice(slice1.start, slice2.stop, step)
+
+    if len(urls) == 0:
+        return urls
+
+    sorted_files_indices = sorted(urls.keys())
+    # if begin > 0:
+    # sorted_files_indices = sorted_files_indices[begin:]
+    idx0 = sorted_files_indices[begin]
+    first_url = urls[idx0]
+
+    merged_indices = [[idx0]]
+    # location = (file_path, data_path, slice)
+    data_location = [[first_url.file_path(), first_url.data_path(), _convert_to_slice(first_url.data_slice())]]
+    pos = 0
+    curr_fp, curr_dp, curr_slice = data_location[pos]
+    skip_next = 0
+    for idx in sorted_files_indices[begin + 1 :]:
+        if skip_next > 1:
+            skip_next -= 1
+            continue
+        url = urls[idx]
+        next_slice = _convert_to_slice(url.data_slice())
+        if (
+            (url.file_path() == curr_fp)
+            and (url.data_path() == curr_dp)
+            and is_contiguous_slice(curr_slice, next_slice, step=subsampling)
+        ):
+            merged_indices[pos].append(idx)
+            merged_slices = merge_slices(curr_slice, next_slice, step=subsampling)
+            data_location[pos][-1] = merged_slices
+            curr_slice = merged_slices
+            skip_next = 0
+        else:  # "jump"
+            if begin > 0 and skip_next == 0:
+                # Skip the "begin" next urls (first of a new block)
+                skip_next = begin
+                continue
+            pos += 1
+            merged_indices.append([idx])
+            data_location.append([url.file_path(), url.data_path(), _convert_to_slice(url.data_slice())])
+            curr_fp, curr_dp, curr_slice = data_location[pos]
+
+    # Format result
+    res = {}
+    for ind, dl in zip(merged_indices, data_location):
+        res.update(dict.fromkeys(ind, DataUrl(file_path=dl[0], data_path=dl[1], data_slice=dl[2])))
+
+    return res
+
+
+def get_first_hdf5_entry(fname):
+    with HDF5File(fname, "r") as fid:
+        entry = list(fid.keys())[0]
+    return entry
+
+
+def hdf5_entry_exists(fname, entry):
+    with HDF5File(fname, "r") as fid:
+        res = fid.get(entry, None) is not None
+    return res
+
+
+def get_h5_value(fname, h5_path, default_ret=None):
+    with HDF5File(fname, "r") as fid:
+        try:
+            val_ptr = fid[h5_path][()]
+        except KeyError:
+            val_ptr = default_ret
+    return val_ptr
+
+
+def get_h5_str_value(dataset_ptr):
+    """
+    Get a HDF5 field which can be bytes or str (depending on h5py version !).
+    """
+    data = dataset_ptr[()]
+    if isinstance(data, str):
+        return data
+    else:
+        return bytes.decode(data)
+
+
+def create_dict_of_indices(images_stack, images_indices):
+    """
+    From an image stack with the images indices, create a dictionary where
+    each index is the image index, and the value is the corresponding image.
+
+    Parameters
+    ----------
+    images_stack: numpy.ndarray
+        A 3D numpy array in the layout (n_images, n_y, n_x)
+    images_indices: array or list of int
+        Array containing the indices of images in the stack
+
+    Examples
+    --------
+    Given a simple array stack:
+
+    >>> images_stack = np.arange(3*4*5).reshape((3,4,5))
+    ... images_indices = [2, 7, 1]
+    ... create_dict_of_indices(images_stack, images_indices)
+    ... # returns {2: array1, 7: array2, 1: array3}
+    """
+    if images_stack.ndim != 3:
+        raise ValueError("Expected a 3D array")
+    if len(images_indices) != images_stack.shape[0]:
+        raise ValueError("images_stack must have as many images as the length of images_indices")
+    res = {}
+    for i in range(len(images_indices)):
+        res[images_indices[i]] = images_stack[i]
+    return res
+
+
+def convert_dict_values(dic, val_replacements, bytes_tostring=False):
+    """
+    Modify a dictionary to be able to export it with silx.io.dicttoh5
+    """
+    modified_dic = {}
+    for key, value in dic.items():
+        if isinstance(key, int):  # np.isscalar ?
+            key = str(key)
+        if isinstance(value, bytes) and bytes_tostring:
+            value = bytes.decode(value.tostring())
+        if isinstance(value, dict):
+            value = convert_dict_values(value, val_replacements, bytes_tostring=bytes_tostring)
+        else:
+            if isinstance(value, DataUrl):
+                value = value.path()
+            elif value.__hash__ is not None and value in val_replacements:
+                value = val_replacements[value]
+        modified_dic[key] = value
+    return modified_dic
+
+
+class _BaseReader(contextlib.AbstractContextManager):
+    def __init__(self, url: DataUrl):
+        if not isinstance(url, DataUrl):
+            raise TypeError("url should be an instance of DataUrl")
+        if url.scheme() not in ("silx", "h5py"):
+            raise ValueError("Valid scheme are silx and h5py")
+        if url.data_slice() is not None:
+            raise ValueError("Data slices are not managed. Data path should " "point to a bliss node (h5py.Group)")
+        self._url = url
+        self._file_handler = None
+
+    def __exit__(self, *exc):
+        return self._file_handler.close()
+
+
+class EntryReader(_BaseReader):
+    """Context manager used to read a bliss node"""
+
+    def __enter__(self):
+        self._file_handler = HDF5File(filename=self._url.file_path(), mode="r")
+        if self._url.data_path() == "":
+            entry = self._file_handler
+        else:
+            entry = self._file_handler[self._url.data_path()]
+        if not isinstance(entry, h5py.Group):
+            raise ValueError("Data path should point to a bliss node (h5py.Group)")
+        return entry
+
+
+class DatasetReader(_BaseReader):
+    """Context manager used to read a bliss node"""
+
+    def __enter__(self):
+        self._file_handler = HDF5File(filename=self._url.file_path(), mode="r")
+        entry = self._file_handler[self._url.data_path()]
+        if not isinstance(entry, h5py.Dataset):
+            raise ValueError("Data path ({}) should point to a dataset (h5py.Dataset)".format(self._url.path()))
+        return entry
+
+
+# TODO: require some utils function to deduce type. And insure homogeneity. Might be moved in tomoscan ?
+def file_format_is_edf(file_format: str):
+    return file_format.lower().lstrip(".") == "edf"
+
+
+def file_format_is_jp2k(file_format: str):
+    return file_format.lower().lstrip(".") in ("jp2k", "jp2")
+
+
+def file_format_is_tiff(file_format: str):
+    return file_format.lower().lstrip(".") in ("tiff", "tif")
+
+
+def file_format_is_hdf5(file_format: str):
+    return file_format.lower().lstrip(".") in ("hdf5", "hdf", "nx", "nexus")
+
+
+def get_output_volume(location: str, file_prefix: Optional[str], file_format: str, multitiff=False) -> VolumeBase:
+    # TODO: see strategy. what if user provide a .nx ... ?
+    # this function should be more generic
+    location, extension = os.path.splitext(location)
+    if extension == "":
+        extension = file_format
+    if file_format_is_edf(extension):
+        return EDFVolume(folder=location, volume_basename=file_prefix)
+    elif file_format_is_jp2k(extension):
+        return JP2KVolume(folder=location, volume_basename=file_prefix)
+    elif file_format_is_hdf5(file_format=extension):
+        if extension is None:
+            if file_prefix is None:
+                location = ".".join([location, extension])
+            else:
+                location = os.path.join(location, ".".join([file_prefix, extension]))
+        return HDF5Volume(file_path=location)
+    elif file_format_is_tiff(extension):
+        if multitiff:
+            return MultiTIFFVolume(file_path=location)
+        else:
+            return TIFFVolume(folder=location, volume_basename=file_prefix)