careamics 0.0.19__py3-none-any.whl

careamics/dataset/dataset_utils/dataset_utils.py

@@ -0,0 +1,118 @@
+"""Dataset utilities."""
+
+import numpy as np
+
+from careamics.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+def get_axes_order(axes_in, ref_axes="STCZYX"):
+    """
+    Get the order of axes based on a reference axes string.
+
+    Parameters
+    ----------
+    axes_in : str
+        Input axes string.
+    ref_axes : str
+        Reference axes string.
+
+    Returns
+    -------
+    list[int]
+        Indices of axes in the reference axes order.
+    """
+    indices = [axes_in.find(k) for k in ref_axes]
+    # remove all non-existing axes (index == -1)
+    new_indices = list(filter(lambda k: k != -1, indices))
+    return new_indices
+
+
+def _get_shape_order(
+    shape_in: tuple[int, ...], axes_in: str, ref_axes: str = "STCZYX"
+) -> tuple[tuple[int, ...], str, list[int]]:
+    """
+    Compute a new shape for the array based on the reference axes.
+
+    Parameters
+    ----------
+    shape_in : tuple[int, ...]
+        Input shape.
+    axes_in : str
+        Input axes.
+    ref_axes : str
+        Reference axes.
+
+    Returns
+    -------
+    tuple[tuple[int, ...], str, list[int]]
+        New shape, new axes, indices of axes in the new axes order.
+    """
+    new_indices = get_axes_order(axes_in, ref_axes)
+
+    # find axes order and get new shape
+    new_axes = [axes_in[ind] for ind in new_indices]
+    new_shape = tuple([shape_in[ind] for ind in new_indices])
+
+    return new_shape, "".join(new_axes), new_indices
+
+
+def reshape_array(x: np.ndarray, axes: str) -> np.ndarray:
+    """Reshape the data to (S, C, (Z), Y, X) by moving axes.
+
+    If the data has both S and T axes, the two axes will be merged. A singleton
+    dimension is added if there are no C axis.
+
+    Parameters
+    ----------
+    x : np.ndarray
+        Input array.
+    axes : str
+        Description of current axes in format `STCZYX`.
+
+    Returns
+    -------
+    np.ndarray
+        Reshaped array with shape (S, C, (Z), Y, X).
+    """
+    _x = x
+    _axes = axes
+
+    # sanity checks
+    if len(_axes) != len(_x.shape):
+        raise ValueError(
+            f"Incompatible data shape ({_x.shape}) and axes ({_axes}). Are the axes "
+            f"correct?"
+        )
+
+    # get new x shape
+    new_x_shape, new_axes, indices = _get_shape_order(_x.shape, _axes)
+
+    # if S is not in the list of axes, then add a singleton S
+    if "S" not in new_axes:
+        new_axes = "S" + new_axes
+        _x = _x[np.newaxis, ...]
+        new_x_shape = (1,) + new_x_shape
+
+        # need to change the array of indices
+        indices = [0] + [1 + i for i in indices]
+
+    # reshape by moving axes
+    destination = list(range(len(indices)))
+    _x = np.moveaxis(_x, indices, destination)
+
+    # remove T if necessary
+    if "T" in new_axes:
+        new_x_shape = (-1,) + new_x_shape[2:]  # remove T and S
+        new_axes = new_axes.replace("T", "")
+
+        # reshape S and T together
+        _x = _x.reshape(new_x_shape)
+
+    # add channel
+    if "C" not in new_axes:
+        # Add channel axis after S
+        _x = np.expand_dims(_x, new_axes.index("S") + 1)
+
+    return _x

careamics/dataset/dataset_utils/file_utils.py

@@ -0,0 +1,141 @@
+"""File utilities."""
+
+from fnmatch import fnmatch
+from pathlib import Path
+from typing import Union
+
+import numpy as np
+
+from careamics.config.support import SupportedData
+from careamics.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+def get_files_size(files: list[Path]) -> float:
+    """Get files size in MB.
+
+    Parameters
+    ----------
+    files : list of pathlib.Path
+        List of files.
+
+    Returns
+    -------
+    float
+        Total size of the files in MB.
+    """
+    return np.sum([f.stat().st_size / 1024**2 for f in files])
+
+
+def list_files(
+    data_path: Union[str, Path],
+    data_type: Union[str, SupportedData],
+    extension_filter: str = "",
+) -> list[Path]:
+    """List recursively files in `data_path` and return a sorted list.
+
+    If `data_path` is a file, its name is validated against the `data_type` using
+    `fnmatch`, and the method returns `data_path` itself.
+
+    By default, if `data_type` is equal to `custom`, all files will be listed. To
+    further filter the files, use `extension_filter`.
+
+    `extension_filter` must be compatible with `fnmatch` and `Path.rglob`, e.g. "*.npy"
+    or "*.czi".
+
+    Parameters
+    ----------
+    data_path : Union[str, Path]
+        Path to the folder containing the data.
+    data_type : Union[str, SupportedData]
+        One of the supported data type (e.g. tif, custom).
+    extension_filter : str, optional
+        Extension filter, by default "".
+
+    Returns
+    -------
+    list[Path]
+        list of pathlib.Path objects.
+
+    Raises
+    ------
+    FileNotFoundError
+        If the data path does not exist.
+    ValueError
+        If the data path is empty or no files with the extension were found.
+    ValueError
+        If the file does not match the requested extension.
+    """
+    # convert to Path
+    data_path = Path(data_path)
+
+    # raise error if does not exists
+    if not data_path.exists():
+        raise FileNotFoundError(f"Data path {data_path} does not exist.")
+
+    # get extension compatible with fnmatch and rglob search
+    extension = SupportedData.get_extension_pattern(data_type)
+
+    if data_type == SupportedData.CUSTOM and extension_filter != "":
+        extension = extension_filter
+
+    # search recurively
+    if data_path.is_dir() and data_path.suffix != ".zarr":
+        # search recursively the path for files with the extension
+        files = sorted(data_path.rglob(extension))
+    else:
+        # raise error if it has the wrong extension
+        if not fnmatch(str(data_path.absolute()), extension):
+            raise ValueError(
+                f"File {data_path} does not match the requested extension "
+                f'"{extension}".'
+            )
+
+        # save in list
+        files = [data_path]
+
+    # raise error if no files were found
+    if len(files) == 0:
+        raise ValueError(
+            f'Data path {data_path} is empty or files with extension "{extension}" '
+            f"were not found."
+        )
+
+    return files
+
+
+def validate_source_target_files(src_files: list[Path], tar_files: list[Path]) -> None:
+    """
+    Validate source and target path lists.
+
+    The two lists should have the same number of files, and the filenames should match.
+
+    Parameters
+    ----------
+    src_files : list of pathlib.Path
+        List of source files.
+    tar_files : list of pathlib.Path
+        List of target files.
+
+    Raises
+    ------
+    ValueError
+        If the number of files in source and target folders is not the same.
+    ValueError
+        If some filenames in Train and target folders are not the same.
+    """
+    # check equal length
+    if len(src_files) != len(tar_files):
+        raise ValueError(
+            f"The number of source files ({len(src_files)}) is not equal to the number "
+            f"of target files ({len(tar_files)})."
+        )
+
+    # check identical names
+    src_names = {f.name for f in src_files}
+    tar_names = {f.name for f in tar_files}
+    difference = src_names.symmetric_difference(tar_names)
+
+    if len(difference) > 0:
+        raise ValueError(f"Source and target files have different names: {difference}.")

careamics/dataset/dataset_utils/iterate_over_files.py

@@ -0,0 +1,84 @@
+"""Function to iterate over files."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Generator
+from pathlib import Path
+from typing import Union
+
+from numpy.typing import NDArray
+from torch.utils.data import get_worker_info
+
+from careamics.config import DataConfig, InferenceConfig
+from careamics.file_io.read import read_tiff
+from careamics.utils.logging import get_logger
+
+from .dataset_utils import reshape_array
+
+logger = get_logger(__name__)
+
+
+def iterate_over_files(
+    data_config: Union[DataConfig, InferenceConfig],
+    data_files: list[Path],
+    target_files: list[Path] | None = None,
+    read_source_func: Callable = read_tiff,
+) -> Generator[tuple[NDArray, NDArray | None], None, None]:
+    """Iterate over data source and yield whole reshaped images.
+
+    Parameters
+    ----------
+    data_config : CAREamics DataConfig or InferenceConfig
+        Configuration.
+    data_files : list of pathlib.Path
+        List of data files.
+    target_files : list of pathlib.Path, optional
+        List of target files, by default None.
+    read_source_func : Callable, optional
+        Function to read the source, by default read_tiff.
+
+    Yields
+    ------
+    NDArray
+        Image.
+    """
+    # When num_workers > 0, each worker process will have a different copy of the
+    # dataset object
+    # Configuring each copy independently to avoid having duplicate data returned
+    # from the workers
+    worker_info = get_worker_info()
+    worker_id = worker_info.id if worker_info is not None else 0
+    num_workers = worker_info.num_workers if worker_info is not None else 1
+
+    # iterate over the files
+    for i, filename in enumerate(data_files):
+        # retrieve file corresponding to the worker id
+        if i % num_workers == worker_id:
+            try:
+                # read data
+                sample = read_source_func(filename, data_config.axes)
+
+                # reshape array
+                reshaped_sample = reshape_array(sample, data_config.axes)
+
+                # read target, if available
+                if target_files is not None:
+                    if filename.name != target_files[i].name:
+                        raise ValueError(
+                            f"File {filename} does not match target file "
+                            f"{target_files[i]}. Have you passed sorted "
+                            f"arrays?"
+                        )
+
+                    # read target
+                    target = read_source_func(target_files[i], data_config.axes)
+
+                    # reshape target
+                    reshaped_target = reshape_array(target, data_config.axes)
+
+                    yield reshaped_sample, reshaped_target
+                else:
+                    yield reshaped_sample, None
+
+            except Exception as e:
+                logger.error(f"Error reading file {filename}: {e}")

careamics/dataset/dataset_utils/running_stats.py

@@ -0,0 +1,189 @@
+"""Computing data statistics."""
+
+import numpy as np
+from numpy.typing import NDArray
+
+
+def compute_normalization_stats(image: NDArray) -> tuple[NDArray, NDArray]:
+    """
+    Compute mean and standard deviation of an array.
+
+    Expected input shape is (S, C, (Z), Y, X). The mean and standard deviation are
+    computed per channel.
+
+    Parameters
+    ----------
+    image : NDArray
+        Input array.
+
+    Returns
+    -------
+    tuple of (list of floats, list of floats)
+        Lists of mean and standard deviation values per channel.
+    """
+    # Define the lists for storing mean and std values
+    means, stds = [], []
+    # Iterate over the channels dimension and compute mean and std
+    for ax in range(image.shape[1]):
+        means.append(image[:, ax, ...].mean())
+        stds.append(image[:, ax, ...].std())
+    return np.stack(means), np.stack(stds)
+
+
+def update_iterative_stats(
+    count: NDArray, mean: NDArray, m2: NDArray, new_values: NDArray
+) -> tuple[NDArray, NDArray, NDArray]:
+    """Update the mean and variance of an array iteratively.
+
+    Parameters
+    ----------
+    count : NDArray
+        Number of elements in the array. Shape: (C,).
+    mean : NDArray
+        Mean of the array. Shape: (C,).
+    m2 : NDArray
+        Variance of the array. Shape: (C,).
+    new_values : NDArray
+        New values to add to the mean and variance. Shape: (C, 1, 1, Z, Y, X).
+
+    Returns
+    -------
+    tuple[NDArray, NDArray, NDArray]
+        Updated count, mean, and variance.
+    """
+    num_channels = len(new_values)
+
+    # --- update channel-wise counts ---
+    count += np.ones_like(count) * np.prod(new_values.shape[1:])
+
+    # --- update channel-wise mean ---
+    # compute (new_values - old_mean) -> shape: (C, Z*Y*X)
+    delta = new_values.reshape(num_channels, -1) - mean.reshape(num_channels, 1)
+    mean += np.sum(delta / count.reshape(num_channels, 1), axis=1)
+
+    # --- update channel-wise SoS ---
+    # compute (new_values - new_mean) -> shape: (C, Z*Y*X)
+    delta2 = new_values.reshape(num_channels, -1) - mean.reshape(num_channels, 1)
+    m2 += np.sum(delta * delta2, axis=1)
+
+    return count, mean, m2
+
+
+def finalize_iterative_stats(
+    count: NDArray, mean: NDArray, m2: NDArray
+) -> tuple[NDArray, NDArray]:
+    """Finalize the mean and variance computation.
+
+    Parameters
+    ----------
+    count : NDArray
+        Number of elements in the array. Shape: (C,).
+    mean : NDArray
+        Mean of the array. Shape: (C,).
+    m2 : NDArray
+        Variance of the array. Shape: (C,).
+
+    Returns
+    -------
+    tuple[NDArray, NDArray]
+        Final channel-wise mean and standard deviation.
+    """
+    std = np.sqrt(m2 / count)
+    if any(c < 2 for c in count):
+        return np.full(mean.shape, np.nan), np.full(std.shape, np.nan)
+    else:
+        return mean, std
+
+
+class WelfordStatistics:
+    """Compute Welford statistics iteratively.
+
+    The Welford algorithm is used to compute the mean and variance of an array
+    iteratively. Based on the implementation from:
+    https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford's_online_algorithm
+    """
+
+    def update(self, array: NDArray, sample_idx: int) -> None:
+        """Update the Welford statistics.
+
+        Parameters
+        ----------
+        array : NDArray
+            Input array.
+        sample_idx : int
+            Current sample number.
+        """
+        self.sample_idx = sample_idx
+        sample_channels = np.array(np.split(array, array.shape[1], axis=1))
+
+        # Initialize the statistics
+        if self.sample_idx == 0:
+            # Compute the mean and standard deviation
+            self.mean, _ = compute_normalization_stats(array)
+            # Initialize the count and m2 with zero-valued arrays of shape (C,)
+            self.count, self.mean, self.m2 = update_iterative_stats(
+                count=np.zeros(array.shape[1]),
+                mean=self.mean,
+                m2=np.zeros(array.shape[1]),
+                new_values=sample_channels,
+            )
+        else:
+            # Update the statistics
+            self.count, self.mean, self.m2 = update_iterative_stats(
+                count=self.count, mean=self.mean, m2=self.m2, new_values=sample_channels
+            )
+
+        self.sample_idx += 1
+
+    def finalize(self) -> tuple[NDArray, NDArray]:
+        """Finalize the Welford statistics.
+
+        Returns
+        -------
+        tuple or numpy arrays
+            Final mean and standard deviation.
+        """
+        return finalize_iterative_stats(self.count, self.mean, self.m2)
+
+
+# from multiprocessing import Value
+# from typing import tuple
+
+# import numpy as np
+
+
+# class RunningStats:
+#     """Calculates running mean and std."""
+
+#     def __init__(self) -> None:
+#         self.reset()
+
+#     def reset(self) -> None:
+#         """Reset the running stats."""
+#         self.avg_mean = Value("d", 0)
+#         self.avg_std = Value("d", 0)
+#         self.m2 = Value("d", 0)
+#         self.count = Value("i", 0)
+
+#     def init(self, mean: float, std: float) -> None:
+#         """Initialize running stats."""
+#         with self.avg_mean.get_lock():
+#             self.avg_mean.value += mean
+#         with self.avg_std.get_lock():
+#             self.avg_std.value = std
+
+#     def compute_std(self) -> tuple[float, float]:
+#         """Compute std."""
+#         if self.count.value >= 2:
+#             self.avg_std.value = np.sqrt(self.m2.value / self.count.value)
+
+#     def update(self, value: float) -> None:
+#         """Update running stats."""
+#         with self.count.get_lock():
+#             self.count.value += 1
+#         delta = value - self.avg_mean.value
+#         with self.avg_mean.get_lock():
+#             self.avg_mean.value += delta / self.count.value
+#         delta2 = value - self.avg_mean.value
+#         with self.m2.get_lock():
+#             self.m2.value += delta * delta2