update data config

README.md

@@ -25,12 +25,9 @@ For image reconstruction code, please refer to https://github.com/gift-surg/Nift
 * Demic (a tool to use NiftyNet). Tested version is v0.1: https://github.com/taigw/Demic/releases/tag/v0.1.
 
 # How to use
-* For fetal brain localization, run `bash/inference_detect.sh`, you need to edit the `PYTHONPATH` environment variable in that file so that it includes the path of NiftyNet and Demic.
+* To get fetal brain detection and segmentation results, run `bash/inference.sh`. You need to edit the `PYTHONPATH` environment variable in that file so that it includes the path of NiftyNet and Demic.
 
-
-* For fetal brain segmentation, run `bash/inference_segment.sh`, you need to edit the `PYTHONPATH` environment variable in that file so that it includes the path of NiftyNet and Demic.
-
-* You can edit `cfg_data_detect.txt` and `cfg_data_segment.txt` to customize the input and output image names.
+* You can edit `cfg_data.txt` to customize the input and output image names.
 
 # Acknowledgement
 This work is part of the GIFT-Surg project (https://www.gift-surg.ac.uk/). It is supported by Wellcome Trust [WT101957; 203145Z/16/Z], EPSRC [EP/L016478/1; NS/A000027/1; NS/A000050/1], and the NIHR UCLH BRC.

bash/inference_detect.sh → bash/inference.sh

@@ -1,3 +1,3 @@
 # edit PYTHONPATH to make sure it includes the path of NiftyNet and Demic
 export PYTHONPATH="/home/guotai/GitHub:/home/guotai/GitHub/NiftyNet-0.2.0:$PYTHONPATH"
-python test.py detect cfg_data_detect.txt
+python test.py cfg_data.txt

cfg_data.txt

@@ -0,0 +1,9 @@
+[image_1]
+input = demo_data/image1.nii.gz
+detect_output  = demo_data/image1_detect.nii.gz
+segment_output = demo_data/image1_segment.nii.gz
+
+[image_2]
+input = demo_data/image2.nii.gz
+detect_output  = demo_data/image2_detect.nii.gz
+segment_output = demo_data/image2_segment.nii.gz

test.py

@@ -17,30 +17,29 @@
 from Demic.util.parse_config import parse_config
 from Demic.util.image_process import *
 
-def model_test(net_config_file, data_config_file, detection_only = False):
-    config = parse_config(net_config_file)
+def model_test(net_config_file, data_config_file):
+    config_net  = parse_config(net_config_file)
+    config_data = parse_config(data_config_file)
     config_detect = {}
-    config_detect['network'] = config['network1']
-    config_detect['network_parameter'] = config['network1_parameter']
-    config_detect['testing'] = config['detect_testing']
+    config_detect['network'] = config_net['network1']
+    config_detect['network_parameter'] = config_net['network1_parameter']
+    config_detect['testing'] = config_net['detect_testing']
     detect_agent = TestAgent(config_detect)
     detect_agent.construct_network()
 
     print('construct network finished')
     config_segment = {}
-    config_segment['network'] = config['network2']
-    config_segment['network_parameter'] = config['network2_parameter']
-    config_segment['testing'] = config['segment_testing']
+    config_segment['network'] = config_net['network2']
+    config_segment['network_parameter'] = config_net['network2_parameter']
+    config_segment['testing'] = config_net['segment_testing']
     segment_agent = TestAgent(config_segment)
     segment_agent.construct_network()
 
     class_num = 2
-    with open(data_config_file) as f:
-        file_names = f.readlines()
-    file_names = [file_name.strip()  for file_name in file_names if file_name[0] != '#']
-    for item in file_names:
-        input_name = item.split(' ')[0]
-        output_name = item.split(' ')[1]
+    for item in config_data:
+        input_name   = config_data[item]['input']
+        detect_name  = config_data[item]['detect_output']
+        segment_name = config_data[item]['segment_output']
         print(input_name)
         # stage 1, detect
         img_dict = load_nifty_volume_as_4d_array(input_name)
@@ -49,39 +48,30 @@ def model_test(net_config_file, data_config_file, detection_only = False):
         outp = detect_agent.test_one_volume(img)
         out = np.asarray(outp > 0.5, np.uint8)
 
-        if(detection_only):
-            margin = [3, 8, 8]
-            out = get_detection_binary_bounding_box(out, margin, None, mode = 0)
-            save_array_as_nifty_volume(out, output_name, input_name)
-            continue
-
-        margin = [3, 20, 20]
-        strt = ndimage.generate_binary_structure(3,2) # iterate structure
-        post = padded_binary_closing(out, strt)
-        post = get_largest_component(post)
-        bb_min, bb_max = get_ND_bounding_box(post, margin)
+        margin = [3, 8, 8]
+        detect_out = get_detection_binary_bounding_box(out, margin, None, mode = 0)
+        save_array_as_nifty_volume(detect_out, detect_name, input_name)
 
         # stage 2, segment
+        margin = [0, 10, 10]
+        bb_min, bb_max = get_ND_bounding_box(detect_out, margin)
+
         img_roi = crop_ND_volume_with_bounding_box(img, bb_min + [0], bb_max + [0])
         outp_roi = segment_agent.test_one_volume(img_roi)
         out_roi = np.asarray(outp_roi > 0.5, np.uint8)
-        post = padded_binary_closing(out_roi, strt)
-        post = get_largest_component(post)
-        out_roi = np.asarray(post*out_roi, np.uint8)
+        strt = ndimage.generate_binary_structure(3,2)
+        out_roi = padded_binary_closing(out_roi, strt)
+        out_roi = get_largest_component(out_roi)
         out = np.zeros(img.shape[:-1], np.uint8)
         out = set_ND_volume_roi_with_bounding_box_range(out, bb_min, bb_max, out_roi)
-
-        save_array_as_nifty_volume(out, output_name, input_name)
+        save_array_as_nifty_volume(out, segment_name, input_name)
 
 if __name__ == '__main__':
-    if(len(sys.argv) != 3):
-        print('Number of arguments should be 3. e.g.')
-        print('    python test.py segment cfg_data_segment.txt')
+    if(len(sys.argv) < 2):
+        print('Number of arguments should be 2. e.g.')
+        print('    python test.py cfg_data.txt')
         exit()
     net_config_file  = 'cfg_net.txt'
-    task = sys.argv[1]
-    assert(task=='segment' or task == 'detect')
-    detection_only = True if task == 'detect' else False
-    data_config_file = sys.argv[2]
+    data_config_file = sys.argv[1]
     assert(os.path.isfile(net_config_file) and os.path.isfile(data_config_file))
-    model_test(net_config_file, data_config_file, detection_only)
+    model_test(net_config_file, data_config_file)