adding prep

httomolibgpu/misc/rescale.py

@@ -19,18 +19,6 @@
     "rescale_to_int",
 ]
 
-rescale_kernel = xp.ElementwiseKernel(
-    "T x, raw T input_min, raw T input_max, raw T factor",
-    "O out",
-    """
-      T x_clean = isnan(x) || isinf(x) ? T(0) : x;
-      T x_clipped = x_clean < input_min ? input_min : (x_clean > input_max ? input_max : x_clean);
-      T x_rebased = x_clipped - input_min;
-      out = O(x_rebased * factor);
-    """,
-    "rescale_to_int",
-)
-
 
 @nvtx.annotate()
 def rescale_to_int(
@@ -97,5 +85,16 @@ def rescale_to_int(
     factor = (output_max - output_min) / (input_max - input_min)
 
     res = xp.empty(data.shape, dtype=output_dtype)
+    rescale_kernel = xp.ElementwiseKernel(
+        "T x, raw T input_min, raw T input_max, raw T factor",
+        "O out",
+        """
+        T x_clean = isnan(x) || isinf(x) ? T(0) : x;
+        T x_clipped = x_clean < input_min ? input_min : (x_clean > input_max ? input_max : x_clean);
+        T x_rebased = x_clipped - input_min;
+        out = O(x_rebased * factor);
+        """,
+        "rescale_to_int",
+    )
     rescale_kernel(data, input_min, input_max, factor, res)
     return res

httomolibgpu/prep/alignment.py

@@ -0,0 +1,175 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# ---------------------------------------------------------------------------
+# Copyright 2022 Diamond Light Source Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ---------------------------------------------------------------------------
+# Created By  : Tomography Team at DLS <[email protected]>
+# Created Date: 01 November 2022
+# ---------------------------------------------------------------------------
+"""Modules for data correction"""
+
+import numpy as xp
+import numpy as np
+
+try:
+    import cupy as xp
+    from cupy import mean
+
+    try:
+        xp.cuda.Device(0).compute_capability
+    except xp.cuda.runtime.CUDARuntimeError:
+        print("CuPy library is a major dependency for HTTomolibgpu, please install")
+        import numpy as np
+except ImportError:
+    import numpy as np
+
+from typing import Dict, List
+import nvtx
+
+__all__ = [
+    "distortion_correction_proj_discorpy",
+]
+
+
+# CuPy implementation of distortion correction from Discorpy
+# https://github.com/DiamondLightSource/discorpy/blob/67743842b60bf5dd45b21b8460e369d4a5e94d67/discorpy/post/postprocessing.py#L111-L148
+# (which is the same as the TomoPy version
+# https://github.com/tomopy/tomopy/blob/c236a2969074f5fc70189fb5545f0a165924f916/source/tomopy/prep/alignment.py#L950-L981
+# but with the additional params `order` and `mode`).
+@nvtx.annotate()
+def distortion_correction_proj_discorpy(
+    data: xp.ndarray,
+    metadata_path: str,
+    preview: Dict[str, List[int]],
+    order: int = 1,
+    mode: str = "reflect",
+):
+    """Unwarp a stack of images using a backward model.
+
+    Parameters
+    ----------
+    data : cp.ndarray
+        3D array.
+
+    metadata_path : str
+        The path to the file containing the distortion coefficients for the
+        data.
+
+    preview : Dict[str, List[int]]
+        A dict containing three key-value pairs:
+        - a list containing the `start` value of each dimension
+        - a list containing the `stop` value of each dimension
+        - a list containing the `step` value of each dimension
+
+    order : int, optional.
+        The order of the spline interpolation.
+
+    mode : {'reflect', 'grid-mirror', 'constant', 'grid-constant', 'nearest',
+           'mirror', 'grid-wrap', 'wrap'}, optional
+        To determine how to handle image boundaries.
+
+    Returns
+    -------
+    cp.ndarray
+        3D array. Distortion-corrected image(s).
+    """
+    from cupyx.scipy.ndimage import map_coordinates
+
+    # Check if it's a stack of 2D images, or only a single 2D image
+    if len(data.shape) == 2:
+        data = xp.expand_dims(data, axis=0)
+
+    # Get info from metadata txt file
+    xcenter, ycenter, list_fact = _load_metadata_txt(metadata_path)
+
+    # Use preview information to offset the x and y coords of the center of
+    # distortion
+    shift = preview["starts"]
+    step = preview["steps"]
+    x_dim = 1
+    y_dim = 0
+    step_check = max([step[i] for i in [x_dim, y_dim]]) > 1
+    if step_check:
+        msg = (
+            "\n***********************************************\n"
+            "!!! ERROR !!! -> Method doesn't work with the step in"
+            " the preview larger than 1 \n"
+            "***********************************************\n"
+        )
+        raise ValueError(msg)
+
+    x_offset = shift[x_dim]
+    y_offset = shift[y_dim]
+    xcenter = xcenter - x_offset
+    ycenter = ycenter - y_offset
+
+    height, width = data.shape[y_dim + 1], data.shape[x_dim + 1]
+    xu_list = xp.arange(width) - xcenter
+    yu_list = xp.arange(height) - ycenter
+    xu_mat, yu_mat = xp.meshgrid(xu_list, yu_list)
+    ru_mat = xp.sqrt(xu_mat**2 + yu_mat**2)
+    fact_mat = xp.sum(
+        xp.asarray([factor * ru_mat**i for i, factor in enumerate(list_fact)]), axis=0
+    )
+    xd_mat = xp.asarray(
+        xp.clip(xcenter + fact_mat * xu_mat, 0, width - 1), dtype=xp.float32
+    )
+    yd_mat = xp.asarray(
+        xp.clip(ycenter + fact_mat * yu_mat, 0, height - 1), dtype=xp.float32
+    )
+    indices = [xp.reshape(yd_mat, (-1, 1)), xp.reshape(xd_mat, (-1, 1))]
+    indices = xp.asarray(indices, dtype=xp.float32)
+
+    # Loop over images and unwarp them
+    for i in range(data.shape[0]):
+        mat = map_coordinates(data[i], indices, order=order, mode=mode)
+        mat = xp.reshape(mat, (height, width))
+        data[i] = mat
+
+    return data
+
+
+def _load_metadata_txt(file_path):
+    """
+    Load distortion coefficients from a text file.
+    Order of the infor in the text file:
+    xcenter
+    ycenter
+    factor_0
+    factor_1
+    factor_2
+    ...
+    Parameters
+    ----------
+    file_path : str
+        Path to the file
+    Returns
+    -------
+    tuple of float and list of floats
+        Tuple of (xcenter, ycenter, list_fact).
+    """
+    with open(file_path, "r") as f:
+        x = f.read().splitlines()
+        list_data = []
+        for i in x:
+            list_data.append(float(i.split()[-1]))
+    xcenter = list_data[0]
+    ycenter = list_data[1]
+    list_fact = list_data[2:]
+
+    return xcenter, ycenter, list_fact
+
+
+## %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%  ##

httomolibgpu/prep/normalize.py

@@ -0,0 +1,140 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# ---------------------------------------------------------------------------
+# Copyright 2022 Diamond Light Source Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ---------------------------------------------------------------------------
+# Created By  : Tomography Team at DLS <[email protected]>
+# Created Date: 01 November 2022
+# ---------------------------------------------------------------------------
+"""Modules for raw projection data normalization"""
+
+import numpy as xp
+import numpy as np
+
+try:
+    import cupy as xp
+    from cupy import mean
+
+    try:
+        xp.cuda.Device(0).compute_capability
+    except xp.cuda.runtime.CUDARuntimeError:
+        print("CuPy library is a major dependency for HTTomolibgpu, please install")
+        import numpy as np
+except ImportError:
+    import numpy as np
+
+import nvtx
+from numpy import float32
+from typing import Tuple
+
+__all__ = ["normalize"]
+
+
+@nvtx.annotate()
+def normalize(
+    data: xp.ndarray,
+    flats: xp.ndarray,
+    darks: xp.ndarray,
+    cutoff: float = 10.0,
+    minus_log: bool = True,
+    nonnegativity: bool = False,
+    remove_nans: bool = False,
+) -> xp.ndarray:
+    """
+    Normalize raw projection data using the flat and dark field projections.
+    This is a raw CUDA kernel implementation with CuPy wrappers.
+
+    Parameters
+    ----------
+    data : cp.ndarray
+        Projection data as a CuPy array.
+    flats : cp.ndarray
+        3D flat field data as a CuPy array.
+    darks : cp.ndarray
+        3D dark field data as a CuPy array.
+    cutoff : float, optional
+        Permitted maximum value for the normalised data.
+    minus_log : bool, optional
+        Apply negative log to the normalised data.
+    nonnegativity : bool, optional
+        Remove negative values in the normalised data.
+    remove_nans : bool, optional
+        Remove NaN and Inf values in the normalised data.
+
+    Returns
+    -------
+    cp.ndarray
+        Normalised 3D tomographic data as a CuPy array.
+    """
+
+    _check_valid_input(data, flats, darks)
+
+    dark0 = xp.empty(darks.shape[1:], dtype=float32)
+    flat0 = xp.empty(flats.shape[1:], dtype=float32)
+    out = xp.empty(data.shape, dtype=float32)
+    mean(darks, axis=0, dtype=float32, out=dark0)
+    mean(flats, axis=0, dtype=float32, out=flat0)
+
+    kernel_name = "normalisation"
+    kernel = r"""
+        float denom = float(flats) - float(darks);
+        if (denom < eps) {
+            denom = eps;
+        }
+        float v = (float(data) - float(darks))/denom;
+        """
+    if minus_log:
+        kernel += "v = -log(v);\n"
+        kernel_name += "_mlog"
+    if nonnegativity:
+        kernel += "if (v < 0.0f) v = 0.0f;\n"
+        kernel_name += "_nneg"
+    if remove_nans:
+        kernel += "if (isnan(v)) v = 0.0f;\n"
+        kernel += "if (isinf(v)) v = 0.0f;\n"
+        kernel_name += "_remnan"
+    kernel += "if (v > cutoff) v = cutoff;\n"
+    kernel += "out = v;\n"
+
+    normalisation_kernel = xp.ElementwiseKernel(
+        "T data, U flats, U darks, raw float32 cutoff",
+        "float32 out",
+        kernel,
+        kernel_name,
+        options=("-std=c++11",),
+        loop_prep="constexpr float eps = 1.0e-07;",
+        no_return=True,
+    )
+
+    normalisation_kernel(data, flat0, dark0, float32(cutoff), out)
+
+    return out
+
+
+def _check_valid_input(data, flats, darks) -> None:
+    """Helper function to check the validity of inputs to normalisation functions"""
+    if data.ndim != 3:
+        raise ValueError("Input data must be a 3D stack of projections")
+
+    if flats.ndim not in (2, 3):
+        raise ValueError("Input flats must be 2D or 3D data only")
+
+    if darks.ndim not in (2, 3):
+        raise ValueError("Input darks must be 2D or 3D data only")
+
+    if flats.ndim == 2:
+        flats = flats[xp.newaxis, :, :]
+    if darks.ndim == 2:
+        darks = darks[xp.newaxis, :, :]