From f042f701ea45a59d2809dd0ecf44e66d0082382e Mon Sep 17 00:00:00 2001 From: Guotai Wang Date: Mon, 5 Mar 2018 16:33:50 +0000 Subject: [PATCH] update test, save spacing for segmentation result --- test.py | 4 ++-- util/data_loader.py | 11 ++++++++--- util/data_process.py | 25 ++++++++++++++++++++----- 3 files changed, 30 insertions(+), 10 deletions(-) diff --git a/test.py b/test.py index c2a7958..6081f42 100644 --- a/test.py +++ b/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()) diff --git a/util/data_loader.py b/util/data_loader.py index 34a666c..2d08570 100644 --- a/util/data_loader.py +++ b/util/data_loader.py @@ -79,7 +79,7 @@ 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): """ @@ -87,6 +87,7 @@ def load_data(self): """ 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]] diff --git a/util/data_process.py b/util/data_process.py index 7d478cd..724b45b 100644 --- a/util/data_process.py +++ b/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): """