update test, save spacing for segmentation result

test.py

@@ -307,7 +307,7 @@ def test(config_file):
     margin = config_test.get('roi_patch_margin', 5)
 
     for i in range(image_num):
-        [temp_imgs, temp_weight, temp_name, temp_bbox, temp_size] = dataloader.get_image_data_with_name(i)
+        [temp_imgs, temp_weight, temp_name, img_names, temp_bbox, temp_size] = dataloader.get_image_data_with_name(i)
         t0 = time.time()
         # 5.1, test of 1st network
         if(config_net1):
@@ -442,7 +442,7 @@ def test(config_file):
         test_time.append(time.time() - t0)
         final_label = np.zeros(temp_size, np.int16)
         final_label = set_ND_volume_roi_with_bounding_box_range(final_label, temp_bbox[0], temp_bbox[1], out_label)
-        save_array_as_nifty_volume(final_label, save_folder+"/{0:}.nii.gz".format(temp_name))
+        save_array_as_nifty_volume(final_label, save_folder+"/{0:}.nii.gz".format(temp_name), img_names[0])
         print(temp_name)
     test_time = np.asarray(test_time)
     print('test time', test_time.mean())

util/data_loader.py

@@ -79,14 +79,15 @@ def __load_one_volume(self, patient_name, mod):
         assert(volume_name is not None)
         volume_name = os.path.join(patient_dir, volume_name)
         volume = load_3d_volume_as_array(volume_name)
-        return volume
+        return volume, volume_name
 
     def load_data(self):
         """
         load all the training/testing data
         """
         self.patient_names = self.__get_patient_names()
         assert(len(self.patient_names)  > 0)
+        ImageNames = []
         X = []
         W = []
         Y = []
@@ -95,8 +96,9 @@ def load_data(self):
         data_num = self.data_num if (self.data_num is not None) else len(self.patient_names)
         for i in range(data_num):
             volume_list = []
+            volume_name_list = []
             for mod_idx in range(len(self.modality_postfix)):
-                volume = self.__load_one_volume(self.patient_names[i], self.modality_postfix[mod_idx])
+                volume, volume_name = self.__load_one_volume(self.patient_names[i], self.modality_postfix[mod_idx])
                 if(mod_idx == 0):
                     margin = 5
                     bbmin, bbmax = get_ND_bounding_box(volume, margin)
@@ -109,6 +111,8 @@ def load_data(self):
                 if(self.intensity_normalize[mod_idx]):
                     volume = itensity_normalize_one_volume(volume)
                 volume_list.append(volume)
+                volume_name_list.append(volume_name)
+            ImageNames.append(volume_name_list)
             X.append(volume_list)
             W.append(weight)
             bbox.append([bbmin, bbmax])
@@ -121,6 +125,7 @@ def load_data(self):
                 Y.append(label)
             if((i+1)%50 == 0 or (i+1) == data_num):
                 print('Data load, {0:}% finished'.format((i+1)*100.0/data_num))
+        self.image_names = ImageNames
         self.data   = X
         self.weight = W
         self.label  = Y
@@ -266,4 +271,4 @@ def get_image_data_with_name(self, i):
         """
         Used for testing, get one image data and patient name
         """
-        return [self.data[i], self.weight[i], self.patient_names[i], self.bbox[i], self.in_size[i]]
+        return [self.data[i], self.weight[i], self.patient_names[i], self.image_names[i], self.bbox[i], self.in_size[i]]

util/data_process.py

@@ -46,31 +46,46 @@ def load_mha_volume_as_array(filename):
     nda = sitk.GetArrayFromImage(img)
     return nda
 
-def load_nifty_volume_as_array(filename):
+def load_nifty_volume_as_array(filename, with_header = False):
     """
     load nifty image into numpy array, and transpose it based on the [z,y,x] axis order
     The output array shape is like [Depth, Height, Width]
     inputs:
         filename: the input file name, should be *.nii or *.nii.gz
+        with_header: return affine and hearder infomation
     outputs:
         data: a numpy data array
     """
     img = nibabel.load(filename)
     data = img.get_data()
     data = np.transpose(data, [2,1,0])
-    return data
+    if(with_header):
+        return data, img.affine, img.header
+    else:
+        return data
 
-def save_array_as_nifty_volume(data, filename):
+def save_array_as_nifty_volume(data, filename, reference_name = None):
     """
     save a numpy array as nifty image
     inputs:
         data: a numpy array with shape [Depth, Height, Width]
         filename: the ouput file name
+        reference_name: file name of the reference image of which affine and header are used
     outputs: None
     """
+
     data = np.transpose(data, [2, 1, 0])
-    img = nibabel.Nifti1Image(data, np.eye(4))
-    nibabel.save(img, filename)
+    if(reference_name is None):
+        affine = np.eye(4)
+        header = None
+        img = nibabel.Nifti1Image(data, affine, header)
+        nibabel.save(img, filename)
+    else:
+        img_ref = sitk.ReadImage(reference_name)
+        img = sitk.Image(data.shape[0], data.shape[1], data.shape[2], sitk.sitkInt16)
+        img.copyInformation(img_ref)
+        sitk.WriteImage(img, filename)
+
 
 def itensity_normalize_one_volume(volume):
     """