PSF and attenuation

.gitignore

@@ -171,6 +171,10 @@ workflows/work/
 # Macos
 .DS_Store
 
-# files
+# Data files
 *.zarr
 *.nii.gz
+*.sif
+
+# Config files
+*.ini

Dockerfile

@@ -19,5 +19,5 @@ RUN pip install --upgrade pip setuptools wheel build
 # Install with verbose output
 COPY linumpy ./linumpy
 COPY scripts ./scripts
-COPY pyproject.toml .
+COPY pyproject.toml requirements.txt README.md setup.py ./
 RUN pip install --no-cache-dir -v -e .

linumpy/preproc/icorr.py

@@ -516,6 +516,7 @@ def getSmoothIntensityTransition(vol, slicesStart):
 
 def getAttenuation_Vermeer2013(vol, dz=6.5e-6, mask=None, C=None):
     """Estimates the attenuation coefficient using the Vermeer2013 model.
+
     Parameters
     ----------
     vol : ndarray
@@ -593,6 +594,95 @@ def getAttenuation_Vermeer2013(vol, dz=6.5e-6, mask=None, C=None):
     return mu
 
 
+def get_extendedAttenuation_Vermeer2013(vol, mask=None, k=10, sigma=5,
+                                        sigma_bottom=3, dz=1, res=6.5,
+                                        zshift=3, fillHoles=False):
+    """Compute the local effective tissue attenuation using
+    the extended Vermeer model.
+
+    Parameters
+    ----------
+    vol: ndarray
+        OCT volume to process
+    mask: ndarray
+        Optional tissue mask. If none is given the water/tissue
+        interface will be detected from the data.
+    k: int
+        Median filter kernel size (px) applied before the attenuation
+        coefficient computation (applied in the XY direction).
+    sigma: int
+        Gaussian filter kernel size (px) applied axially before the
+        exponential signal fit used to extend the Alines for the extended
+        Vermeer signal evalution.
+    dz: int
+        Number of axial pixel to consider when computing the bottom slice
+        signal for the signal extension.
+    res: float
+        Axial resolution in micron / pixel
+    zshift: int
+        Number of pixel under the water-tissue interface to ignore while
+        fitting the exponential function for signal extension.
+
+    Returns
+    -------
+    ndarray
+        Computed attenuation coefficients.
+    """
+    # First the slice is denoised with a small median filter
+    if k > 0:
+        vol = sitk.GetArrayFromImage(
+            sitk.Median(sitk.GetImageFromArray(vol), (0, k, k))
+        )
+
+    # Computing tissue mask
+    if mask is None:
+        # Detecting the water / tissue interface
+        interface = xyzcorr.findTissueInterface(
+            vol, s_xy=3, s_z=1, order=1, useLog=True
+        )
+        mask = xyzcorr.maskUnderInterface(vol, interface + zshift, returnMask=True)
+
+    # Lets fit an exponential function on each Aline to extend the tissue slice.
+    exp_fit = get_gradientAttenuation(gaussian_filter(vol, (0, 0, sigma)))
+    exp_fit = np.ma.masked_array(exp_fit, ~mask).mean(axis=2)
+
+    # Fill holes left by NaN values
+    exp_fit[np.isnan(exp_fit)] = 0
+    # exp_fit = gaussian_filter(exp_fit, sigma_bottom);
+
+    # Get the signal at the interface for each Aline
+    interface_bottom = (
+        vol.shape[2] - xyzcorr.getInterfaceDepthFromMask(mask[:, :, ::-1]) - 1 - dz
+    )
+    mask_bottom = xyzcorr.maskUnderInterface(vol, interface_bottom, returnMask=True)
+    mask_bottom = (mask_bottom * mask).astype(bool)
+    i0 = np.ma.masked_array(vol, ~mask_bottom).mean(axis=2)
+    # i0 = gaussian_filter(i0, sigma_bottom)
+
+    # Compute the end-of-scan bias
+    epsilon = 1e-3
+    C = np.zeros_like(i0)
+    C[exp_fit > epsilon] = i0[exp_fit > epsilon] / exp_fit[exp_fit > epsilon]
+    C = gaussian_filter(C, sigma_bottom)
+
+    # Compute the attenuation
+    attn_cropped = getAttenuation_Vermeer2013(vol, dz=res * 1e-06, mask=mask, C=C)
+
+    # Remove NaN
+    attn_cropped[np.isnan(attn_cropped)] = 0
+
+    # Only keep attn within the mask
+    attn_cropped[~mask.astype(bool)] = 0
+
+    # Fill holes
+    if fillHoles:
+        attn_cropped = sitk.GetArrayFromImage(
+            sitk.GrayscaleFillhole(sitk.GetImageFromArray(attn_cropped))
+        )
+
+    return attn_cropped
+
+
 def getAttenuation_Faber2004(vol, mask=None, dz=6.5e-6, N=4):
     """Estimates the attenuation coefficient using the Faber2004 model.
     Parameters

linumpy/psf/psf_estimator.py

@@ -0,0 +1,145 @@
+import numpy as np
+from linumpy.preproc.xyzcorr import findTissueInterface
+from linumpy.preproc.icorr import confocalPSF, fit_TissueConfocalModel
+
+from scipy.ndimage import binary_dilation, binary_fill_holes, gaussian_filter
+from scipy.stats import zscore
+from skimage.filters import threshold_li
+from skimage.morphology import disk
+
+
+# TODO: Fine-tune default values for 10x microscope or give heuristic
+# for fixing them.
+def extract_psfParametersFromMosaic(
+    vol, f=0.01, nProfiles=10, zr_0=610.0, res=6.5, nIterations=15
+):
+    """Computes the confocal PSF from a slice
+
+    Parameters
+    ----------
+    vol : ndarray
+        A stitched tissue slice with axes in order (x, y, z).
+    f : float
+        Smoothing factor (in fraction of image size).
+    nProfiles : int
+        Number of intensity profile to use.
+    zr_0 : float
+        Initial Rayleigh length to use in micron (default=%(default)s for a 3X objective)
+    res : float
+        Z resolution (in micron).
+
+    Returns
+    -------
+    (2,) tuple
+        Focal depth (zf) and Rayleigh length (zr) in micron
+
+    """
+
+    nx, ny, nz = vol.shape
+    k = int(0.5 * f * (nx + ny))
+    aip = vol.mean(axis=2)
+
+    # Compute water-tissue interface
+    interface = findTissueInterface(vol).astype(int)
+
+    # Compute the agarose mask with the li thresholding method
+    thresh = threshold_li(aip)
+    mask_tissue = binary_fill_holes(aip > thresh)
+    mask_agarose = ~binary_fill_holes(binary_dilation(mask_tissue, disk(k)))
+    mask_agarose[aip == 0] = 0
+    del mask_tissue
+
+    # Get min and max interface depth for the agarose
+    zmin = np.percentile(interface[mask_agarose], 2.5)
+
+    # Get the average iProfile / interface depth
+    profilePerInterfaceDepth = np.zeros((nProfiles, nz))
+    for ii in range(nProfiles):
+        for z in range(nz):
+            profilePerInterfaceDepth[ii, z] = np.mean(
+                vol[:, :, z][mask_agarose * (interface == zmin + ii)]
+            )
+
+    # Detect outliers
+    iProfile_gradient = np.abs(
+        gaussian_filter(profilePerInterfaceDepth, sigma=(0, 2), order=1)
+    )
+    profile_mask = np.abs(zscore(iProfile_gradient, axis=1)) <= 1.0
+    for ii in range(nProfiles):
+        profile_mask[ii, 0:int(zmin + ii)] = 0
+
+    z = np.linspace(0, nz * res, nz)
+    zf_list = list()
+    zr_list = list()
+    total_err = list()
+    for z0 in range(nProfiles):
+        # Find the coarse alignment of the focus based on
+        # pre-established Rayleigh length from thorlab
+        errList = list()
+        for zf in range(nz):
+            a = profilePerInterfaceDepth[z0, zf]
+            synthetic_signal = confocalPSF(z, zf, zr_0, a)
+            err = np.abs(synthetic_signal - profilePerInterfaceDepth[z0, :])
+            err = np.mean(err[profile_mask[z0, :]])
+            errList.append(err)
+
+        errList = np.array(errList)
+        zf = np.argmin(errList) * res
+        a = profilePerInterfaceDepth[z0, int(zf / res)]
+
+        if not (np.isnan(a)):
+            last_zr = zr_0
+            for _ in range(nIterations):
+                # Optimize the model (without using attenuation)
+                iProfile = profilePerInterfaceDepth[z0, :]
+                output = fit_TissueConfocalModel(
+                    iProfile,
+                    int(z0 + zmin),
+                    last_zr,
+                    res,
+                    returnParameters=True,
+                    return_fullModel=True,
+                    useBumpModel=True,
+                )
+                zf = output["parameters"]["zf"]
+                zr = output["parameters"]["zr"]
+                last_zr = zr
+
+            zf_list.append(zf)
+            zr_list.append(zr)
+            err_fit = (output["tissue_psf"] - profilePerInterfaceDepth[z0, :]) ** 2.0
+            total_err.append(np.mean(err_fit))
+
+    min_err = np.argmin(total_err)
+    zf_final = zf_list[min_err]
+    zr_final = zr_list[min_err]
+
+    return zf_final, zr_final
+
+
+def get_3dPSF(zf, zr, res, volshape):
+    """Generate a 3D PSF based on Gaussian beam parameters.
+
+    Parameters
+    ----------
+    zf : float
+        Focal depth in microns
+    zr : float
+        Rayleigh length in microns
+    res : float
+        Axial resolution in micron / pixel
+    volshape : (3,) list of int
+        Output volume shape in pixel
+
+    Returns
+    -------
+    ndarray
+        3D PSF of shape 'volshape'
+    """
+    # TODO: Invert axes to agree with OME-zarr convention?
+    nx, ny, nz = volshape[0:3]
+    z = np.linspace(0, res * nz, nz)
+    psf = confocalPSF(z, zf, zr)
+    psf = np.tile(np.reshape(psf, (1, 1, nz)), (nx, ny, 1))
+
+    return psf

scripts/linum_compensate_attenuation.py

@@ -0,0 +1,48 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+Compensate the tissue attenuation using a precomputed attenuation
+bias field.
+"""
+
+import argparse
+from pathlib import Path
+import dask.array as da
+
+from linumpy.io.zarr import save_zarr, read_omezarr
+
+
+def _build_arg_parser():
+    p = argparse.ArgumentParser(
+        description=__doc__, formatter_class=argparse.RawTextHelpFormatter)
+    p.add_argument("input",
+                   help="Input volume (.ome.zarr)")
+    p.add_argument("bias",
+                   help="Attenuation bias field (.ome.zarr)")
+    p.add_argument("output",
+                   help="Compensated volume (.ome.zarr)")
+
+    return p
+
+
+def main():
+    # Parse arguments
+    p = _build_arg_parser()
+    args = p.parse_args()
+
+    # Load volume and bias field
+    vol, res = read_omezarr(args.input, level=0)
+    bias, _ = read_omezarr(args.bias, level=0)
+    chunks = vol.chunks
+
+    # Apply correction
+    vol_dask = da.from_zarr(vol)
+    bias_dask = da.from_zarr(bias)
+    vol_dask /= bias_dask
+
+    # Save the output
+    save_zarr(vol_dask.astype(da.float32), args.output, scales=res, chunks=chunks)
+
+
+if __name__ == "__main__":
+    main()

scripts/linum_compensate_for_psf.py

@@ -0,0 +1,65 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+import argparse
+import numpy as np
+from linumpy.io.zarr import read_omezarr, save_zarr
+from linumpy.psf.psf_estimator import extract_psfParametersFromMosaic, get_3dPSF
+
+
+def _build_arg_parser():
+    p = argparse.ArgumentParser()
+    p.add_argument('in_zarr',
+                   help='Input stitched 3D slice (OME-zarr).')
+    p.add_argument('out_zarr',
+                   help='Output volume corrected for beam PSF (OME-zarr).')
+    p.add_argument('--out_psf',
+                   help='Optional output PSF filename.')
+    p.add_argument('--nz', type=int, default=25,
+                   help='The "nz" first voxels belonging to background [%(default)s].')
+    p.add_argument('--n_profiles', type=int, default=10,
+                   help='Number of intensity profiles to use [%(default)s].')
+    p.add_argument('--n_iterations', type=int, default=15,
+                   help='Number of iterations [%(default)s].')
+    p.add_argument('--smooth', type=float, default=0.01,
+                   help='Smoothing factor as a fraction of volume depth [%(default)s].')
+    return p
+
+
+def main():
+    parser = _build_arg_parser()
+    args = parser.parse_args()
+
+    # 1. load stitched tissue slice
+    vol, res = read_omezarr(args.in_zarr, level=0)
+    chunks = vol.chunks
+    vol = np.moveaxis(vol, (0, 1, 2), (2, 1, 0))
+    res = res[::-1]
+    res_axial_microns = res[2] * 1000
+
+    # 2. estimate psf
+    zf, zr = extract_psfParametersFromMosaic(vol, nProfiles=args.n_profiles,
+                                             res=res_axial_microns, f=args.smooth,
+                                             nIterations=args.n_iterations)
+    psf_3d = get_3dPSF(zf, zr, res_axial_microns, vol.shape)
+
+    # Compensate by the PSF
+    background = np.mean(vol[..., :args.nz])
+    output = (vol - background) / psf_3d + background
+
+    # remove negative values
+    output -= output.min()
+
+    # TODO: Use dask arrays
+    output = np.moveaxis(output, (0, 1, 2), (2, 1, 0))
+    res = res[::-1]
+
+    if args.out_psf:
+        psf_3d = np.moveaxis(psf_3d, (0, 1, 2), (2, 1, 0))
+        # when there are too many levels it'll break the viewer for some reason
+        save_zarr(psf_3d.astype(np.float32), args.out_psf, scales=res, chunks=chunks)
+
+    save_zarr(output.astype(np.float32), args.out_zarr, scales=res, chunks=chunks)
+
+
+if __name__ == '__main__':
+    main()