aaa_web_app / Comp2Comp-main /comp2comp /visualization /linear_planar_reformation.py
AdritRao's picture
Upload 62 files
a3290d1
raw
history blame
3.45 kB
"""
@author: louisblankemeier
"""
import numpy as np
def linear_planar_reformation(
medical_volume: np.ndarray, segmentation: np.ndarray, centroids, dimension="axial"
):
if dimension == "sagittal" or dimension == "coronal":
centroids = sorted(centroids, key=lambda x: x[2])
elif dimension == "axial":
centroids = sorted(centroids, key=lambda x: x[0])
centroids = [(int(x[0]), int(x[1]), int(x[2])) for x in centroids]
sagittal_centroids = [centroids[i][0] for i in range(0, len(centroids))]
coronal_centroids = [centroids[i][1] for i in range(0, len(centroids))]
axial_centroids = [centroids[i][2] for i in range(0, len(centroids))]
sagittal_vals, coronal_vals, axial_vals = [], [], []
if dimension == "sagittal":
sagittal_vals = [sagittal_centroids[0]] * axial_centroids[0]
if dimension == "coronal":
coronal_vals = [coronal_centroids[0]] * axial_centroids[0]
if dimension == "axial":
axial_vals = [axial_centroids[0]] * sagittal_centroids[0]
for i in range(1, len(axial_centroids)):
if dimension == "sagittal" or dimension == "coronal":
num = axial_centroids[i] - axial_centroids[i - 1]
elif dimension == "axial":
num = sagittal_centroids[i] - sagittal_centroids[i - 1]
if dimension == "sagittal":
interp = list(
np.linspace(sagittal_centroids[i - 1], sagittal_centroids[i], num=num)
)
sagittal_vals.extend(interp)
if dimension == "coronal":
interp = list(
np.linspace(coronal_centroids[i - 1], coronal_centroids[i], num=num)
)
coronal_vals.extend(interp)
if dimension == "axial":
interp = list(
np.linspace(axial_centroids[i - 1], axial_centroids[i], num=num)
)
axial_vals.extend(interp)
if dimension == "sagittal":
sagittal_vals.extend(
[sagittal_centroids[-1]] * (medical_volume.shape[2] - len(sagittal_vals))
)
sagittal_vals = np.array(sagittal_vals)
sagittal_vals = sagittal_vals.astype(int)
if dimension == "coronal":
coronal_vals.extend(
[coronal_centroids[-1]] * (medical_volume.shape[2] - len(coronal_vals))
)
coronal_vals = np.array(coronal_vals)
coronal_vals = coronal_vals.astype(int)
if dimension == "axial":
axial_vals.extend(
[axial_centroids[-1]] * (medical_volume.shape[0] - len(axial_vals))
)
axial_vals = np.array(axial_vals)
axial_vals = axial_vals.astype(int)
if dimension == "sagittal":
sagittal_image = medical_volume[sagittal_vals, :, range(len(sagittal_vals))]
sagittal_label = segmentation[sagittal_vals, :, range(len(sagittal_vals))]
if dimension == "coronal":
coronal_image = medical_volume[:, coronal_vals, range(len(coronal_vals))]
coronal_label = segmentation[:, coronal_vals, range(len(coronal_vals))]
if dimension == "axial":
axial_image = medical_volume[range(len(axial_vals)), :, axial_vals]
axial_label = segmentation[range(len(axial_vals)), :, axial_vals]
if dimension == "sagittal":
return sagittal_image, sagittal_label
if dimension == "coronal":
return coronal_image, coronal_label
if dimension == "axial":
return axial_image, axial_label