refactor get_dicoms_path and add get_assets_dir

drivendataorg · Oct 4, 2017 · 28660bf · 28660bf
1 parent 59c029f
commit 28660bf
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 11 deletions.
diff --git a/prediction/src/algorithms/segment/src/data_generation.py b/prediction/src/algorithms/segment/src/data_generation.py
@@ -14,9 +14,16 @@ def get_dicom_paths(in_docker=True):
         in_docker: whether this method is invoked from within docker or from the prediction directory
     """
     if in_docker:
-        return glob.glob(os.path.join('..', 'images_full', 'LIDC-IDRI-*', '**', '**'))
+        paths = glob.glob(os.path.join('..', 'images_full', 'LIDC-IDRI-*', '**', '**'))
     else:
-        return glob.glob(os.path.join('..', 'tests', 'assets', 'test_image_data', 'full', 'LIDC-IDRI-*', '**', '**'))
+        paths = glob.glob(os.path.join('..', 'tests', 'assets', 'test_image_data', 'full', 'LIDC-IDRI-*', '**', '**'))
+    return paths
+
+
+def get_assets_dir():
+    current_dir = os.path.dirname(os.path.realpath(__file__))
+    parent_dir = os.path.dirname(current_dir)
+    return os.path.abspath(os.path.join(parent_dir, 'assets'))
 
 
 def prepare_training_data(in_docker=True):
@@ -27,13 +34,11 @@ def prepare_training_data(in_docker=True):
         in_docker: whether this method is invoked from within docker or from the prediction directory
     """
     INTERMEDIATE_MALICIOUS = 3
-
-    current_dir = os.path.dirname(os.path.realpath(__file__))
-    assets_dir = os.path.abspath(os.path.join(current_dir, '..', 'assets'))
+    assets_dir = get_assets_dir()
 
     dicom_paths = sorted(get_dicom_paths(in_docker=in_docker))
     for path in dicom_paths:
-        directories = path.split('/')
+        directories = path.split(os.path.sep)
         lidc_id_path_index = 2 if in_docker else 5
         lidc_id = directories[lidc_id_path_index]
         lung_patient_file = os.path.join(assets_dir, "segmented_lung_patient_{}".format(lidc_id))

diff --git a/prediction/src/algorithms/segment/src/training.py b/prediction/src/algorithms/segment/src/training.py
@@ -8,7 +8,7 @@
 from keras.models import load_model
 from src.preprocess.lung_segmentation import save_lung_segments
 
-from .data_generation import get_dicom_paths
+from .data_generation import get_dicom_paths, get_assets_dir
 from .model import simple_model_3d
 
 
@@ -52,8 +52,7 @@ def train(load_checkpoint=False):
     CUBOID_IMAGE_SHAPE = get_data_shape()
     CUBOID_BATCH = 1  # How many training pairs should be passed to model.fit in one batch
 
-    current_dir = os.path.dirname(os.path.realpath(__file__))
-    assets_dir = os.path.abspath(os.path.join(current_dir, '../assets'))
+    assets_dir = get_assets_dir()
     dicom_paths = get_dicom_paths()
 
     labels = glob.glob(os.path.join(assets_dir, "segmented_lung_patient_*.npy"))
@@ -70,8 +69,9 @@ def train(load_checkpoint=False):
     for batch_index in range(0, len(labels), CUBOID_BATCH):
 
         for index, path in enumerate(dicom_paths[batch_index:batch_index + CUBOID_BATCH]):
-            lidc_id = path.split('/')[5]
-            patient_id = path.split('/')[-1]
+            directories = path.split(os.path.sep)
+            lidc_id = directories[5]
+            patient_id = directories[-1]
             _, input_img = save_lung_segments(path, patient_id)
             output_img = np.load(os.path.join(assets_dir, "segmented_lung_patient_{}.npy").format(lidc_id))