Utils.py

 # %% Import modules
import os, cv2, h5py, sys, random, matplotlib, natsort, tifffile, json
matplotlib.use('module://backend_interagg')
import numpy as np
import pandas as pd
import nibabel as nib
import matplotlib.pyplot as plt
import SimpleITK as sitk

from glob import glob
from PIL import Image
from tqdm import tqdm
from pathlib import Path
from matplotlib import cm
from skimage.color import rgb2hed, hed2rgb

sys.path.append('/home/choiyoonho/Documents/HT/3D tissue visualization/co-registration') # <- Superglue
from models.matching import Matching
from models.utils import (compute_pose_error, compute_epipolar_error, estimate_pose, make_matching_plot,
                          error_colormap, AverageTimer, pose_auc, read_image, rotate_intrinsics, rotate_pose_inplane,
                          scale_intrinsics)
from histolab.filters.image_filters import Compose, OtsuThreshold
from histolab.filters.morphological_filters import BinaryClosing, RemoveSmallObjects, RemoveSmallHoles

sys.path.append('/home/choiyoonho/Documents/HT/hovernet2')
from misc.wsi_handler import get_file_handler
from misc.viz_utils import visualize_instances_dict


# %% Functions
def show(img):
    plt.imshow(img, cmap='gray')
    plt.axis('off')
    plt.show()
    plt.clf()

def save_png_from_3DHT(data, thumbnail_dir, roi, slicenum=4):
    data_ = data[:,:,slicenum]
    data_ = (data_-data_.min())/(data_.max()-data_.min())
    image = Image.fromarray(data_*255)
    image = image.convert('L')
    image.save(os.path.join(thumbnail_dir, roi+'_3D.png'))

def save_nii_from_3DHT(data, nii_dir, roi):
    nifti_img = nib.Nifti1Image(data, affine=np.eye(4))
    nib.save(nifti_img, os.path.join(nii_dir, roi + '.nii'))

def pop_insert(list, pop_idx, insert_idx):
    element_to_move = list.pop(pop_idx)  # 0번째 원소를 추출하여 저장
    list.insert(insert_idx, element_to_move)  # 2번째 위치에 원소 삽입
    return list

def convert_value(value):
     if isinstance(value, np.integer):
         return int(value)
     elif isinstance(value, np.floating):
         return float(value)
     elif isinstance(value, np.ndarray):
         return value.tolist()
     else:
         return value

def convert_attributes(attributes):
     converted_attributes = {}
     for key, value in attributes.items():
         converted_value = convert_value(value)
         converted_attributes[key] = converted_value
     return converted_attributes

class HDF5Wrapper:
    def __init__(self, filepath):
        self.filepath = filepath
        self.file = h5py.File(filepath, 'a')
    def __setitem__(self, name, val):
        if isinstance(val, np.ndarray):
            self.file.create_dataset(name, data=val)
        else:
            self.file.attrs[name] = str(val)
    def close(self):
        self.file.close()

def save_hdf5_with_attributes(data, filepath, attributes):
    hdf5 = HDF5Wrapper(filepath)
    hdf5['data'] = data
    # 속성 저장
    for key, value in attributes.items():
        hdf5[key] = value
    hdf5.close()
    print('Complete to save {}'.format(filepath))

def load_hdf5_with_attributes(filepath):
    file = h5py.File(filepath, 'r')
    attributes = {}
    for key in file.attrs.keys():
        if isinstance(file.attrs[key], np.ndarray):
            attributes[key] = file.attrs[key]
        else:
            attributes[key] = str(file.attrs[key])
    data = file['data'][()]
    file.close()
    return data, attributes

def composed_filters(image_rgb):
    filters = Compose([OtsuThreshold(), BinaryClosing(9), RemoveSmallHoles(), RemoveSmallObjects()])
    return filters(image_rgb)

def composed_filters_he(image_rgb):
    filters = Compose([OtsuThreshold(), BinaryClosing(1), RemoveSmallHoles(), RemoveSmallObjects()])
    return filters(image_rgb)
# %%
path_to_tmp_files_superglue = '../tmp_files_superglue'
os.makedirs(path_to_tmp_files_superglue, exist_ok=True)
def generate_superglue_keypoints():
    input_pairs = '{}/scannet_sample_pairs_with_gt.txt'.format(path_to_tmp_files_superglue)
    input_dir = '{}'.format(path_to_tmp_files_superglue)
    output_dir = '{}'.format(path_to_tmp_files_superglue)
    max_length = -1
    resize = [-1]
    superglue = 'outdoor'
    max_keypoints = 4096 #4096
    keypoint_threshold = 0.1
    nms_radius = 4
    sinkhorn_iterations=20
    match_threshold=0.2
    viz = False
    viz_extension='png'
    cache= False
    shuffle = False
    eval_=False
    force_cpu = False
    resize_float = False
    if len(resize) == 2 and resize[1] == -1:
        resize = resize[0:1]
    if len(resize) == 2:
        print('Will resize to {}x{} (WxH)'.format(
            resize[0], resize[1]))
    elif len(resize) == 1 and resize[0] > 0:
        print('Will resize max dimension to {}'.format(resize[0]))
    elif len(resize) == 1:
        print('Will not resize images')
    else:
        raise ValueError('Cannot specify more than two integers for --resize')

    with open(input_pairs, 'r') as f:
        pairs = [l.split() for l in f.readlines()]

    if max_length > -1:
        pairs = pairs[0:np.min([len(pairs), max_length])]

    if shuffle:
        random.Random(0).shuffle(pairs)

    if eval_:
        if not all([len(p) == 38 for p in pairs]):
            raise ValueError(
                'All pairs should have ground truth info for evaluation.'
                'File \"{}\" needs 38 valid entries per row'.format(input_pairs))

    # Load the SuperPoint and SuperGlue models.
    os.environ["CUDA_VISIBLE_DEVICES"] = "0, 1"
    device = 'cpu' #'cuda' if torch.cuda.is_available() and not force_cpu else 'cpu'
    print('Running inference on device \"{}\"'.format(device))
    config = {
        'superpoint': {
            'nms_radius': nms_radius,
            'keypoint_threshold': keypoint_threshold,
            'max_keypoints': max_keypoints
        },
        'superglue': {
            'weights': superglue,
            'sinkhorn_iterations': sinkhorn_iterations,
            'match_threshold': match_threshold,
        }
    }
    matching = Matching(config).eval().to(device)

    # Create the output directories if they do not exist already.
    input_dir = Path(input_dir)
    print('Looking for data in directory \"{}\"'.format(input_dir))
    output_dir = Path(output_dir)
    output_dir.mkdir(exist_ok=True, parents=True)
    print('Will write matches to directory \"{}\"'.format(output_dir))
    if eval:
        print('Will write evaluation results',
              'to directory \"{}\"'.format(output_dir))
    if viz:
        print('Will write visualization images to',
              'directory \"{}\"'.format(output_dir))

    timer = AverageTimer(newline=True)
    for i, pair in enumerate(pairs):
        name0, name1 = pair[:2]
        stem0, stem1 = Path(name0).stem, Path(name1).stem
        matches_path = output_dir / '{}_{}_matches.npz'.format(stem0, stem1)
        print('matches_path', matches_path)
        eval_path = output_dir / '{}_{}_evaluation.npz'.format(stem0, stem1)
        viz_path = output_dir / '{}_{}_matches.{}'.format(stem0, stem1, viz_extension)
        viz_eval_path = output_dir / \
            '{}_{}_evaluation.{}'.format(stem0, stem1, viz_extension)

        # Handle --cache logic.
        do_match = True
        do_eval = eval_
        do_viz = viz
        do_viz_eval = eval_ and viz
        if cache:
            if matches_path.exists():
                try:
                    results = np.load(matches_path)
                except:
                    raise IOError('Cannot load matches .npz file: %s' %
                                  matches_path)

                kpts0, kpts1 = results['keypoints0'], results['keypoints1']
                matches, conf = results['matches'], results['match_confidence']
                do_match = False
            if eval_ and eval_path.exists():
                try:
                    results = np.load(eval_path)
                except:
                    raise IOError('Cannot load eval .npz file: %s' % eval_path)
                err_R, err_t = results['error_R'], results['error_t']
                precision = results['precision']
                matching_score = results['matching_score']
                num_correct = results['num_correct']
                epi_errs = results['epipolar_errors']
                do_eval = False
            if viz and viz_path.exists():
                do_viz = False
            if viz and eval_ and viz_eval_path.exists():
                do_viz_eval = False
            timer.update('load_cache')

        if not (do_match or do_eval or do_viz or do_viz_eval):
            timer.print('Finished pair {:5} of {:5}'.format(i, len(pairs)))
            continue

        # If a rotation integer is provided (e.g. from EXIF data), use it:
        if len(pair) >= 5:
            rot0, rot1 = int(pair[2]), int(pair[3])
        else:
            rot0, rot1 = 0, 0

        # Load the image pair.
        image0, inp0, scales0 = read_image(
            input_dir / name0, device, resize, rot0, resize_float)
        image1, inp1, scales1 = read_image(
            input_dir / name1, device, resize, rot1, resize_float)
        if image0 is None or image1 is None:
            print('Problem reading image pair: {} {}'.format(
                input_dir/name0, input_dir/name1))
            exit(1)
        timer.update('load_image')

        if do_match:
            # Perform the matching.
            pred = matching({'image0': inp0, 'image1': inp1})
            pred = {k: v[0].cpu().detach().numpy() for k, v in pred.items()}
            kpts0, kpts1 = pred['keypoints0'], pred['keypoints1']
            matches, conf = pred['matches0'], pred['matching_scores0']
            timer.update('matcher')

            # Write the matches to disk.
            out_matches = {'keypoints0': kpts0, 'keypoints1': kpts1,
                           'matches': matches, 'match_confidence': conf}
            np.savez(str(matches_path), **out_matches)

        # Keep the matching keypoints.
        valid = matches > -1
        mkpts0 = kpts0[valid]
        mkpts1 = kpts1[matches[valid]]
        mconf = conf[valid]

        if do_eval:
            # Estimate the pose and compute the pose error.
            assert len(pair) == 38, 'Pair does not have ground truth info'
            K0 = np.array(pair[4:13]).astype(float).reshape(3, 3)
            K1 = np.array(pair[13:22]).astype(float).reshape(3, 3)
            T_0to1 = np.array(pair[22:]).astype(float).reshape(4, 4)

            # Scale the intrinsics to resized image.
            K0 = scale_intrinsics(K0, scales0)
            K1 = scale_intrinsics(K1, scales1)

            # Update the intrinsics + extrinsics if EXIF rotation was found.
            if rot0 != 0 or rot1 != 0:
                cam0_T_w = np.eye(4)
                cam1_T_w = T_0to1
                if rot0 != 0:
                    K0 = rotate_intrinsics(K0, image0.shape, rot0)
                    cam0_T_w = rotate_pose_inplane(cam0_T_w, rot0)
                if rot1 != 0:
                    K1 = rotate_intrinsics(K1, image1.shape, rot1)
                    cam1_T_w = rotate_pose_inplane(cam1_T_w, rot1)
                cam1_T_cam0 = cam1_T_w @ np.linalg.inv(cam0_T_w)
                T_0to1 = cam1_T_cam0

            epi_errs = compute_epipolar_error(mkpts0, mkpts1, T_0to1, K0, K1)
            correct = epi_errs < 5e-4
            num_correct = np.sum(correct)
            precision = np.mean(correct) if len(correct) > 0 else 0
            matching_score = num_correct / len(kpts0) if len(kpts0) > 0 else 0

            thresh = 1.  # In pixels relative to resized image size.
            ret = estimate_pose(mkpts0, mkpts1, K0, K1, thresh)
            if ret is None:
                err_t, err_R = np.inf, np.inf
            else:
                R, t, inliers = ret
                err_t, err_R = compute_pose_error(T_0to1, R, t)

            # Write the evaluation results to disk.
            out_eval = {'error_t': err_t,
                        'error_R': err_R,
                        'precision': precision,
                        'matching_score': matching_score,
                        'num_correct': num_correct,
                        'epipolar_errors': epi_errs}
            np.savez(str(eval_path), **out_eval)
            timer.update('eval')

        if do_viz:
            # Visualize the matches.
            color = cm.jet(mconf)
            text = [
                'SuperGlue',
                'Keypoints: {}:{}'.format(len(kpts0), len(kpts1)),
                'Matches: {}'.format(len(mkpts0)),
            ]
            if rot0 != 0 or rot1 != 0:
                text.append('Rotation: {}:{}'.format(rot0, rot1))

            # Display extra parameter info.
            k_thresh = matching.superpoint.config['keypoint_threshold']
            m_thresh = matching.superglue.config['match_threshold']
            small_text = [
                'Keypoint Threshold: {:.4f}'.format(k_thresh),
                'Match Threshold: {:.2f}'.format(m_thresh),
                'Image Pair: {}:{}'.format(stem0, stem1),
            ]

            make_matching_plot(
                image0, image1, kpts0, kpts1, mkpts0, mkpts1, color,
                text, viz_path,show_keypoints,
                fast_viz, opencv_display, 'Matches', small_text)

            timer.update('viz_match')

        if do_viz_eval:
            # Visualize the evaluation results for the image pair.
            color = np.clip((epi_errs - 0) / (1e-3 - 0), 0, 1)
            color = error_colormap(1 - color)
            deg, delta = ' deg', 'Delta '
            if not fast_viz:
                deg, delta = '°', '$\\Delta$'
            e_t = 'FAIL' if np.isinf(err_t) else '{:.1f}{}'.format(err_t, deg)
            e_R = 'FAIL' if np.isinf(err_R) else '{:.1f}{}'.format(err_R, deg)
            text = [
                'SuperGlue',
                '{}R: {}'.format(delta, e_R), '{}t: {}'.format(delta, e_t),
                'inliers: {}/{}'.format(num_correct, (matches > -1).sum()),
            ]
            if rot0 != 0 or rot1 != 0:
                text.append('Rotation: {}:{}'.format(rot0, rot1))

            # Display extra parameter info (only works with --fast_viz).
            k_thresh = matching.superpoint.config['keypoint_threshold']
            m_thresh = matching.superglue.config['match_threshold']
            small_text = [
                'Keypoint Threshold: {:.4f}'.format(k_thresh),
                'Match Threshold: {:.2f}'.format(m_thresh),
                'Image Pair: {}:{}'.format(stem0, stem1),
            ]

            make_matching_plot(
                image0, image1, kpts0, kpts1, mkpts0,
                mkpts1, color, text, viz_eval_path,
                show_keypoints, fast_viz,
                opencv_display, 'Relative Pose', small_text)

            timer.update('viz_eval')

        timer.print('Finished pair {:5} of {:5}'.format(i, len(pairs)))

    if eval_:
        # Collate the results into a final table and print to terminal.
        pose_errors = []
        precisions = []
        matching_scores = []
        for pair in pairs:
            name0, name1 = pair[:2]
            stem0, stem1 = Path(name0).stem, Path(name1).stem
            eval_path = output_dir / \
                '{}_{}_evaluation.npz'.format(stem0, stem1)
            results = np.load(eval_path)
            pose_error = np.maximum(results['error_t'], results['error_R'])
            pose_errors.append(pose_error)
            precisions.append(results['precision'])
            matching_scores.append(results['matching_score'])
        thresholds = [5, 10, 20]
        aucs = pose_auc(pose_errors, thresholds)
        aucs = [100.*yy for yy in aucs]
        prec = 100.*np.mean(precisions)
        ms = 100.*np.mean(matching_scores)

def get_rgb2hed(thumbnail, write_basepath=None) -> list:
    'rgb -> [hematoxylin, Eosin, DAB]'

    ihc_hed = rgb2hed(np.array(thumbnail))

    null = np.zeros_like(ihc_hed[:, :, 0])
    ihc_h = hed2rgb(np.stack((ihc_hed[:, :, 0], null, null), axis=-1))  # hematoxylin
    ihc_e = hed2rgb(np.stack((null, ihc_hed[:, :, 1], null), axis=-1))  # Eosin
    ihc_d = hed2rgb(np.stack((null, null, ihc_hed[:, :, 2]), axis=-1))  # DAB

    if write_basepath:
        # thumbnail.save(write_basepath+'/'+'hed.png')
        cv2.imwrite(write_basepath + '/' + 'ihc_h.png', ihc_h * 255)
        cv2.imwrite(write_basepath + '/' + 'ihc_e.png', ihc_e * 255)
        cv2.imwrite(write_basepath + '/' + 'ihc_d.png', ihc_d * 255)

    return [ihc_h, ihc_e, ihc_d]

def unmixHE(img, saveFile=None, Io=240, alpha=1, beta=0.15):
    HERef = np.array([[0.5626, 0.2159],
                      [0.7201, 0.8012],
                      [0.4062, 0.5581]])

    maxCRef = np.array([1.9705, 1.0308])

    # define height and width of image
    h, w, c = img.shape

    # reshape image
    img = img.reshape((-1, 3))

    # calculate optical density
    OD = -np.log((img.astype(np.float) + 1) / Io)

    # remove transparent pixels
    ODhat = OD[~np.any(OD < beta, axis=1)]

    # compute eigenvectors
    eigvals, eigvecs = np.linalg.eigh(np.cov(ODhat.T))

    # eigvecs *= -1

    # project on the plane spanned by the eigenvectors corresponding to the two
    # largest eigenvalues
    That = ODhat.dot(eigvecs[:, 1:3])

    phi = np.arctan2(That[:, 1], That[:, 0])

    minPhi = np.percentile(phi, alpha)
    maxPhi = np.percentile(phi, 100 - alpha)

    vMin = eigvecs[:, 1:3].dot(np.array([(np.cos(minPhi), np.sin(minPhi))]).T)
    vMax = eigvecs[:, 1:3].dot(np.array([(np.cos(maxPhi), np.sin(maxPhi))]).T)

    # a heuristic to make the vector corresponding to hematoxylin first and the
    # one corresponding to eosin second
    if vMin[0] > vMax[0]:
        HE = np.array((vMin[:, 0], vMax[:, 0])).T
    else:
        HE = np.array((vMax[:, 0], vMin[:, 0])).T

    # rows correspond to channels (RGB), columns to OD values
    Y = np.reshape(OD, (-1, 3)).T

    # determine concentrations of the individual stains
    C = np.linalg.lstsq(HE, Y, rcond=None)[0]

    # normalize stain concentrations
    maxC = np.array([np.percentile(C[0, :], 99), np.percentile(C[1, :], 99)])
    tmp = np.divide(maxC, maxCRef)
    C2 = np.divide(C, tmp[:, np.newaxis])

    # recreate the image using reference mixing matrix
    Inorm = np.multiply(Io, np.exp(-HERef.dot(C2)))
    Inorm[Inorm > 255] = 254
    Inorm = np.reshape(Inorm.T, (h, w, 3)).astype(np.uint8)

    # unmix hematoxylin and eosin
    H = np.multiply(Io, np.exp(np.expand_dims(-HERef[:, 0], axis=1).dot(np.expand_dims(C2[0, :], axis=0))))
    H[H > 255] = 254
    H = np.reshape(H.T, (h, w, 3)).astype(np.uint8)

    E = np.multiply(Io, np.exp(np.expand_dims(-HERef[:, 1], axis=1).dot(np.expand_dims(C2[1, :], axis=0))))
    E[E > 255] = 254
    E = np.reshape(E.T, (h, w, 3)).astype(np.uint8)

    # if saveFile is not None:
    # Image.fromarray(Inorm).save(saveFile+'.png')
    # Image.fromarray(H).save(saveFile+'_H.png')
    # Image.fromarray(E).save(saveFile+'_E.png')
    return Inorm, H, E

def initial_transform(fixed_image, moving_image):
    initial_transform = sitk.CenteredTransformInitializer(fixed_image, moving_image, sitk.AffineTransform(2),
                                                          sitk.CenteredTransformInitializerFilter.GEOMETRY)
    # moving_resampled = sitk.Resample(moving_image, fixed_image, initial_transform, sitk.sitkLinear, 0.0, moving_image.GetPixelID())
    registration_method = sitk.ImageRegistrationMethod()
    registration_method.SetMetricAsCorrelation()
    registration_method.SetMetricSamplingStrategy(registration_method.RANDOM)
    registration_method.SetMetricSamplingPercentage(0.01)
    registration_method.SetInterpolator(sitk.sitkLinear)
    registration_method.SetOptimizerAsGradientDescent(learningRate=1.0, numberOfIterations=100,
                                                      convergenceMinimumValue=1e-6, convergenceWindowSize=10)
    registration_method.SetOptimizerScalesFromPhysicalShift()
    registration_method.SetShrinkFactorsPerLevel(shrinkFactors=[2, 1])  # [8, 4, 2, 1]
    registration_method.SetSmoothingSigmasPerLevel(smoothingSigmas=[1, 0])  # [3, 2, 1, 0]
    registration_method.SmoothingSigmasAreSpecifiedInPhysicalUnitsOn()
    registration_method.SetInitialTransform(initial_transform, inPlace=False)

    final_transform = registration_method.Execute(sitk.Cast(fixed_image, sitk.sitkFloat32),
                                                  sitk.Cast(moving_image, sitk.sitkFloat32))

    return final_transform

def bspline_registration(fixed, moving):
    transformDomainMeshSize = [8] * moving.GetDimension()
    tx = sitk.BSplineTransformInitializer(fixed,
                                          transformDomainMeshSize)

    print("Initial Parameters:");
    print(tx.GetParameters())

    R = sitk.ImageRegistrationMethod()
    R.SetMetricAsCorrelation()

    R.SetOptimizerAsLBFGSB(gradientConvergenceTolerance=1e-5,
                           numberOfIterations=100,
                           maximumNumberOfCorrections=5,
                           maximumNumberOfFunctionEvaluations=1000,
                           costFunctionConvergenceFactor=1e+7)
    R.SetInitialTransform(tx, True)
    R.SetInterpolator(sitk.sitkLinear)

    # R.AddCommand( sitk.sitkIterationEvent, lambda: command_iteration(R) )

    outTx = R.Execute(fixed, moving)
    # a=outTx.TransformPoint((0,0))

    print("-------")
    print(outTx)
    print("Optimizer stop condition: {0}".format(R.GetOptimizerStopConditionDescription()))
    print(" Iteration: {0}".format(R.GetOptimizerIteration()))
    print(" Metric value: {0}".format(R.GetMetricValue()))

    resampler = sitk.ResampleImageFilter()
    resampler.SetReferenceImage(fixed)
    resampler.SetInterpolator(sitk.sitkLinear)
    # resampler.SetDefaultPixelValue(100)
    resampler.SetTransform(outTx)

    out = resampler.Execute(moving)
    return out, outTx

def non_rigid_registration(fixedImage, movingImage, default_tranform='bspline', grid_size=16, NumberOfResolutions=4, MaximumNumberOfIterations=500):
    elastixImageFilter = sitk.ElastixImageFilter()
    elastixImageFilter.SetFixedImage(fixedImage)
    elastixImageFilter.SetMovingImage(movingImage)

    parameterMapVector = sitk.VectorOfParameterMap()
    parameterMapVector.append(sitk.GetDefaultParameterMap("bspline"))
    elastixImageFilter.SetParameterMap(parameterMapVector)

    sitk.PrintParameterMap(elastixImageFilter.GetParameterMap())
    elastixImageFilter.LogToConsoleOn()
    elastixImageFilter.Execute()
    transformParameterMap = elastixImageFilter.GetTransformParameterMap()
    outGrayArray = sitk.GetArrayFromImage(elastixImageFilter.GetResultImage())

    return outGrayArray, transformParameterMap  # return deformationField

def deform_array(init_tx, transformParameterMap, movingArray: np.array, refImage: sitk.Image):
    if len(movingArray.shape) == 2:
        movingArray = np.expand_dims(movingArray, axis=2)

    initArray, outArray = [], []
    for c in range(movingArray.shape[2]):
        movingChannelImage = sitk.GetImageFromArray(movingArray[:, :, c])
        movingResampled = sitk.Resample(movingChannelImage, refImage, init_tx, sitk.sitkLinear,
                                        movingChannelImage[0, 0],
                                        movingChannelImage.GetPixelID())  # default pixel=moving_image[0,0]
        initArray.append(sitk.GetArrayFromImage(movingResampled))

        transformixImageFilter = sitk.TransformixImageFilter()
        transformixImageFilter.SetTransformParameterMap(transformParameterMap)

        transformixImageFilter.SetMovingImage(movingResampled)
        transformixImageFilter.ComputeDeformationFieldOn()
        transformixImageFilter.Execute()
        out = transformixImageFilter.GetResultImage()

        outArray.append(sitk.GetArrayFromImage(out))
    transformedArray = np.dstack(initArray)  # intial rigid transformed image, before deformation
    deformedArray = np.dstack(outArray)  # final deformed image, after rigid, norigid transformation
    return np.squeeze(transformedArray), np.squeeze(deformedArray), transformixImageFilter.GetDeformationField()

def deform_array1(init_tx, outTx, movingArray: np.array, refImage: sitk.Image):
    if len(movingArray.shape) == 2:
        movingArray = np.expand_dims(movingArray, axis=2)

    initArray, outArray = [], []
    for c in range(movingArray.shape[2]):
        movingChannelImage = sitk.GetImageFromArray(movingArray[:, :, c])
        movingResampled = sitk.Resample(movingChannelImage, refImage, init_tx, sitk.sitkLinear,
                                        movingChannelImage[0, 0],
                                        movingChannelImage.GetPixelID())  # default pixel=moving_image[0,0]
        initArray.append(sitk.GetArrayFromImage(movingResampled))

        resampler = sitk.ResampleImageFilter()
        resampler.SetReferenceImage(refImage)
        resampler.SetInterpolator(sitk.sitkLinear)
        # resampler.SetDefaultPixelValue(100)
        resampler.SetTransform(outTx)

        out = resampler.Execute(movingResampled)
        outArray.append(sitk.GetArrayFromImage(out))
    transformedArray = np.dstack(initArray)  # intial rigid transformed image, before deformation
    deformedArray = np.dstack(outArray)  # final deformed image, after rigid, norigid transformation
    return np.squeeze(transformedArray), np.squeeze(deformedArray)

def write_deform_field1(init_trans, dis_tx, prefix, fixedImage):

    grid_image = sitk.GridSource(outputPixelType=sitk.sitkUInt16,
                                 size=fixedImage.GetSize(),
                                 sigma=(0.1, 0.1), gridSpacing=(32.0, 32.0))
    # grid_image.CopyInformation(deformationField)\
    gridArr = np.zeros((grid_image.GetSize()[1], grid_image.GetSize()[0]))

    deformArray = sitk.GetArrayFromImage(fixedImage)

    # grid_image = sitk.ReadImage('Z:/PUBLIC/lab_members/inyeop_jang/data/organized_datasets/Van_Abel_HE_thumbnails/1664369/1664369_MR16-1693 I4_LN_HE.png')
    SRCtoTRG = np.zeros_like(sitk.GetArrayFromImage(fixedImage))

    df = pd.DataFrame(index=range(deformArray.shape[0] * deformArray.shape[1]),
                      columns=['source_x', 'source_y', 'target_x', 'target_y'])
    print(prefix)
    index = 0
    for y in tqdm(range(deformArray.shape[0])):
        for x in range(deformArray.shape[1]):
            tx, ty = init_trans.TransformPoint((x, y))
            tx, ty = int(np.floor(tx)), int(np.floor(ty))
            if tx < 0 or tx >= deformArray.shape[1] or ty < 0 or ty >= deformArray.shape[0]: continue

            nx, ny = dis_tx.TransformPoint(
                (tx, ty))  # (x - deformArray[(y,x)][0], y -deformArray[(y,x)][1]) #dis_tx.TransformPoint((x,y))
            if nx >= 0 and nx < SRCtoTRG.shape[1] and ny >= 0 and ny < SRCtoTRG.shape[0]:
                new_x, new_y = int(np.floor(nx)), int(np.floor(ny))
                SRCtoTRG[(new_y, new_x)] = fixedImage[(x, y)]
                gridArr[(new_y, new_x)] = grid_image[(x, y)]
                df.loc[index, ['source_x', 'source_y', 'target_x', 'target_y']] = [x, y, new_x,
                                                                                   new_y]  # fixed -> moved
                index += 1

    df.to_csv(prefix + '/deformField.csv', index=False)
    cv2.imwrite(prefix + '/displaceFixedImage.jpg', SRCtoTRG)
    cv2.imwrite(prefix + '/deformGrid.jpg', gridArr.astype(np.uint8))

    # resampler = sitk.ResampleImageFilter()
    # resampler.SetReferenceImage(deformationField)  # Or any target geometry
    # resampler.SetTransform(sitk.DisplacementFieldTransform(
    #     sitk.Cast(deformationField, sitk.sitkVectorFloat64)))
    # warped = resampler.Execute(movingR)
    # cv2.imwrite(prefix+'warped.jpg', sitk.GetArrayFromImage(warped))

def write_deform_field(init_trans, deformationField, prefix, fixedImage):
    # afine_tx =sitk.AffineTransform(sitk.Cast(deformationField, sitk.sitkVectorFloat64))
    dis_tx = sitk.DisplacementFieldTransform(sitk.Cast(deformationField, sitk.sitkVectorFloat64))

    grid_image = sitk.GridSource(outputPixelType=sitk.sitkUInt16,
                                 size=deformationField.GetSize(),
                                 sigma=(0.1, 0.1), gridSpacing=(32.0, 32.0))
    # grid_image.CopyInformation(deformationField)\
    gridArr = np.zeros((grid_image.GetSize()[1], grid_image.GetSize()[0]))

    deformArray = sitk.GetArrayFromImage(deformationField)

    # grid_image = sitk.ReadImage('Z:/PUBLIC/lab_members/inyeop_jang/data/organized_datasets/Van_Abel_HE_thumbnails/1664369/1664369_MR16-1693 I4_LN_HE.png')
    SRCtoTRG = np.zeros_like(sitk.GetArrayFromImage(fixedImage))

    df = pd.DataFrame(index=range(deformArray.shape[0] * deformArray.shape[1]),
                      columns=['source_x', 'source_y', 'target_x', 'target_y'])
    print(prefix)
    index = 0
    for y in tqdm(range(deformArray.shape[0])):
        for x in range(deformArray.shape[1]):
            tx, ty = init_trans.TransformPoint((x, y))
            tx, ty = int(np.floor(tx)), int(np.floor(ty))
            if tx < 0 or tx >= deformArray.shape[1] or ty < 0 or ty >= deformArray.shape[0]: continue

            nx, ny = dis_tx.TransformPoint(
                (tx, ty))  # (x - deformArray[(y,x)][0], y -deformArray[(y,x)][1]) #dis_tx.TransformPoint((x,y))
            if nx >= 0 and nx < SRCtoTRG.shape[1] and ny >= 0 and ny < SRCtoTRG.shape[0]:
                new_x, new_y = int(np.floor(nx)), int(np.floor(ny))
                SRCtoTRG[(new_y, new_x)] = fixedImage[(x, y)]
                gridArr[(new_y, new_x)] = grid_image[(x, y)]
                df.loc[index, ['source_x', 'source_y', 'target_x', 'target_y']] = [x, y, new_x, new_y]
                index += 1

    df.to_csv(prefix + '/deformField.csv', index=False)
    cv2.imwrite(prefix + '/displaceFixedImage.jpg', SRCtoTRG)
    cv2.imwrite(prefix + '/deformGrid.jpg', gridArr.astype(np.uint8))

    # resampler = sitk.ResampleImageFilter()
    # resampler.SetReferenceImage(deformationField)  # Or any target geometry
    # resampler.SetTransform(sitk.DisplacementFieldTransform(
    #     sitk.Cast(deformationField, sitk.sitkVectorFloat64)))
    # warped = resampler.Execute(movingR)
    # cv2.imwrite(prefix+'warped.jpg', sitk.GetArrayFromImage(warped))

def inverse_deformationfield(deformationField):
    for y in range(deformationField.GetSize()[1]):
        for x in range(deformationField.GetSize()[0]):
            deformationField[x, y] = (-deformationField[x, y][0], -deformationField[x, y][1])
    return deformationField

def deform_by_deformatinfield(deformationField: sitk.Image, sourceArray: np.array, refImage: sitk.Image = None):
    resampler = sitk.ResampleImageFilter()

    # deformationField = sitk.GetImageFromArray(deformationArray)
    if refImage == None: refImage = deformationField
    resampler.SetReferenceImage(refImage)  # Or any target geometry
    resampler.SetTransform(sitk.DisplacementFieldTransform(
        sitk.Cast(deformationField, sitk.sitkVectorFloat64)))
    resampler.SetInterpolator(sitk.sitkNearestNeighbor)

    if len(sourceArray.shape) == 2:  # if grayimage
        sourceArray = np.expand_dims(sourceArray, axis=2)

    outChannel = []
    for C in range(sourceArray.shape[2]):
        deformed = sitk.GetArrayFromImage(resampler.Execute(sitk.GetImageFromArray(sourceArray[:, :, C])))
        outChannel.append(deformed)
    outArray = np.dstack(outChannel).squeeze()
    # cv2.imwrite('warped1664369.jpg', RGB)
    return outArray

def load_IHC_HE(fixed_path, moving_path):
    fixedArray = cv2.imread(fixed_path, cv2.IMREAD_COLOR)
    # movingImage = sitk.ReadImage(moving_path, sitk.sitkFloat32)
    movingArray = cv2.imread(moving_path, cv2.IMREAD_COLOR)  # IHC

    if 'HE' in moving_path:
        IHC = fixedArray
        HE = movingArray
    else:
        IHC = movingArray
        HE = fixedArray

    # Separate the stains from the IHC image
    ihc_hed = rgb2hed(IHC)
    # Create an RGB image for each of the stains
    null = np.zeros_like(ihc_hed[:, :, 0])
    ihc_h = hed2rgb(np.stack((ihc_hed[:, :, 0], null, null), axis=-1))  # hematoxylin
    fixedIHC = (ihc_h[:, :, 0] * 255.0).astype(np.uint8)

    _, movingH, _ = unmixHE(HE)
    movingHE = cv2.cvtColor(movingH, cv2.COLOR_RGB2GRAY)

    return fixedIHC, movingHE

def read_json(json_path):
    bbox_list = []
    centroid_list = []
    contour_list = []
    type_list = []
    with open(json_path) as json_file:
        data = json.load(json_file)
        mag_info = data['mag']
        nuc_info = data['nuc']
        for inst in nuc_info:
            inst_info = nuc_info[inst]
            inst_centroid = inst_info['centroid']
            centroid_list.append(inst_centroid)
            inst_contour = inst_info['contour']
            contour_list.append(inst_contour)
            inst_bbox = inst_info['bbox']
            bbox_list.append(inst_bbox)
            inst_type = inst_info['type']
            type_list.append(inst_type)

    print('Number of centroids', len(centroid_list))
    print('Number of contours', len(contour_list))
    print('Number of bounding boxes', len(bbox_list))

    # each item is a list of coordinates - let's take a look!
    print('-'*60)
    print(centroid_list[0])
    print('-'*60)
    print(contour_list[0])
    print('-'*60)
    print(bbox_list[0])
    return bbox_list, centroid_list, contour_list, type_list

def get_info_dict(contour_list, coords_xmin, coords_ymin, coords_xmax, coords_ymax):
    info_dict = {}
    count = 0
    for idx, cnt in enumerate(contour_list):
        cnt_tmp = np.array(cnt)
        cnt_tmp = cnt_tmp[(cnt_tmp[:,0] >= coords_xmin) & (cnt_tmp[:,0] <= coords_xmax) & (cnt_tmp[:,1] >= coords_ymin) & (cnt_tmp[:,1] <= coords_ymax)]
        label = str(type_list[idx])
        if cnt_tmp.shape[0] > 0:
            cnt_adj = cnt_tmp.astype('int')
            info_dict[idx] = {'contour': cnt_adj, 'type':label}
            count += 1
    return info_dict

def plot_2d_images_at_z_indices(arr_3d, z_indices, key, savepath=None):
    fig, axs = plt.subplots(1, len(z_indices), figsize=(15, 4))
    for i, z_index in enumerate(z_indices):
        axs[i].imshow(arr_3d[:, :, z_index], cmap='gray')
        axs[i].set_title(f"Z Index {z_index}")
        axs[i].axis('off')
    fig.suptitle(key, fontsize=16)
    plt.subplots_adjust(wspace=0.05)
    plt.tight_layout()
    if savepath !=None:
        plt.savefig(savepath)
    else:
        plt.show()
def save_to_nii(data, path, resolution_x=0.15543286502361298, resolution_y=0.15543286502361298, resolution_z=0.949573814868927):
    nifti_img = nib.Nifti1Image(data.astype(np.int16), affine=np.eye(4))
    nifti_img.header['pixdim'][1] = resolution_x
    nifti_img.header['pixdim'][2] = resolution_y
    nifti_img.header['pixdim'][3] = resolution_z
    nib.save(nifti_img, path)

def post_process_cell_mask(cell_mask, collagen_mask):
    if cell_mask.shape == collagen_mask.shape:
        non_overlap_mask = np.zeros_like(cell_mask)
        non_overlap_mask[np.logical_and(collagen_mask == 0, cell_mask != 0)] = 1
        new_cell_mask = non_overlap_mask * cell_mask
        return new_cell_mask
    else:
        raise ('The shape of cell mask and collagen mask are different')


# %% Model functions
from tensorflow.keras.layers import (Input, Dense, Conv2D, MaxPooling2D, Dropout, Add, GlobalAveragePooling2D,
                                     Activation, BatchNormalization, UpSampling2D, Cropping2D, Concatenate)
from tensorflow.keras.models import Sequential, Model
def _relu(inputs):
    return Activation('relu')(inputs)
def _bn_relu(inputs):
    bn = BatchNormalization()(inputs)
    relu = Activation('relu')(bn)
    return relu
def _conv_relu(out_dim, ks, strides, padding):
    def f(inputs):
        conv = Conv2D(out_dim, ks, strides=strides, padding=padding)(inputs)
        conv = _relu(conv)
        return conv
    return f

def _conv_bn_relu(out_dim, ks, strides, padding):
    def f(inputs):
        conv = Conv2D(out_dim, ks, strides=strides, padding=padding)(inputs)
        conv = _bn_relu(conv)
        return conv
    return f
def _res_unit(n_in, n_mid, n_out, s, inputs, _conv_blk=_conv_relu, _skip_act=_relu):
    conv = _conv_blk(n_mid, 1, strides=(1, 1), padding='same')(inputs)
    conv = _conv_blk(n_mid, 3, strides=(s, s), padding='same')(conv)
    conv = Conv2D(n_out, 1, strides=(1, 1), padding='same')(conv)
    if n_in != n_out:
        shortcut = Conv2D(n_out, 1, strides=(s, s), padding='same')(inputs)
    else:
        shortcut = inputs
    conv = Add()([conv, shortcut])  # all the operation should be implemented as layers
    conv = _skip_act(conv)
    return conv
def _res_stage(n_in, n_mid, n_out, first_strides, inputs, n_units, _conv_blk=_conv_relu, _skip_act=_relu):
    unit = _res_unit(n_in, n_mid, n_out, first_strides, inputs, _conv_blk, _skip_act)
    for i in range(n_units - 1):
        #         print('stage %s' % str(i))
        unit = _res_unit(n_out, n_mid, n_out, 1, unit, _conv_blk, _skip_act)
    return unit
def _dense_blk(inputs, n_units, act=_relu, _con_blk=_conv_relu):
	for i in range(n_units):
		conv = act(inputs)
		conv = _con_blk(128, 1, strides=(1,1), padding='same')(conv)
		conv = Conv2D(32, 5, strides=(1,1), padding='same')(conv)
		# inputs = Cropping2D((2,2))(inputs)
		inputs = Concatenate(axis=-1)([inputs, conv])
	return inputs

def encoder(inputs):
    base_chnl = 64
    stage_chls = 2**np.arange(0,4) * base_chnl # 64, 128, 256, 512
    stage_units = [3,4,6,3]
    # conv1
    conv1 = _conv_relu(base_chnl, 7, strides=(1,1), padding='same')(inputs)
    # pool1
    # pool1 = MaxPooling2D(pool_size=(3,3), strides=(2,2), padding='same')(conv1)
    # stage 2
    stage1 = _res_stage(stage_chls[0],   stage_chls[0], stage_chls[0]*4, 1, conv1, stage_units[0])
    stage2 = _res_stage(stage_chls[0]*4, stage_chls[1], stage_chls[1]*4, 2, stage1, stage_units[1])
    stage3 = _res_stage(stage_chls[1]*4, stage_chls[2], stage_chls[2]*4, 2, stage2, stage_units[2])
    stage4 = _res_stage(stage_chls[2]*4, stage_chls[3], stage_chls[3]*4, 2, stage3, stage_units[3])
    stage4 = Conv2D(1024, 1, strides=(1,1), padding='same')(stage4)   # 这里不带activation
    return [stage1, stage2, stage3, stage4]

def decoder(encoder):
	chnls = 2**np.arange(1,5) * 64
	d1, d2, d3, d4 = encoder
	conv1 = UpSampling2D(size=(2,2))(d4)  # 66, 1024
	conv1 = Add()([d3, conv1])
	conv1 = Conv2D(256, 5, strides=(1,1), padding='same')(conv1)  # 62, 256, 5*5
	dense1 = _dense_blk(conv1, 8, _relu, _conv_relu)
	conv1 = Conv2D(512, 1, strides=(1,1), padding='same')(dense1) # no activation

	conv2 = UpSampling2D(size=(2,2))(conv1)
	conv2 = Add()([d2, conv2])
	conv2 = Conv2D(128, 5, strides=(1,1), padding='same')(conv2)  # no activation
	dense2 = _dense_blk(conv2, 4, _relu, _conv_relu)
	conv2 = Conv2D(256, 1, strides=(1,1), padding='same')(dense2) # no activation

	conv3 = UpSampling2D(size=(2,2))(conv2)
	conv3 = Add()([d1, conv3])
	conv3 = Conv2D(64, 5, strides=(1,1), padding='same')(conv3)    # no activation
	return conv3

def hvnet(input_shape, n_chnl):
    inputs = Input(input_shape+(n_chnl,))
    enc = encoder(inputs)

    np_decoder = decoder(enc)
    np_decoder = _relu(np_decoder)
    np_decoder = Conv2D(1, 1, strides=(1,1), padding='same', activation='sigmoid', name='np')(np_decoder)

    hv_decoder = decoder(enc)
    hv_decoder = _relu(hv_decoder)
    hv_decoder = Conv2D(2, 1, strides=(1,1), padding='same', name='hv')(hv_decoder)

    clg_decoder = decoder(enc)
    clg_decoder = _relu(clg_decoder)
    clg_decoder = Conv2D(1, 1, strides=(1, 1), padding='same', activation='sigmoid', name='clg')(clg_decoder)
    model = Model(inputs, outputs=[np_decoder, hv_decoder, clg_decoder])
    return model

def normalization(img):
    norm_img = (img-img.min())/(img.max()-img.min())
    return norm_img

def std(x):
    t = (x-x.mean())/(x.std())
    return t

def histogram_equalization(image):
    # Calculate the histogram
    hist, bins = np.histogram(image.flatten(), bins=1000000, range=[0, 1])
    # Calculate the cumulative distribution function (CDF)
    cdf = hist.cumsum()
    # Normalize the CDF to the range [0, 1]
    cdf_normalized = cdf / float(cdf[-1])
    # Perform histogram equalization
    equalized_image = np.interp(image.flatten(), bins[:-1], cdf_normalized)
    equalized_image = equalized_image.reshape(image.shape)
    # Normalize the equalized image to the range [0, 1]
    equalized_image = (equalized_image - np.min(equalized_image)) / (np.max(equalized_image) - np.min(equalized_image))
    return equalized_image

def get_data_from_tilenames(root, tilenames, cellmask = None):
    x = []
    y = []
    # Iterate over each PNG file in png_files
    for item in tilenames:
        if len(item.split('_')) != 3:
            y.append(int(item.split('_')[0]))
            x.append(int(item.split('_')[1]))
        else:
            y.append(int(item.split('.')[0].split('_')[1]))
            x.append(int(item.split('.')[0].split('_')[2]))
    # return x, y, num_channels
    if cellmask != None:
        num_channels = np.load(os.path.join(root, tilenames[0])).shape[-1]
        empty_arr = np.zeros((5, (np.max(x)+1)*256, (np.max(y)+1)*256, num_channels)).astype(np.uint8)
        for i in range(len(x)):
            tile_img = np.load(os.path.join(root, tilenames[i]))
            empty_arr[:, x[i]*256 :x[i]*256 + 256, y[i]*256: y[i]*256 +256, :] = tile_img
    else:
        num_channels = np.load(os.path.join(root, tilenames[0])).shape[2]
        empty_arr = np.zeros(((np.max(x)+1)*256, (np.max(y)+1)*256, num_channels)).astype(np.float32)
        for i in range(len(x)):
            tile_img = np.load(os.path.join(root, tilenames[i]))
            print(tile_img.shape)
            empty_arr[x[i]*256 :x[i]*256 + 256, y[i]*256: y[i]*256 +256, :] = tile_img
    return empty_arr