Diff of /patient_loader.py [000000] .. [4f54f1]

Switch to side-by-side view

--- a
+++ b/patient_loader.py
@@ -0,0 +1,110 @@
+import os
+import cv2
+import numpy as np
+from random import randrange
+
+import utils
+import config
+import lung_segmentation as ls
+
+
+segmentation_algo = ls.get_segmentation_algorithm()
+
+
+class PatientImageLoader(object):
+    def __init__(self, images_dir):
+        self._images_input = images_dir or config.SEGMENTED_LUNGS_DIR
+        self._augment = False
+
+    def load_scans(self, patient):
+        if 'augm' in patient:
+            self._augment = True
+            patient = patient.split('-')[0]
+        return utils.load_patient_image(self._images_input, patient)
+
+    @property
+    def images_input(self):
+        return self._images_input
+
+    @property
+    def name(self):
+        return 'base_image_loader'
+
+
+# Tests with the mean scans loader are not using
+# lung segmentation, only compressed sorted slices in HU units.
+class MeanScansLoader(PatientImageLoader):
+    def __init__(self, images_dir=None):
+        super(MeanScansLoader, self).__init__(images_dir)
+
+    def load_scans(self, patient):
+        image = utils.load_patient_image(self._images_input, patient)
+        image = utils.resize(image)
+
+        return utils.get_mean_chunk_slices(image)
+
+    @property
+    def name(self):
+        return 'mean_scans_loader'
+
+
+class SegmentedGaussianLungsLoader(PatientImageLoader):
+    def __init__(self, images_dir=config.SEGMENTED_LUNGS_DIR):
+        super(SegmentedGaussianLungsLoader, self).__init__(images_dir)
+
+    def process_scans(self, image):
+        image = np.stack([cv2.GaussianBlur(scan, (5, 5), 0) for scan in image])
+        image = utils.resize(image)
+
+        return utils.trim_pad_slices(image, pad_with_existing=False)
+
+    def load_scans(self, patient):
+        image = utils.load_patient_image(self._images_input, patient)
+        return self.process_scans(image)
+
+    @property
+    def name(self):
+        return 'segmented_gaussian_lungs_loader'
+
+
+# Default loader
+class SegmentedLungsScansLoader(PatientImageLoader):
+    def __init__(self, images_dir=config.SEGMENTED_LUNGS_DIR):
+        super(SegmentedLungsScansLoader, self).__init__(images_dir)
+
+    def process_scans(self, image):
+        image = segmentation_algo.get_slices_with_nodules(image)
+        image = utils.resize(image)
+
+        if self._augment:
+            angle = randrange(-15, 15)
+            image = utils.rotate_scans(image, angle)
+
+        return utils.trim_pad_slices(image, pad_with_existing=True)
+
+    def load_scans(self, patient):
+        image = super(SegmentedLungsScansLoader, self).load_scans(patient)
+        return self.process_scans(image)
+
+    @property
+    def name(self):
+        return 'segmented_lungs_loader'
+
+
+class NodulesScansLoader(PatientImageLoader):
+    def __init__(self, images_dir=config.SEGMENTED_LUNGS_DIR):
+        super(NodulesScansLoader, self).__init__(images_dir)
+
+    def process_scans(self, patient):
+        image = utils.load_patient_image(self._images_input, patient)
+        nodules = segmentation_algo.get_lung_nodules_candidates(image)
+        nodules = utils.resize(nodules)
+
+        return utils.trim_pad_slices(nodules)
+
+    def load_scans(self, patient):
+        return self.process_scans(patient)
+
+    @property
+    def name(self):
+        return 'nodules_scans_loader'
\ No newline at end of file