add new dependency for reading segmentation masks, add segmentation i…

…mage loading

environment-CPU.yml

@@ -15,3 +15,4 @@ dependencies:
         - kaggle
         - pydicom
         - opencv-python
+        - nibabel

environment-CUDA.yml

@@ -16,3 +16,4 @@ dependencies:
         - kaggle
         - pydicom
         - opencv-python
+        - nibabel

environment-WIN.yml

@@ -15,3 +15,4 @@ dependencies:
         - kaggle
         - pydicom
         - opencv-python
+        - nibabel

kaggleDataLoader.py

@@ -3,6 +3,7 @@
 
 import pandas as pd
 import pydicom
+import nibabel as nib
 import numpy as np
 import cv2
 import os
@@ -11,7 +12,7 @@
 TRAIN_IMAGES_PATH = f'{RSNA_2022_PATH}/train_images'
 TEST_IMAGES_PATH  = f'{RSNA_2022_PATH}/test_images'
 
-class kaggleDataLoader:
+class KaggleDataLoader:
 
     def __init__(self):
         # Load metadata
@@ -26,13 +27,13 @@ def __init__(self):
         # https://www.kaggle.com/competitions/rsna-2022-cervical-spine-fracture-detection/discussion/344862
         bad_scans           = ['1.2.826.0.1.3680043.20574', '1.2.826.0.1.3680043.29952']
         for uid in bad_scans:
-            self.train_df.drop(self.train_df[self.train_df['StudyInstanceUID'] == uid].index, axis = 0, inplace = True)
+            self.trainDf.drop(self.trainDf[self.trainDf['StudyInstanceUID'] == uid].index, axis = 0, inplace = True)
 
-    def BBoxFromIndex(self, id, sliceNum):
+    def bboxFromIndex(self, id, sliceNum):
         box = self.trainBbox.loc[(self.trainBbox.StudyInstanceUID == id & self.trainBbox.slice_number == sliceNum), :]
         return list(box.values[0])
 
-    def get_fractured_bones(self, patientID):
+    def fracturedBones(self, patientID):
         fractured_bones = []
         temp = self.trainDf.loc[self.trainDf.StudyInstanceUID == patientID, ['C1', 'C2', 'C3', 'C4', 'C5', 'C6', 'C7']]
         temp = list(temp.values[0])  # there is one row per id
@@ -45,7 +46,7 @@ def get_fractured_bones(self, patientID):
     image loading function for rsna competition images, stored as 'dicom' images
     """
 
-    def load_dicom(self, path):
+    def loadDicom(self, path):
         img = pydicom.dcmread(path)
         img.PhotometricInterpretation = 'YBR_FULL'
         data = img.pixel_array
@@ -55,11 +56,17 @@ def load_dicom(self, path):
         data = (data * 255).astype(np.uint8)
         return cv2.cvtColor(data, cv2.COLOR_GRAY2RGB), img
 
-    def load_slice_from_id(self, patientID, sliceIndex):
+    def loadSliceFromId(self, patientID, sliceIndex):
         targetPath = os.path.join(self.trainImagePath, patientID, str(sliceIndex), ".dcm")
-        return self.load_dicom(targetPath)
+        return self.loadDicom(targetPath)
 
-    def load_dataset(self):
+    def loadSegmentationsForPatient(self, patientID):
+        segmentations = nib.load(os.path.join(self.segPath, patientID, '.nii')).get_fdata()
+        segmentations = segmentations[:, ::-1, ::-1]
+        segmentations = segmentations.transpose(2, 1, 0)
+        return segmentations
+
+    def loadDataset(self):
         """
         prepare full dataset for training
         """