AvtomatskaPoravnavaMarkerjev/FirstTryLinear.txt

import os
import numpy as np
import scipy
from scipy.spatial.distance import cdist
from scipy.spatial.transform import Rotation as R
import slicer
from DICOMLib import DICOMUtils
from collections import deque
import SimpleITK as sitk
import sitkUtils

#exec(open("C:/Users/lkomar/Documents/Prostata/FirstTryRegister.py").read())

# Define a threshold for grouping nearby points (in voxel space)
#distance_threshold = 4  # This can be adjusted based on your dataset

# Function to group points that are close to each other
def group_points(points, threshold):
    grouped_points = []
    while points:
        point = points.pop()  # Take one point from the list
        group = [point]  # Start a new group
        
        # Find all points close to this one
        distances = cdist([point], points)  # Calculate distances from this point to others
        close_points = [i for i, dist in enumerate(distances[0]) if dist < threshold]
        
        # Add the close points to the group
        group.extend([points[i] for i in close_points])
        
        # Remove the grouped points from the list
        points = [point for i, point in enumerate(points) if i not in close_points]
        
        # Add the group to the result
        grouped_points.append(group)
    
    return grouped_points


def region_growing(image_data, seed, intensity_threshold, max_distance):
    dimensions = image_data.GetDimensions()
    visited = set()
    region = []
    queue = deque([seed])

    while queue:
        x, y, z = queue.popleft()
        if (x, y, z) in visited:
            continue

        visited.add((x, y, z))
        voxel_value = image_data.GetScalarComponentAsDouble(x, y, z, 0)
        
        if voxel_value >= intensity_threshold:
            region.append((x, y, z))
            # Add neighbors within bounds
            for dx, dy, dz in [(1, 0, 0), (-1, 0, 0), (0, 1, 0), (0, -1, 0), (0, 0, 1), (0, 0, -1)]:
                nx, ny, nz = x + dx, y + dy, z + dz
                if 0 <= nx < dimensions[0] and 0 <= ny < dimensions[1] and 0 <= nz < dimensions[2]:
                    if (nx, ny, nz) not in visited:
                        queue.append((nx, ny, nz))

    return region


def detect_points_region_growing(volume_name, intensity_threshold=3000, x_min=90, x_max=380, y_min=190, y_max=380, z_min=80, z_max=120, max_distance=9, centroid_merge_threshold=5):
    volume_node = slicer.util.getNode(volume_name)
    if not volume_node:
        raise RuntimeError(f"Volume {volume_name} not found.")
    
    image_data = volume_node.GetImageData()
    matrix = vtk.vtkMatrix4x4()
    volume_node.GetIJKToRASMatrix(matrix)

    dimensions = image_data.GetDimensions()
    detected_regions = []

    # Check if it's CT or CBCT
    is_cbct = "cbct" in volume_name.lower()

    if is_cbct:
        valid_x_min, valid_x_max = 0, dimensions[0] - 1
        valid_y_min, valid_y_max = 0, dimensions[1] - 1
        valid_z_min, valid_z_max = 0, dimensions[2] - 1
    else:
        valid_x_min, valid_x_max = max(x_min, 0), min(x_max, dimensions[0] - 1)
        valid_y_min, valid_y_max = max(y_min, 0), min(y_max, dimensions[1] - 1)
        valid_z_min, valid_z_max = max(z_min, 0), min(z_max, dimensions[2] - 1)

    visited = set()

    def grow_region(x, y, z):
        if (x, y, z) in visited:
            return None

        voxel_value = image_data.GetScalarComponentAsDouble(x, y, z, 0)
        if voxel_value < intensity_threshold:
            return None

        region = region_growing(image_data, (x, y, z), intensity_threshold, max_distance=max_distance)
        if region:
            for point in region:
                visited.add(tuple(point))
            return region
        return None

    regions = []
    for z in range(valid_z_min, valid_z_max + 1):
        for y in range(valid_y_min, valid_y_max + 1):
            for x in range(valid_x_min, valid_x_max + 1):
                region = grow_region(x, y, z)
                if region:
                    regions.append(region)

    # Collect centroids using intensity-weighted average
    centroids = []
    for region in regions:
        points = np.array([matrix.MultiplyPoint([*point, 1])[:3] for point in region])
        intensities = np.array([image_data.GetScalarComponentAsDouble(*point, 0) for point in region])
        
        if intensities.sum() > 0:
            weighted_centroid = np.average(points, axis=0, weights=intensities)
            max_intensity = intensities.max()
            centroids.append((np.round(weighted_centroid, 2), max_intensity))

    unique_centroids = []
    for centroid, intensity in centroids:
        if not any(np.linalg.norm(centroid - existing_centroid) < centroid_merge_threshold for existing_centroid, _ in unique_centroids):
            unique_centroids.append((centroid, intensity))
            
    markups_node = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLMarkupsFiducialNode", f"Markers_{volume_name}")
    for centroid, intensity in unique_centroids:
        markups_node.AddControlPoint(*centroid)
        #print(f"Detected Centroid (RAS): {centroid}, Max Intensity: {intensity}")

    return unique_centroids



def compute_Kabsch_rotation(moving_points, fixed_points):
    """
    Computes the optimal rotation matrix to align moving_points to fixed_points.
    
    Parameters:
    moving_points (list or ndarray): List of points to be rotated CBCT
    fixed_points (list or ndarray): List of reference points CT

    Returns:
    ndarray: Optimal rotation matrix.
    """
    assert len(moving_points) == len(fixed_points), "Point lists must be the same length."

    # Convert to numpy arrays
    moving = np.array(moving_points)
    fixed = np.array(fixed_points)

    # Compute centroids
    centroid_moving = np.mean(moving, axis=0)
    centroid_fixed = np.mean(fixed, axis=0)

    # Center the points
    moving_centered = moving - centroid_moving
    fixed_centered = fixed - centroid_fixed

    # Compute covariance matrix
    H = np.dot(moving_centered.T, fixed_centered)

    # SVD decomposition
    U, _, Vt = np.linalg.svd(H)
    Rotate_optimal = np.dot(Vt.T, U.T)

    # Correct improper rotation (reflection)
    if np.linalg.det(Rotate_optimal) < 0:
        Vt[-1, :] *= -1
        Rotate_optimal = np.dot(Vt.T, U.T)

    return Rotate_optimal

def compute_Horn_rotation(moving_points, fixed_points):
    """
    Computes the optimal rotation matrix using Horn's method.

    Parameters:
    moving_points (list or ndarray): List of points to be rotated, CBCT
    fixed_points (list or ndarray): List of reference points, CT

    Returns:
    ndarray: Optimal rotation matrix.
    """
    assert len(moving_points) == len(fixed_points), "Point lists must be the same length."
    
    moving = np.array(moving_points)
    fixed = np.array(fixed_points)
    
    # Compute centroids
    centroid_moving = np.mean(moving, axis=0)
    centroid_fixed = np.mean(fixed, axis=0)
    
    # Center the points
    moving_centered = moving - centroid_moving
    fixed_centered = fixed - centroid_fixed
    
    # Compute cross-dispersion matrix
    H = np.dot(moving_centered.T, fixed_centered)
    
    # Compute SVD of H
    U, _, Vt = np.linalg.svd(H)
    
    # Compute rotation matrix
    R = np.dot(Vt.T, U.T)
    
    # Ensure a proper rotation (avoid reflection)
    if np.linalg.det(R) < 0:
        Vt[-1, :] *= -1
        R = np.dot(Vt.T, U.T)
    
    return R

def compute_quaternion_rotation(moving_points, fixed_points):
    """
    Computes the optimal rotation matrix using quaternions.

    Parameters:
    moving_points (list or ndarray): List of points to be rotated.
    fixed_points (list or ndarray): List of reference points.

    Returns:
    ndarray: Optimal rotation matrix.
    """
    assert len(moving_points) == len(fixed_points), "Point lists must be the same length."
    
    moving = np.array(moving_points)
    fixed = np.array(fixed_points)
    
    # Compute centroids
    centroid_moving = np.mean(moving, axis=0)
    centroid_fixed = np.mean(fixed, axis=0)
    
    # Center the points
    moving_centered = moving - centroid_moving
    fixed_centered = fixed - centroid_fixed
    
    # Construct the cross-dispersion matrix
    M = np.dot(moving_centered.T, fixed_centered)
    
    # Construct the N matrix for quaternion solution
    A = M - M.T
    delta = np.array([A[1, 2], A[2, 0], A[0, 1]])
    trace = np.trace(M)
    
    N = np.zeros((4, 4))
    N[0, 0] = trace
    N[1:, 0] = delta
    N[0, 1:] = delta
    N[1:, 1:] = M + M.T - np.eye(3) * trace
    
    # Compute the eigenvector corresponding to the maximum eigenvalue
    eigvals, eigvecs = np.linalg.eigh(N)
    q_optimal = eigvecs[:, np.argmax(eigvals)]  # Optimal quaternion
    
    # Convert quaternion to rotation matrix
    w, x, y, z = q_optimal
    R = np.array([
        [1 - 2*(y**2 + z**2), 2*(x*y - z*w), 2*(x*z + y*w)],
        [2*(x*y + z*w), 1 - 2*(x**2 + z**2), 2*(y*z - x*w)],
        [2*(x*z - y*w), 2*(y*z + x*w), 1 - 2*(x**2 + y**2)]
    ])
    
    return R

def compute_translation(moving_points, fixed_points, rotation_matrix):
    """
    Computes the translation vector to align moving_points to fixed_points given a rotation matrix.
    
    Parameters:
    moving_points (list or ndarray): List of points to be translated.
    fixed_points (list or ndarray): List of reference points.
    rotation_matrix (ndarray): Rotation matrix.

    Returns:
    ndarray: Translation vector.
    """
    # Convert to numpy arrays
    moving = np.array(moving_points)
    fixed = np.array(fixed_points)

    # Compute centroids
    centroid_moving = np.mean(moving, axis=0)
    centroid_fixed = np.mean(fixed, axis=0)

    # Compute translation
    translation = centroid_fixed - np.dot(centroid_moving, rotation_matrix)

    return translation

# def apply_transform(points, rotation_matrix, translation_vector):
#     points = np.array(points)
#     transformed_points = np.dot(points, rotation_matrix.T) + translation_vector
#     return transformed_points

def apply_transform(volume, rotation_matrix, translation_vector):
    """
    Transforms a 3D volume using a given rotation matrix and translation vector.

    Parameters:
    volume (sitk.Image): Input 3D volume to be transformed.
    rotation_matrix (ndarray): 3x3 rotation matrix.
    translation_vector (ndarray): 1x3 translation vector.

    Returns:
    sitk.Image: Transformed 3D volume.
    """
    # Create an affine transform
    transform = sitk.AffineTransform(3)
    transform.SetMatrix(rotation_matrix.flatten())
    translation_vector = translation_vector.tolist()
    transform.SetTranslation(translation_vector)

    # Resample the volume
    resampler = sitk.ResampleImageFilter()
    resampler.SetReferenceImage(volume)
    resampler.SetInterpolator(sitk.sitkLinear)  # Linear interpolation
    resampler.SetTransform(transform)
    resampler.SetDefaultPixelValue(0)  # Background value for areas outside the original volume

    transformed_volume = resampler.Execute(volume)
    return transformed_volume



# Initialize lists and dictionary
cbct_list = []
ct_list = []
volume_points_dict = {}

# Process loaded volumes
for volumeNode in slicer.util.getNodesByClass("vtkMRMLScalarVolumeNode"):
    volumeName = volumeNode.GetName()
    shNode = slicer.vtkMRMLSubjectHierarchyNode.GetSubjectHierarchyNode(slicer.mrmlScene)
    imageItem = shNode.GetItemByDataNode(volumeNode)
    
    modality = shNode.GetItemAttribute(imageItem, 'DICOM.Modality')
    #print(modality)
    
    # Check if the volume is loaded into the scene
    if not slicer.mrmlScene.IsNodePresent(volumeNode):
        print(f"Volume {volumeName} not present in the scene.")
        continue
    
    # Determine scan type
    if "cbct" in volumeName.lower():
        cbct_list.append(volumeName)
        scan_type = "CBCT"
    else:
        ct_list.append(volumeName)
        scan_type = "CT"
    
    # Detect points using region growing
    grouped_points = detect_points_region_growing(volumeName, intensity_threshold=3000)
    volume_points_dict[(scan_type, volumeName)] = grouped_points

# Print the results
# print(f"\nCBCT Volumes: {cbct_list}")
# print(f"CT Volumes: {ct_list}")
# print("\nDetected Points by Volume:")
# for (scan_type, vol_name), points in volume_points_dict.items():
#     print(f"{scan_type} Volume '{vol_name}': {len(points)} points detected.")


if cbct_list and ct_list:
    # Izberi prvi CT volumen kot referenco
    ct_volume_name = ct_list[0]
    ct_points = [centroid for centroid, _ in volume_points_dict[("CT", ct_volume_name)]]

    if len(ct_points) < 3:
        print("CT volumen nima dovolj točk za registracijo.")
    else:
        print("CT points: ", np.array(ct_points))
        
        for cbct_volume_name in cbct_list:
            # Izvleci točke za trenutni CBCT volumen
            cbct_points = [centroid for centroid, _ in volume_points_dict[("CBCT", cbct_volume_name)]]

            print(f"\nProcessing CBCT Volume: {cbct_volume_name}")
            if len(cbct_points) < 3:
                print(f"CBCT Volume '{cbct_volume_name}' nima dovolj točk za registracijo.")
                continue

            #print("CBCT points: ", np.array(cbct_points))

            # Compute rotation matrices using different methods
            svd_rotation_matrix = compute_Kabsch_rotation(cbct_points, ct_points)
            horn_rotation_matrix = compute_Horn_rotation(cbct_points, ct_points)
            quaternion_rotation_matrix = compute_quaternion_rotation(cbct_points, ct_points)

            # Compute translations for each method
            svd_translation_vector = compute_translation(cbct_points, ct_points, svd_rotation_matrix)
            horn_translation_vector = compute_translation(cbct_points, ct_points, horn_rotation_matrix)
            quaternion_translation_vector = compute_translation(cbct_points, ct_points, quaternion_rotation_matrix)

            # Display the results for the current CBCT volume
            print("\nSVD Method:")
            print("Rotation Matrix:\n", svd_rotation_matrix)
            print("Translation Vector:\n", svd_translation_vector)

            print("\nHorn Method:")
            print("Rotation Matrix:\n", horn_rotation_matrix)
            print("Translation Vector:\n", horn_translation_vector)

            print("\nQuaternion Method:")
            print("Rotation Matrix:\n", quaternion_rotation_matrix)
            print("Translation Vector:\n", quaternion_translation_vector)

            # Apply chosen transformation (select one method manually)
            chosen_rotation_matrix = svd_rotation_matrix  # Change to horn_rotation_matrix or quaternion_rotation_matrix if needed
            chosen_translation_vector = svd_translation_vector #also change here

            # Pridobitev CBCT slike kot SimpleITK
            cbct_volume_node = slicer.util.getNode(cbct_volume_name)
            cbct_image_sitk = sitkUtils.PullVolumeFromSlicer(cbct_volume_node)

            transformed_cbct_image = apply_transform(cbct_image_sitk, chosen_rotation_matrix, chosen_translation_vector)

            print(f"\nTransformed Volume '{cbct_volume_name}' using SVD Method")

            # Shranimo nazaj v Slicer
            sitkUtils.PushVolumeToSlicer(transformed_cbct_image, name=f"Transformed_{cbct_volume_name}")
else:
    print("CBCT ali CT volumen ni bil najden.")


# def compute_rigid_transform(moving_points, fixed_points):
#     assert len(moving_points) == len(fixed_points), "Point lists must be the same length."

#     # Convert to numpy arrays
#     moving = np.array(moving_points)
#     fixed = np.array(fixed_points)

#     # Compute centroids
#     centroid_moving = np.mean(moving, axis=0)
#     centroid_fixed = np.mean(fixed, axis=0)

#     # Center the points
#     moving_centered = moving - centroid_moving
#     fixed_centered = fixed - centroid_fixed

#     # Compute covariance matrix
#     H = np.dot(moving_centered.T, fixed_centered)

#     # SVD decomposition
#     U, _, Vt = np.linalg.svd(H)
#     Rotate_optimal = np.dot(Vt.T, U.T)

#     # Correct improper rotation (reflection)
#     if np.linalg.det(Rotate_optimal) < 0:
#         Vt[-1, :] *= -1
#         Rotate_optimal = np.dot(Vt.T, U.T)

#     # Compute translation
#     translation = centroid_fixed - np.dot(centroid_moving, Rotate_optimal)

#     return Rotate_optimal, translation