import subprocess
required_libraries = [
"numpy", "os", "pydicom", "SimpleITK", "pandas", "trimesh",
"pyntcloud", "scikit-image", "matplotlib", "tqdm", "nibabel",
"re", "scipy", "opencv-python", 'segmentation_models', 'random',
'trimesh'
]
for lib in required_libraries:
try:
__import__(lib)
except ImportError:
print(f"{lib} is not installed. Installing...")
subprocess.check_call(["pip", "install", lib])
print(f"{lib} has been installed.")
from preprocessing import preprocessing, VisualValidationMSK, transform_string
from image_preprocessing import image_cropping, uniform_cropping, uniform_resizing
from writeout_dataset import Export2CompressedNifiti
from MSKMulticlass import CreateMasks4MulticlassMSK
from createDirectories import createDirectoriesfunc
from preparingTestInstance import preprocessTestScans
from data_augmentation import DataAugmentation
import os
# import tensorflow
# from tensorflow import keras
# from keras.models import load_model
# from keras.utils import to_categorical
import nibabel as nib
import numpy as np
# import segmentation_models as sm
# from keras.metrics import MeanIoU
from skimage.measure import marching_cubes
import matplotlib.pyplot as plt
import trimesh
import SimpleITK as sitk
import nibabel as nib
import numpy as np
createDirectoriesfunc()
def Preprocessing(scans_path, scan_data_folders, Cropping):
print('-'*30)
print('Loading and preprocessing training data...')
print('-'*30)
# /research\resmed202100086-tws ----> Address for raw data
# /research\resabi202200010-Friedlander
total_slices_raw_data = 0
DataOnlyAOI = False
ExportDatasets = True
multiclass_mask_output_dir = None
# Only set multiclassSegmentation to True once all the AOI masks are preprocessed or its onto the final AOI
if (Cropping == "256x256"):
Cropping = True
else:
Cropping = False
scan_data_folders = [scan_data_folders]
# Creates the preprocessed scan and mask data and exports to Data folder
# Formats binary masks and applies preprocessing steps
segmasks = []
paitent_id = (scan_data_folders[0].split('_'))[-1]
files = os.listdir(('{}/Raw NIFITI Segmentation Masks (3D Slicer Output)').format(scans_path))
for file in files:
file_id = file.split('_')[-1]
file_id = file_id.split('.')[0]
if (file_id == paitent_id):
segmasks.append(file)
print(segmasks)
for segmask in segmasks:
print(segmask)
imgs_train, imgs_mask_train, median_aoi_index = preprocessing(scans_path, segmask, scan_data_folders, total_slices_raw_data, DataOnlyAOI, Cropping)
orientation = (segmask.split('_'))[1]
colab_fname = [transform_string(segmask)]
Export2CompressedNifiti(imgs_train, scans_path, colab_fname, imgs_mask_train, orientation)
# User Input
# Multi-class mask creation
individual_mask_directory = ('{}/nnUNet Data/masks').format(scans_path)
multiclass_mask_output_dir = ('{}/nnUNet Data/multiclass_masks').format(scans_path)
scan_dir = ('{}/nnUNet Data/scans').format(scans_path)
input_scan_dir = ('{}/nnUNet Data/unprocessed_scans').format(scans_path)
TIBIA_encoding = 1
FEMUR_encoding = 2
FIBULA_encoding = 3
PELVIS_encoding = 4
mask_index = int((scan_data_folders[0].split('_'))[0])
AOIThresholding = True
FriedLanderDataset = False
print('Multi-Class Segmentation Task Data Preparation!')
CreateMasks4MulticlassMSK(input_scan_dir, scan_dir, individual_mask_directory, mask_index, TIBIA_encoding, FEMUR_encoding, FIBULA_encoding, PELVIS_encoding, multiclass_mask_output_dir, AOIThresholding, FriedLanderDataset)
print('-'*30)
print('Completed Preprocessing Stage!')
print('-'*30)
def preprocessTestScansMain(cutoffSlice, seg_scan_dir, folders):
preprocessTestScans(cutoffSlice, seg_scan_dir, [folders])
def VisualiseSlice(basedir, colab_fname, slice_idx_bin):
import nibabel as nib
import numpy as np
import matplotlib.pyplot as plt
def superimpose_images(image1, image2):
image1 = image1 / np.max(image1)
image2 = image2 / np.max(image2)
alpha = 0.5
superimposed_image = alpha * image1 + (1 - alpha) * image2
return superimposed_image
if ((colab_fname.split('_'))[0] != 'msk'):
orientation = (colab_fname.split('_'))[1]
colab_fname = transform_string(colab_fname)
mask_index = '{:03d}'.format(int(((colab_fname).split('_'))[1]))
nii_img_scan = nib.load(('{}/nnUNet Data/scans/msk_{}.nii.gz').format(basedir, mask_index))
nii_img_mask = nib.load(('{}/nnUNet Data/masks/{}/{}_{}.nii.gz').format(basedir, ((colab_fname).split('_'))[0],colab_fname, orientation))
if ((colab_fname.split('_'))[0] == 'msk'):
nii_img_scan = nib.load(('{}/nnUNet Data/scans/{}.nii.gz').format(basedir, colab_fname))
nii_img_mask = nib.load(('{}/nnUNet Data/multiclass_masks/{}.nii.gz').format(basedir, colab_fname))
orientation = '_'
mask_data = nii_img_mask.get_fdata()
scan_data = nii_img_scan.get_fdata()
image1 = scan_data[int(slice_idx_bin), :, :, 0]
image2 = mask_data[int(slice_idx_bin), :, :, 0]
superimposed_image = superimpose_images(image1, image2)
plt.imshow(superimposed_image, cmap='gray')
plt.title(('Validating Scan & Mask on Slice {} of {}_{}').format(slice_idx_bin, colab_fname, orientation))
plt.axis('off')
plt.show()
plt.imshow(image2, cmap='gray')
plt.title(('Validating Mask on Slice {} of {}_{}').format(slice_idx_bin, colab_fname, orientation))
plt.axis('off')
plt.show()
plt.imshow(image1, cmap='gray')
plt.title(('Validating Scan on Slice {} of {}_{}').format(slice_idx_bin, colab_fname, orientation))
plt.axis('off')
plt.show()
print('-'*30)
print('Visulisations Generated!')
print('-'*30)
def DataAug(basedir, aug_fname, num_augs):
imgs_train = nib.load(('{}/nnUNet Data/scans/{}.nii.gz').format(basedir, aug_fname))
imgs_mask_train = nib.load(('{}/nnUNet Data/multiclass_masks/{}.nii.gz').format(basedir, aug_fname))
imgs_train = imgs_train.get_fdata()
imgs_mask_train = imgs_mask_train.get_fdata()
num_train = len(imgs_train)
num_augs = int(num_augs)
augmented_images_train, augmented_masks_train = DataAugmentation(imgs_train, imgs_mask_train, num_augs)
augmented_images_train_sorted, augmented_masks_train_sorted = [], []
for i in range (int(num_augs)):
augmented_images_train_temp = augmented_images_train[i::num_augs]
augmented_masks_train_temp = augmented_masks_train[i::num_augs]
augmented_images_train_sorted.append(augmented_images_train_temp)
augmented_masks_train_sorted.append(augmented_masks_train_temp)
augmented_images_train_sorted = np.array(augmented_images_train_sorted)
augmented_masks_train_sorted = np.array(augmented_masks_train_sorted)
print('Augmented Training Scan Shape: ', augmented_images_train_sorted.shape)
print('Augmented Training Mask Shape: ',augmented_masks_train_sorted.shape)
for i in range (len(augmented_images_train_sorted)):
temp = np.expand_dims(augmented_images_train_sorted[i], axis=-1)
temp = temp.astype('float32')
np.savetxt('D:\MRI - Tairawhiti (User POV)/train.txt', temp[100,:,:,0], fmt="%d", delimiter=",")
# temp /= 255. # scale scans to [0, 1]
nii_img_train = nib.Nifti1Image(temp, affine=np.eye(4))
output_file_path = ('{}/nnUNet Data/scans/{}_aug{}.nii.gz').format(basedir, aug_fname, i)
nib.save(nii_img_train, output_file_path)
temp = np.expand_dims(augmented_masks_train_sorted[i], axis=-1)
temp = temp.astype(int)
np.savetxt('D:\MRI - Tairawhiti (User POV)/train_mask.txt', temp[100,:,:,0], fmt="%d", delimiter=",")
nii_img_mask = nib.Nifti1Image(temp, affine=np.eye(4))
output_file_path = ('{}/nnUNet Data/multiclass_masks/{}_aug{}.nii.gz').format(basedir, aug_fname, i)
nib.save(nii_img_mask, output_file_path)
print('-'*30)
print('Data Augmentation Completed & Exported!')
print('-'*30)
def AutoSegModel(subjectfname, base_dir):
model_dir = ('{}/Pre-Trained Models (Google Colab)/res34_backbone_20_epochs_dicefocal_256_4P_12batch_maxF1_pelvis_aug_best.hdf5').format(base_dir)
model = load_model(model_dir, compile=False)
n_classes = 5
BACKBONE = 'resnet34'
model_used = 'U-Net(resnet34)'
img_dir = ('{}/nnUNet Data/scans/{}.nii.gz').format(base_dir, subjectfname)
img = nib.load(img_dir)
img_data = img.get_fdata()
X_test = np.repeat(img_data, 3, axis=3)
print(('Fined-Tuned Model: {}').format((model_dir.split('/'))[-1]))
print('Prediction Scan Stack: ', subjectfname)
print('Number of Segmentation Classes: ', n_classes)
print('\n')
print("Test Images Shape: ", X_test.shape)
print('\n')
preprocess_input = sm.get_preprocessing(BACKBONE)
X_test_processed = preprocess_input(X_test)
# Prediction
y_pred=model.predict(X_test_processed)
y_pred_argmax=np.argmax(y_pred, axis=3)
y_pred_argmax = np.expand_dims(y_pred_argmax, axis = -1)
print('Pred Mask Shape: ', y_pred_argmax.shape)
print("Pred Mask Labels: ", np.unique(y_pred_argmax))
print('\n')
combined_mask = y_pred_argmax.astype(np.int32)
combined_img = nib.Nifti1Image(combined_mask, affine=np.eye(4), dtype=np.int32)
nib.save(combined_img, ("{}/{}_pred.nii.gz").format(model_dir, subjectfname), dtype = np.uint8)
print('Exported Prediction Segmentation: ', (("{}/{}_pred.nii.gz").format(model_dir, subjectfname)))
print('\n')
def Export3DStructure(segmentation_data, predfname, class_msk, model_used):
# Generate a surface mesh using marching cubes
vertices, faces, normals, _ = marching_cubes(segmentation_data, level=0)
# Create a Trimesh object
mesh = trimesh.Trimesh(vertices=vertices, faces=faces, vertex_normals=normals)
# Save the mesh as a PLY file
ply_path = ('{}/{}_{}_{}.ply').format(model_dir, predfname, class_msk, model_used)
mesh.export(ply_path)
segmentation_data_all = y_pred_argmax
segmentation_data_all = segmentation_data_all[:,:,:,0]
tibia_seg_data = np.where(segmentation_data_all != 1, 0, 1)
femur_seg_data = np.where(segmentation_data_all != 2, 0, 2)
fibula_seg_data = np.where(segmentation_data_all != 3, 0, 3)
pelvis_seg_data = np.where(segmentation_data_all != 4, 0, 4)
segmentation_data = [segmentation_data_all, tibia_seg_data, femur_seg_data, fibula_seg_data, pelvis_seg_data]
class_msk = ['ALL', 'TIBIA', 'FEMUR', 'FIBULA', 'PELVIS']
for i in range (len(segmentation_data)):
Export3DStructure(segmentation_data[i], subjectfname, class_msk[i], model_used)
Datasets
Models
