Moved EXTRACTED_IMAGE_DIR to prediction/config.py and updated the rel…

…evant files to reflect that change. (Closes: drivendataorg#208)
Serhiy-Shekhovtsov · Nov 14, 2017 · ca2358a · ca2358a
1 parent 052ae01
commit ca2358a
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 2 deletions.
diff --git a/prediction/config.py b/prediction/config.py
@@ -22,6 +22,8 @@ class Config(object):
     SMALL_DICOM_PATHS = path.join(PARENT_DIR, 'images')
     FULL_DICOM_PATHS_WILDCARD = path.join(FULL_DICOM_PATHS, *LIDC_WILDCARD)
     SMALL_DICOM_PATHS_WILDCARD = path.join(FULL_DICOM_PATHS, *LIDC_WILDCARD)
+    DATA_DIR = path.abspath(path.join(CURRENT_DIR, 'data'))
+    EXTRACTED_IMAGE_DIR = path.abspath(path.join(CURRENT_DIR, 'extracted'))
 
 
 class Production(Config):

diff --git a/prediction/src/algorithms/identify/prediction.py b/prediction/src/algorithms/identify/prediction.py
@@ -15,10 +15,15 @@
 
 from ...preprocess.lung_segmentation import rescale_patient_images
 
+try:
+    from ....config import Config
+except ValueError:
+    from config import Config
+
 CUBE_SIZE = 32
 CROP_SIZE = 32
 MEAN_PIXEL_VALUE = 41
-EXTRACTED_IMAGE_DIR = "data/extracted/"
+EXTRACTED_IMAGE_DIR = Config.EXTRACTED_IMAGE_DIR
 NODULE_DETECTION_DIR = "data/detections/"
 K.set_image_dim_ordering("tf")
 POS_WEIGHT = 2

diff --git a/prediction/src/preprocess/lung_segmentation.py b/prediction/src/preprocess/lung_segmentation.py
@@ -14,6 +14,11 @@
 from skimage.morphology import disk, binary_erosion, binary_closing
 from skimage.segmentation import clear_border
 
+try:
+    from ...config import Config
+except ValueError:
+    from config import Config
+
 
 def get_z_range(dicom_path):
     """Return the Z range of the images in the DICOM path
@@ -47,7 +52,7 @@ def save_lung_segments(dicom_path, patient_id):
         Original patient images (z, x, y),
         Rescaled mask images (z, x, y)
     """
-    EXTRACTED_IMAGE_DIR = "data/extracted/"
+    EXTRACTED_IMAGE_DIR = Config.EXTRACTED_IMAGE_DIR
     TARGET_VOXEL_MM = 1.00
     target_dir = os.path.join(os.getcwd(), EXTRACTED_IMAGE_DIR, patient_id)
     os.makedirs(target_dir, exist_ok=True)