preprocess_script.py

import os
import glob
import warnings
warnings.filterwarnings("ignore")
import numpy as np
from PIL import Image
import nibabel as nib
import pandas as pd
from skimage.exposure import rescale_intensity

from tqdm import tqdm
import multiprocessing as mp
from skimage.transform import resize
from treelib import Tree


from utils.data_io import check_folder,read_original_cd_file


def combine_slices(config):
    """
    Combine slices into stack images
    :param config: parameters
    :return:
    """
    # signal.emit(True,'sss')
    num_slice = config["num_slice"]
    embryo_names = config["embryo_names"]
    max_times = config["max_times"]
    xy_res = config["xy_resolution"]
    z_res = config["z_resolution"]
    out_size = config["out_size"]
    raw_folder = config["raw_folder"]
    stack_folder = config["target_folder"]
    is_save_nuc = config["is_save_nuc"]
    is_save_memb = config["is_save_memb"]
    is_save_labelled_nuc=config['is_save_labelled_nuc']

    # -------------------get output size depend on x y z resolution and input shape-----------
    raw_memb_files = glob.glob(os.path.join(raw_folder, embryo_names[0], "tifR", "*.tif"))
    # print(raw_memb_files)
    raw_size = list(np.asarray(Image.open(raw_memb_files[0])).shape) + [int(num_slice * z_res / xy_res)]
    out_res = [res * in_scale / out_scale for res, in_scale, out_scale in zip([xy_res, xy_res, xy_res], raw_size, out_size)]
    print('composed nii.gz given shape',out_size, 'calculating shape ratio should be :::',raw_size[:2]+[int(num_slice * z_res / xy_res)])

    print('composed nii.gz resolution should be :::',out_res)
    # ----------------------------------------------------------------------------------------

    # multiprocessing
    mpPool = mp.Pool(mp.cpu_count() - 1)

    for idx_embryo,embryo_name in enumerate(embryo_names):

        # # get lineage file
        # if is_save_labelled_nuc:
        # else:
        #     lineage_file_path = None

        # ------------------ compose tiff and save raw nucleus as nii.gz--------------------------------
        if is_save_nuc:
            origin_folder = os.path.join(raw_folder, embryo_name, "tif")
            target_folder = os.path.join(stack_folder, embryo_name, "RawNuc")
            if not os.path.isdir(target_folder):
                os.makedirs(target_folder)

            configs = []
            for tp in range(1, max_times[idx_embryo] + 1):
                configs.append((origin_folder, target_folder, embryo_name, tp, out_size, num_slice, out_res))

            for idx, _ in enumerate(tqdm(mpPool.imap_unordered(stack_nuc_slices, configs), total=len(configs),
                                         desc="1/3 Stack nucleus of {}".format(embryo_name))):
                pass
        # ----------------------------------------------------------------

        # ------------------ compose tiff and save raw memb as nii.gz--------------------------------
        if is_save_memb:
            origin_folder = os.path.join(raw_folder, embryo_name, "tifR")
            target_folder = os.path.join(stack_folder, embryo_name, "RawMemb")
            if not os.path.isdir(target_folder):
                os.makedirs(target_folder)

            configs = []
            for tp in range(1, max_times[idx_embryo] + 1):
                configs.append((origin_folder, target_folder, embryo_name, tp, out_size, num_slice, out_res))
                # stack_memb_slices(configs[0])
            for idx, _ in enumerate(tqdm(mpPool.imap_unordered(stack_memb_slices, configs), total=len(configs),
                                         desc="2/3 Stack membrane of {}".format(embryo_name))):
                # TODO: Process Name: `2/3 Stack membrane`; Current status: `idx`; Final status: max_time
                pass
        # ----------------------------------------------------------------------------------------
        # save nucleus
        if is_save_labelled_nuc:
            lineage_file_path = os.path.join(config["raw_folder"], embryo_name, "aceNuc", "CD{}.csv".format(embryo_name))
            number_dictionary_path = config["number_dictionary_path"]
            target_folder = os.path.join(stack_folder, embryo_name, "SegNuc")
            if not os.path.isdir(target_folder):
                os.makedirs(target_folder)
            pd_lineage = read_original_cd_file(lineage_file_path)
            print('cd file looks like:::',pd_lineage)

            pd_number = pd.read_csv(number_dictionary_path, names=["name", "label"])
            number_dict = pd.Series(pd_number.label.values, index=pd_number.name).to_dict()

            configs = []
            for tp in range(1, max_times[idx_embryo] + 1):
                configs.append((embryo_name, number_dict, pd_lineage, tp, raw_size, out_size, out_res,
                                xy_res / z_res, target_folder))
                # save_nuc_seg(configs[0])
            for idx, _ in enumerate(tqdm(mpPool.imap_unordered(save_nuc_seg, configs), total=len(configs),
                                         desc="3/3 Construct nucleus location of {}".format(embryo_name))):
                # TODO: Process Name: `3/3 Construct nucleus location`; Current status: `idx`; Final status: max_time
                pass
            # for tp in range(1, max_time+1):
            #     save_nuc_seg(embryo_name=embryo_name,
            #                  name_dict=name_dict,
            #                  pd_lineage=pd_lineage,
            #                  tp=tp,
            #                  raw_size=raw_size,
            #                  out_size=out_size,
            #                  out_res=out_res,
            #                  dif_res=xy_res/z_res,
            #                  save_folder=target_folder)
            # shutil.copy(lineage_file, os.path.join(stack_folder, embryo_name))

# ============================================
# save raw nucleus stack
# ============================================
def stack_nuc_slices(para):
    [raw_folder, save_folder, embryo_name, tp, out_size, num_slice, res] = para

    out_stack = []
    save_file_name = "{}_{}_rawNuc.nii.gz".format(embryo_name, str(tp).zfill(3))
    for i_slice in range(1, num_slice + 1):
        raw_file_name = "{}_L1-t{}-p{}.tif".format(embryo_name, str(tp).zfill(3), str(i_slice).zfill(2))

        img = np.asanyarray(Image.open(os.path.join(raw_folder, raw_file_name)))
        out_stack.insert(0, img)
    img_stack = np.transpose(np.stack(out_stack), axes=(1, 2, 0))
    img_stack = resize(image=img_stack, output_shape=out_size, preserve_range=True, order=1).astype(np.uint8)
    nib_stack = nib.Nifti1Image(img_stack, np.eye(4))
    nib_stack.header.set_xyzt_units(xyz=3, t=8)
    nib_stack.header["pixdim"] = [1.0, res[0], res[1], res[2], 0., 0., 0., 0.]
    save_file = os.path.join(save_folder, save_file_name)
    check_folder(save_file)
    nib.save(nib_stack, save_file)

# ============================================
# save raw membrane stack
# ============================================
def stack_memb_slices(para):
    [raw_folder, save_folder, embryo_name, tp, out_size, num_slice, res] = para

    out_stack = []
    save_file_name = "{}_{}_rawMemb.nii.gz".format(embryo_name, str(tp).zfill(3))
    for i_slice in range(1, num_slice+1):
        #r"D:\TemDownload\201112plc1_late_Lng\tifR\c elegans 3.lif_Series001_Lng_001_t00_z08_ch01.tif"
        raw_file_name = "{}_L1-t{}-p{}.tif".format(embryo_name, str(tp).zfill(3), str(i_slice).zfill(2))
        # transform the image to array and short them in a list
        img = np.asanyarray(Image.open(os.path.join(raw_folder, raw_file_name)))
        out_stack.insert(0, img)
    img_stack = np.transpose(np.stack(out_stack), axes=(1, 2, 0)) # trasnpose the image from zxy to xyz
    v_min, v_max = np.percentile(img_stack, (0.2, 99.9)) # erase the outrange grayscale
    img_stack = rescale_intensity(img_stack, in_range=(v_min, v_max), out_range=(0, 255.0))
    # cut xy, interpolate z
    img_stack = resize(image=img_stack, output_shape=out_size, preserve_range=True, order=1).astype(np.uint8)
    nib_stack = nib.Nifti1Image(img_stack, np.eye(4))
    nib_stack.header.set_xyzt_units(xyz=3, t=8)
    nib_stack.header["pixdim"] = [1.0, res[0], res[1], res[2], 0., 0., 0., 0.]
    save_file = os.path.join(save_folder, save_file_name)
    check_folder(save_file)
    nib.save(nib_stack, save_file)

# =============================================
# save nucleus segmentation
# =============================================
def save_nuc_seg(para):
    [embryo_name, name_dict, pd_lineage, tp, raw_size, out_size, out_res, dif_res, save_folder] = para

    zoom_ratio = [y / x for x, y in zip(raw_size, out_size)] # raw and out size is x, y , z
    tp_lineage = pd_lineage[pd_lineage["time"] == tp]
    tp_lineage.loc[:, "x"] = (tp_lineage["x"] * zoom_ratio[1]).astype(np.int16) # the 'x' in cd file is the y in input and output
    tp_lineage.loc[:, "y"] = (np.floor(tp_lineage["y"] * zoom_ratio[0])).astype(np.int16) # the 'y' in cd file is the x in the input and output
    tp_lineage.loc[:, "z"] = (out_size[2] - np.floor(tp_lineage["z"] * (zoom_ratio[2] / dif_res))).astype(np.int16)

    # !!!! x <--> y
    nuc_dict = dict(
        zip(tp_lineage["cell"], zip(tp_lineage["y"].values, tp_lineage["x"].values, tp_lineage["z"].values)))
    labels = [name_dict[name] for name in list(nuc_dict.keys())]
    locs = list(nuc_dict.values())
    out_seg = np.zeros(out_size, dtype=np.uint16)
    out_seg[tuple(zip(*locs))] = labels

    save_file_name = "_".join([embryo_name, str(tp).zfill(3), "segNuc.nii.gz"])
    nib_stack = nib.Nifti1Image(out_seg, np.eye(4))
    nib_stack.header.set_xyzt_units(xyz=3, t=8)
    nib_stack.header["pixdim"] = [1.0, out_res[1], out_res[0], out_res[2], 0., 0., 0., 0.]
    save_file = os.path.join(save_folder, save_file_name)
    check_folder(save_file)
    nib.save(nib_stack, save_file)


def add_number_dict(configs):
    '''
    Construct cell tree structure with cell names
    :param nucleus_file:  the name list file to the tree initilization
    :param max_time: the maximum time point to be considered
    :return cell_tree: cell tree structure where each time corresponds to one cell (with specific name)
    '''

    ##  Construct cell
    #  Add unregulized naming
    raw_folder=configs['raw_folder']
    embryo_names=configs['embryo_names']
    for idx, embryo_name in enumerate(embryo_names):

        cell_tree = Tree()
        cell_tree.create_node('P0', 'P0')
        cell_tree.create_node('AB', 'AB', parent='P0')
        cell_tree.create_node('P1', 'P1', parent='P0')
        cell_tree.create_node('EMS', 'EMS', parent='P1')
        cell_tree.create_node('P2', 'P2', parent='P1')
        cell_tree.create_node('P3', 'P3', parent='P2')
        cell_tree.create_node('C', 'C', parent='P2')
        cell_tree.create_node('P4', 'P4', parent='P3')
        cell_tree.create_node('D', 'D', parent='P3')
        cell_tree.create_node('Z2', 'Z2', parent='P4')
        cell_tree.create_node('Z3', 'Z3', parent='P4')

        # EMS
        cell_tree.create_node('E', 'E', parent='EMS')
        cell_tree.create_node('MS', 'MS', parent='EMS')


        # read and combine all names from different acetrees
        ## Get cell number
        try:
            pd_number = pd.read_csv('./tem_important_files/number_dictionary.csv', names=["name", "label"])
            number_dictionary = pd.Series(pd_number.label.values, index=pd_number.name).to_dict()
        except:
            number_dictionary = {}

        # =====================================
        # dynamic update the name dictionary
        # =====================================
        cell_in_dictionary = list(number_dictionary.keys())
        # Read the name excel and construct the tree with complete segCell
        ace_pd = read_original_cd_file(os.path.join(raw_folder,embryo_name,'aceNuc','CD{}.csv'.format(embryo_name)))

        # ace_pd = ace_pd[ace_pd.time <= max_time]
        cell_list = list(ace_pd.cell.unique())
        add_cell_list = list(set(cell_list) - set(cell_in_dictionary))
        add_cell_list.sort()
        if len(add_cell_list) > 0:
            print("Name dictionary updated !!!")
            add_number_dictionary = dict(zip(add_cell_list, range(len(cell_in_dictionary) + 1, len(cell_in_dictionary) + len(add_cell_list) + 1)))
            number_dictionary.update(add_number_dictionary)
            pd_number_dictionary = pd.DataFrame.from_dict(number_dictionary, orient="index")
            pd_number_dictionary.to_csv('./tem_important_files/number_dictionary.csv', header=False)



if __name__ == "__main__":
    # -----------------generate the number_dictionary----------------------------
    config_generating_number_dict=dict(raw_folder=r"F:\ProjectData\MembraneProject\AllRawData",
                embryo_names=["221017plc1p2"])
    add_number_dict(config_generating_number_dict)
    # ---------------------------------------------------------------------------

    # ------------------generate nii.gz from raw tiff grayscale images------------
    config_composing_tiffs = dict(num_slice=94,
                  embryo_names=["221017plc1p2"],
                  max_times = [240],
                  xy_resolution = 0.09,
                  z_resolution = 0.43,
                  # 94  *   0.43/0.09  *  356/712
                  out_size = [256, 356, 224], # todo: need to be calculated with the vertical image amount
                  raw_folder=r"F:\ProjectData\MembraneProject\AllRawData",
                  target_folder=r"D:\ProjectData\AllDataPacked",
                  is_save_nuc = False,
                  is_save_memb = False,
                  is_save_labelled_nuc = True,
                  number_dictionary_path = r'./tem_important_files/number_dictionary.csv'
                  )

    combine_slices(config_composing_tiffs)
    # ---------------------------------------------------------------------------