Python SimpleITK.WriteImage() Examples

def _transform(dcm_dir,save_dir):
###   from dcm to nii
    reader = sitk.ImageSeriesReader()
    dicom_names = reader.GetGDCMSeriesFileNames(dcm_dir)
    reader.SetFileNames(dicom_names)
    image2 = reader.Execute()
    
###   transform 3D image to array
    image_array = sitk.GetArrayFromImage(image2)  #  z,y,x

###   crop the dark voxel
#    new_array,range_list = get_bound(image_array)


###   transform array to 3D image
    image3 = sitk.GetImageFromArray(image_array)

###   save 3D image
    name = dcm_dir.split('/')[-1] + '.nii'
    save_path = os.path.join(save_dir,name)   #   get the save path
    sitk.WriteImage(image3,save_path) 

def N4BiasFieldCorrection(src_path, dst_path):
        '''
        This function carry out BiasFieldCorrection for the files in a specific directory
        :param src_path: path of the source file
        :param dst_path: path of the target file
        :return:
        '''
        print("N4 bias correction runs.")
        inputImage = sitk.ReadImage(src_path)

        maskImage = sitk.OtsuThreshold(inputImage, 0, 1, 200)
        sitk.WriteImage(maskImage, dst_path)

        inputImage = sitk.Cast(inputImage, sitk.sitkFloat32)

        corrector = sitk.N4BiasFieldCorrectionImageFilter()

        # corrector.SetMaximumNumberOfIterations(10)

        output = corrector.Execute(inputImage, maskImage)
        sitk.WriteImage(output, dst_path)
        print("Finished N4 Bias Field Correction.....")

    # normalize the data(zero mean and unit variance) 

def correct_bias(in_file, out_file, image_type=sitk.sitkFloat64):
    """
    Corrects the bias using ANTs N4BiasFieldCorrection. If this fails, will then attempt to correct bias using SimpleITK
    :param in_file: input file path
    :param out_file: output file path
    :return: file path to the bias corrected image
    """
    correct = N4BiasFieldCorrection()
    correct.inputs.input_image = in_file
    correct.inputs.output_image = out_file
    try:
        done = correct.run()
        return done.outputs.output_image
    except IOError:
        warnings.warn(RuntimeWarning("ANTs N4BIasFieldCorrection could not be found."
                                     "Will try using SimpleITK for bias field correction"
                                     " which will take much longer. To fix this problem, add N4BiasFieldCorrection"
                                     " to your PATH system variable. (example: EXPORT PATH=${PATH}:/path/to/ants/bin)"))
        input_image = sitk.ReadImage(in_file, image_type)
        output_image = sitk.N4BiasFieldCorrection(input_image, input_image > 0)
        sitk.WriteImage(output_image, out_file)
        return os.path.abspath(out_file) 

def ReadMhd_WriteNifty(impath,impath_new):

    '''
    input:param impath & impsth_new
    return: imsave which aim to save the format changed image to a new designated path;
    all functions used in this function are from SimpleITK
    sitk.ReadImage():aim to read the pending image in
    sitk.GetArrayFromImage():aim to get array from the image
    sitk.GetImageFromArray()&sitk.WriteImage():both achieve to write and save the new image as what we want

    '''

# main part of the function
    image=sitk.ReadImage(impath)
    image_arr=sitk.GetArrayFromImage(image)
    imnew=sitk.GetImageFromArray(image_arr)
    imsave=sitk.WriteImage(imnew,impath_new)
    return imsave

# an example of using the function
# read the new image in the format of Nifity 

def correct_bias(in_file, out_file, image_type=sitk.sitkFloat64):
    """
    Corrects the bias using ANTs N4BiasFieldCorrection. If this fails, will then attempt to correct bias using SimpleITK
    :param in_file: input file path
    :param out_file: output file path
    :return: file path to the bias corrected image
    """
    correct = N4BiasFieldCorrection()
    correct.inputs.input_image = in_file
    correct.inputs.output_image = out_file
    try:
        done = correct.run()
        return done.outputs.output_image
    except IOError:
        warnings.warn(RuntimeWarning("ANTs N4BIasFieldCorrection could not be found."
                                     "Will try using SimpleITK for bias field correction"
                                     " which will take much longer. To fix this problem, add N4BiasFieldCorrection"
                                     " to your PATH system variable. (example: EXPORT PATH=${PATH}:/path/to/ants/bin)"))
        input_image = sitk.ReadImage(in_file, image_type)
        output_image = sitk.N4BiasFieldCorrection(input_image, input_image > 0)
        sitk.WriteImage(output_image, out_file)
        return os.path.abspath(out_file) 

def eyesize_basic_test():

  downloader = imagedownloader.ImageDownloader()
  estimator = eyesize.EyeSize()

  image_name = "TralitusSaltrator.jpg"

  downloader.set_figshare_id("1066744")
  downloader.set_image_name(image_name)
  downloader.download()

  input_image = sitk.ReadImage(image_name)

  estimator.set_image(input_image)
  estimator.set_seed_point([204,400])

  eyes_segmented,radius_estimate = estimator.estimate()

  sitk.WriteImage(eyes_segmented,'SegmentedEye.png')

  assert radius_estimate == 85 

def loop(self):
    if self.trigger == "internal":
      if self.fps_label:
        while self.inputs[0].poll():
          self.camera.max_fps = self.inputs[0].recv()[self.fps_label]
      t,img = self.camera.read_image()
    elif self.trigger == "external":
      data = self.inputs[0].recv()  # wait for a signal
      if data is None:
        return
      t,img = self.camera.get_image()
    self.timer = time.time()
    if self.save_folder:
      if not sitk:
        raise IOError("[Camera] Cannot save image, sitk is not installed !")
      image = sitk.GetImageFromArray(img)
      sitk.WriteImage(image,
               self.save_folder + "img_%.6d_%.5f.tiff" % (
               self.loops, t-self.t0))
    self.loops += 1
    self.send([t-self.t0,img]) 

def get_background_mask(in_folder, out_file, truth_name="GlistrBoost_ManuallyCorrected"):
    """
    This function computes a common background mask for all of the data in a subject folder.
    :param in_folder: a subject folder from the BRATS dataset.
    :param out_file: an image containing a mask that is 1 where the image data for that subject contains the background.
    :param truth_name: how the truth file is labeled int he subject folder
    :return: the path to the out_file
    """
    background_image = None
    for name in config["all_modalities"] + [truth_name]:
        image = sitk.ReadImage(get_image(in_folder, name))
        if background_image:
            if name == truth_name and not (image.GetOrigin() == background_image.GetOrigin()):
                image.SetOrigin(background_image.GetOrigin())
            background_image = sitk.And(image == 0, background_image)
        else:
            background_image = image == 0
    sitk.WriteImage(background_image, out_file)
    return os.path.abspath(out_file) 

def eyesize_noisy_test():

  downloader = imagedownloader.ImageDownloader()
  estimator = eyesize.EyeSize()

  image_name = "TralitusSaltrator.jpg"

  downloader.set_figshare_id("1066744")
  downloader.set_image_name(image_name)
  downloader.download()

  input_image = sitk.ReadImage(image_name)

  estimator.set_image(input_image)
  estimator.set_seed_point([204,400])

  eyes_segmented,radius_estimate = estimator.estimate()

  sitk.WriteImage(eyes_segmented,'SegmentedEye.png')

  expected_value = 85
  assert radius_estimate == expected_value, \
    "Problem with estimating radius: current: %s, expected: %s" % (radius_estimate, expected_value) 

def nib_resize1(save_dir1,name1,image_shape):

    save_path = os.path.join(save_dir1,name1)
    # load
    image1 = nib.load(save_path)
    image_data1 = image1.get_data()
    # print('before resize :',image_data1.shape)
    # resize
    image1 = resize(image1, image_shape)

    ###   transform to array
    image_data1 = image1.get_data()
    # image_data1 = image_data1[:,:,::-1]

    # print('after resize : ',image_data1.shape)
    image1 = sitk.GetImageFromArray(image_data1)
    sitk.WriteImage(image1,save_path) 

def split_4d_nifti(filename, output_folder):
    img_itk = sitk.ReadImage(filename)
    dim = img_itk.GetDimension()
    file_base = filename.split("/")[-1]
    if dim == 3:
        shutil.copy(filename, join(output_folder, file_base[:-7] + "_0000.nii.gz"))
        return
    elif dim != 4:
        raise RuntimeError("Unexpected dimensionality: %d of file %s, cannot split" % (dim, filename))
    else:
        img_npy = sitk.GetArrayFromImage(img_itk)
        spacing = img_itk.GetSpacing()
        origin = img_itk.GetOrigin()
        direction = np.array(img_itk.GetDirection()).reshape(4,4)
        # now modify these to remove the fourth dimension
        spacing = tuple(list(spacing[:-1]))
        origin = tuple(list(origin[:-1]))
        direction = tuple(direction[:-1, :-1].reshape(-1))
        for i, t in enumerate(range(img_npy.shape[0])):
            img = img_npy[t]
            img_itk_new = sitk.GetImageFromArray(img)
            img_itk_new.SetSpacing(spacing)
            img_itk_new.SetOrigin(origin)
            img_itk_new.SetDirection(direction)
            sitk.WriteImage(img_itk_new, join(output_folder, file_base[:-7] + "_%04.0d.nii.gz" % i)) 

def save_train_dataset_as_nifti(results_dir=os.path.join(paths.results_folder, "final"),
            out_dir=os.path.join(paths.results_folder, "training_set_results")):
    if not os.path.isdir(out_dir):
        os.mkdir(out_dir)
    a = load_dataset()
    for fold in range(5):
        working_dir = os.path.join(results_dir, "fold%d"%fold, "validation")
        ids_in_fold = os.listdir(working_dir)
        ids_in_fold.sort()
        ids_in_fold = [i for i in ids_in_fold if os.path.isdir(os.path.join(working_dir, i))]
        ids_in_fold_as_int = [int(i) for i in ids_in_fold]
        for pat_id in ids_in_fold_as_int:
            pat_in_dataset = a[pat_id]
            seg_pred = np.load(os.path.join(working_dir, "%03.0d"%pat_id, "segs.npz"))['seg_pred']
            b = convert_to_original_coord_system(seg_pred, pat_in_dataset)
            sitk_img = sitk.GetImageFromArray(b)
            sitk_img.SetSpacing(pat_in_dataset['spacing'])
            sitk_img.SetDirection(pat_in_dataset['direction'])
            sitk_img.SetOrigin(pat_in_dataset['origin'])
            sitk.WriteImage(sitk_img, os.path.join(out_dir, pat_in_dataset['name'] + ".nii.gz")) 

def load_remove_save(input_file: str, output_file: str, for_which_classes: list,
                     minimum_valid_object_size: dict = None):
    # Only objects larger than minimum_valid_object_size will be removed. Keys in minimum_valid_object_size must
    # match entries in for_which_classes
    img_in = sitk.ReadImage(input_file)
    img_npy = sitk.GetArrayFromImage(img_in)
    volume_per_voxel = float(np.prod(img_in.GetSpacing(), dtype=np.float64))

    image, largest_removed, kept_size = remove_all_but_the_largest_connected_component(img_npy, for_which_classes,
                                                                                       volume_per_voxel,
                                                                                       minimum_valid_object_size)
    # print(input_file, "kept:", kept_size)
    img_out_itk = sitk.GetImageFromArray(image)
    img_out_itk = copy_geometry(img_out_itk, img_in)
    sitk.WriteImage(img_out_itk, output_file)
    return largest_removed, kept_size 

def convert_labels_back_to_BraTS_2018_2019_convention(input_folder: str, output_folder: str):
    """
    reads all prediction files (nifti) in the input folder, converts the labels back to BraTS convention and saves the
    result in output_folder
    :param input_folder:
    :param output_folder:
    :return:
    """
    maybe_mkdir_p(output_folder)
    nii = subfiles(input_folder, suffix='.nii.gz', join=False)
    for n in nii:
        a = sitk.ReadImage(join(input_folder, n))
        b = sitk.GetArrayFromImage(a)
        c = convert_labels_back_to_BraTS(b)
        d = sitk.GetImageFromArray(c)
        d.CopyInformation(a)
        sitk.WriteImage(d, join(output_folder, n)) 

def remove_all_but_the_two_largest_conn_comp(img_itk_file: str, file_out: str):
    """
    This was not used. I was just curious because others used this. Turns out this is not necessary for my networks
    """
    img_itk = sitk.ReadImage(img_itk_file)
    img_npy = sitk.GetArrayFromImage(img_itk)

    labelmap, num_labels = label((img_npy > 0).astype(int))

    if num_labels > 2:
        label_sizes = []
        for i in range(1, num_labels + 1):
            label_sizes.append(np.sum(labelmap == i))
        argsrt = np.argsort(label_sizes)[::-1] # two largest are now argsrt[0] and argsrt[1]
        keep_mask = (labelmap == argsrt[0] + 1) | (labelmap == argsrt[1] + 1)
        img_npy[~keep_mask] = 0
        new = sitk.GetImageFromArray(img_npy)
        new.CopyInformation(img_itk)
        sitk.WriteImage(new, file_out)
        print(os.path.basename(img_itk_file), num_labels, label_sizes)
    else:
        shutil.copy(img_itk_file, file_out) 

def copy_BraTS_segmentation_and_convert_labels(in_file, out_file):
    # use this for segmentation only!!!
    # nnUNet wants the labels to be continuous. BraTS is 0, 1, 2, 4 -> we make that into 0, 1, 2, 3
    img = sitk.ReadImage(in_file)
    img_npy = sitk.GetArrayFromImage(img)

    uniques = np.unique(img_npy)
    for u in uniques:
        if u not in [0, 1, 2, 4]:
            raise RuntimeError('unexpected label')

    seg_new = np.zeros_like(img_npy)
    seg_new[img_npy == 4] = 3
    seg_new[img_npy == 2] = 1
    seg_new[img_npy == 1] = 2
    img_corr = sitk.GetImageFromArray(seg_new)
    img_corr.CopyInformation(img)
    sitk.WriteImage(img_corr, out_file) 

def export_segmentations_postprocess(indir, outdir):
    maybe_mkdir_p(outdir)
    niftis = subfiles(indir, suffix='nii.gz', join=False)
    for n in niftis:
        print("\n", n)
        identifier = str(n.split("_")[-1][:-7])
        outfname = join(outdir, "test-segmentation-%s.nii" % identifier)
        img = sitk.ReadImage(join(indir, n))
        img_npy = sitk.GetArrayFromImage(img)
        lmap, num_objects = label((img_npy > 0).astype(int))
        sizes = []
        for o in range(1, num_objects + 1):
            sizes.append((lmap == o).sum())
        mx = np.argmax(sizes) + 1
        print(sizes)
        img_npy[lmap != mx] = 0
        img_new = sitk.GetImageFromArray(img_npy)
        img_new.CopyInformation(img)
        sitk.WriteImage(img_new, outfname) 

def copy_BraTS_segmentation_and_convert_labels(in_file, out_file):
    # use this for segmentation only!!!
    # nnUNet wants the labels to be continuous. BraTS is 0, 1, 2, 4 -> we make that into 0, 1, 2, 3
    img = sitk.ReadImage(in_file)
    img_npy = sitk.GetArrayFromImage(img)

    uniques = np.unique(img_npy)
    for u in uniques:
        if u not in [0, 1, 2, 4]:
            raise RuntimeError('unexpected label')

    seg_new = np.zeros_like(img_npy)
    seg_new[img_npy == 4] = 3
    seg_new[img_npy == 2] = 1
    seg_new[img_npy == 1] = 2
    img_corr = sitk.GetImageFromArray(seg_new)
    img_corr.CopyInformation(img)
    sitk.WriteImage(img_corr, out_file) 

def get_background_mask(in_folder, out_file, truth_name="GlistrBoost_ManuallyCorrected"):
    """
    This function computes a common background mask for all of the data in a subject folder.
    :param in_folder: a subject folder from the BRATS dataset.
    :param out_file: an image containing a mask that is 1 where the image data for that subject contains the background.
    :param truth_name: how the truth file is labeled int he subject folder
    :return: the path to the out_file
    """
    background_image = None
    for name in config["all_modalities"] + [truth_name]:
        image = sitk.ReadImage(get_image(in_folder, name))
        if background_image:
            if name == truth_name and not (image.GetOrigin() == background_image.GetOrigin()):
                image.SetOrigin(background_image.GetOrigin())
            background_image = sitk.And(image == 0, background_image)
        else:
            background_image = image == 0
    sitk.WriteImage(background_image, out_file)
    return os.path.abspath(out_file) 

def convert_for_submission(source_dir, target_dir):
    files = subfiles(source_dir, suffix=".nii.gz", join=False)
    maybe_mkdir_p(target_dir)
    for f in files:
        splitted = f.split("__")
        case_id = int(splitted[1])
        timestep = int(splitted[2][:-7])
        t = join(target_dir, "test%02d_%02d_nnUNet.nii" % (case_id, timestep))
        img = sitk.ReadImage(join(source_dir, f))
        sitk.WriteImage(img, t) 

def convert_to_nii_gz(filename):
    f = sitk.ReadImage(filename)
    sitk.WriteImage(f, os.path.splitext(filename)[0] + ".nii.gz")
    os.remove(filename) 

def save_image(img_arr, path):
    itk_img = sitk.GetImageFromArray(img_arr, isVector=False)
    sitk.WriteImage(itk_img, path) 

def check_origin(in_file, in_file2):
    image = sitk.ReadImage(in_file)
    image2 = sitk.ReadImage(in_file2)
    if not image.GetOrigin() == image2.GetOrigin():
        image.SetOrigin(image2.GetOrigin())
        sitk.WriteImage(image, in_file) 

def ReadMhd_WriteNifty(impath, impath_new):
    '''
    input:param impath & impsth_new
    return: imsave which aim to save the format changed image to a new designated path;
    all functions used in this function are from SimpleITK
    sitk.ReadImage():aim to read the pending image in
    sitk.GetArrayFromImage():aim to get array from the image
    sitk.GetImageFromArray()&sitk.WriteImage():both achieve to write and save the new image as what we want

    '''

# main part of the function
    # src : target image
    # img_array : the array form image
    # img : the new created image what we want
    # imresult : return value
    src = sitk.ReadImage(impath)
    img_array = sitk.GetArrayFromImage(src)
    img = sitk.GetImageFromArray(img_array)
    imresult = sitk.WriteImage(img, impath_new)

    return imresult


# an example of using the function
# read the new image in the format of nrry 

def ReadMhd_Write2Nrrd(mhd_full_path_name, nrrd_full_path_name):
    Image = sitk.ReadImage(mhd_full_path_name)
    image_array = sitk.GetArrayFromImage(Image)
    img = sitk.GetImageFromArray(image_array)
    sitk.WriteImage(img, nrrd_full_path_name)
    print('The picture has been transferred.')

# ReadMhd_Write2Nrrd('E:\IMAGE-BME\yanghelin\PATIENT_DICOM.mhd','E:\IMAGE-BME\yanghelin\PATIENT_DICOM.nrrd') 

def eyesize_cleanplate_test():

  downloader = imagedownloader.ImageDownloader()
  estimator = eyesize.EyeSize()

  image_name = "TralitusSaltrator.jpg"

  inputs_dir = os.path.join(os.getcwd(), "inputs")
  outputs_dir = os.path.join(os.getcwd(), "outputs")
  
  image_path = os.path.join(inputs_dir, image_name)
  
  # Clear inputs folder
  if os.path.isdir(inputs_dir):
    shutil.rmtree(inputs_dir)
  os.mkdir(inputs_dir)
  
  # Clear outputs folder
  if os.path.isdir(outputs_dir):
    shutil.rmtree(outputs_dir)
  os.mkdir(outputs_dir)

  downloader.set_figshare_id("1066744")
  downloader.set_image_name(image_path)
  downloader.download()

  input_image = sitk.ReadImage(image_path)

  estimator.set_image(input_image)
  estimator.set_seed_point([204,400])

  eyes_segmented,radius_estimate = estimator.estimate()

  sitk.WriteImage(eyes_segmented,'SegmentedEye.png')

  expected_value = 85
  assert radius_estimate == expected_value, \
    "Problem with estimating radius: current: %s, expected: %s" % (radius_estimate, expected_value) 

def ConvertToFormat(filePath, finalPath):


'''
;PURPOSE:
;	Definite a function that can input a img file and convert to an other curtain format img file.
;
;INPUT:
;	A img file path:     (*.mhd)
;	Other format path:   (*.nrrd, *nii)
;
;OUTPUT:
;	Other format in the input path.
;
;RETURN VALUE:
;	void
;
;AUTHER:
;	Hang Pan, Dalian University of Technology
;
;CREATION DATE:
;	2018-04-06
'''
   try:
        img = sitk.GetImageFromArray(
    sitk.GetArrayFromImage(
        sitk.ReadImage(filePath)))
    except Exception as e:
        print("ReadError: This file cannot be read!")

    try:
        sitk.WriteImage(img, finalPath)
    except Exception as e:
        print("FormatError: This format is not Support!") 

def export_for_submission(source_dir, target_dir):
    """
    promise wants mhd :-/
    :param source_dir:
    :param target_dir:
    :return:
    """
    files = subfiles(source_dir, suffix=".nii.gz", join=False)
    target_files = [join(target_dir, i[:-7] + ".mhd") for i in files]
    maybe_mkdir_p(target_dir)
    for f, t in zip(files, target_files):
        img = sitk.ReadImage(join(source_dir, f))
        sitk.WriteImage(img, t) 

def format_conv(path, my_format):
    # Define function format_conv
    # Call SimpleITK to read the image
    src = sitk.ReadImage(path)
    img_array = sitk.GetArrayFromImage(src)     # Convert image to array
    img = sitk.GetImageFromArray(img_array)     # Convert array to image

    # Get the path the file located
    l = len(path)
    for i in range(0, l):
        if path[l - 1 - i] == '\\':       # Find the last \
            end = l - i
            break

    self_path = path[0:end]              # Get the path of file

    if my_format == 'nii':                                  # chose which format to save
        # save the file named result in the same path
        sitk.WriteImage(img, self_path + 'result.nii')
        print('format_conv')                                # Tell the user function worked successfully
    elif my_format == 'nrrd':
        sitk.WriteImage(img, self_path + 'result.nrrd')
        print('format_conv')
    else:
        # If there is an unexpected format, inform the error
        print('Sorry, We do not support this format.') 

def save_updated_image(img_arr, path, origin, spacing):
    itk_scaled_img = sitk.GetImageFromArray(img_arr, isVector=False)
    itk_scaled_img.SetSpacing(spacing)
    itk_scaled_img.SetOrigin(origin)
    sitk.WriteImage(itk_scaled_img, path)