nnUNet_Brain_EPTN / libraries /utils_preprocess.py

Upload Process library

b51cd04 verified over 2 years ago

5.17 kB

#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Created on Sat Sep  5 20:34:46 2020
@author: ana
"""


# import general libraries
import os
import numpy as np
import pandas as pd
import scipy
import nibabel as nib
import pydicom
import glob
import warnings
from copy import deepcopy
from scipy.ndimage import find_objects
from scipy.ndimage.morphology import binary_fill_holes
from skimage import measure
from matplotlib.patches import Polygon
from libraries.utils_nii_dicom import set_header_info


###############################################################################
###################################   FUNCTIONS   #############################
###############################################################################

def delete_all(obj):
    for i in vars(obj):
        del obj.i

def overwrite_ct_threshold(ct_image, body = None, artefact = None, contrast = None):
    # Change the HU out of the body to air: -1000    
    if body is not None:
        # Change the HU outside the body to -1000
        ct_image[body==0]=-1000
    if artefact is not None:
        # Change the HU to muscle: 14
        ct_image[artefact==1]=14
    if contrast is not None:
        # Change the HU to water: 0 Houndsfield Unit: CT unit
        ct_image[contrast==1]=0
    # Threshold above 1560HU
    ct_image[ct_image > 1560] = 1560
    return ct_image

def remove_dose_outside_mask(dose, mask = None):
    # Put dose = 0 outside the mask (typically the body)   
    if mask is not None:
        dose[mask==0]=0
    return dose

def get_and_save_tv(roi, nib_header, tv_names, dst_dir):
    tv = np.zeros(nib_header['dim'][1:4])
    roi_names = list(roi.keys())
    if tv_names is not None:
        for c in tv_names:
            if c in roi_names:
                tv[roi[c]>0]=int(c.split('_')[1])/100
                
    # create nii, update header and save
    save_nii_image(tv, nib_header,dst_dir, 'struct_tv')

    return tv

def get_and_save_oars(roi, nib_header, oar_names, dst_dir):
    oars = np.zeros(nib_header['dim'][1:4])
    
    roi_names = list(roi.keys())
    contoursexist = []
    
    if oar_names is not None:
        for i in range(len(oar_names)):
            if oar_names[i] in roi_names:
                if roi[oar_names[i]] is not None:
                    contoursexist.append(1)
                    oars += (2**i)*roi[oar_names[i]]
                else:
                    contoursexist.append(0)        
                    
    # create nii, update header and save 
    save_nii_image(oars, nib_header,dst_dir, 'struct_oar', contours_exist = contoursexist)

        
def get_and_save_sample_probability(tv,dst_dir):
    # get sample probability for patch-based training
    bufftv=np.zeros_like(tv)
    bufftv[tv!=0]=1
    m=bufftv.sum(axis=0).sum(axis=0).sum()/1000
    sample_probability_slc=(bufftv.sum(axis=0).sum(axis=0)+m)/(bufftv.sum(axis=0).sum(axis=0)+m).sum()
    sample_probability_row=(bufftv.sum(axis=1).sum(axis=1)+m)/(bufftv.sum(axis=1).sum(axis=1)+m).sum()
    sample_probability_col=(bufftv.sum(axis=0).sum(axis=-1)+m)/(bufftv.sum(axis=0).sum(axis=-1)+m).sum()
    
    np.savez_compressed(os.path.join(dst_dir,'sample_probability_row.npz'),sample_probability_row)
    np.savez_compressed(os.path.join(dst_dir,'sample_probability_col.npz'),sample_probability_col)
    np.savez_compressed(os.path.join(dst_dir,'sample_probability_slc.npz'),sample_probability_slc)
    

def decompress_struct(struct_nib,struct_list):
    
    struct_data = struct_nib.get_fdata()
    roi = dict.fromkeys(struct_list, None) 
    # get contourexists from header
    contoursexist = list(struct_nib.header.extensions[0].get_content())
    for i in range(len(contoursexist)):
        if contoursexist[i] == 1:
            roi[struct_list[i]] = np.bitwise_and(struct_data.astype(int), 2 ** i).astype(bool)

    return roi

def binary_to_integers(struct,select_list):
        
    roi_names = list(struct.keys())
    # get shape
    for k in roi_names:
        if struct[k] is not None:
            struct_shape = struct[k].shape
    # initialize roi
    roi = np.zeros(struct_shape)     
    
    # initialize label
    label = 0
    for i in range(len(select_list)):
        label = label + 1
        if select_list[i] in roi_names:
            if struct[select_list[i]] is not None:
                # populate roi matrix with selected rois in select_list
                roi[struct[select_list[i]]] = label
        
    return roi

def save_nii_image(nib_img, nib_header,dst_dir, image_name, contours_exist = None):
    
    image_nii = nib.Nifti1Image(nib_img, affine=np.eye(4)) # for Nifti1 header, change for a Nifti2 type of header
    # Update header fields
    if contours_exist is not None:
        image_nii = set_header_info(image_nii, nib_header['pixdim'][1:4], [nib_header['qoffset_x'],nib_header['qoffset_y'],nib_header['qoffset_z']], contours_exist = contours_exist)
    else:
        image_nii = set_header_info(image_nii, nib_header['pixdim'][1:4], [nib_header['qoffset_x'],nib_header['qoffset_y'],nib_header['qoffset_z']])
    # Save  nii 
    nib.save(image_nii, os.path.join(dst_dir,image_name +'.nii.gz'))



      

	#!/usr/bin/env python2
	# -- coding: utf-8 --
	"""
	Created on Sat Sep 5 20:34:46 2020

	@author: ana
	"""


	# import general libraries
	import os
	import numpy as np
	import pandas as pd
	import scipy
	import nibabel as nib
	import pydicom
	import glob
	import warnings
	from copy import deepcopy
	from scipy.ndimage import find_objects
	from scipy.ndimage.morphology import binary_fill_holes
	from skimage import measure
	from matplotlib.patches import Polygon
	from libraries.utils_nii_dicom import set_header_info


	###############################################################################
	################################### FUNCTIONS #############################
	###############################################################################

	def delete_all(obj):
	for i in vars(obj):
	del obj.i

	def overwrite_ct_threshold(ct_image, body = None, artefact = None, contrast = None):
	# Change the HU out of the body to air: -1000
	if body is not None:
	# Change the HU outside the body to -1000
	ct_image[body==0]=-1000
	if artefact is not None:
	# Change the HU to muscle: 14
	ct_image[artefact==1]=14
	if contrast is not None:
	# Change the HU to water: 0 Houndsfield Unit: CT unit
	ct_image[contrast==1]=0
	# Threshold above 1560HU
	ct_image[ct_image > 1560] = 1560
	return ct_image

	def remove_dose_outside_mask(dose, mask = None):
	# Put dose = 0 outside the mask (typically the body)
	if mask is not None:
	dose[mask==0]=0
	return dose

	def get_and_save_tv(roi, nib_header, tv_names, dst_dir):
	tv = np.zeros(nib_header['dim'][1:4])
	roi_names = list(roi.keys())
	if tv_names is not None:
	for c in tv_names:
	if c in roi_names:
	tv[roi[c]>0]=int(c.split('_')[1])/100

	# create nii, update header and save
	save_nii_image(tv, nib_header,dst_dir, 'struct_tv')

	return tv

	def get_and_save_oars(roi, nib_header, oar_names, dst_dir):
	oars = np.zeros(nib_header['dim'][1:4])

	roi_names = list(roi.keys())
	contoursexist = []

	if oar_names is not None:
	for i in range(len(oar_names)):
	if oar_names[i] in roi_names:
	if roi[oar_names[i]] is not None:
	contoursexist.append(1)
	oars += (2*i)roi[oar_names[i]]
	else:
	contoursexist.append(0)

	# create nii, update header and save
	save_nii_image(oars, nib_header,dst_dir, 'struct_oar', contours_exist = contoursexist)


	def get_and_save_sample_probability(tv,dst_dir):
	# get sample probability for patch-based training
	bufftv=np.zeros_like(tv)
	bufftv[tv!=0]=1
	m=bufftv.sum(axis=0).sum(axis=0).sum()/1000
	sample_probability_slc=(bufftv.sum(axis=0).sum(axis=0)+m)/(bufftv.sum(axis=0).sum(axis=0)+m).sum()
	sample_probability_row=(bufftv.sum(axis=1).sum(axis=1)+m)/(bufftv.sum(axis=1).sum(axis=1)+m).sum()
	sample_probability_col=(bufftv.sum(axis=0).sum(axis=-1)+m)/(bufftv.sum(axis=0).sum(axis=-1)+m).sum()

	np.savez_compressed(os.path.join(dst_dir,'sample_probability_row.npz'),sample_probability_row)
	np.savez_compressed(os.path.join(dst_dir,'sample_probability_col.npz'),sample_probability_col)
	np.savez_compressed(os.path.join(dst_dir,'sample_probability_slc.npz'),sample_probability_slc)


	def decompress_struct(struct_nib,struct_list):

	struct_data = struct_nib.get_fdata()
	roi = dict.fromkeys(struct_list, None)
	# get contourexists from header
	contoursexist = list(struct_nib.header.extensions[0].get_content())
	for i in range(len(contoursexist)):
	if contoursexist[i] == 1:
	roi[struct_list[i]] = np.bitwise_and(struct_data.astype(int), 2 ** i).astype(bool)

	return roi

	def binary_to_integers(struct,select_list):

	roi_names = list(struct.keys())
	# get shape
	for k in roi_names:
	if struct[k] is not None:
	struct_shape = struct[k].shape
	# initialize roi
	roi = np.zeros(struct_shape)

	# initialize label
	label = 0
	for i in range(len(select_list)):
	label = label + 1
	if select_list[i] in roi_names:
	if struct[select_list[i]] is not None:
	# populate roi matrix with selected rois in select_list
	roi[struct[select_list[i]]] = label

	return roi

	def save_nii_image(nib_img, nib_header,dst_dir, image_name, contours_exist = None):

	image_nii = nib.Nifti1Image(nib_img, affine=np.eye(4)) # for Nifti1 header, change for a Nifti2 type of header
	# Update header fields
	if contours_exist is not None:
	image_nii = set_header_info(image_nii, nib_header['pixdim'][1:4], [nib_header['qoffset_x'],nib_header['qoffset_y'],nib_header['qoffset_z']], contours_exist = contours_exist)
	else:
	image_nii = set_header_info(image_nii, nib_header['pixdim'][1:4], [nib_header['qoffset_x'],nib_header['qoffset_y'],nib_header['qoffset_z']])
	# Save nii
	nib.save(image_nii, os.path.join(dst_dir,image_name +'.nii.gz'))