Comp2Comp-main/comp2comp/hip/hip_utils.py

"""
@author: louisblankemeier
"""

import math
import os
import shutil

import cv2
import nibabel as nib
import numpy as np
import scipy.ndimage as ndi
from scipy.ndimage import zoom
from skimage.morphology import ball, binary_erosion

from comp2comp.hip.hip_visualization import method_visualizer


def compute_rois(medical_volume, segmentation, model, output_dir, save=False):
    left_femur_mask = segmentation.get_fdata() == model.categories["femur_left"]
    left_femur_mask = left_femur_mask.astype(np.uint8)
    right_femur_mask = segmentation.get_fdata() == model.categories["femur_right"]
    right_femur_mask = right_femur_mask.astype(np.uint8)
    left_head_roi, left_head_centroid, left_head_hu = get_femural_head_roi(
        left_femur_mask, medical_volume, output_dir, "left_head"
    )
    right_head_roi, right_head_centroid, right_head_hu = get_femural_head_roi(
        right_femur_mask, medical_volume, output_dir, "right_head"
    )
    (
        left_intertrochanter_roi,
        left_intertrochanter_centroid,
        left_intertrochanter_hu,
    ) = get_femural_head_roi(
        left_femur_mask, medical_volume, output_dir, "left_intertrochanter"
    )
    (
        right_intertrochanter_roi,
        right_intertrochanter_centroid,
        right_intertrochanter_hu,
    ) = get_femural_head_roi(
        right_femur_mask, medical_volume, output_dir, "right_intertrochanter"
    )
    (
        left_neck_roi,
        left_neck_centroid,
        left_neck_hu,
    ) = get_femural_neck_roi(
        left_femur_mask,
        medical_volume,
        left_intertrochanter_roi,
        left_intertrochanter_centroid,
        left_head_roi,
        left_head_centroid,
        output_dir,
    )
    (
        right_neck_roi,
        right_neck_centroid,
        right_neck_hu,
    ) = get_femural_neck_roi(
        right_femur_mask,
        medical_volume,
        right_intertrochanter_roi,
        right_intertrochanter_centroid,
        right_head_roi,
        right_head_centroid,
        output_dir,
    )
    combined_roi = (
        left_head_roi
        + (right_head_roi)  # * 2)
        + (left_intertrochanter_roi)  # * 3)
        + (right_intertrochanter_roi)  # * 4)
        + (left_neck_roi)  # * 5)
        + (right_neck_roi)  # * 6)
    )

    if save:
        # make roi directory if it doesn't exist
        parent_output_dir = os.path.dirname(output_dir)
        roi_output_dir = os.path.join(parent_output_dir, "rois")
        if not os.path.exists(roi_output_dir):
            os.makedirs(roi_output_dir)

        # Convert left ROI to NIfTI
        left_roi_nifti = nib.Nifti1Image(combined_roi, medical_volume.affine)
        left_roi_path = os.path.join(roi_output_dir, "roi.nii.gz")
        nib.save(left_roi_nifti, left_roi_path)
        shutil.copy(
            os.path.join(
                os.path.dirname(os.path.abspath(__file__)),
                "tunnelvision.ipynb",
            ),
            parent_output_dir,
        )

    return {
        "left_head": {
            "roi": left_head_roi,
            "centroid": left_head_centroid,
            "hu": left_head_hu,
        },
        "right_head": {
            "roi": right_head_roi,
            "centroid": right_head_centroid,
            "hu": right_head_hu,
        },
        "left_intertrochanter": {
            "roi": left_intertrochanter_roi,
            "centroid": left_intertrochanter_centroid,
            "hu": left_intertrochanter_hu,
        },
        "right_intertrochanter": {
            "roi": right_intertrochanter_roi,
            "centroid": right_intertrochanter_centroid,
            "hu": right_intertrochanter_hu,
        },
        "left_neck": {
            "roi": left_neck_roi,
            "centroid": left_neck_centroid,
            "hu": left_neck_hu,
        },
        "right_neck": {
            "roi": right_neck_roi,
            "centroid": right_neck_centroid,
            "hu": right_neck_hu,
        },
    }


def get_femural_head_roi(
    femur_mask,
    medical_volume,
    output_dir,
    anatomy,
    visualize_method=False,
    min_pixel_count=20,
):
    top = np.where(femur_mask.sum(axis=(0, 1)) != 0)[0].max()
    top_mask = femur_mask[:, :, top]

    print(f"======== Computing {anatomy} femur ROIs ========")

    while True:
        labeled, num_features = ndi.label(top_mask)

        component_sizes = np.bincount(labeled.ravel())
        valid_components = np.where(component_sizes >= min_pixel_count)[0][1:]

        if len(valid_components) == 2:
            break

        top -= 1
        if top < 0:
            print("Two connected components not found in the femur mask.")
            break
        top_mask = femur_mask[:, :, top]

    if len(valid_components) == 2:
        # Find the center of mass for each connected component
        center_of_mass_1 = list(
            ndi.center_of_mass(top_mask, labeled, valid_components[0])
        )
        center_of_mass_2 = list(
            ndi.center_of_mass(top_mask, labeled, valid_components[1])
        )

        # Assign left_center_of_mass to be the center of mass with lowest value in the first dimension
        if center_of_mass_1[0] < center_of_mass_2[0]:
            left_center_of_mass = center_of_mass_1
            right_center_of_mass = center_of_mass_2
        else:
            left_center_of_mass = center_of_mass_2
            right_center_of_mass = center_of_mass_1

        print(f"Left center of mass: {left_center_of_mass}")
        print(f"Right center of mass: {right_center_of_mass}")

    if anatomy == "left_intertrochanter" or anatomy == "right_head":
        center_of_mass = left_center_of_mass
    elif anatomy == "right_intertrochanter" or anatomy == "left_head":
        center_of_mass = right_center_of_mass

    coronal_slice = femur_mask[:, round(center_of_mass[1]), :]
    coronal_image = medical_volume.get_fdata()[:, round(center_of_mass[1]), :]
    sagittal_slice = femur_mask[round(center_of_mass[0]), :, :]
    sagittal_image = medical_volume.get_fdata()[round(center_of_mass[0]), :, :]

    zooms = medical_volume.header.get_zooms()
    zoom_factor = zooms[2] / zooms[1]

    coronal_slice = zoom(coronal_slice, (1, zoom_factor), order=1).round()
    coronal_image = zoom(coronal_image, (1, zoom_factor), order=3).round()
    sagittal_image = zoom(sagittal_image, (1, zoom_factor), order=3).round()

    centroid = [round(center_of_mass[0]), 0, 0]

    print(f"Starting centroid: {centroid}")

    for _ in range(3):
        sagittal_slice = femur_mask[centroid[0], :, :]
        sagittal_slice = zoom(sagittal_slice, (1, zoom_factor), order=1).round()
        centroid[1], centroid[2], radius_sagittal = inscribe_sagittal(
            sagittal_slice, zoom_factor
        )

        print(f"Centroid after inscribe sagittal: {centroid}")

        axial_slice = femur_mask[:, :, centroid[2]]
        if anatomy == "left_intertrochanter" or anatomy == "right_head":
            axial_slice[round(right_center_of_mass[0]) :, :] = 0
        elif anatomy == "right_intertrochanter" or anatomy == "left_head":
            axial_slice[: round(left_center_of_mass[0]), :] = 0
        centroid[0], centroid[1], radius_axial = inscribe_axial(axial_slice)

        print(f"Centroid after inscribe axial: {centroid}")

    axial_image = medical_volume.get_fdata()[:, :, round(centroid[2])]
    sagittal_image = medical_volume.get_fdata()[round(centroid[0]), :, :]
    sagittal_image = zoom(sagittal_image, (1, zoom_factor), order=3).round()

    if visualize_method:
        method_visualizer(
            sagittal_image,
            axial_image,
            axial_slice,
            sagittal_slice,
            [centroid[2], centroid[1]],
            radius_sagittal,
            [centroid[1], centroid[0]],
            radius_axial,
            output_dir,
            anatomy,
        )

    roi = compute_hip_roi(medical_volume, centroid, radius_sagittal, radius_axial)

    # selem = ndi.generate_binary_structure(3, 1)
    selem = ball(3)
    femur_mask_eroded = binary_erosion(femur_mask, selem)
    roi = roi * femur_mask_eroded
    roi_eroded = roi.astype(np.uint8)

    hu = get_mean_roi_hu(medical_volume, roi_eroded)

    return (roi_eroded, centroid, hu)


def get_femural_neck_roi(
    femur_mask,
    medical_volume,
    intertrochanter_roi,
    intertrochanter_centroid,
    head_roi,
    head_centroid,
    output_dir,
):
    zooms = medical_volume.header.get_zooms()

    direction_vector = np.array(head_centroid) - np.array(intertrochanter_centroid)
    unit_direction_vector = direction_vector / np.linalg.norm(direction_vector)

    z, y, x = np.where(intertrochanter_roi)
    intertrochanter_points = np.column_stack((z, y, x))
    t_start = np.dot(
        intertrochanter_points - intertrochanter_centroid, unit_direction_vector
    ).max()

    z, y, x = np.where(head_roi)
    head_points = np.column_stack((z, y, x))
    t_end = (
        np.linalg.norm(direction_vector)
        + np.dot(head_points - head_centroid, unit_direction_vector).min()
    )

    z, y, x = np.indices(femur_mask.shape)
    coordinates = np.stack((z, y, x), axis=-1)

    distance_to_line_origin = np.dot(
        coordinates - intertrochanter_centroid, unit_direction_vector
    )

    coordinates_zoomed = coordinates * zooms
    intertrochanter_centroid_zoomed = np.array(intertrochanter_centroid) * zooms
    unit_direction_vector_zoomed = unit_direction_vector * zooms

    distance_to_line = np.linalg.norm(
        np.cross(
            coordinates_zoomed - intertrochanter_centroid_zoomed,
            coordinates_zoomed
            - (intertrochanter_centroid_zoomed + unit_direction_vector_zoomed),
        ),
        axis=-1,
    ) / np.linalg.norm(unit_direction_vector_zoomed)

    cylinder_radius = 10

    cylinder_mask = (
        (distance_to_line <= cylinder_radius)
        & (distance_to_line_origin >= t_start)
        & (distance_to_line_origin <= t_end)
    )

    # selem = ndi.generate_binary_structure(3, 1)
    selem = ball(3)
    femur_mask_eroded = binary_erosion(femur_mask, selem)
    roi = cylinder_mask * femur_mask_eroded
    neck_roi = roi.astype(np.uint8)

    hu = get_mean_roi_hu(medical_volume, neck_roi)

    centroid = list(
        intertrochanter_centroid + unit_direction_vector * (t_start + t_end) / 2
    )
    centroid = [round(x) for x in centroid]

    return neck_roi, centroid, hu


def compute_hip_roi(img, centroid, radius_sagittal, radius_axial):
    pixel_spacing = img.header.get_zooms()
    length_i = radius_axial * 0.75 / pixel_spacing[0]
    length_j = radius_axial * 0.75 / pixel_spacing[1]
    length_k = radius_sagittal * 0.75 / pixel_spacing[2]

    roi = np.zeros(img.get_fdata().shape, dtype=np.uint8)
    i_lower = math.floor(centroid[0] - length_i)
    j_lower = math.floor(centroid[1] - length_j)
    k_lower = math.floor(centroid[2] - length_k)
    for i in range(i_lower, i_lower + 2 * math.ceil(length_i) + 1):
        for j in range(j_lower, j_lower + 2 * math.ceil(length_j) + 1):
            for k in range(k_lower, k_lower + 2 * math.ceil(length_k) + 1):
                if (i - centroid[0]) ** 2 / length_i**2 + (
                    j - centroid[1]
                ) ** 2 / length_j**2 + (k - centroid[2]) ** 2 / length_k**2 <= 1:
                    roi[i, j, k] = 1
    return roi


def inscribe_axial(axial_mask):
    dist_map = cv2.distanceTransform(axial_mask, cv2.DIST_L2, cv2.DIST_MASK_PRECISE)
    _, radius_axial, _, center_axial = cv2.minMaxLoc(dist_map)
    center_axial = list(center_axial)
    left_right_center = round(center_axial[1])
    posterior_anterior_center = round(center_axial[0])
    return left_right_center, posterior_anterior_center, radius_axial


def inscribe_sagittal(sagittal_mask, zoom_factor):
    dist_map = cv2.distanceTransform(sagittal_mask, cv2.DIST_L2, cv2.DIST_MASK_PRECISE)
    _, radius_sagittal, _, center_sagittal = cv2.minMaxLoc(dist_map)
    center_sagittal = list(center_sagittal)
    posterior_anterior_center = round(center_sagittal[1])
    inferior_superior_center = round(center_sagittal[0])
    inferior_superior_center = round(inferior_superior_center / zoom_factor)
    return posterior_anterior_center, inferior_superior_center, radius_sagittal


def get_mean_roi_hu(medical_volume, roi):
    masked_medical_volume = medical_volume.get_fdata() * roi
    return np.mean(masked_medical_volume[masked_medical_volume != 0])