import os
import numpy as np
import nibabel as nib
from PIL import Image
from util import fileutil
class Normalizer:
@staticmethod
def dispersed(min=None, max=None):
return lambda x: (x - (min if min else x.min())) / \
((max if max else x.max()) - (min if min else x.min())) if \
((max if max else x.max()) - (min if min else x.min())) != 0 else x
@staticmethod
def z_score(mean=None, std=None):
return lambda x: (x - (mean if mean else np.mean(x, axis=0))) / (std if std else np.std(x, axis=0))
def resize_nii(data, size):
original_img_dtype = data.dtype
if original_img_dtype == np.int32 or original_img_dtype == np.int64: data = data.astype(np.int16)
new_img = np.zeros(shape=list(size) + [data.shape[2]], dtype=data.dtype)
for z in range(data.shape[2]):
new_img[:, :, z] = np.array(Image.fromarray(data[:, :, z]).resize(tuple(size), resample=0))
return new_img.astype(original_img_dtype)
if __name__ == '__main__':
in_dir = '/data/Data/NII/WFY/bladder/Raw512/Train/img'
out_dir = '/data/Data/NII/WFY/bladder/After256/Train/img_zscore'
resize_xy = [256, 256]
save_label_val = 1
normalizer = Normalizer.z_score()
# normalizer = Normalizer.dispersed(min=-1024, max=1024)
paths = fileutil.list_all_file_path(in_dir)
for path in paths:
print(path)
pre_nii = nib.load(path)
nii_data = np.array(pre_nii.dataobj)
is_ct = nii_data[nii_data < 0].any()
if is_ct:
nii_data[nii_data < -1024] = -1024
nii_data[nii_data > 1024] = 1024
if normalizer: nii_data = normalizer(nii_data)
else:
if save_label_val: nii_data[nii_data != save_label_val] = 0
if resize_xy: nii_data = resize_nii(nii_data, resize_xy)
post_nii = nib.Nifti1Image(nii_data, pre_nii.affine)
out_mask_path = f'{out_dir}/{os.path.basename(path)}'
fileutil.check_and_create_dir_for_new_file(out_mask_path)
nib.save(post_nii, out_mask_path)