<reponame>bf777/MesoNet<filename>mesonet/atlas_brain_matching.py
"""
MesoNet
Authors: <NAME> and <NAME>, <NAME>
https://github.com/bf777/MesoNet
Licensed under the Creative Commons Attribution 4.0 International License (see LICENSE for details)
"""
from mesonet.utils import natural_sort_key, convert_to_png
from mesonet.mask_functions import atlas_to_mask, applyMask

# from mesonet.niftyreg_align import niftyreg_align
from mesonet.voxelmorph_align import voxelmorph_align, vxm_transform
import numpy as np
import pandas as pd
import cv2
import imutils
import math
import scipy.io
import skimage.io as io
import skimage.transform as trans
from skimage.transform import PiecewiseAffineTransform, warp
import imageio
import os
import fnmatch
import glob
from PIL import Image


def find_peaks(img):
    """
    Locates the peaks of cortical activity in each hemisphere of the brain image.
    :param img: A sensory map (indicating functional activation in the mouse brain).
    :return: An array of the maximum points of activation for the sensory map.
    """
    max_loc_arr = []
    img = cv2.imread(str(img), 0)
    im = img.copy()
    x_min = int(np.around(im.shape[0] / 2))
    im1 = im[:, x_min : im.shape[0]]
    im2 = im[:, 0:x_min]
    (minVal, max_val, minLoc, max_loc) = cv2.minMaxLoc(im1)
    (minVal2, maxVal2, minLoc2, maxLoc2) = cv2.minMaxLoc(im2)
    max_loc = list(max_loc)
    max_loc[0] = max_loc[0] + x_min
    max_loc = tuple(max_loc)
    max_loc_arr.append(max_loc)
    if (max_val - 30) <= maxVal2 <= (max_val + 30):
        max_loc_arr.append(maxLoc2)
    return max_loc_arr


def coords_to_mat(
    sub_dlc_pts, i, output_mask_path, bregma_present, bregma_index, landmark_arr
):
    if bregma_present:
        x_bregma, y_bregma = sub_dlc_pts[bregma_index]
        for pt, landmark in zip(sub_dlc_pts, landmark_arr):
            x_pt = pt[0]
            y_pt = pt[1]
            pt_adj = [landmark, x_pt - x_bregma, y_pt - y_bregma]
            pt_adj_to_mat = np.array(pt_adj, dtype=np.object)
            if not os.path.isdir(os.path.join(output_mask_path, "mat_coords")):
                os.mkdir(os.path.join(output_mask_path, "mat_coords"))
            scipy.io.savemat(
                os.path.join(
                    output_mask_path,
                    "mat_coords/landmarks_{}_{}.mat".format(i, landmark),
                ),
                {"landmark_coords_{}_{}".format(i, landmark): pt_adj_to_mat},
                appendmat=False,
            )


def sensory_to_mat(sub_dlc_pts, bregma_pt, i, output_mask_path):
    x_bregma, y_bregma = bregma_pt
    sensory_names = ["tail_left", "tail_right", "visual", "whisker"]
    for pt, landmark in zip(sub_dlc_pts, sensory_names):
        x_pt = pt[0]
        y_pt = pt[1]
        pt_adj = [landmark, x_pt - x_bregma, y_pt - y_bregma]
        pt_adj_to_mat = np.array(pt_adj, dtype=np.object)
        if not os.path.isdir(os.path.join(output_mask_path, "mat_coords")):
            os.mkdir(os.path.join(output_mask_path, "mat_coords"))
        scipy.io.savemat(
            os.path.join(
                output_mask_path,
                "mat_coords/sensory_peaks_{}_{}.mat".format(i, landmark),
            ),
            {"sensory_peaks_{}_{}".format(i, landmark): pt_adj_to_mat},
            appendmat=False,
        )


def atlas_from_mat(input_file, mat_cnt_list):
    """
    Generates a binary brain atlas from a .mat file.
    :param input_file: The input .mat file representing a brain atlas (with white = 255 and black = 0)
    :param mat_cnt_list: The list to which mat files should be appended
    :return: A thresholded binary brain atlas.
    """
    file = input_file
    atlas_base = np.zeros((512, 512), dtype="uint8")
    if glob.glob(os.path.join(input_file, "*.mat")):
        mat = scipy.io.loadmat(file)
        mat_shape = mat[list(mat.keys())[3]]
        if len(mat_shape.shape) > 2:
            for val in range(0, mat_shape.shape[2]):
                mat_roi = mat_shape[:, :, val]
                mat_resize = cv2.resize(mat_roi, (512, 512))
                mat_resize = np.uint8(mat_resize)
                ret, thresh = cv2.threshold(mat_resize, 5, 255, cv2.THRESH_BINARY_INV)
                mat_roi_cnt = cv2.findContours(
                    thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
                )
                mat_roi_cnt = imutils.grab_contours(mat_roi_cnt)
                c_to_save = max(mat_roi_cnt, key=cv2.contourArea)
                mat_cnt_list.append(c_to_save)
                cv2.drawContours(atlas_base, mat_roi_cnt, -1, (255, 255, 255), 1)
            ret, thresh = cv2.threshold(atlas_base, 5, 255, cv2.THRESH_BINARY_INV)
            io.imsave("atlas_unresized_test.png", thresh)
        else:
            mat = mat["atlas"]
            mat_resize = cv2.resize(mat, (512, 512))
            ret, thresh = cv2.threshold(mat_resize, 5, 255, cv2.THRESH_BINARY_INV)
    else:
        atlas_im = cv2.imread(file, cv2.IMREAD_GRAYSCALE)
        atlas_resize = np.uint8(atlas_im)
        ret, atlas_resize = cv2.threshold(atlas_resize, 127, 255, 0)
        io.imsave("atlas_unresized_test.png", atlas_resize)
        roi_cnt, hierarchy = cv2.findContours(
            atlas_resize, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE
        )[-2:]
        for val in roi_cnt:
            c_to_save = max(val, key=cv2.contourArea)
            mat_cnt_list.append(c_to_save)
            cv2.drawContours(atlas_base, val, -1, (255, 255, 255), 1)
        ret, thresh = cv2.threshold(atlas_base, 5, 255, cv2.THRESH_BINARY_INV)
    return thresh


def atlas_rotate(dlc_pts, im):
    dlc_y_pts = [
        coord if (190 <= coord[0] <= 330) else (1000, 1000) for coord in dlc_pts
    ]
    dlc_y_pts = [coord for coord in dlc_y_pts if coord[0] < 1000]

    # https://stackoverflow.com/questions/42258637/how-to-know-the-angle-between-two-points
    rotate_rad = math.atan2(0, (im.shape[1] / 2) - dlc_y_pts[-1][0])
    rotate_deg = -1 * (abs(math.degrees(rotate_rad)))
    im_rotate_mat = cv2.getRotationMatrix2D(
        (im.shape[1] / 2, im.shape[0] / 2), rotate_deg, 1.0
    )
    im_rotated = cv2.warpAffine(im, im_rotate_mat, (512, 512))
    x_min = int(np.around(im_rotated.shape[0] / 2))
    im_left = im_rotated[:, 0:x_min]
    im_right = im_rotated[:, x_min : im_rotated.shape[0]]
    return im_left, im_right


def getMaskContour(mask_dir, atlas_img, predicted_pts, actual_pts, cwd, n, main_mask):
    """
    Gets the contour of the brain's boundaries and applies a piecewise affine transform to the brain atlas
    based on the cortical landmarks predicted in dlc_predict (and peaks of activity on the sensory map, if available).
    :param mask_dir: The path to the directory containing the U-net masks of the brain's boundaries.
    :param atlas_img: The brain atlas to be transformed.
    :param predicted_pts: The coordinates of the cortical landmarks predicted in dlc_predict (or, for the second run
    of this function, the coordinates of the peaks of activity in the sensory map).
    :param actual_pts: The fixed coordinates of the cortical landmarks on the brain atlas (or, for the second run of
    this function, the fixed coordinates of the peaks of sensory activity on the brain atlas).
    :param cwd: The path to the current working directory.
    :param n: The number of the current image in the directory.
    """
    c_landmarks = np.empty([0, 2])
    c_atlas_landmarks = np.empty([0, 2])
    mask = cv2.imread(mask_dir, cv2.IMREAD_GRAYSCALE)
    atlas_to_warp = atlas_img
    mask = np.uint8(mask)
    cnts, hierarchy = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)[-2:]
    for cnt in cnts:
        cnt = cnt[:, 0, :]
        cnt = np.asarray(cnt).astype("float32")
        c_landmarks = np.concatenate((c_landmarks, cnt))
        c_atlas_landmarks = np.concatenate((c_atlas_landmarks, cnt))
    c_landmarks = np.concatenate((c_landmarks, predicted_pts))
    c_atlas_landmarks = np.concatenate((c_atlas_landmarks, actual_pts))
    tform = PiecewiseAffineTransform()
    tform.estimate(c_atlas_landmarks, c_landmarks)
    dst = warp(atlas_to_warp, tform, output_shape=(512, 512))
    if main_mask:
        io.imsave(os.path.join(cwd, "mask_{}.png".format(n)), mask)
    return dst


def atlasBrainMatch(
    brain_img_dir,
    sensory_img_dir,
    coords_input,
    sensory_match,
    mat_save,
    threshold,
    git_repo_base,
    region_labels,
    landmark_arr_orig,
    use_unet,
    use_dlc,
    atlas_to_brain_align,
    model,
    olfactory_check,
    plot_landmarks,
    align_once,
    original_label,
    use_voxelmorph,
    exist_transform,
    voxelmorph_model="motif_model_atlas.h5",
    vxm_template_path="templates",
    dlc_template_path="dlc_templates",
    flow_path="",
):
    """
    Align and overlap brain atlas onto brain image based on four landmark locations in the brain image and the atlas.
    :param brain_img_dir: The directory containing each brain image to be used.
    :param sensory_img_dir: The directory containing each sensory image to be used (if you are aligning each brain
    image using a sensory map).
    :param coords_input: Predicted locations of the four landmarks on the brain image from the file generated by
    DeepLabCut.
    :param sensory_match: Whether or not a sensory map is to be used.
    :param mat_save: Whether or not to export each brain region to a .mat file in applyMask, which is called at the end
    of this function.
    :param threshold: The threshold for the cv2.opening operation carried out in applyMask, which is called at the end
    of this function.
    :param git_repo_base: The path to the base git repository containing necessary resources for MesoNet (reference
    atlases, DeepLabCut config files, etc.)
    :param region_labels: Choose whether or not to attempt to label each region with its name from the Allen Institute
    Mouse Brain Atlas.
    :param landmark_arr_orig: The original array of landmarks from DeepLabCut (to be distinguished from any automatic
    exclusions to landmark array based on prediction quality).
    :param use_unet: Choose whether or not to identify the borders of the cortex using a U-net model.
    :param atlas_to_brain_align: If True, registers the atlas to each brain image. If False, registers each brain image
    to the atlas.
    :param model: The name of the U-net model (for passthrough to mask_functions.py)
    :param olfactory_check: If True, draws olfactory bulb contours on the brain image.
    :param plot_landmarks: If True, plots DeepLabCut landmarks (large circles) and original alignment landmarks (small
    circles) on final brain image.
    :param align_once: if True, carries out all alignments based on the alignment of the first atlas and brain. This can
    save time if you have many frames of the same brain with a fixed camera position.
    :param original_label: if True, uses a brain region labelling approach that attempts to automatically sort brain
    regions in a consistent order (left to right by hemisphere, then top to bottom for vertically aligned regions). This
    approach may be more flexible if you're using a custom brain atlas (i.e. not one in which region is filled with a
    unique number).
    :param exist_transform: if True, uses an existing voxelmorph transformation field for all data instead of predicting
    a new transformation.
    :param voxelmorph_model: the name of a .h5 model located in the models folder of the git repository for MesoNet,
    generated using voxelmorph and containing weights for a voxelmorph local deformation model.
    :param vxm_template_path: the path to a template atlas (.npy or .mat( to which the brain image will be aligned in
    voxelmorph.
    :param flow_path: the path to a voxelmorph transformation field that will be used to transform all data instead of
    predicting a new transformation if exist_transform is True.
    """
    # load brain images folder
    brain_img_arr = []
    dlc_img_arr = []
    peak_arr = []
    atlas_label_list = []
    dst_list = []
    vxm_template_list = []
    br_list = []
    olfactory_bulbs_to_use_list = []
    olfactory_bulbs_to_use_pre_align_list = []

    voxelmorph_model_path = os.path.join(
        git_repo_base, "models", "voxelmorph", voxelmorph_model
    )

    # Prepare template for VoxelMorph
    convert_to_png(vxm_template_path)
    vxm_template_orig = cv2.imread(
        glob.glob(os.path.join(git_repo_base, "atlases", vxm_template_path, "*.png"))[0]
    )

    # Prepare output folder
    cwd = os.getcwd()
    output_mask_path = os.path.join(cwd, "../output_mask")
    # Output folder for transparent masks and masks overlaid onto brain image
    output_overlay_path = os.path.join(cwd, "../output_overlay")
    if not os.path.isdir(output_mask_path):
        os.mkdir(output_mask_path)
    if not os.path.isdir(output_overlay_path):
        os.mkdir(output_overlay_path)

    if not atlas_to_brain_align:
        im = cv2.imread(
            os.path.join(git_repo_base, "atlases/Atlas_workflow2_binary.png")
        )
    else:
        if use_voxelmorph and not use_dlc:
            im = cv2.imread(
                os.path.join(git_repo_base, "atlases/Atlas_for_Voxelmorph_binary.png")
            )
        else:
            im = cv2.imread(
                os.path.join(git_repo_base, "atlases/Atlas_workflow1_binary.png")
            )
        im_left = cv2.imread(os.path.join(git_repo_base, "atlases/left_hemi.png"))
        ret, im_left = cv2.threshold(im_left, 5, 255, cv2.THRESH_BINARY_INV)
        im_right = cv2.imread(os.path.join(git_repo_base, "atlases/right_hemi.png"))
        ret, im_right = cv2.threshold(im_right, 5, 255, cv2.THRESH_BINARY_INV)
        im_left = np.uint8(im_left)
        im_right = np.uint8(im_right)
        im = np.uint8(im)
    # im = atlas_from_mat(os.path.join(git_repo_base, 'atlases/atlas_ROIs.mat'))
    atlas = im
    # FOR ALIGNING BRAIN TO ATLAS
    for num, file in enumerate(os.listdir(cwd)):
        if fnmatch.fnmatch(file, "*.png") and "mask" not in file:
            dlc_img_arr.append(os.path.join(cwd, file))
    for num, file in enumerate(os.listdir(brain_img_dir)):
        if fnmatch.fnmatch(file, "*.png"):
            brain_img_arr.append(os.path.join(brain_img_dir, file))
            brain_img_arr.sort(key=natural_sort_key)
        elif fnmatch.fnmatch(file, "*.tif"):
            tif_stack = imageio.mimread(os.path.join(brain_img_dir, file))
            for tif_im in tif_stack:
                brain_img_arr.append(tif_im)
    # i_coord, j_coord = np.array([(100, 256, 413, 256), (148, 254, 148, 446)])

    # https://www.pyimagesearch.com/2014/07/21/detecting-circles-images-using-opencv-hough-circles/
    coord_circles_img = cv2.imread(
        os.path.join(
            git_repo_base, "atlases", "multi_landmark", "landmarks_new_binary.png"
        ),
        cv2.IMREAD_GRAYSCALE,
    )
    coord_circles_img = np.uint8(coord_circles_img)
    # detect circles in the image
    circles, hierarchy = cv2.findContours(
        coord_circles_img, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE
    )[-2:]
    # ensure at least some circles were found
    if circles is not None:
        # convert the (x, y) coordinates and radius of the circles to integers
        atlas_arr = np.array(
            [
                (
                    int(cv2.moments(circle)["m10"] / cv2.moments(circle)["m00"]),
                    int(cv2.moments(circle)["m01"] / cv2.moments(circle)["m00"]),
                )
                for circle in circles
            ]
        )

    atlas_arr = np.array(
        [
            (102, 148),
            (166, 88),
            (214, 454),
            (256, 88),
            (256, 256),
            (256, 428),
            (410, 148),
            (346, 88),
            (298, 454),
        ]
    )

    peak_arr_flat = []
    peak_arr_total = []

    if sensory_match:
        for num, file in enumerate(brain_img_arr):
            img_name = str(os.path.splitext(os.path.basename(file))[0])
            sensory_img_for_brain = os.path.join(sensory_img_dir, img_name)
            if glob.glob(sensory_img_for_brain):
                sensory_img_for_brain_dir = os.listdir(sensory_img_for_brain)
                sensory_img_for_brain_dir.sort(key=natural_sort_key)
                for num_im, file_im in enumerate(sensory_img_for_brain_dir):
                    sensory_im = io.imread(
                        os.path.join(sensory_img_dir, img_name, file_im)
                    )
                    sensory_im = trans.resize(sensory_im, (512, 512))
                    io.imsave(
                        os.path.join(sensory_img_dir, img_name, file_im), sensory_im
                    )
                    peak = find_peaks(os.path.join(sensory_img_dir, img_name, file_im))
                    peak_arr.append(peak)
            for x in peak_arr:
                for y in x:
                    peak_arr_flat.append(y)
            peak_arr_total.append(peak_arr_flat)
            peak_arr_flat = []
            peak_arr = []

    dlc_pts = []
    atlas_pts = []
    sensory_peak_pts = []
    sensory_atlas_pts = []
    sub_dlc_pts = []
    sub_atlas_pts = []
    sub_sensory_peak_pts = []
    sub_sensory_atlas_pts = []

    bregma_index_list = []
    bregma_list = []
    if use_dlc:
        bregma_present = True
    else:
        bregma_present = False

    coords = pd.read_csv(coords_input)
    x_coord = coords.iloc[2:, 1::3]
    y_coord = coords.iloc[2:, 2::3]
    accuracy = coords.iloc[2:, 3::3]
    acc_left_total = accuracy.iloc[:, 0:5]
    acc_right_total = accuracy.iloc[:, 3:8]
    landmark_indices = [0, 1, 2, 3, 4, 5, 6, 7, 8]  # [0, 3, 2, 1]
    for arr_index, i in enumerate(range(0, len(x_coord))):
        landmark_arr = landmark_arr_orig
        x_coord_flat = x_coord.iloc[i].values.astype("float32")
        y_coord_flat = y_coord.iloc[i].values.astype("float32")
        x_coord_flat = x_coord_flat[landmark_arr]
        y_coord_flat = y_coord_flat[landmark_arr]
        dlc_list = []
        atlas_list = []
        for (coord_x, coord_y) in zip(x_coord_flat, y_coord_flat):
            dlc_coord = (coord_x, coord_y)
            dlc_list.append(dlc_coord)
        for coord_atlas in atlas_arr:
            atlas_coord = (coord_atlas[0], coord_atlas[1])
            atlas_list.append(atlas_coord)
        atlas_list = [atlas_list[i] for i in landmark_arr]

        # Initialize result as max value
        landmark_indices = landmark_indices[0 : len(landmark_arr)]
        atlas_indices = landmark_arr

        pts_dist = np.absolute(
            np.asarray(atlas_list) - np.asarray((im.shape[0] / 2, im.shape[1] / 2))
        )
        pts_avg_dist = [np.mean(v) for v in pts_dist]
        bregma_index = np.argmin(np.asarray(pts_avg_dist))

        for j in landmark_indices:
            sub_dlc_pts.append([x_coord_flat[j], y_coord_flat[j]])
        for j in atlas_indices:
            sub_atlas_pts.append([atlas_arr[j][0], atlas_arr[j][1]])

        dlc_pts.append(sub_dlc_pts)
        atlas_pts.append(sub_atlas_pts)
        coords_to_mat(
            sub_dlc_pts, i, output_mask_path, bregma_present, bregma_index, landmark_arr
        )
        bregma_index_list.append(bregma_index)
        sub_dlc_pts = []
        sub_atlas_pts = []
    if sensory_match:
        k_coord, m_coord = np.array([(189, 323, 435, 348), (315, 315, 350, 460)])
        coords_peak = peak_arr_total
        for img_num, img in enumerate(brain_img_arr):
            for j in [1, 0, 3, 2]:  # Get peak values from heatmaps
                sub_sensory_peak_pts.append(
                    [coords_peak[img_num][j][0], coords_peak[img_num][j][1]]
                )
            for j in [0, 1, 2, 3]:  # Get circle locations
                sub_sensory_atlas_pts.append([k_coord[j], m_coord[j]])
            sensory_peak_pts.append(sub_sensory_peak_pts)
            sensory_atlas_pts.append(sub_sensory_atlas_pts)
            sensory_to_mat(
                sub_sensory_peak_pts, dlc_pts[img_num][3], img_num, output_mask_path
            )
            sub_sensory_peak_pts = []
            sub_sensory_atlas_pts = []
        sensory_peak_pts, sensory_atlas_pts = (
            np.asarray(sensory_peak_pts).astype("float32"),
            np.asarray(sensory_atlas_pts).astype("float32"),
        )

    for (n, br) in enumerate(brain_img_arr):
        vxm_template = np.uint8(vxm_template_orig)
        vxm_template = cv2.resize(vxm_template, (512, 512))

        align_val = n
        if atlas_to_brain_align:
            im = np.uint8(im)
            br = cv2.imread(br)
            br = np.uint8(br)
            br = cv2.resize(br, (512, 512))
        else:
            # FOR ALIGNING BRAIN TO ATLAS
            if ".png" in br:
                im = cv2.imread(br)
            else:
                im = br
            im = np.uint8(im)
            im = cv2.resize(im, (512, 512))

        if atlas_to_brain_align:
            if use_voxelmorph and not use_dlc:
                atlas_mask_dir = os.path.join(
                    git_repo_base, "atlases/Atlas_for_Voxelmorph_border.png"
                )
            else:
                atlas_mask_dir = os.path.join(
                    git_repo_base, "atlases/atlas_smooth2_binary.png"
                )
        else:
            atlas_mask_dir = os.path.join(
                git_repo_base, "atlases/atlas_smooth2_binary.png"
            )
        atlas_mask_dir_left = os.path.join(
            git_repo_base, "atlases/left_hemisphere_smooth.png"
        )
        atlas_mask_dir_right = os.path.join(
            git_repo_base, "atlases/right_hemisphere_smooth.png"
        )
        atlas_label_mask_dir = os.path.join(
            git_repo_base, "atlases/diff_colour_regions/Common_atlas.mat"
        )
        atlas_label_mask_dir_left = os.path.join(
            git_repo_base, "atlases/diff_colour_regions/atlas_left_hemisphere.csv"
        )
        atlas_label_mask_dir_right = os.path.join(
            git_repo_base, "atlases/diff_colour_regions/atlas_right_hemisphere.csv"
        )
        atlas_label_mask_left = np.genfromtxt(
            atlas_label_mask_dir_left, delimiter=","
        )
        atlas_label_mask_right = np.genfromtxt(
            atlas_label_mask_dir_right, delimiter=","
        )
        atlas_mask_left = cv2.imread(atlas_mask_dir_left, cv2.IMREAD_UNCHANGED)
        atlas_mask_left = cv2.resize(atlas_mask_left, (im.shape[0], im.shape[1]))
        atlas_mask_left = np.uint8(atlas_mask_left)
        atlas_mask_right = cv2.imread(atlas_mask_dir_right, cv2.IMREAD_UNCHANGED)
        atlas_mask_right = cv2.resize(atlas_mask_right, (im.shape[0], im.shape[1]))
        atlas_mask_right = np.uint8(atlas_mask_right)

        atlas_mask = cv2.imread(atlas_mask_dir, cv2.IMREAD_UNCHANGED)
        atlas_mask = cv2.resize(atlas_mask, (im.shape[0], im.shape[1]))
        atlas_mask = np.uint8(atlas_mask)
        mask_dir = os.path.join(cwd, "../output_mask/{}.png".format(n))
        if use_voxelmorph and n == 1:
            try:
                os.remove(mask_dir)
            except:
                print("Cannot remove second U-Net output for VoxelMorph!")

        print("Performing first transformation of atlas {}...".format(n))

        mask_warped_path = os.path.join(
            output_mask_path, "{}_mask_warped.png".format(str(n))
        )
        mask_warped_path_alt_left = os.path.join(
            output_mask_path, "{}_mask_warped_left.png".format(str(n))
        )
        mask_warped_path_alt_right = os.path.join(
            output_mask_path, "{}_mask_warped_right.png".format(str(n))
        )
        brain_to_atlas_mask_path = os.path.join(
            output_mask_path, "{}_brain_to_atlas_mask.png".format(str(n))
        )
        if use_dlc:
            # First alignment of brain atlas using three cortical landmarks and standard affine transform
            atlas_pts_for_input = np.array([atlas_pts[n][0 : len(dlc_pts[n])]]).astype(
                "float32"
            )
            pts_for_input = np.array([dlc_pts[n]]).astype("float32")

            if align_once:
                align_val = 0
            else:
                align_val = n

            if len(atlas_pts_for_input[0]) == 2:
                atlas_pts_for_input = np.append(atlas_pts_for_input[0], [[0, 0]], axis=0)
                pts_for_input = np.append(pts_for_input[0], [[0, 0]], axis=0)
            if len(atlas_pts_for_input[0]) <= 2:
                warp_coords = cv2.estimateAffinePartial2D(
                    atlas_pts_for_input, pts_for_input
                )[0]
                if atlas_to_brain_align:
                    atlas_warped_left = cv2.warpAffine(im_left, warp_coords, (512, 512))
                    atlas_warped_right = cv2.warpAffine(im_right, warp_coords, (512, 512))
                    atlas_warped = cv2.bitwise_or(atlas_warped_left, atlas_warped_right)
                    ret, atlas_warped = cv2.threshold(
                        atlas_warped, 5, 255, cv2.THRESH_BINARY_INV
                    )
                    atlas_left_transform_path = os.path.join(
                        output_mask_path, "{}_atlas_left_transform.png".format(str(n))
                    )
                    atlas_right_transform_path = os.path.join(
                        output_mask_path, "{}_atlas_right_transform.png".format(str(n))
                    )
                    io.imsave(atlas_left_transform_path, atlas_warped_left)
                    io.imsave(atlas_right_transform_path, atlas_warped_right)
                else:
                    atlas_warped = cv2.warpAffine(im, warp_coords, (512, 512))
            elif len(atlas_pts_for_input[0]) == 3:
                warp_coords = cv2.getAffineTransform(atlas_pts_for_input, pts_for_input)
                atlas_warped = cv2.warpAffine(im, warp_coords, (512, 512))
            elif len(atlas_pts_for_input[0]) >= 4:
                im_final_size = (512, 512)

                left = acc_left_total.iloc[n, :].values.astype("float32").tolist()
                right = acc_right_total.iloc[n, :].values.astype("float32").tolist()
                left = np.argsort(left).tolist()
                right = np.argsort(right).tolist()
                right = [x + 1 for x in right]
                # if set([1, 3, 5, 7]).issubset(landmark_arr):
                if set([0, 3, 5, 6]).issubset(landmark_arr) and len(landmark_arr) >= 7:
                    left = [0, 3, 5]
                    right = [3, 5, 6]
                    # OLD METHOD FOR PREFERENTIALLY USING TOP CENTRE AND LAMBDA
                    # elif 4 in landmark_arr and 6 in landmark_arr:
                    #     # If top centre and lambda landmarks are selected,
                    #     # add the first landmark from the left hemisphere and the
                    #     # last one from the right hemisphere.
                    #     centreline_pts = [landmark_arr.index(x) for x in landmark_arr if x in [4, 6]]
                    #     left = [x for x in landmark_arr if x in [0, 1, 2, 3, 5]][0]
                    #     right = [landmark_arr.index(x) for x in reversed(landmark_arr) if x in [8, 7, 5, 4, 3]][0]
                    #     left = centreline_pts.extend([left])
                    #     right = [right].extend(centreline_pts)
                    #     print(landmark_indices)
                    #     print(landmark_arr)
                    #     print(left)
                    #     print(right)
                else:
                    # left = [x for x in landmark_indices if x in range(0, 6)][0:2]
                    # right = [x for x in landmark_indices if x in range(3, 9)][0:2]
                    left = [landmark_arr.index(x) for x in landmark_arr if x in [0, 1, 2, 3, 4, 5]][0:3]
                    # right = [landmark_arr.index(x) for x in reversed(landmark_arr) if x in [8, 7, 6, 5, 4, 3]][0:3]
                    right =  [landmark_arr.index(x) for x in landmark_arr if x in [3, 4, 5, 6, 7, 8]][-3:]
                    print(landmark_indices)
                    print(landmark_arr)
                    print(left)
                    print(right)

                # try:
                atlas_pts_left = np.array(
                    [
                        atlas_pts[align_val][left[0]],
                        atlas_pts[align_val][left[1]],
                        atlas_pts[align_val][left[2]],
                    ],
                    dtype=np.float32,
                )
                atlas_pts_right = np.array(
                    [
                        atlas_pts[align_val][right[0]],
                        atlas_pts[align_val][right[1]],
                        atlas_pts[align_val][right[2]],
                    ],
                    dtype=np.float32,
                )
                dlc_pts_left = np.array(
                    [
                        dlc_pts[align_val][left[0]],
                        dlc_pts[align_val][left[1]],
                        dlc_pts[align_val][left[2]],
                    ],
                    dtype=np.float32,
                )
                dlc_pts_right = np.array(
                    [
                        dlc_pts[align_val][right[0]],
                        dlc_pts[align_val][right[1]],
                        dlc_pts[align_val][right[2]],
                    ],
                    dtype=np.float32,
                )

                # except:
                #     # Old: left: 0, 2, 3; right: 1, 2, 3
                #     atlas_pts_left = np.array(
                #         [
                #             atlas_pts[align_val][0],
                #             atlas_pts[align_val][1],
                #             atlas_pts[align_val][2],
                #         ],
                #         dtype=np.float32,
                #     )
                #     atlas_pts_right = np.array(
                #         [
                #             atlas_pts[align_val][5],
                #             atlas_pts[align_val][6],
                #             atlas_pts[align_val][7],
                #         ],
                #         dtype=np.float32,
                #     )
                #     dlc_pts_left = np.array(
                #         [
                #             dlc_pts[align_val][0],
                #             dlc_pts[align_val][1],
                #             dlc_pts[align_val][2],
                #         ],
                #         dtype=np.float32,
                #     )
                #     dlc_pts_right = np.array(
                #         [
                #             dlc_pts[align_val][5],
                #             dlc_pts[align_val][6],
                #             dlc_pts[align_val][7],
                #         ],
                #         dtype=np.float32,
                #     )

                warp_coords_left = cv2.getAffineTransform(atlas_pts_left, dlc_pts_left)
                warp_coords_right = cv2.getAffineTransform(atlas_pts_right, dlc_pts_right)
                warp_coords_brain_atlas_left = cv2.getAffineTransform(dlc_pts_left, atlas_pts_left)
                warp_coords_brain_atlas_right = cv2.getAffineTransform(dlc_pts_right, atlas_pts_right)
                if atlas_to_brain_align:
                    atlas_warped_left = cv2.warpAffine(
                        im_left, warp_coords_left, im_final_size
                    )
                    atlas_warped_right = cv2.warpAffine(
                        im_right, warp_coords_right, im_final_size
                    )
                    atlas_warped = cv2.bitwise_or(atlas_warped_left, atlas_warped_right)
                    ret, atlas_warped = cv2.threshold(
                        atlas_warped, 5, 255, cv2.THRESH_BINARY_INV
                    )
                    if not original_label:
                        atlas_label_left = cv2.warpAffine(
                            atlas_label_mask_left, warp_coords_left, im_final_size
                        )
                        atlas_label_right = cv2.warpAffine(
                            atlas_label_mask_right, warp_coords_right, im_final_size
                        )
                        atlas_label = cv2.bitwise_or(atlas_label_left, atlas_label_right)

                else:
                    pts_np = np.array(
                        [
                            dlc_pts[align_val][0],
                            dlc_pts[align_val][1],
                            dlc_pts[align_val][2],
                        ],
                        dtype=np.float32,
                    )
                    atlas_pts_np = np.array(
                        [
                            atlas_pts[align_val][0],
                            atlas_pts[align_val][1],
                            atlas_pts[align_val][2],
                        ],
                        dtype=np.float32,
                    )
                    warp_coords = cv2.getAffineTransform(pts_np, atlas_pts_np)
                    atlas_warped = cv2.warpAffine(im, warp_coords, (512, 512))

            # if atlas_to_brain_align:
            if len(atlas_pts_for_input[0]) == 2:
                atlas_mask_left_warped = cv2.warpAffine(
                    atlas_mask_left, warp_coords, (512, 512)
                )
                atlas_mask_right_warped = cv2.warpAffine(
                    atlas_mask_right, warp_coords, (512, 512)
                )
                atlas_mask_warped = cv2.bitwise_or(
                    atlas_mask_left_warped, atlas_mask_right_warped
                )
            if len(atlas_pts_for_input[0]) == 3:
                atlas_mask_warped = cv2.warpAffine(atlas_mask, warp_coords, (512, 512))
            if len(atlas_pts_for_input[0]) >= 4:
                atlas_mask_left_warped = cv2.warpAffine(
                    atlas_mask_left, warp_coords_left, (512, 512)
                )
                atlas_mask_right_warped = cv2.warpAffine(
                    atlas_mask_right, warp_coords_right, (512, 512)
                )
                atlas_mask_warped = cv2.bitwise_or(
                    atlas_mask_left_warped, atlas_mask_right_warped
                )
            atlas_mask_warped = np.uint8(atlas_mask_warped)
            io.imsave(mask_warped_path_alt_left, atlas_mask_left_warped)
            io.imsave(mask_warped_path_alt_right, atlas_mask_right_warped)
            # brain_to_atlas_mask = cv2.bitwise_or(
            #     atlas_mask_warped, im
            # )
            # io.imsave(brain_to_atlas_mask_path, brain_to_atlas_mask)
            hemispheres = ['left', 'right']
            olfactory_bulbs_to_use = []
            if olfactory_check and use_unet:
                if align_once and n != 0:
                    olfactory_bulbs = olfactory_bulbs_to_use_pre_align_list[0]
                else:
                    mask = cv2.imread(mask_dir, cv2.IMREAD_GRAYSCALE)
                    cnts_for_olfactory = cv2.findContours(
                        mask.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
                    )
                    cnts_for_olfactory = imutils.grab_contours(cnts_for_olfactory)
                    if len(cnts_for_olfactory) == 3:
                        olfactory_bulbs = sorted(
                            cnts_for_olfactory, key=cv2.contourArea, reverse=True
                        )[1:3]
                    else:
                        olfactory_bulbs = sorted(
                            cnts_for_olfactory, key=cv2.contourArea, reverse=True
                        )[2:4]
                if align_once and n == 0:
                    olfactory_bulbs_to_use_pre_align_list.append(olfactory_bulbs)
                if len(olfactory_bulbs) == 0:
                    # If there's no olfactory bulbs found, disable future olfactory
                    # bulb checking
                    olfactory_check = False
            for hemisphere in hemispheres:
                new_data = []
                if hemisphere == 'left':
                    mask_path = mask_warped_path_alt_left
                    mask_warped_to_use = atlas_mask_left_warped
                    if olfactory_check and use_unet:
                        if len(olfactory_bulbs) >= 1:
                            bulb = olfactory_bulbs[0]
                        bulb_fill = 300
                else:
                    mask_path = mask_warped_path_alt_right
                    mask_warped_to_use = atlas_mask_right_warped
                    if olfactory_check and use_unet:
                        if len(olfactory_bulbs) > 1:
                            bulb = olfactory_bulbs[1]
                        bulb_fill = 400
                if olfactory_check and use_unet:
                    olfactory_bulbs_to_use = olfactory_bulbs
                    try:
                        cv2.fillPoly(mask_warped_to_use, pts=[bulb], color=[255, 255, 255])
                        io.imsave(mask_path, mask_warped_to_use)
                    except:
                        print('No olfactory bulb found!')
                    mask_warped_to_use = cv2.cvtColor(mask_warped_to_use, cv2.COLOR_BGR2GRAY)
                img_edited = Image.open(mask_path)
                # Generates a transparent version of the brain atlas.
                img_rgba = img_edited.convert("RGBA")
                data = img_rgba.getdata()
                for pixel in data:
                    if pixel[0] == 255 and pixel[1] == 255 and pixel[2] == 255:
                        new_data.append((pixel[0], pixel[1], pixel[2], 0))
                    else:
                        new_data.append(pixel)
                img_rgba.putdata(new_data)
                img_rgba.save(os.path.join(output_mask_path, "{}_brain_atlas_mask_transparent_{}.png".format(n, hemisphere)))
                img_transparent = cv2.imread(
                    os.path.join(output_mask_path, "{}_brain_atlas_mask_transparent_{}.png".format(n, hemisphere))
                )
                brain_atlas_transparent = cv2.bitwise_and(im, img_transparent)

                io.imsave(os.path.join(output_mask_path, "{}_brain_atlas_transparent_{}.png".format(n, hemisphere)),
                          brain_atlas_transparent)
                if hemisphere == 'left':
                    if atlas_to_brain_align:
                        brain_to_atlas_warped_left = cv2.warpAffine(brain_atlas_transparent, warp_coords_left, (512, 512))
                        if olfactory_check and use_unet:
                            olfactory_warped_left = cv2.warpAffine(mask_warped_to_use, warp_coords_left, (512, 512))
                            olfactory_warped_cnts_left = cv2.findContours(
                                olfactory_warped_left.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
                            )
                            olfactory_warped_cnts_left = imutils.grab_contours(olfactory_warped_cnts_left)
                            olfactory_warped_left = sorted(
                                olfactory_warped_cnts_left, key=cv2.contourArea, reverse=True
                            )[-1]
                    else:
                        brain_to_atlas_warped_left = cv2.warpAffine(brain_atlas_transparent, warp_coords_brain_atlas_left, (512, 512))
                        if olfactory_check and use_unet:
                            olfactory_warped_left = cv2.warpAffine(mask_warped_to_use, warp_coords_brain_atlas_left, (512, 512))
                            olfactory_warped_cnts_left = cv2.findContours(
                                olfactory_warped_left.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
                            )
                            olfactory_warped_cnts_left = imutils.grab_contours(olfactory_warped_cnts_left)
                            olfactory_warped_left = sorted(
                                olfactory_warped_cnts_left, key=cv2.contourArea, reverse=True
                            )[-1]
                else:
                    if atlas_to_brain_align:
                        brain_to_atlas_warped_right = cv2.warpAffine(brain_atlas_transparent, warp_coords_right, (512, 512))
                        olfactory_warped_right = cv2.warpAffine(mask_warped_to_use, warp_coords_right, (512, 512))
                        if olfactory_check and use_unet:
                            olfactory_warped_right = cv2.warpAffine(mask_warped_to_use, warp_coords_right, (512, 512))
                            olfactory_warped_cnts_right = cv2.findContours(
                                olfactory_warped_right.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
                            )
                            olfactory_warped_cnts_right = imutils.grab_contours(olfactory_warped_cnts_right)
                            olfactory_warped_right = sorted(
                                olfactory_warped_cnts_right, key=cv2.contourArea, reverse=True
                            )[-1]
                    else:
                        brain_to_atlas_warped_right = cv2.warpAffine(brain_atlas_transparent, warp_coords_brain_atlas_right, (512, 512))
                        olfactory_warped_right = cv2.warpAffine(mask_warped_to_use, warp_coords_brain_atlas_right, (512, 512))
                        if olfactory_check and use_unet:
                            olfactory_warped_right = cv2.warpAffine(mask_warped_to_use, warp_coords_brain_atlas_right, (512, 512))
                            olfactory_warped_cnts_right = cv2.findContours(
                                olfactory_warped_right.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
                            )
                            olfactory_warped_cnts_right = imutils.grab_contours(olfactory_warped_cnts_right)
                            olfactory_warped_right = sorted(
                                olfactory_warped_cnts_right, key=cv2.contourArea, reverse=True
                            )[-1]
            if olfactory_check and use_unet:
                olfactory_bulbs_to_use_list.append([olfactory_warped_left, olfactory_warped_right])
            brain_to_atlas_mask = cv2.bitwise_or(
                brain_to_atlas_warped_left, brain_to_atlas_warped_right
            )
            io.imsave(brain_to_atlas_mask_path, brain_to_atlas_mask)


            # Second alignment of brain atlas using cortical landmarks and piecewise affine transform
            print("Performing second transformation of atlas {}...".format(n))

            atlas_first_transform_path = os.path.join(
                output_mask_path, "{}_atlas_first_transform.png".format(str(n))
            )

            if atlas_to_brain_align:
                dst = atlas_warped
            else:
                dst = brain_to_atlas_mask

            io.imsave(atlas_first_transform_path, dst)

            # If a sensory map of the brain is provided, do a third alignment of the brain atlas using up to
            # four peaks of sensory activity
            if sensory_match:
                original_label = True
                if atlas_to_brain_align:
                    dst = getMaskContour(
                        mask_dir,
                        atlas_warped,
                        sensory_peak_pts[align_val],
                        sensory_atlas_pts[align_val],
                        cwd,
                        align_val,
                        False,
                    )
                    atlas_mask_warped = getMaskContour(
                        atlas_first_transform_path,
                        atlas_mask_warped,
                        sensory_peak_pts[align_val],
                        sensory_atlas_pts[align_val],
                        cwd,
                        align_val,
                        False,
                    )
                    atlas_mask_warped = cv2.resize(
                        atlas_mask_warped, (im.shape[0], im.shape[1])
                    )
                else:
                    dst = atlas_warped
        else:
            # If we're not using DeepLabCut...
            if atlas_to_brain_align and not use_voxelmorph:
                dst = cv2.bitwise_or(im_left, im_right)
                ret, dst = cv2.threshold(
                    dst, 5, 255, cv2.THRESH_BINARY_INV
                )
            else:
                dst = im
            dst = np.uint8(dst)

        if use_dlc:
            if atlas_to_brain_align:
                io.imsave(mask_warped_path, atlas_mask_warped)
            else:
                io.imsave(mask_warped_path, atlas_mask)
        else:
            io.imsave(mask_warped_path, dst)
            # if atlas_to_brain_align:
            if use_voxelmorph:
                atlas_mask_warped = atlas_mask
            else:
                atlas_mask_warped = cv2.bitwise_or(
                    atlas_mask_left, atlas_mask_right
                )
            atlas_mask_warped = cv2.cvtColor(atlas_mask_warped, cv2.COLOR_BGR2GRAY)
            ret, atlas_mask_warped = cv2.threshold(
                atlas_mask_warped, 5, 255, cv2.THRESH_BINARY
            )
            atlas_mask_warped = np.uint8(atlas_mask_warped)
            original_label = True
            io.imsave(mask_warped_path, atlas_mask_warped)
        # Resize images back to 512x512
        dst = cv2.resize(dst, (im.shape[0], im.shape[1]))
        atlas_path = os.path.join(output_mask_path, "{}_atlas.png".format(str(n)))

        vxm_template_output_path = os.path.join(
            output_mask_path, "{}_vxm_template.png".format(str(n))
        )

        dst_list.append(dst)
        if use_voxelmorph:
            vxm_template_list.append(vxm_template)
            io.imsave(vxm_template_output_path, vxm_template_list[n])

        if atlas_to_brain_align:
            io.imsave(atlas_path, dst)
            br_list.append(br)
        else:
            brain_warped_path = os.path.join(
                output_mask_path, "{}_brain_warp.png".format(str(n))
            )
            io.imsave(brain_warped_path, dst)
            io.imsave(atlas_path, atlas)

        if atlas_to_brain_align:
            if original_label:
                atlas_label = []
            atlas_label = atlas_to_mask(
                atlas_path,
                mask_dir,
                mask_warped_path,
                output_mask_path,
                n,
                use_unet,
                use_voxelmorph,
                atlas_to_brain_align,
                git_repo_base,
                olfactory_check,
                [],
                atlas_label
            )
            atlas_label_list.append(atlas_label)
        elif not use_dlc:
            io.imsave(os.path.join(output_mask_path, "{}.png".format(n)), dst)
        if bregma_present:
            bregma_val = int(bregma_index_list[n])
            bregma_list.append(dlc_pts[n][bregma_val])

    # Carries out VoxelMorph on each motif-based functional map (MBFM) that has been aligned to a raw brain image
    if use_voxelmorph:
        align_once = True
        if len(dst_list) == 1:
            n_to_use = 0
        else:
            n_to_use = 1
        for (n_post, dst_post), vxm_template_post in zip(enumerate([dst_list[n_to_use]]), [vxm_template_list[n_to_use]]):
            output_img, flow_post = voxelmorph_align(
               voxelmorph_model_path, dst_post, vxm_template_post, exist_transform, flow_path
            )
            flow_path_after = os.path.join(output_mask_path, "{}_flow.npy".format(str(n_post)))
            np.save(flow_path_after, flow_post)
            output_path_after = os.path.join(output_mask_path, "{}_output_img.png".format(str(n_post)))
            io.imsave(output_path_after, output_img)
            if not exist_transform:
                dst_gray = cv2.cvtColor(atlas, cv2.COLOR_BGR2GRAY)
                dst_post = vxm_transform(dst_gray, flow_path_after)
                ret, dst_post = cv2.threshold(
                    dst_post, 1, 255, cv2.THRESH_BINARY
                ) # 5, 255
                dst_post = np.uint8(dst_post)

            mask_warped_path = os.path.join(
                output_mask_path, "{}_mask_warped.png".format(str(n_post))
            )

            atlas_first_transform_path_post = os.path.join(
                output_mask_path, "{}_atlas_first_transform.png".format(str(n_post))
            )

            io.imsave(atlas_first_transform_path_post, dst_post)
            if align_once and use_voxelmorph:
                atlas_path = os.path.join(output_mask_path, "{}_atlas.png".format(str(0)))

            brain_warped_path = os.path.join(
                output_mask_path, "{}_brain_warp.png".format(str(n_post))
            )
            mask_dir = os.path.join(cwd, "../output_mask/{}.png".format(n_post))
            dst_post = cv2.resize(dst_post, (im.shape[0], im.shape[1]))
            if not atlas_to_brain_align:
                atlas_to_brain_align = True
                original_label = True
            if atlas_to_brain_align:
                io.imsave(atlas_path, dst_post)
            else:
                io.imsave(brain_warped_path, dst_post)
            if atlas_to_brain_align:
                if original_label:
                    atlas_label = []
                if use_voxelmorph and olfactory_check and use_unet and use_dlc:
                    if align_once:
                        olfactory_bulbs_to_use_check = olfactory_bulbs_to_use_list[0]
                    else:
                        olfactory_bulbs_to_use_check = olfactory_bulbs_to_use_list[n_post]
                else:
                    olfactory_bulbs_to_use_check = []
                atlas_label = atlas_to_mask(
                    atlas_path,
                    mask_dir,
                    mask_warped_path,
                    output_mask_path,
                    n_post,
                    use_unet,
                    use_voxelmorph,
                    atlas_to_brain_align,
                    git_repo_base,
                    olfactory_check,
                    olfactory_bulbs_to_use_check,
                    atlas_label
                )
                atlas_label_list.append(atlas_label)


    # Converts the transformed brain atlas into a segmentation template for the original brain image
    applyMask(
        brain_img_dir,
        output_mask_path,
        output_overlay_path,
        output_overlay_path,
        mat_save,
        threshold,
        git_repo_base,
        bregma_list,
        atlas_to_brain_align,
        model,
        dlc_pts,
        atlas_pts,
        olfactory_check,
        use_unet,
        use_dlc,
        use_voxelmorph,
        plot_landmarks,
        align_once,
        atlas_label_list,
        olfactory_bulbs_to_use_list,
        region_labels,
        original_label
    )

"""
MesoNet
Authors: <NAME> and <NAME>, <NAME>
https://github.com/bf777/MesoNet
Licensed under the Creative Commons Attribution 4.0 International License (see LICENSE for details)
"""
from mesonet.utils import natural_sort_key, convert_to_png
from mesonet.mask_functions import atlas_to_mask, applyMask

# from mesonet.niftyreg_align import niftyreg_align
from mesonet.voxelmorph_align import voxelmorph_align, vxm_transform
import numpy as np
import pandas as pd
import cv2
import imutils
import math
import scipy.io
import skimage.io as io
import skimage.transform as trans
from skimage.transform import PiecewiseAffineTransform, warp
import imageio
import os
import fnmatch
import glob
from PIL import Image


def find_peaks(img):
    """
    Locates the peaks of cortical activity in each hemisphere of the brain image.
    :param img: A sensory map (indicating functional activation in the mouse brain).
    :return: An array of the maximum points of activation for the sensory map.
    """
    max_loc_arr = []
    img = cv2.imread(str(img), 0)
    im = img.copy()
    x_min = int(np.around(im.shape[0] / 2))
    im1 = im[:, x_min : im.shape[0]]
    im2 = im[:, 0:x_min]
    (minVal, max_val, minLoc, max_loc) = cv2.minMaxLoc(im1)
    (minVal2, maxVal2, minLoc2, maxLoc2) = cv2.minMaxLoc(im2)
    max_loc = list(max_loc)
    max_loc[0] = max_loc[0] + x_min
    max_loc = tuple(max_loc)
    max_loc_arr.append(max_loc)
    if (max_val - 30) <= maxVal2 <= (max_val + 30):
        max_loc_arr.append(maxLoc2)
    return max_loc_arr


def coords_to_mat(
    sub_dlc_pts, i, output_mask_path, bregma_present, bregma_index, landmark_arr
):
    if bregma_present:
        x_bregma, y_bregma = sub_dlc_pts[bregma_index]
        for pt, landmark in zip(sub_dlc_pts, landmark_arr):
            x_pt = pt[0]
            y_pt = pt[1]
            pt_adj = [landmark, x_pt - x_bregma, y_pt - y_bregma]
            pt_adj_to_mat = np.array(pt_adj, dtype=np.object)
            if not os.path.isdir(os.path.join(output_mask_path, "mat_coords")):
                os.mkdir(os.path.join(output_mask_path, "mat_coords"))
            scipy.io.savemat(
                os.path.join(
                    output_mask_path,
                    "mat_coords/landmarks_{}_{}.mat".format(i, landmark),
                ),
                {"landmark_coords_{}_{}".format(i, landmark): pt_adj_to_mat},
                appendmat=False,
            )


def sensory_to_mat(sub_dlc_pts, bregma_pt, i, output_mask_path):
    x_bregma, y_bregma = bregma_pt
    sensory_names = ["tail_left", "tail_right", "visual", "whisker"]
    for pt, landmark in zip(sub_dlc_pts, sensory_names):
        x_pt = pt[0]
        y_pt = pt[1]
        pt_adj = [landmark, x_pt - x_bregma, y_pt - y_bregma]
        pt_adj_to_mat = np.array(pt_adj, dtype=np.object)
        if not os.path.isdir(os.path.join(output_mask_path, "mat_coords")):
            os.mkdir(os.path.join(output_mask_path, "mat_coords"))
        scipy.io.savemat(
            os.path.join(
                output_mask_path,
                "mat_coords/sensory_peaks_{}_{}.mat".format(i, landmark),
            ),
            {"sensory_peaks_{}_{}".format(i, landmark): pt_adj_to_mat},
            appendmat=False,
        )


def atlas_from_mat(input_file, mat_cnt_list):
    """
    Generates a binary brain atlas from a .mat file.
    :param input_file: The input .mat file representing a brain atlas (with white = 255 and black = 0)
    :param mat_cnt_list: The list to which mat files should be appended
    :return: A thresholded binary brain atlas.
    """
    file = input_file
    atlas_base = np.zeros((512, 512), dtype="uint8")
    if glob.glob(os.path.join(input_file, "*.mat")):
        mat = scipy.io.loadmat(file)
        mat_shape = mat[list(mat.keys())[3]]
        if len(mat_shape.shape) > 2:
            for val in range(0, mat_shape.shape[2]):
                mat_roi = mat_shape[:, :, val]
                mat_resize = cv2.resize(mat_roi, (512, 512))
                mat_resize = np.uint8(mat_resize)
                ret, thresh = cv2.threshold(mat_resize, 5, 255, cv2.THRESH_BINARY_INV)
                mat_roi_cnt = cv2.findContours(
                    thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
                )
                mat_roi_cnt = imutils.grab_contours(mat_roi_cnt)
                c_to_save = max(mat_roi_cnt, key=cv2.contourArea)
                mat_cnt_list.append(c_to_save)
                cv2.drawContours(atlas_base, mat_roi_cnt, -1, (255, 255, 255), 1)
            ret, thresh = cv2.threshold(atlas_base, 5, 255, cv2.THRESH_BINARY_INV)
            io.imsave("atlas_unresized_test.png", thresh)
        else:
            mat = mat["atlas"]
            mat_resize = cv2.resize(mat, (512, 512))
            ret, thresh = cv2.threshold(mat_resize, 5, 255, cv2.THRESH_BINARY_INV)
    else:
        atlas_im = cv2.imread(file, cv2.IMREAD_GRAYSCALE)
        atlas_resize = np.uint8(atlas_im)
        ret, atlas_resize = cv2.threshold(atlas_resize, 127, 255, 0)
        io.imsave("atlas_unresized_test.png", atlas_resize)
        roi_cnt, hierarchy = cv2.findContours(
            atlas_resize, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE
        )[-2:]
        for val in roi_cnt:
            c_to_save = max(val, key=cv2.contourArea)
            mat_cnt_list.append(c_to_save)
            cv2.drawContours(atlas_base, val, -1, (255, 255, 255), 1)
        ret, thresh = cv2.threshold(atlas_base, 5, 255, cv2.THRESH_BINARY_INV)
    return thresh


def atlas_rotate(dlc_pts, im):
    dlc_y_pts = [
        coord if (190 <= coord[0] <= 330) else (1000, 1000) for coord in dlc_pts
    ]
    dlc_y_pts = [coord for coord in dlc_y_pts if coord[0] < 1000]

    # https://stackoverflow.com/questions/42258637/how-to-know-the-angle-between-two-points
    rotate_rad = math.atan2(0, (im.shape[1] / 2) - dlc_y_pts[-1][0])
    rotate_deg = -1 * (abs(math.degrees(rotate_rad)))
    im_rotate_mat = cv2.getRotationMatrix2D(
        (im.shape[1] / 2, im.shape[0] / 2), rotate_deg, 1.0
    )
    im_rotated = cv2.warpAffine(im, im_rotate_mat, (512, 512))
    x_min = int(np.around(im_rotated.shape[0] / 2))
    im_left = im_rotated[:, 0:x_min]
    im_right = im_rotated[:, x_min : im_rotated.shape[0]]
    return im_left, im_right


def getMaskContour(mask_dir, atlas_img, predicted_pts, actual_pts, cwd, n, main_mask):
    """
    Gets the contour of the brain's boundaries and applies a piecewise affine transform to the brain atlas
    based on the cortical landmarks predicted in dlc_predict (and peaks of activity on the sensory map, if available).
    :param mask_dir: The path to the directory containing the U-net masks of the brain's boundaries.
    :param atlas_img: The brain atlas to be transformed.
    :param predicted_pts: The coordinates of the cortical landmarks predicted in dlc_predict (or, for the second run
    of this function, the coordinates of the peaks of activity in the sensory map).
    :param actual_pts: The fixed coordinates of the cortical landmarks on the brain atlas (or, for the second run of
    this function, the fixed coordinates of the peaks of sensory activity on the brain atlas).
    :param cwd: The path to the current working directory.
    :param n: The number of the current image in the directory.
    """
    c_landmarks = np.empty([0, 2])
    c_atlas_landmarks = np.empty([0, 2])
    mask = cv2.imread(mask_dir, cv2.IMREAD_GRAYSCALE)
    atlas_to_warp = atlas_img
    mask = np.uint8(mask)
    cnts, hierarchy = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)[-2:]
    for cnt in cnts:
        cnt = cnt[:, 0, :]
        cnt = np.asarray(cnt).astype("float32")
        c_landmarks = np.concatenate((c_landmarks, cnt))
        c_atlas_landmarks = np.concatenate((c_atlas_landmarks, cnt))
    c_landmarks = np.concatenate((c_landmarks, predicted_pts))
    c_atlas_landmarks = np.concatenate((c_atlas_landmarks, actual_pts))
    tform = PiecewiseAffineTransform()
    tform.estimate(c_atlas_landmarks, c_landmarks)
    dst = warp(atlas_to_warp, tform, output_shape=(512, 512))
    if main_mask:
        io.imsave(os.path.join(cwd, "mask_{}.png".format(n)), mask)
    return dst


def atlasBrainMatch(
    brain_img_dir,
    sensory_img_dir,
    coords_input,
    sensory_match,
    mat_save,
    threshold,
    git_repo_base,
    region_labels,
    landmark_arr_orig,
    use_unet,
    use_dlc,
    atlas_to_brain_align,
    model,
    olfactory_check,
    plot_landmarks,
    align_once,
    original_label,
    use_voxelmorph,
    exist_transform,
    voxelmorph_model="motif_model_atlas.h5",
    vxm_template_path="templates",
    dlc_template_path="dlc_templates",
    flow_path="",
):
    """
    Align and overlap brain atlas onto brain image based on four landmark locations in the brain image and the atlas.
    :param brain_img_dir: The directory containing each brain image to be used.
    :param sensory_img_dir: The directory containing each sensory image to be used (if you are aligning each brain
    image using a sensory map).
    :param coords_input: Predicted locations of the four landmarks on the brain image from the file generated by
    DeepLabCut.
    :param sensory_match: Whether or not a sensory map is to be used.
    :param mat_save: Whether or not to export each brain region to a .mat file in applyMask, which is called at the end
    of this function.
    :param threshold: The threshold for the cv2.opening operation carried out in applyMask, which is called at the end
    of this function.
    :param git_repo_base: The path to the base git repository containing necessary resources for MesoNet (reference
    atlases, DeepLabCut config files, etc.)
    :param region_labels: Choose whether or not to attempt to label each region with its name from the Allen Institute
    Mouse Brain Atlas.
    :param landmark_arr_orig: The original array of landmarks from DeepLabCut (to be distinguished from any automatic
    exclusions to landmark array based on prediction quality).
    :param use_unet: Choose whether or not to identify the borders of the cortex using a U-net model.
    :param atlas_to_brain_align: If True, registers the atlas to each brain image. If False, registers each brain image
    to the atlas.
    :param model: The name of the U-net model (for passthrough to mask_functions.py)
    :param olfactory_check: If True, draws olfactory bulb contours on the brain image.
    :param plot_landmarks: If True, plots DeepLabCut landmarks (large circles) and original alignment landmarks (small
    circles) on final brain image.
    :param align_once: if True, carries out all alignments based on the alignment of the first atlas and brain. This can
    save time if you have many frames of the same brain with a fixed camera position.
    :param original_label: if True, uses a brain region labelling approach that attempts to automatically sort brain
    regions in a consistent order (left to right by hemisphere, then top to bottom for vertically aligned regions). This
    approach may be more flexible if you're using a custom brain atlas (i.e. not one in which region is filled with a
    unique number).
    :param exist_transform: if True, uses an existing voxelmorph transformation field for all data instead of predicting
    a new transformation.
    :param voxelmorph_model: the name of a .h5 model located in the models folder of the git repository for MesoNet,
    generated using voxelmorph and containing weights for a voxelmorph local deformation model.
    :param vxm_template_path: the path to a template atlas (.npy or .mat( to which the brain image will be aligned in
    voxelmorph.
    :param flow_path: the path to a voxelmorph transformation field that will be used to transform all data instead of
    predicting a new transformation if exist_transform is True.
    """
    # load brain images folder
    brain_img_arr = []
    dlc_img_arr = []
    peak_arr = []
    atlas_label_list = []
    dst_list = []
    vxm_template_list = []
    br_list = []
    olfactory_bulbs_to_use_list = []
    olfactory_bulbs_to_use_pre_align_list = []

    voxelmorph_model_path = os.path.join(
        git_repo_base, "models", "voxelmorph", voxelmorph_model
    )

    # Prepare template for VoxelMorph
    convert_to_png(vxm_template_path)
    vxm_template_orig = cv2.imread(
        glob.glob(os.path.join(git_repo_base, "atlases", vxm_template_path, "*.png"))[0]
    )

    # Prepare output folder
    cwd = os.getcwd()
    output_mask_path = os.path.join(cwd, "../output_mask")
    # Output folder for transparent masks and masks overlaid onto brain image
    output_overlay_path = os.path.join(cwd, "../output_overlay")
    if not os.path.isdir(output_mask_path):
        os.mkdir(output_mask_path)
    if not os.path.isdir(output_overlay_path):
        os.mkdir(output_overlay_path)

    if not atlas_to_brain_align:
        im = cv2.imread(
            os.path.join(git_repo_base, "atlases/Atlas_workflow2_binary.png")
        )
    else:
        if use_voxelmorph and not use_dlc:
            im = cv2.imread(
                os.path.join(git_repo_base, "atlases/Atlas_for_Voxelmorph_binary.png")
            )
        else:
            im = cv2.imread(
                os.path.join(git_repo_base, "atlases/Atlas_workflow1_binary.png")
            )
        im_left = cv2.imread(os.path.join(git_repo_base, "atlases/left_hemi.png"))
        ret, im_left = cv2.threshold(im_left, 5, 255, cv2.THRESH_BINARY_INV)
        im_right = cv2.imread(os.path.join(git_repo_base, "atlases/right_hemi.png"))
        ret, im_right = cv2.threshold(im_right, 5, 255, cv2.THRESH_BINARY_INV)
        im_left = np.uint8(im_left)
        im_right = np.uint8(im_right)
        im = np.uint8(im)
    # im = atlas_from_mat(os.path.join(git_repo_base, 'atlases/atlas_ROIs.mat'))
    atlas = im
    # FOR ALIGNING BRAIN TO ATLAS
    for num, file in enumerate(os.listdir(cwd)):
        if fnmatch.fnmatch(file, "*.png") and "mask" not in file:
            dlc_img_arr.append(os.path.join(cwd, file))
    for num, file in enumerate(os.listdir(brain_img_dir)):
        if fnmatch.fnmatch(file, "*.png"):
            brain_img_arr.append(os.path.join(brain_img_dir, file))
            brain_img_arr.sort(key=natural_sort_key)
        elif fnmatch.fnmatch(file, "*.tif"):
            tif_stack = imageio.mimread(os.path.join(brain_img_dir, file))
            for tif_im in tif_stack:
                brain_img_arr.append(tif_im)
    # i_coord, j_coord = np.array([(100, 256, 413, 256), (148, 254, 148, 446)])

    # https://www.pyimagesearch.com/2014/07/21/detecting-circles-images-using-opencv-hough-circles/
    coord_circles_img = cv2.imread(
        os.path.join(
            git_repo_base, "atlases", "multi_landmark", "landmarks_new_binary.png"
        ),
        cv2.IMREAD_GRAYSCALE,
    )
    coord_circles_img = np.uint8(coord_circles_img)
    # detect circles in the image
    circles, hierarchy = cv2.findContours(
        coord_circles_img, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE
    )[-2:]
    # ensure at least some circles were found
    if circles is not None:
        # convert the (x, y) coordinates and radius of the circles to integers
        atlas_arr = np.array(
            [
                (
                    int(cv2.moments(circle)["m10"] / cv2.moments(circle)["m00"]),
                    int(cv2.moments(circle)["m01"] / cv2.moments(circle)["m00"]),
                )
                for circle in circles
            ]
        )

    atlas_arr = np.array(
        [
            (102, 148),
            (166, 88),
            (214, 454),
            (256, 88),
            (256, 256),
            (256, 428),
            (410, 148),
            (346, 88),
            (298, 454),
        ]
    )

    peak_arr_flat = []
    peak_arr_total = []

    if sensory_match:
        for num, file in enumerate(brain_img_arr):
            img_name = str(os.path.splitext(os.path.basename(file))[0])
            sensory_img_for_brain = os.path.join(sensory_img_dir, img_name)
            if glob.glob(sensory_img_for_brain):
                sensory_img_for_brain_dir = os.listdir(sensory_img_for_brain)
                sensory_img_for_brain_dir.sort(key=natural_sort_key)
                for num_im, file_im in enumerate(sensory_img_for_brain_dir):
                    sensory_im = io.imread(
                        os.path.join(sensory_img_dir, img_name, file_im)
                    )
                    sensory_im = trans.resize(sensory_im, (512, 512))
                    io.imsave(
                        os.path.join(sensory_img_dir, img_name, file_im), sensory_im
                    )
                    peak = find_peaks(os.path.join(sensory_img_dir, img_name, file_im))
                    peak_arr.append(peak)
            for x in peak_arr:
                for y in x:
                    peak_arr_flat.append(y)
            peak_arr_total.append(peak_arr_flat)
            peak_arr_flat = []
            peak_arr = []

    dlc_pts = []
    atlas_pts = []
    sensory_peak_pts = []
    sensory_atlas_pts = []
    sub_dlc_pts = []
    sub_atlas_pts = []
    sub_sensory_peak_pts = []
    sub_sensory_atlas_pts = []

    bregma_index_list = []
    bregma_list = []
    if use_dlc:
        bregma_present = True
    else:
        bregma_present = False

    coords = pd.read_csv(coords_input)
    x_coord = coords.iloc[2:, 1::3]
    y_coord = coords.iloc[2:, 2::3]
    accuracy = coords.iloc[2:, 3::3]
    acc_left_total = accuracy.iloc[:, 0:5]
    acc_right_total = accuracy.iloc[:, 3:8]
    landmark_indices = [0, 1, 2, 3, 4, 5, 6, 7, 8]  # [0, 3, 2, 1]
    for arr_index, i in enumerate(range(0, len(x_coord))):
        landmark_arr = landmark_arr_orig
        x_coord_flat = x_coord.iloc[i].values.astype("float32")
        y_coord_flat = y_coord.iloc[i].values.astype("float32")
        x_coord_flat = x_coord_flat[landmark_arr]
        y_coord_flat = y_coord_flat[landmark_arr]
        dlc_list = []
        atlas_list = []
        for (coord_x, coord_y) in zip(x_coord_flat, y_coord_flat):
            dlc_coord = (coord_x, coord_y)
            dlc_list.append(dlc_coord)
        for coord_atlas in atlas_arr:
            atlas_coord = (coord_atlas[0], coord_atlas[1])
            atlas_list.append(atlas_coord)
        atlas_list = [atlas_list[i] for i in landmark_arr]

        # Initialize result as max value
        landmark_indices = landmark_indices[0 : len(landmark_arr)]
        atlas_indices = landmark_arr

        pts_dist = np.absolute(
            np.asarray(atlas_list) - np.asarray((im.shape[0] / 2, im.shape[1] / 2))
        )
        pts_avg_dist = [np.mean(v) for v in pts_dist]
        bregma_index = np.argmin(np.asarray(pts_avg_dist))

        for j in landmark_indices:
            sub_dlc_pts.append([x_coord_flat[j], y_coord_flat[j]])
        for j in atlas_indices:
            sub_atlas_pts.append([atlas_arr[j][0], atlas_arr[j][1]])

        dlc_pts.append(sub_dlc_pts)
        atlas_pts.append(sub_atlas_pts)
        coords_to_mat(
            sub_dlc_pts, i, output_mask_path, bregma_present, bregma_index, landmark_arr
        )
        bregma_index_list.append(bregma_index)
        sub_dlc_pts = []
        sub_atlas_pts = []
    if sensory_match:
        k_coord, m_coord = np.array([(189, 323, 435, 348), (315, 315, 350, 460)])
        coords_peak = peak_arr_total
        for img_num, img in enumerate(brain_img_arr):
            for j in [1, 0, 3, 2]:  # Get peak values from heatmaps
                sub_sensory_peak_pts.append(
                    [coords_peak[img_num][j][0], coords_peak[img_num][j][1]]
                )
            for j in [0, 1, 2, 3]:  # Get circle locations
                sub_sensory_atlas_pts.append([k_coord[j], m_coord[j]])
            sensory_peak_pts.append(sub_sensory_peak_pts)
            sensory_atlas_pts.append(sub_sensory_atlas_pts)
            sensory_to_mat(
                sub_sensory_peak_pts, dlc_pts[img_num][3], img_num, output_mask_path
            )
            sub_sensory_peak_pts = []
            sub_sensory_atlas_pts = []
        sensory_peak_pts, sensory_atlas_pts = (
            np.asarray(sensory_peak_pts).astype("float32"),
            np.asarray(sensory_atlas_pts).astype("float32"),
        )

    for (n, br) in enumerate(brain_img_arr):
        vxm_template = np.uint8(vxm_template_orig)
        vxm_template = cv2.resize(vxm_template, (512, 512))

        align_val = n
        if atlas_to_brain_align:
            im = np.uint8(im)
            br = cv2.imread(br)
            br = np.uint8(br)
            br = cv2.resize(br, (512, 512))
        else:
            # FOR ALIGNING BRAIN TO ATLAS
            if ".png" in br:
                im = cv2.imread(br)
            else:
                im = br
            im = np.uint8(im)
            im = cv2.resize(im, (512, 512))

        if atlas_to_brain_align:
            if use_voxelmorph and not use_dlc:
                atlas_mask_dir = os.path.join(
                    git_repo_base, "atlases/Atlas_for_Voxelmorph_border.png"
                )
            else:
                atlas_mask_dir = os.path.join(
                    git_repo_base, "atlases/atlas_smooth2_binary.png"
                )
        else:
            atlas_mask_dir = os.path.join(
                git_repo_base, "atlases/atlas_smooth2_binary.png"
            )
        atlas_mask_dir_left = os.path.join(
            git_repo_base, "atlases/left_hemisphere_smooth.png"
        )
        atlas_mask_dir_right = os.path.join(
            git_repo_base, "atlases/right_hemisphere_smooth.png"
        )
        atlas_label_mask_dir = os.path.join(
            git_repo_base, "atlases/diff_colour_regions/Common_atlas.mat"
        )
        atlas_label_mask_dir_left = os.path.join(
            git_repo_base, "atlases/diff_colour_regions/atlas_left_hemisphere.csv"
        )
        atlas_label_mask_dir_right = os.path.join(
            git_repo_base, "atlases/diff_colour_regions/atlas_right_hemisphere.csv"
        )
        atlas_label_mask_left = np.genfromtxt(
            atlas_label_mask_dir_left, delimiter=","
        )
        atlas_label_mask_right = np.genfromtxt(
            atlas_label_mask_dir_right, delimiter=","
        )
        atlas_mask_left = cv2.imread(atlas_mask_dir_left, cv2.IMREAD_UNCHANGED)
        atlas_mask_left = cv2.resize(atlas_mask_left, (im.shape[0], im.shape[1]))
        atlas_mask_left = np.uint8(atlas_mask_left)
        atlas_mask_right = cv2.imread(atlas_mask_dir_right, cv2.IMREAD_UNCHANGED)
        atlas_mask_right = cv2.resize(atlas_mask_right, (im.shape[0], im.shape[1]))
        atlas_mask_right = np.uint8(atlas_mask_right)

        atlas_mask = cv2.imread(atlas_mask_dir, cv2.IMREAD_UNCHANGED)
        atlas_mask = cv2.resize(atlas_mask, (im.shape[0], im.shape[1]))
        atlas_mask = np.uint8(atlas_mask)
        mask_dir = os.path.join(cwd, "../output_mask/{}.png".format(n))
        if use_voxelmorph and n == 1:
            try:
                os.remove(mask_dir)
            except:
                print("Cannot remove second U-Net output for VoxelMorph!")

        print("Performing first transformation of atlas {}...".format(n))

        mask_warped_path = os.path.join(
            output_mask_path, "{}_mask_warped.png".format(str(n))
        )
        mask_warped_path_alt_left = os.path.join(
            output_mask_path, "{}_mask_warped_left.png".format(str(n))
        )
        mask_warped_path_alt_right = os.path.join(
            output_mask_path, "{}_mask_warped_right.png".format(str(n))
        )
        brain_to_atlas_mask_path = os.path.join(
            output_mask_path, "{}_brain_to_atlas_mask.png".format(str(n))
        )
        if use_dlc:
            # First alignment of brain atlas using three cortical landmarks and standard affine transform
            atlas_pts_for_input = np.array([atlas_pts[n][0 : len(dlc_pts[n])]]).astype(
                "float32"
            )
            pts_for_input = np.array([dlc_pts[n]]).astype("float32")

            if align_once:
                align_val = 0
            else:
                align_val = n

            if len(atlas_pts_for_input[0]) == 2:
                atlas_pts_for_input = np.append(atlas_pts_for_input[0], [[0, 0]], axis=0)
                pts_for_input = np.append(pts_for_input[0], [[0, 0]], axis=0)
            if len(atlas_pts_for_input[0]) <= 2:
                warp_coords = cv2.estimateAffinePartial2D(
                    atlas_pts_for_input, pts_for_input
                )[0]
                if atlas_to_brain_align:
                    atlas_warped_left = cv2.warpAffine(im_left, warp_coords, (512, 512))
                    atlas_warped_right = cv2.warpAffine(im_right, warp_coords, (512, 512))
                    atlas_warped = cv2.bitwise_or(atlas_warped_left, atlas_warped_right)
                    ret, atlas_warped = cv2.threshold(
                        atlas_warped, 5, 255, cv2.THRESH_BINARY_INV
                    )
                    atlas_left_transform_path = os.path.join(
                        output_mask_path, "{}_atlas_left_transform.png".format(str(n))
                    )
                    atlas_right_transform_path = os.path.join(
                        output_mask_path, "{}_atlas_right_transform.png".format(str(n))
                    )
                    io.imsave(atlas_left_transform_path, atlas_warped_left)
                    io.imsave(atlas_right_transform_path, atlas_warped_right)
                else:
                    atlas_warped = cv2.warpAffine(im, warp_coords, (512, 512))
            elif len(atlas_pts_for_input[0]) == 3:
                warp_coords = cv2.getAffineTransform(atlas_pts_for_input, pts_for_input)
                atlas_warped = cv2.warpAffine(im, warp_coords, (512, 512))
            elif len(atlas_pts_for_input[0]) >= 4:
                im_final_size = (512, 512)

                left = acc_left_total.iloc[n, :].values.astype("float32").tolist()
                right = acc_right_total.iloc[n, :].values.astype("float32").tolist()
                left = np.argsort(left).tolist()
                right = np.argsort(right).tolist()
                right = [x + 1 for x in right]
                # if set([1, 3, 5, 7]).issubset(landmark_arr):
                if set([0, 3, 5, 6]).issubset(landmark_arr) and len(landmark_arr) >= 7:
                    left = [0, 3, 5]
                    right = [3, 5, 6]
                    # OLD METHOD FOR PREFERENTIALLY USING TOP CENTRE AND LAMBDA
                    # elif 4 in landmark_arr and 6 in landmark_arr:
                    #     # If top centre and lambda landmarks are selected,
                    #     # add the first landmark from the left hemisphere and the
                    #     # last one from the right hemisphere.
                    #     centreline_pts = [landmark_arr.index(x) for x in landmark_arr if x in [4, 6]]
                    #     left = [x for x in landmark_arr if x in [0, 1, 2, 3, 5]][0]
                    #     right = [landmark_arr.index(x) for x in reversed(landmark_arr) if x in [8, 7, 5, 4, 3]][0]
                    #     left = centreline_pts.extend([left])
                    #     right = [right].extend(centreline_pts)
                    #     print(landmark_indices)
                    #     print(landmark_arr)
                    #     print(left)
                    #     print(right)
                else:
                    # left = [x for x in landmark_indices if x in range(0, 6)][0:2]
                    # right = [x for x in landmark_indices if x in range(3, 9)][0:2]
                    left = [landmark_arr.index(x) for x in landmark_arr if x in [0, 1, 2, 3, 4, 5]][0:3]
                    # right = [landmark_arr.index(x) for x in reversed(landmark_arr) if x in [8, 7, 6, 5, 4, 3]][0:3]
                    right =  [landmark_arr.index(x) for x in landmark_arr if x in [3, 4, 5, 6, 7, 8]][-3:]
                    print(landmark_indices)
                    print(landmark_arr)
                    print(left)
                    print(right)

                # try:
                atlas_pts_left = np.array(
                    [
                        atlas_pts[align_val][left[0]],
                        atlas_pts[align_val][left[1]],
                        atlas_pts[align_val][left[2]],
                    ],
                    dtype=np.float32,
                )
                atlas_pts_right = np.array(
                    [
                        atlas_pts[align_val][right[0]],
                        atlas_pts[align_val][right[1]],
                        atlas_pts[align_val][right[2]],
                    ],
                    dtype=np.float32,
                )
                dlc_pts_left = np.array(
                    [
                        dlc_pts[align_val][left[0]],
                        dlc_pts[align_val][left[1]],
                        dlc_pts[align_val][left[2]],
                    ],
                    dtype=np.float32,
                )
                dlc_pts_right = np.array(
                    [
                        dlc_pts[align_val][right[0]],
                        dlc_pts[align_val][right[1]],
                        dlc_pts[align_val][right[2]],
                    ],
                    dtype=np.float32,
                )

                # except:
                #     # Old: left: 0, 2, 3; right: 1, 2, 3
                #     atlas_pts_left = np.array(
                #         [
                #             atlas_pts[align_val][0],
                #             atlas_pts[align_val][1],
                #             atlas_pts[align_val][2],
                #         ],
                #         dtype=np.float32,
                #     )
                #     atlas_pts_right = np.array(
                #         [
                #             atlas_pts[align_val][5],
                #             atlas_pts[align_val][6],
                #             atlas_pts[align_val][7],
                #         ],
                #         dtype=np.float32,
                #     )
                #     dlc_pts_left = np.array(
                #         [
                #             dlc_pts[align_val][0],
                #             dlc_pts[align_val][1],
                #             dlc_pts[align_val][2],
                #         ],
                #         dtype=np.float32,
                #     )
                #     dlc_pts_right = np.array(
                #         [
                #             dlc_pts[align_val][5],
                #             dlc_pts[align_val][6],
                #             dlc_pts[align_val][7],
                #         ],
                #         dtype=np.float32,
                #     )

                warp_coords_left = cv2.getAffineTransform(atlas_pts_left, dlc_pts_left)
                warp_coords_right = cv2.getAffineTransform(atlas_pts_right, dlc_pts_right)
                warp_coords_brain_atlas_left = cv2.getAffineTransform(dlc_pts_left, atlas_pts_left)
                warp_coords_brain_atlas_right = cv2.getAffineTransform(dlc_pts_right, atlas_pts_right)
                if atlas_to_brain_align:
                    atlas_warped_left = cv2.warpAffine(
                        im_left, warp_coords_left, im_final_size
                    )
                    atlas_warped_right = cv2.warpAffine(
                        im_right, warp_coords_right, im_final_size
                    )
                    atlas_warped = cv2.bitwise_or(atlas_warped_left, atlas_warped_right)
                    ret, atlas_warped = cv2.threshold(
                        atlas_warped, 5, 255, cv2.THRESH_BINARY_INV
                    )
                    if not original_label:
                        atlas_label_left = cv2.warpAffine(
                            atlas_label_mask_left, warp_coords_left, im_final_size
                        )
                        atlas_label_right = cv2.warpAffine(
                            atlas_label_mask_right, warp_coords_right, im_final_size
                        )
                        atlas_label = cv2.bitwise_or(atlas_label_left, atlas_label_right)

                else:
                    pts_np = np.array(
                        [
                            dlc_pts[align_val][0],
                            dlc_pts[align_val][1],
                            dlc_pts[align_val][2],
                        ],
                        dtype=np.float32,
                    )
                    atlas_pts_np = np.array(
                        [
                            atlas_pts[align_val][0],
                            atlas_pts[align_val][1],
                            atlas_pts[align_val][2],
                        ],
                        dtype=np.float32,
                    )
                    warp_coords = cv2.getAffineTransform(pts_np, atlas_pts_np)
                    atlas_warped = cv2.warpAffine(im, warp_coords, (512, 512))

            # if atlas_to_brain_align:
            if len(atlas_pts_for_input[0]) == 2:
                atlas_mask_left_warped = cv2.warpAffine(
                    atlas_mask_left, warp_coords, (512, 512)
                )
                atlas_mask_right_warped = cv2.warpAffine(
                    atlas_mask_right, warp_coords, (512, 512)
                )
                atlas_mask_warped = cv2.bitwise_or(
                    atlas_mask_left_warped, atlas_mask_right_warped
                )
            if len(atlas_pts_for_input[0]) == 3:
                atlas_mask_warped = cv2.warpAffine(atlas_mask, warp_coords, (512, 512))
            if len(atlas_pts_for_input[0]) >= 4:
                atlas_mask_left_warped = cv2.warpAffine(
                    atlas_mask_left, warp_coords_left, (512, 512)
                )
                atlas_mask_right_warped = cv2.warpAffine(
                    atlas_mask_right, warp_coords_right, (512, 512)
                )
                atlas_mask_warped = cv2.bitwise_or(
                    atlas_mask_left_warped, atlas_mask_right_warped
                )
            atlas_mask_warped = np.uint8(atlas_mask_warped)
            io.imsave(mask_warped_path_alt_left, atlas_mask_left_warped)
            io.imsave(mask_warped_path_alt_right, atlas_mask_right_warped)
            # brain_to_atlas_mask = cv2.bitwise_or(
            #     atlas_mask_warped, im
            # )
            # io.imsave(brain_to_atlas_mask_path, brain_to_atlas_mask)
            hemispheres = ['left', 'right']
            olfactory_bulbs_to_use = []
            if olfactory_check and use_unet:
                if align_once and n != 0:
                    olfactory_bulbs = olfactory_bulbs_to_use_pre_align_list[0]
                else:
                    mask = cv2.imread(mask_dir, cv2.IMREAD_GRAYSCALE)
                    cnts_for_olfactory = cv2.findContours(
                        mask.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
                    )
                    cnts_for_olfactory = imutils.grab_contours(cnts_for_olfactory)
                    if len(cnts_for_olfactory) == 3:
                        olfactory_bulbs = sorted(
                            cnts_for_olfactory, key=cv2.contourArea, reverse=True
                        )[1:3]
                    else:
                        olfactory_bulbs = sorted(
                            cnts_for_olfactory, key=cv2.contourArea, reverse=True
                        )[2:4]
                if align_once and n == 0:
                    olfactory_bulbs_to_use_pre_align_list.append(olfactory_bulbs)
                if len(olfactory_bulbs) == 0:
                    # If there's no olfactory bulbs found, disable future olfactory
                    # bulb checking
                    olfactory_check = False
            for hemisphere in hemispheres:
                new_data = []
                if hemisphere == 'left':
                    mask_path = mask_warped_path_alt_left
                    mask_warped_to_use = atlas_mask_left_warped
                    if olfactory_check and use_unet:
                        if len(olfactory_bulbs) >= 1:
                            bulb = olfactory_bulbs[0]
                        bulb_fill = 300
                else:
                    mask_path = mask_warped_path_alt_right
                    mask_warped_to_use = atlas_mask_right_warped
                    if olfactory_check and use_unet:
                        if len(olfactory_bulbs) > 1:
                            bulb = olfactory_bulbs[1]
                        bulb_fill = 400
                if olfactory_check and use_unet:
                    olfactory_bulbs_to_use = olfactory_bulbs
                    try:
                        cv2.fillPoly(mask_warped_to_use, pts=[bulb], color=[255, 255, 255])
                        io.imsave(mask_path, mask_warped_to_use)
                    except:
                        print('No olfactory bulb found!')
                    mask_warped_to_use = cv2.cvtColor(mask_warped_to_use, cv2.COLOR_BGR2GRAY)
                img_edited = Image.open(mask_path)
                # Generates a transparent version of the brain atlas.
                img_rgba = img_edited.convert("RGBA")
                data = img_rgba.getdata()
                for pixel in data:
                    if pixel[0] == 255 and pixel[1] == 255 and pixel[2] == 255:
                        new_data.append((pixel[0], pixel[1], pixel[2], 0))
                    else:
                        new_data.append(pixel)
                img_rgba.putdata(new_data)
                img_rgba.save(os.path.join(output_mask_path, "{}_brain_atlas_mask_transparent_{}.png".format(n, hemisphere)))
                img_transparent = cv2.imread(
                    os.path.join(output_mask_path, "{}_brain_atlas_mask_transparent_{}.png".format(n, hemisphere))
                )
                brain_atlas_transparent = cv2.bitwise_and(im, img_transparent)

                io.imsave(os.path.join(output_mask_path, "{}_brain_atlas_transparent_{}.png".format(n, hemisphere)),
                          brain_atlas_transparent)
                if hemisphere == 'left':
                    if atlas_to_brain_align:
                        brain_to_atlas_warped_left = cv2.warpAffine(brain_atlas_transparent, warp_coords_left, (512, 512))
                        if olfactory_check and use_unet:
                            olfactory_warped_left = cv2.warpAffine(mask_warped_to_use, warp_coords_left, (512, 512))
                            olfactory_warped_cnts_left = cv2.findContours(
                                olfactory_warped_left.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
                            )
                            olfactory_warped_cnts_left = imutils.grab_contours(olfactory_warped_cnts_left)
                            olfactory_warped_left = sorted(
                                olfactory_warped_cnts_left, key=cv2.contourArea, reverse=True
                            )[-1]
                    else:
                        brain_to_atlas_warped_left = cv2.warpAffine(brain_atlas_transparent, warp_coords_brain_atlas_left, (512, 512))
                        if olfactory_check and use_unet:
                            olfactory_warped_left = cv2.warpAffine(mask_warped_to_use, warp_coords_brain_atlas_left, (512, 512))
                            olfactory_warped_cnts_left = cv2.findContours(
                                olfactory_warped_left.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
                            )
                            olfactory_warped_cnts_left = imutils.grab_contours(olfactory_warped_cnts_left)
                            olfactory_warped_left = sorted(
                                olfactory_warped_cnts_left, key=cv2.contourArea, reverse=True
                            )[-1]
                else:
                    if atlas_to_brain_align:
                        brain_to_atlas_warped_right = cv2.warpAffine(brain_atlas_transparent, warp_coords_right, (512, 512))
                        olfactory_warped_right = cv2.warpAffine(mask_warped_to_use, warp_coords_right, (512, 512))
                        if olfactory_check and use_unet:
                            olfactory_warped_right = cv2.warpAffine(mask_warped_to_use, warp_coords_right, (512, 512))
                            olfactory_warped_cnts_right = cv2.findContours(
                                olfactory_warped_right.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
                            )
                            olfactory_warped_cnts_right = imutils.grab_contours(olfactory_warped_cnts_right)
                            olfactory_warped_right = sorted(
                                olfactory_warped_cnts_right, key=cv2.contourArea, reverse=True
                            )[-1]
                    else:
                        brain_to_atlas_warped_right = cv2.warpAffine(brain_atlas_transparent, warp_coords_brain_atlas_right, (512, 512))
                        olfactory_warped_right = cv2.warpAffine(mask_warped_to_use, warp_coords_brain_atlas_right, (512, 512))
                        if olfactory_check and use_unet:
                            olfactory_warped_right = cv2.warpAffine(mask_warped_to_use, warp_coords_brain_atlas_right, (512, 512))
                            olfactory_warped_cnts_right = cv2.findContours(
                                olfactory_warped_right.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
                            )
                            olfactory_warped_cnts_right = imutils.grab_contours(olfactory_warped_cnts_right)
                            olfactory_warped_right = sorted(
                                olfactory_warped_cnts_right, key=cv2.contourArea, reverse=True
                            )[-1]
            if olfactory_check and use_unet:
                olfactory_bulbs_to_use_list.append([olfactory_warped_left, olfactory_warped_right])
            brain_to_atlas_mask = cv2.bitwise_or(
                brain_to_atlas_warped_left, brain_to_atlas_warped_right
            )
            io.imsave(brain_to_atlas_mask_path, brain_to_atlas_mask)


            # Second alignment of brain atlas using cortical landmarks and piecewise affine transform
            print("Performing second transformation of atlas {}...".format(n))

            atlas_first_transform_path = os.path.join(
                output_mask_path, "{}_atlas_first_transform.png".format(str(n))
            )

            if atlas_to_brain_align:
                dst = atlas_warped
            else:
                dst = brain_to_atlas_mask

            io.imsave(atlas_first_transform_path, dst)

            # If a sensory map of the brain is provided, do a third alignment of the brain atlas using up to
            # four peaks of sensory activity
            if sensory_match:
                original_label = True
                if atlas_to_brain_align:
                    dst = getMaskContour(
                        mask_dir,
                        atlas_warped,
                        sensory_peak_pts[align_val],
                        sensory_atlas_pts[align_val],
                        cwd,
                        align_val,
                        False,
                    )
                    atlas_mask_warped = getMaskContour(
                        atlas_first_transform_path,
                        atlas_mask_warped,
                        sensory_peak_pts[align_val],
                        sensory_atlas_pts[align_val],
                        cwd,
                        align_val,
                        False,
                    )
                    atlas_mask_warped = cv2.resize(
                        atlas_mask_warped, (im.shape[0], im.shape[1])
                    )
                else:
                    dst = atlas_warped
        else:
            # If we're not using DeepLabCut...
            if atlas_to_brain_align and not use_voxelmorph:
                dst = cv2.bitwise_or(im_left, im_right)
                ret, dst = cv2.threshold(
                    dst, 5, 255, cv2.THRESH_BINARY_INV
                )
            else:
                dst = im
            dst = np.uint8(dst)

        if use_dlc:
            if atlas_to_brain_align:
                io.imsave(mask_warped_path, atlas_mask_warped)
            else:
                io.imsave(mask_warped_path, atlas_mask)
        else:
            io.imsave(mask_warped_path, dst)
            # if atlas_to_brain_align:
            if use_voxelmorph:
                atlas_mask_warped = atlas_mask
            else:
                atlas_mask_warped = cv2.bitwise_or(
                    atlas_mask_left, atlas_mask_right
                )
            atlas_mask_warped = cv2.cvtColor(atlas_mask_warped, cv2.COLOR_BGR2GRAY)
            ret, atlas_mask_warped = cv2.threshold(
                atlas_mask_warped, 5, 255, cv2.THRESH_BINARY
            )
            atlas_mask_warped = np.uint8(atlas_mask_warped)
            original_label = True
            io.imsave(mask_warped_path, atlas_mask_warped)
        # Resize images back to 512x512
        dst = cv2.resize(dst, (im.shape[0], im.shape[1]))
        atlas_path = os.path.join(output_mask_path, "{}_atlas.png".format(str(n)))

        vxm_template_output_path = os.path.join(
            output_mask_path, "{}_vxm_template.png".format(str(n))
        )

        dst_list.append(dst)
        if use_voxelmorph:
            vxm_template_list.append(vxm_template)
            io.imsave(vxm_template_output_path, vxm_template_list[n])

        if atlas_to_brain_align:
            io.imsave(atlas_path, dst)
            br_list.append(br)
        else:
            brain_warped_path = os.path.join(
                output_mask_path, "{}_brain_warp.png".format(str(n))
            )
            io.imsave(brain_warped_path, dst)
            io.imsave(atlas_path, atlas)

        if atlas_to_brain_align:
            if original_label:
                atlas_label = []
            atlas_label = atlas_to_mask(
                atlas_path,
                mask_dir,
                mask_warped_path,
                output_mask_path,
                n,
                use_unet,
                use_voxelmorph,
                atlas_to_brain_align,
                git_repo_base,
                olfactory_check,
                [],
                atlas_label
            )
            atlas_label_list.append(atlas_label)
        elif not use_dlc:
            io.imsave(os.path.join(output_mask_path, "{}.png".format(n)), dst)
        if bregma_present:
            bregma_val = int(bregma_index_list[n])
            bregma_list.append(dlc_pts[n][bregma_val])

    # Carries out VoxelMorph on each motif-based functional map (MBFM) that has been aligned to a raw brain image
    if use_voxelmorph:
        align_once = True
        if len(dst_list) == 1:
            n_to_use = 0
        else:
            n_to_use = 1
        for (n_post, dst_post), vxm_template_post in zip(enumerate([dst_list[n_to_use]]), [vxm_template_list[n_to_use]]):
            output_img, flow_post = voxelmorph_align(
               voxelmorph_model_path, dst_post, vxm_template_post, exist_transform, flow_path
            )
            flow_path_after = os.path.join(output_mask_path, "{}_flow.npy".format(str(n_post)))
            np.save(flow_path_after, flow_post)
            output_path_after = os.path.join(output_mask_path, "{}_output_img.png".format(str(n_post)))
            io.imsave(output_path_after, output_img)
            if not exist_transform:
                dst_gray = cv2.cvtColor(atlas, cv2.COLOR_BGR2GRAY)
                dst_post = vxm_transform(dst_gray, flow_path_after)
                ret, dst_post = cv2.threshold(
                    dst_post, 1, 255, cv2.THRESH_BINARY
                ) # 5, 255
                dst_post = np.uint8(dst_post)

            mask_warped_path = os.path.join(
                output_mask_path, "{}_mask_warped.png".format(str(n_post))
            )

            atlas_first_transform_path_post = os.path.join(
                output_mask_path, "{}_atlas_first_transform.png".format(str(n_post))
            )

            io.imsave(atlas_first_transform_path_post, dst_post)
            if align_once and use_voxelmorph:
                atlas_path = os.path.join(output_mask_path, "{}_atlas.png".format(str(0)))

            brain_warped_path = os.path.join(
                output_mask_path, "{}_brain_warp.png".format(str(n_post))
            )
            mask_dir = os.path.join(cwd, "../output_mask/{}.png".format(n_post))
            dst_post = cv2.resize(dst_post, (im.shape[0], im.shape[1]))
            if not atlas_to_brain_align:
                atlas_to_brain_align = True
                original_label = True
            if atlas_to_brain_align:
                io.imsave(atlas_path, dst_post)
            else:
                io.imsave(brain_warped_path, dst_post)
            if atlas_to_brain_align:
                if original_label:
                    atlas_label = []
                if use_voxelmorph and olfactory_check and use_unet and use_dlc:
                    if align_once:
                        olfactory_bulbs_to_use_check = olfactory_bulbs_to_use_list[0]
                    else:
                        olfactory_bulbs_to_use_check = olfactory_bulbs_to_use_list[n_post]
                else:
                    olfactory_bulbs_to_use_check = []
                atlas_label = atlas_to_mask(
                    atlas_path,
                    mask_dir,
                    mask_warped_path,
                    output_mask_path,
                    n_post,
                    use_unet,
                    use_voxelmorph,
                    atlas_to_brain_align,
                    git_repo_base,
                    olfactory_check,
                    olfactory_bulbs_to_use_check,
                    atlas_label
                )
                atlas_label_list.append(atlas_label)


    # Converts the transformed brain atlas into a segmentation template for the original brain image
    applyMask(
        brain_img_dir,
        output_mask_path,
        output_overlay_path,
        output_overlay_path,
        mat_save,
        threshold,
        git_repo_base,
        bregma_list,
        atlas_to_brain_align,
        model,
        dlc_pts,
        atlas_pts,
        olfactory_check,
        use_unet,
        use_dlc,
        use_voxelmorph,
        plot_landmarks,
        align_once,
        atlas_label_list,
        olfactory_bulbs_to_use_list,
        region_labels,
        original_label
    )
MesoNet Authors: <NAME> and <NAME>, <NAME> https://github.com/bf777/MesoNet Licensed under the Creative Commons Attribution 4.0 International License (see LICENSE for details) # from mesonet.niftyreg_align import niftyreg_align Locates the peaks of cortical activity in each hemisphere of the brain image.     :param img: A sensory map (indicating functional activation in the mouse brain).     :return: An array of the maximum points of activation for the sensory map. Generates a binary brain atlas from a .mat file.     :param input_file: The input .mat file representing a brain atlas (with white = 255 and black = 0)     :param mat_cnt_list: The list to which mat files should be appended     :return: A thresholded binary brain atlas. # https://stackoverflow.com/questions/42258637/how-to-know-the-angle-between-two-points Gets the contour of the brain's boundaries and applies a piecewise affine transform to the brain atlas     based on the cortical landmarks predicted in dlc_predict (and peaks of activity on the sensory map, if available).     :param mask_dir: The path to the directory containing the U-net masks of the brain's boundaries.     :param atlas_img: The brain atlas to be transformed.     :param predicted_pts: The coordinates of the cortical landmarks predicted in dlc_predict (or, for the second run     of this function, the coordinates of the peaks of activity in the sensory map).     :param actual_pts: The fixed coordinates of the cortical landmarks on the brain atlas (or, for the second run of     this function, the fixed coordinates of the peaks of sensory activity on the brain atlas).     :param cwd: The path to the current working directory.     :param n: The number of the current image in the directory. Align and overlap brain atlas onto brain image based on four landmark locations in the brain image and the atlas.     :param brain_img_dir: The directory containing each brain image to be used.     :param sensory_img_dir: The directory containing each sensory image to be used (if you are aligning each brain     image using a sensory map).     :param coords_input: Predicted locations of the four landmarks on the brain image from the file generated by     DeepLabCut.     :param sensory_match: Whether or not a sensory map is to be used.     :param mat_save: Whether or not to export each brain region to a .mat file in applyMask, which is called at the end     of this function.     :param threshold: The threshold for the cv2.opening operation carried out in applyMask, which is called at the end     of this function.     :param git_repo_base: The path to the base git repository containing necessary resources for MesoNet (reference     atlases, DeepLabCut config files, etc.)     :param region_labels: Choose whether or not to attempt to label each region with its name from the Allen Institute     Mouse Brain Atlas.     :param landmark_arr_orig: The original array of landmarks from DeepLabCut (to be distinguished from any automatic     exclusions to landmark array based on prediction quality).     :param use_unet: Choose whether or not to identify the borders of the cortex using a U-net model.     :param atlas_to_brain_align: If True, registers the atlas to each brain image. If False, registers each brain image     to the atlas.     :param model: The name of the U-net model (for passthrough to mask_functions.py)     :param olfactory_check: If True, draws olfactory bulb contours on the brain image.     :param plot_landmarks: If True, plots DeepLabCut landmarks (large circles) and original alignment landmarks (small     circles) on final brain image.     :param align_once: if True, carries out all alignments based on the alignment of the first atlas and brain. This can     save time if you have many frames of the same brain with a fixed camera position.     :param original_label: if True, uses a brain region labelling approach that attempts to automatically sort brain     regions in a consistent order (left to right by hemisphere, then top to bottom for vertically aligned regions). This     approach may be more flexible if you're using a custom brain atlas (i.e. not one in which region is filled with a     unique number).     :param exist_transform: if True, uses an existing voxelmorph transformation field for all data instead of predicting     a new transformation.     :param voxelmorph_model: the name of a .h5 model located in the models folder of the git repository for MesoNet,     generated using voxelmorph and containing weights for a voxelmorph local deformation model.     :param vxm_template_path: the path to a template atlas (.npy or .mat( to which the brain image will be aligned in     voxelmorph.     :param flow_path: the path to a voxelmorph transformation field that will be used to transform all data instead of     predicting a new transformation if exist_transform is True. # load brain images folder # Prepare template for VoxelMorph # Prepare output folder # Output folder for transparent masks and masks overlaid onto brain image # im = atlas_from_mat(os.path.join(git_repo_base, 'atlases/atlas_ROIs.mat')) # FOR ALIGNING BRAIN TO ATLAS # i_coord, j_coord = np.array([(100, 256, 413, 256), (148, 254, 148, 446)]) # https://www.pyimagesearch.com/2014/07/21/detecting-circles-images-using-opencv-hough-circles/ # detect circles in the image # ensure at least some circles were found # convert the (x, y) coordinates and radius of the circles to integers # [0, 3, 2, 1] # Initialize result as max value # Get peak values from heatmaps # Get circle locations # FOR ALIGNING BRAIN TO ATLAS # First alignment of brain atlas using three cortical landmarks and standard affine transform # if set([1, 3, 5, 7]).issubset(landmark_arr): # OLD METHOD FOR PREFERENTIALLY USING TOP CENTRE AND LAMBDA # elif 4 in landmark_arr and 6 in landmark_arr: #     # If top centre and lambda landmarks are selected, #     # add the first landmark from the left hemisphere and the #     # last one from the right hemisphere. #     centreline_pts = [landmark_arr.index(x) for x in landmark_arr if x in [4, 6]] #     left = [x for x in landmark_arr if x in [0, 1, 2, 3, 5]][0] #     right = [landmark_arr.index(x) for x in reversed(landmark_arr) if x in [8, 7, 5, 4, 3]][0] #     left = centreline_pts.extend([left]) #     right = [right].extend(centreline_pts) #     print(landmark_indices) #     print(landmark_arr) #     print(left) #     print(right) # left = [x for x in landmark_indices if x in range(0, 6)][0:2] # right = [x for x in landmark_indices if x in range(3, 9)][0:2] # right = [landmark_arr.index(x) for x in reversed(landmark_arr) if x in [8, 7, 6, 5, 4, 3]][0:3] # try: # except: #     # Old: left: 0, 2, 3; right: 1, 2, 3 #     atlas_pts_left = np.array( #         [ #             atlas_pts[align_val][0], #             atlas_pts[align_val][1], #             atlas_pts[align_val][2], #         ], #         dtype=np.float32, #     ) #     atlas_pts_right = np.array( #         [ #             atlas_pts[align_val][5], #             atlas_pts[align_val][6], #             atlas_pts[align_val][7], #         ], #         dtype=np.float32, #     ) #     dlc_pts_left = np.array( #         [ #             dlc_pts[align_val][0], #             dlc_pts[align_val][1], #             dlc_pts[align_val][2], #         ], #         dtype=np.float32, #     ) #     dlc_pts_right = np.array( #         [ #             dlc_pts[align_val][5], #             dlc_pts[align_val][6], #             dlc_pts[align_val][7], #         ], #         dtype=np.float32, #     ) # if atlas_to_brain_align: # brain_to_atlas_mask = cv2.bitwise_or( #     atlas_mask_warped, im # ) # io.imsave(brain_to_atlas_mask_path, brain_to_atlas_mask) # If there's no olfactory bulbs found, disable future olfactory # bulb checking # Generates a transparent version of the brain atlas. # Second alignment of brain atlas using cortical landmarks and piecewise affine transform # If a sensory map of the brain is provided, do a third alignment of the brain atlas using up to # four peaks of sensory activity # If we're not using DeepLabCut... # if atlas_to_brain_align: # Resize images back to 512x512 # Carries out VoxelMorph on each motif-based functional map (MBFM) that has been aligned to a raw brain image # 5, 255 # Converts the transformed brain atlas into a segmentation template for the original brain image