--- a +++ b/Preprocessing Medical Data Pipeline/MSKMulticlass.py @@ -0,0 +1,123 @@ +import numpy as np +import nibabel as nib +import matplotlib.pyplot as plt +import os +from image_preprocessing import uniform_resizing + +def CreateMasks4MulticlassMSK(input_scan_dir, scan_dir, directory, mask_index, TIBIA_encoding, FEMUR_encoding, FIBULA_encoding, PELVIS_encoding, output_dir, AOIThresholding, FriedLanderDataset): + aoi_fnames = os.listdir(directory) + resizing_size = 256 + print('Regions of Interest for Segmentation: ', aoi_fnames) + + fnames = [] + for aoi_fname in aoi_fnames: + directory_aoi = (directory + ('/{}').format(aoi_fname)) + aoi_fname = os.listdir(directory_aoi) + aoi_fname = sorted(aoi_fname) + fnames.append(aoi_fname) + + n_segmentation_classes = len(fnames) + print('Number of Segmentation Classes: ', n_segmentation_classes) + + + + suffix_index_dict = {} + for idx, file_list in enumerate(fnames): + for filename in file_list: + # Extract the suffix index from the filename + parts = filename.split('_') + if len(parts) > 1: + suffix_index = parts[1].split('.')[0] # Extracting the index before '.' + suffix_index = int(suffix_index) # Convert index to integer + # Add the filename to the corresponding list in the dictionary + if suffix_index in suffix_index_dict: + suffix_index_dict[suffix_index].append(filename) + else: + suffix_index_dict[suffix_index] = [filename] + + # for index, filenames in suffix_index_dict.items(): + # print(f"Segmentation Mask {index}: {filenames}") + + combined_mask = None + for fname_mask in suffix_index_dict[mask_index]: + print(fname_mask) + mask = nib.load(('{}/{}/{}').format(directory, (fname_mask.split('_'))[0], fname_mask)) + mask_data = mask.get_fdata().astype(np.uint8) + aoi = (fname_mask.split('_'))[0] + + if (FriedLanderDataset == True): + mask_data = np.expand_dims(mask_data, axis = -1) + mask_data = mask_data.transpose(1, 0, 2, 3) + print(('Unprocessed {} Mask Shape: {}').format(aoi, mask_data.shape)) + mask_data = uniform_resizing(mask_data, resizing_size) + + + if (aoi == 'TIBIA'): + mask_data = np.where(mask_data != 0, TIBIA_encoding, 0) + if (aoi == 'FEMUR'): + mask_data = np.where(mask_data != 0, FEMUR_encoding, 0) + if (aoi == 'FIBULA'): + mask_data = np.where(mask_data != 0, FIBULA_encoding, 0) + if (aoi == 'PELVIS'): + mask_data = np.where(mask_data != 0, PELVIS_encoding, 0) + if combined_mask is None: + combined_mask = mask_data + else: + combined_mask += mask_data # Combine the pixel arrays by adding them element-wise + combined_mask = np.array(combined_mask) + combined_mask[combined_mask > n_segmentation_classes] = 0 + + if (AOIThresholding == True): + threshold = 0 + + combined_mask = np.array(combined_mask) + indices = np.transpose(np.nonzero(combined_mask != 0)) + if indices.size > 0: + first_non_zero_index_2d = tuple(indices[0]) + else: + first_non_zero_index_2d = None + + if indices.size > 0: + last_non_zero_index_2d = tuple(indices[-1]) + else: + last_non_zero_index_2d = None + + + print(first_non_zero_index_2d) + print(last_non_zero_index_2d) + first_slice_aoi = (first_non_zero_index_2d[0]) - int(threshold) + last_slice_aoi = (last_non_zero_index_2d[0]) + int(threshold) + print("AOI Slice Start (Final with Thresholding): ", first_slice_aoi) + print("AOI Slice End (Final with Thresholding): ", last_slice_aoi) + + mask_index = '{:03d}'.format(mask_index) + combined_mask = combined_mask[first_slice_aoi:last_slice_aoi, :, :, :] + scan = nib.load(('{}/msk_{}.nii.gz').format(input_scan_dir, mask_index)) + scan_data = scan.get_fdata() + # scan_og = nib.Nifti1Image(scan_data, scan.affine) + # nib.save(scan_og, ('{}/msk_{}_og.nii.gz').format(output_dir, mask_index)) + + if (FriedLanderDataset == True): + scan_data = np.expand_dims(scan_data, axis = -1) + scan_data = scan_data.transpose(1, 0, 2, 3) + print('Unprocessed Scan Shape: ', scan_data.shape) + scan_data = uniform_resizing(scan_data, resizing_size) + scan_data = scan_data.astype('float32') + scan_data /= np.max(scan_data) # scale scans to [0, 1] + print('Max Pixel Value in Scan: ', np.max(scan_data)) + + scan_data = scan_data[first_slice_aoi:last_slice_aoi, :, :, :] + print('Final Training Scans Input Shape: ', scan_data.shape) + print('Final Training Masks Input Shape: ', combined_mask.shape) + + # np.savetxt('output.txt', scan_data[100,:,:,0], fmt="%d", delimiter=",") + combined_mask = combined_mask.astype(int) + combined_img = nib.Nifti1Image(combined_mask, mask.affine) + print('Multi-class Labels: ', np.unique(combined_mask)) + nib.save(combined_img, ('{}/msk_{}.nii.gz').format(output_dir, mask_index)) + + combined_img = nib.Nifti1Image(scan_data, scan.affine) + nib.save(combined_img, ('{}/msk_{}.nii.gz').format(scan_dir, mask_index)) + + print(('MSK Multiclass Mask Made Using {}, Saved To {}, and Region of Interest Slice Thresholding = {} !').format(aoi_fnames, ('{}/msk_{}.nii.gz').format(output_dir, mask_index), AOIThresholding)) + print('\n') \ No newline at end of file