B6_dapi_template/preprocess.py

from pathlib import Path
import xmltodict
import argparse

import javabridge
import bioformats
import numpy as np
import nibabel as nib
from PIL import Image, ImageSequence
from scipy import ndimage
from skimage import morphology, filters
from nipype.interfaces.ants.segmentation import N4BiasFieldCorrection


def get_out_paths(out_dir, stem):
    out_path_scaled = Path(out_dir, f'{stem}.nii')
    out_path_unscaled = Path(out_dir, f'{stem}_us.nii')
    return out_path_scaled, out_path_unscaled


def bias_field_correction(image, output, dim=2):
    N4BiasFieldCorrection(
        input_image=image,
        output_image=output,
        dimension=dim
    ).run()


def image_data_to_nii(pixdim, image_data, shrink, out_dir, file_path, save_unscaled=False):
    """
    Saves an array of image data as a nifti file, with correct pixel dimensions in
    :param save_unscaled: boolean whether to save unscaled, as well as scaled nifti file
    :param shrink: factor by which to shrink image dimensions by
    :return: location of scaled and possibly unscaled nifti files
    """
    image_dim = (image_data.shape[0], image_data.shape[1])
    scale = 1/shrink
    new_dim = (round(image_dim[1] * scale), round(image_dim[0] * scale))
    new_arr = np.ndarray(new_dim + (image_data.shape[2],))

    for i in range(0,image_data.shape[2]):
        cur_channel = image_data[:, :, i]
        resized = np.array(Image.fromarray(cur_channel).resize(new_dim)).transpose()
        new_arr[:, :, i] = resized
    path_scaled, path_unscaled = get_out_paths(out_dir, Path(file_path).stem)
    nii_scaled = nib.Nifti1Image(new_arr, np.eye(4))
    nii_scaled.header['xyzt_units'] = 3
    nii_scaled.header['pixdim'][1:3] = pixdim * shrink, pixdim * shrink
    nib.save(nii_scaled, str(path_scaled))

    if save_unscaled:
        nii_unscaled = nib.Nifti1Image(image_data, np.eye(4))
        nii_unscaled.header['xyzt_units'] = 3
        nii_unscaled.header['pixdim'][1:3] = pixdim, pixdim
        nib.save(nii_unscaled, str(path_unscaled))
    print(f'Preprocessed:  {path_scaled}\n')
    return path_scaled, path_unscaled


def nd2_to_nii(nd2_path, out_dir, series, shrink=10):
    """
    Wrapper function for image_data_to_nii, for converting nd2 files to nifti
    """
    javabridge.start_vm(bioformats.JARS)
    image = np.array(bioformats.load_image(str(nd2_path), series=series-1, rescale=False))
    meta_dict = xmltodict.parse(bioformats.get_omexml_metadata(str(nd2_path)))
    vox_size_x = float(meta_dict['OME']['Image'][series-1]['Pixels']['@PhysicalSizeX'])
    javabridge.kill_vm()
    return image_data_to_nii(vox_size_x, image, shrink, out_dir, nd2_path)


def tiff_to_nii(tif_path, out_dir, pixdim=None, shrink=10):
    """
    Wrapper function for image_data_to_nii, for converting tiff files to nifti
    """
    tif_image = Image.open(tif_path)
    tif_header = dict(tif_image.tag)
    output = np.empty(np.array(tif_image).shape + (0,))
    if not pixdim:
        pixdim = 10e6/tif_header[282][0][0]
    for i, page in enumerate(ImageSequence.Iterator(tif_image)):
        page_data = np.expand_dims(np.array(page), 2)
        output = np.concatenate((output, page_data), 2)
    return image_data_to_nii(pixdim, output, shrink, out_dir, tif_path)


def split_nii_channels(nii_path, out_dir=None, flip=False, mask_index=-1, bias=False):
    """
    Converts a single multi-channel nifti file to multiple single-channel nifti files, and masks foreground
    :param flip: Whether to vertically flip the image, in order to properly align nifti file with template
    :param mask_index: index of DAPI stained channel, on which to mask other channels on
    :param bias: whether to bias-field correct the image
    :return: Location of multiple single-channel nifti files
    """
    if out_dir is None:
        out_dir = nii_path.parent
    nii = nib.load(str(nii_path))
    nii_data = nii.get_data()
    nii_header = nii.header

    if mask_index == -1:
        mask_index = nii_data.shape[2] - 1
    paths = []

    for i in range(0, nii_data.shape[2]):
        out_path = out_dir / f'im_c{i+1}.nii'
        channel_data = nii_data[:, :, i]

        if flip:
            channel_data = np.flip(channel_data, 1)

        if i == mask_index:
            channel_data = mask_foreground(channel_data)

        new_header = nii_header
        new_header['dim'][0] = 2
        nii = nib.Nifti1Image(channel_data, np.eye(4), header=new_header)
        nib.save(nii, str(out_path))

        if i == mask_index and bias:
            bias_field_correction(str(out_path), str(out_path))
            corrected = nib.load(str(out_path))
            corrected_data = corrected.get_data()
            corrected_normalized = corrected_data / np.mean(corrected_data[corrected_data != 0])
            nii_corrected = nib.Nifti1Image(corrected_normalized, corrected.affine, corrected.header)
            nib.save(nii_corrected, str(out_path))
        paths.append(out_path)
    return paths


def mask_foreground(raw_data):
    """
        Mask the foreground of an image, using otsu threshold and connected components to remove background noise
    """
    raw_max = raw_data.max()
    raw_data = raw_data / raw_max
    blurred_data = ndimage.gaussian_filter(raw_data, 4)

    threshold = filters.threshold_otsu(raw_data) / 2
    threshold_data = blurred_data > threshold
    connected_structure = ndimage.generate_binary_structure(2, 2)  # Connects adjacent and diagonal.
    padded_comp, padded_nr = ndimage.label(threshold_data, structure=connected_structure)

    comps, comps_count = np.unique(padded_comp, return_counts=True)
    comps_count, comps = zip(*sorted(zip(comps_count, comps), reverse=True))

    two_biggest_cc = ((comps[0], np.average(comps[0])), (comps[1], np.average(comps[1])))
    biggest_cc = max(two_biggest_cc, key=lambda a: a[1])[0]

    foreground_mask = np.where(padded_comp == biggest_cc, True, False)
    closed = morphology.binary_closing(foreground_mask, selem=morphology.square(30))
    raw_data = np.where(closed, raw_data, 0)
    return raw_data * raw_max


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("file", help="Location of file to process")
    parser.add_argument("dir", help="Directory for preprocessed Files")
    parser.add_argument("-s", type=int, help="Series to extract")
    parser.add_argument('-b', type=bool, default= True, help='Bias field correct image')
    parser.add_argument('--pdim', type=float, help='Pixel dimensions (Retrieved from Tiff file if not set)')
    args = parser.parse_args()

    filetype = Path(args.file).suffix.lower()
    if filetype == '.tif' or filetype == '.tiff':
        n_path = tiff_to_nii(Path(args.file), Path(args.dir))
    elif filetype == '.nd2':
        n_path = nd2_to_nii(Path(args.file), Path(args.dir), args.series)