Switch to side-by-side view

--- 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