Centerline/centerline.py

import os, sys

currentdir = os.path.dirname(os.path.realpath(__file__))
parentdir = os.path.dirname(currentdir)
sys.path.append(parentdir)  # PYTHON > 3.3 does not allow relative referencing

# import tensorflow as tf
# tf.enable_eager_execution()
# import neurite.py.utils as neurite_utils

from skimage.morphology import skeletonize_3d, ball
from skimage.morphology import binary_closing, binary_opening
from skimage.filters import median
from skimage.measure import regionprops, label
from skimage.transform import warp

from scipy.ndimage import zoom
from scipy.interpolate import LinearNDInterpolator, Rbf

import h5py
import numpy as np
from tqdm import tqdm
import re
import nibabel as nib
from nilearn.image import resample_img

from Centerline.graph_utils import graph_to_ndarray, deform_graph, get_bifurcation_nodes, subsample_graph, \
    apply_displacement
from Centerline.skeleton_to_graph import get_graph_from_skeleton
from Centerline.visualization_utils import plot_skeleton, compare_graphs

from ddmr.utils.operators import min_max_norm
from ddmr.utils import constants as C

import cupy
from cupyx.scipy.ndimage import zoom as zoom_gpu
from cupyx.scipy.ndimage import map_coordinates

DATASET_LOCATION = '/mnt/EncryptedData1/Users/javier/vessel_registration/3Dirca/dataset/EVAL'
DATASET_NAMES = ['Affine', 'None', 'Translation']
DATASET_FILENAME = 'volume'
IMGS_FOLDER = '/home/jpdefrutos/workspace/ddmr/Centerline/centerlines'

DATASTE_RAW_FILES = '/mnt/EncryptedData1/Users/javier/vessel_registration/3Dirca/nifti3'
LITS_SEGMENTATION_FILE = 'segmentation'
LITS_CT_FILE = 'volume'


def warp_volume(volume, disp_map, indexing='ij'):
    assert indexing is 'ij' or 'xy', 'Invalid indexing option. Only "ij" or "xy"'
    grid_i = np.linspace(0, disp_map.shape[0], disp_map.shape[0], endpoint=False)
    grid_j = np.linspace(0, disp_map.shape[1], disp_map.shape[1], endpoint=False)
    grid_k = np.linspace(0, disp_map.shape[2], disp_map.shape[2], endpoint=False)
    grid_i, grid_j, grid_k = np.meshgrid(grid_i, grid_j, grid_k, indexing=indexing)
    grid_i = (grid_i.flatten() + disp_map[..., 0].flatten())[..., np.newaxis]
    grid_j = (grid_j.flatten() + disp_map[..., 1].flatten())[..., np.newaxis]
    grid_k = (grid_k.flatten() + disp_map[..., 2].flatten())[..., np.newaxis]
    coords = np.hstack([grid_i, grid_j, grid_k]).reshape([*disp_map.shape[:-1], -1])
    coords = coords.transpose((-1, 0, 1, 2))
    # The returned volume has indexing xy
    return warp(volume, coords)


def keep_largest_segmentation(img):
    label_img = label(img)
    rp = regionprops(label_img)  # Regions labeled with 0 (bg) are ignored
    biggest_area = (0, 0)
    for l in range(0, label_img.max()):
        if rp[l].area > biggest_area[1]:
            biggest_area = (l + 1, rp[l].area)
    img[label_img != biggest_area[0]] = 0.
    return img


def preprocess_image(img, keep_largest=False):
    ret = binary_closing(img, ball(1))
    ret = binary_opening(ret, ball(1))
    #ret = median(ret, ball(1), mode='constant')
    if keep_largest:
        ret = keep_largest_segmentation(ret)
    return ret.astype(np.float)


def build_displacement_map_interpolator(disp_map, backwards=False, indexing='ij'):
    grid_i = np.linspace(0, disp_map.shape[0], disp_map.shape[0], endpoint=False)
    grid_j = np.linspace(0, disp_map.shape[1], disp_map.shape[1], endpoint=False)
    grid_k = np.linspace(0, disp_map.shape[2], disp_map.shape[2], endpoint=False)
    grid_i, grid_j, grid_k = np.meshgrid(grid_i, grid_j, grid_k, indexing=indexing)
    grid_i = grid_i.flatten()
    grid_j = grid_j.flatten()
    grid_k = grid_k.flatten()
    # To generate the moving image, we used backwards mapping were the input was the fix image
    # Now we are doing direct mapping from the fix graph coordinates to the moving coordinates
    # The application points of the displacement map are thus the transformed "moving image"-grid
    #   and the displacement vectors are reversed
    if backwards:
        coords = np.hstack([grid_i[..., np.newaxis], grid_j[..., np.newaxis], grid_k[..., np.newaxis]])
        return LinearNDInterpolator(coords, np.reshape(disp_map, [-1, 3]))
    else:
        grid_i = (grid_i + disp_map[..., 0].flatten())
        grid_j = (grid_j + disp_map[..., 1].flatten())
        grid_k = (grid_k + disp_map[..., 2].flatten())

        coords = np.hstack([grid_i[..., np.newaxis], grid_j[..., np.newaxis], grid_k[..., np.newaxis]])
        return LinearNDInterpolator(coords, -np.reshape(disp_map, [-1, 3]))


def resample_segmentation(img, output_shape, preserve_range, threshold=None, gpu=True):
    # Preserve range can be a bool (keep or not the original dyn. range) or a list with a new dyn. range
    zoom_f = np.divide(np.asarray(output_shape), np.asarray(img.shape))

    if gpu:
        out_img = zoom_gpu(cupy.asarray(img), zoom_f, order=1)  # order = 0 or 1
    else:
        out_img = zoom(img, zoom_f)
    if isinstance(preserve_range, bool):
        if preserve_range:
            range_min, range_max = np.amin(img), np.amax(img)
            out_img = min_max_norm(out_img)
            out_img = out_img * (range_max - range_min) + range_min
    elif isinstance(preserve_range, list):
        range_min, range_max = preserve_range
        out_img = min_max_norm(out_img)
        out_img = out_img * (range_max - range_min) + range_min

    if threshold is not None and out_img.min() < threshold < out_img.max():
        range_min, range_max = np.amin(out_img), np.amax(out_img)
        out_img[out_img > threshold] = range_max
        out_img[out_img < range_max] = range_min
    return cupy.asnumpy(out_img) if gpu else out_img


if __name__ == '__main__':
    for dataset_name in DATASET_NAMES:
        dataset_loc = os.path.join(DATASET_LOCATION, dataset_name)
        dataset_files = os.listdir(dataset_loc)
        dataset_files.sort()
        dataset_files = [os.path.join(dataset_loc, f) for f in dataset_files if DATASET_FILENAME in f]

        iterator = tqdm(dataset_files)
        for file_path in iterator:
            file_num = int(re.findall('(\d+)', os.path.split(file_path)[-1])[0])

            iterator.set_description('{} ({}): laoding data'.format(file_num, dataset_name))
            vol_file = h5py.File(file_path, 'r')
            # fix_vessels = vol_file[C.H5_FIX_VESSELS_MASK][..., 0]
            disp_map = vol_file[C.H5_GT_DISP][:]
            bbox = vol_file['parameters/bbox'][:]
            bbox_min = bbox[:3]
            bbox_max = bbox[3:] + bbox_min

            # Load vessel segmentation mask and resize to 64^3
            fix_labels = nib.load(os.path.join(DATASTE_RAW_FILES, 'segmentation-{:04d}.nii.gz'.format(file_num)))
            fix_vessels = fix_labels.slicer[..., 1]
            fix_vessels = resample_img(fix_vessels, np.eye(3))
            fix_vessels = np.asarray(fix_vessels.dataobj)
            fix_vessels = preprocess_image(fix_vessels)
            fix_vessels = resample_segmentation(fix_vessels, vol_file['parameters/first_reshape'][:], [0, 1], 0.3,
                                                gpu=True)
            fix_vessels = fix_vessels[bbox_min[0]:bbox_max[0], bbox_min[1]:bbox_max[1], bbox_min[2]:bbox_max[2]]
            fix_vessels = resample_segmentation(fix_vessels, [64] * 3, [0, 1], 0.3, gpu=True)
            fix_vessels = preprocess_image(fix_vessels)

            mov_vessels = preprocess_image(warp_volume(fix_vessels, disp_map))
            mov_skel = skeletonize_3d(mov_vessels)
            ### Fix the incorrect scaling ###
            disp_map *= 2
            bbox_size = np.asarray(bbox[3:])  # Only load the bbox size
            rescale_factors = 64 / bbox_size

            disp_map[..., 0] = np.multiply(disp_map[..., 0], rescale_factors[0])
            disp_map[..., 1] = np.multiply(disp_map[..., 1], rescale_factors[1])
            disp_map[..., 2] = np.multiply(disp_map[..., 2], rescale_factors[2])
            #################################

            iterator.set_description('{} ({}): getting graphs'.format(file_num, dataset_name))
            # Prepare displacement map
            disp_map_interpolator = build_displacement_map_interpolator(disp_map, backwards=False)

            # Get skeleton and graph
            fix_skel = skeletonize_3d(fix_vessels)
            fix_graph = get_graph_from_skeleton(fix_skel, subsample=True)
            mov_graph = get_graph_from_skeleton(mov_skel, subsample=True) # deform_graph(fix_graph, disp_map_interpolator)

            ##### TODO: ERASE Check the mov graph ######
            # check_mov_vessels = vol_file[C.H5_MOV_VESSELS_MASK][..., 0]
            # check_mov_vessels = preprocess_image(check_mov_vessels)
            # check_mov_skel = skeletonize_3d(check_mov_vessels)
            # check_mov_graph = get_graph_from_skeleton(check_mov_skel, subsample=True)
            ###########
            fix_pts, fix_nodes, fix_edges = graph_to_ndarray(fix_graph)
            mov_pts, mov_nodes, mov_edges = graph_to_ndarray(mov_graph)

            fix_bifur_loc, fix_bifur_id = get_bifurcation_nodes(fix_graph)
            mov_bifur_loc, mov_bifur_id = get_bifurcation_nodes(mov_graph)

            iterator.set_description('{} ({}): saving data'.format(file_num, dataset_name))
            pts_file_path, pts_file_name = os.path.split(file_path)
            pts_file_name = pts_file_name.replace(DATASET_FILENAME, 'points')
            pts_file_path = os.path.join(pts_file_path, pts_file_name)
            pts_file = h5py.File(pts_file_path, 'w')

            pts_file.create_dataset('fix/points', data=fix_pts)
            pts_file.create_dataset('fix/nodes', data=fix_nodes)
            pts_file.create_dataset('fix/edges', data=fix_edges)
            pts_file.create_dataset('fix/bifurcations', data=fix_bifur_loc)
            pts_file.create_dataset('fix/graph', data=fix_graph)
            pts_file.create_dataset('fix/img', data=fix_vessels)
            pts_file.create_dataset('fix/skeleton', data=fix_skel)
            pts_file.create_dataset('fix/centroid', data=vol_file[C.H5_FIX_CENTROID][:])

            pts_file.create_dataset('mov/points', data=mov_pts)
            pts_file.create_dataset('mov/nodes', data=mov_nodes)
            pts_file.create_dataset('mov/edges', data=mov_edges)
            pts_file.create_dataset('mov/bifurcations', data=mov_bifur_loc)
            pts_file.create_dataset('mov/graph', data=mov_graph)
            pts_file.create_dataset('mov/img', data=mov_vessels)
            pts_file.create_dataset('mov/skeleton', data=mov_skel)
            pts_file.create_dataset('mov/centroid', data=vol_file[C.H5_MOV_CENTROID][:])

            pts_file.create_dataset('parameters/voxel_size', data=vol_file['parameters/voxel_size'][:])
            pts_file.create_dataset('parameters/original_affine', data=vol_file['parameters/original_affine'][:])
            pts_file.create_dataset('parameters/isotropic_affine', data=vol_file['parameters/isotropic_affine'][:])
            pts_file.create_dataset('parameters/original_shape', data=vol_file['parameters/original_shape'][:])
            pts_file.create_dataset('parameters/isotropic_shape', data=vol_file['parameters/isotropic_shape'][:])
            pts_file.create_dataset('parameters/first_reshape', data=vol_file['parameters/first_reshape'][:])
            pts_file.create_dataset('parameters/bbox', data=vol_file['parameters/bbox'][:])
            pts_file.create_dataset('parameters/last_reshape', data=vol_file['parameters/last_reshape'][:])

            pts_file.create_dataset('displacement_map', data=disp_map)

            vol_file.close()
            pts_file.close()

            iterator.set_description('{} ({}): drawing plots'.format(file_num, dataset_name))
            num = pts_file_name.split('-')[-1].split('.hd5')[0]
            imgs_folder = os.path.join(IMGS_FOLDER, dataset_name, num)
            os.makedirs(imgs_folder, exist_ok=True)
            plot_skeleton(fix_vessels, fix_skel, fix_graph, os.path.join(imgs_folder, 'fix'), ['.pdf', '.png'])
            plot_skeleton(mov_vessels, mov_skel, mov_graph, os.path.join(imgs_folder, 'mov'), ['.pdf', '.png'])
            iterator.set_description('{} ({})'.format(file_num, dataset_name))