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