prepare.py

import os
import shutil
import numpy as np
from copy import deepcopy as copy
import cv2
import math
import datetime
import time
from scipy import ndimage as ndi
from scipy.io import loadmat
import numpy as npS
import scipy
from scipy.ndimage.interpolation import zoom
from glob import glob
from skimage import measure, morphology
import SimpleITK as sitk
from scipy.ndimage.morphology import binary_dilation, generate_binary_structure, distance_transform_edt
from skimage.morphology import convex_hull_image, disk, binary_closing
from multiprocessing import Pool
from functools import partial
import sys
sys.path.append('preprocessing')
# from preprocessing.step1b import step1_python
from preprocessing.step1 import step1_python,load_itk_image
# from preprocessing.step1 import load_itk_image
import warnings
import pydicom as dicom


def load_scan(path):
	"""
	:param path: input CT path, dicom format
	:return: CT image
	"""
	#slices = [dicom.read_file(path + '/' + s, force = True) for s in os.listdir(path)]
	slices = []
	for s in os.listdir(path):
		if s.endswith('.dcm'):
			slices.append(dicom.read_file(path + '/' + s, force=True))

	slices.sort(key=lambda x: float(x.ImagePositionPatient[2]))
	if slices[0].ImagePositionPatient[2] == slices[1].ImagePositionPatient[2]:
		sec_num = 2
		while slices[0].ImagePositionPatient[2] == slices[sec_num].ImagePositionPatient[2]:
			sec_num = sec_num + 1
		slice_num = int(len(slices) / sec_num)
		slices.sort(key=lambda x:float(x.InstanceNumber))
		slices = slices[0:slice_num]
		slices.sort(key=lambda x:float(x.ImagePositionPatient[2]))
	try:
		slice_thickness = np.abs(slices[0].ImagePositionPatient[2] - slices[1].ImagePositionPatient[2])
	except:
		slice_thickness = np.abs(slices[0].SliceLocation - slices[1].SliceLocation)
		
	for s in slices:
		s.SliceThickness = slice_thickness
		
	return slices


def resample(imgs, spacing, new_spacing, order=1):
	"""
	:param imgs: CT image
	:param spacing: CT original voxel spacing
	:param new_spacing: target voxel spacing
	:param order: interpolation choice
	:return: CT image with new voxel spacing, new voxel spacing
	"""
	if len(imgs.shape) == 3:
		new_shape = np.round(imgs.shape * spacing / new_spacing)
		true_spacing = spacing * imgs.shape / new_shape
		resize_factor = new_shape / imgs.shape
		imgs = zoom(imgs, resize_factor, mode='nearest', order=order)
		return imgs, true_spacing
	elif len(imgs.shape) == 4:
		n = imgs.shape[-1]
		newimg = []
		for i in range(n):
			slice = imgs[:,:,:,i]
			newslice, true_spacing = resample(slice, spacing, new_spacing)
			newimg.append(newslice)
		newimg = np.transpose(np.array(newimg), [1, 2, 3, 0])
		return newimg, true_spacing
	else:
		raise ValueError('wrong shape')


def process_mask(mask):
	"""
	:param mask: input lung mask
	:return: convex hull processing on lung mask to avoid over-segmentation
	"""
	convex_mask = np.copy(mask)
	for i_layer in range(convex_mask.shape[0]):
		mask1 = np.ascontiguousarray(mask[i_layer])
		if np.sum(mask1) > 0:
			mask2 = convex_hull_image(mask1)
			if np.sum(mask2) > 1.5*np.sum(mask1):
				mask2 = mask1
		else:
			mask2 = mask1
		convex_mask[i_layer] = mask2
	struct = generate_binary_structure(3, 1)
	dilatedMask = binary_dilation(convex_mask, structure=struct, iterations=10)
	return dilatedMask


def lumTrans_hu(img):
	"""
	:param img: CT image
	:return: Hounsfield Unit window clipped and normalized
	"""
	lungwin = np.array([-1000.,400.])
	newimg = (img-lungwin[0])/(lungwin[1]-lungwin[0])
	newimg[newimg < 0] = 0
	newimg[newimg > 1] = 1
	newimg = (newimg*255).astype('uint8')
	return newimg


def save_itk(image, origin, spacing, filename):
	"""
	:param image: images to be saved
	:param origin: CT origin
	:param spacing: CT spacing
	:param filename: save name
	:return: None
	"""
	if type(origin) != tuple:
		if type(origin) == list:
			origin = tuple(reversed(origin))
		else:
			origin = tuple(reversed(origin.tolist()))
	if type(spacing) != tuple:
		if type(spacing) == list:
			spacing = tuple(reversed(spacing))
		else:
			spacing = tuple(reversed(spacing.tolist()))
	itkimage = sitk.GetImageFromArray(image, isVector=False)
	itkimage.SetSpacing(spacing)
	itkimage.SetOrigin(origin)
	sitk.WriteImage(itkimage, filename, True)


def savenpy(data_path, prep_folder):
	"""
	:param data_path: input CT data path
	:param prep_folder:
	:return: None
	"""
	resolution = np.array([1, 1, 1])

	#windows系统添加的是反斜杠\
	name = data_path.split('\\')[-1].split('.nii')[0]#获得文件名
	# name = data_path.split('/')[-1].split('.nii')[0]

	assert (os.path.exists(data_path) is True)
	im, m1, m2, mtotal, origin, spacing = step1_python(data_path)
	# print('Origin: ', origin, ' Spacing: ', spacing, 'img shape: ', im.shape)
	Mask = m1 + m2
	xx, yy, zz = np.where(Mask)
	box = np.array([[np.min(xx), np.max(xx)], [np.min(yy), np.max(yy)], [np.min(zz),np.max(zz)]])
	margin = 5
	box = np.vstack([np.max([[0, 0, 0], box[:, 0] - margin], 0), np.min([np.array(Mask.shape), box[:, 1] + margin], axis=0).T]).T

	# save the lung mask
	data_savepath = os.path.join(prep_folder, name+'_lung_mask.nii.gz')
	Mask_crop = Mask[box[0, 0]:box[0, 1],
				  box[1,0]:box[1,1],
				  box[2,0]:box[2,1]]
	save_itk(Mask_crop.astype(dtype='uint8'), origin, spacing, data_savepath)
	# convex_mask = m1
	# dm1 = process_mask(m1)
	# dm2 = process_mask(m2)
	# dilatedMask = (dm1 + dm2) | mtotal
	# Mask = m1+m2
	# dilatedMask = dilatedMask.astype('uint8')
	# Mask = Mask.astype('uint8')

	# save the CT image
	im[np.isnan(im)] = -2000
	sliceim_hu = lumTrans_hu(im)
	shapeorg = sliceim_hu.shape
	box_shape = np.array([[0, shapeorg[0]], [0, shapeorg[1]], [0, shapeorg[2]]])
	sliceim2_hu = sliceim_hu[box[0, 0]:box[0, 1],
				  box[1,0]:box[1,1],
				  box[2,0]:box[2,1]]
	# save box (image original shape and cropped window region)
	box = np.concatenate([box, box_shape], axis=0)
	np.save(os.path.join(prep_folder, name+'_box.npy'), box)
	# save processed image
	data_savepath = os.path.join(prep_folder, name+'_clean_hu.nii.gz')
	# data_savepath = os.path.join(prep_folder, name + '_clean_hu.nii.gz')
	save_itk(sliceim2_hu.astype(dtype='uint8'), origin, spacing, data_savepath)
	print(name)
	
	return


def preprocess_CT(inputpath=None, savepath=None):
	"""
	Preprocess the CT images in the input path to extract lung field
	Save the processed images in the save path
	:param inputpath: input data path
	:param savepath: output save path
	:return: save directory path
	"""
	warnings.filterwarnings("ignore")
	warnings.simplefilter(action='ignore', category=FutureWarning)
	if not os.path.exists(savepath):
		os.mkdir(savepath)

	filelist = glob(os.path.join(inputpath, '*.nii*'))  # default nifty format
	print(inputpath, filelist)
	for curfilepath in filelist:
		start_time = time.time()
		print('starting preprocessing lung CT')
		savenpy(data_path=curfilepath, prep_folder=savepath)
		end_time = time.time()
		print('end preprocessing lung CT, time %d seconds'%(end_time-start_time))

	return savepath

def preprocess_CTlabel(labelInputpath,boxpath,savepath):
	"""
	将CT图像的标签图像裁剪为clean_hu文件大小

	"""
	labelfilelist = glob(os.path.join(labelInputpath,'*.gz'))
	boxfilelist = glob(os.path.join(boxpath,'*.npy'))

	i=0;
	for curfilepath in labelfilelist:
		curboxpath = boxfilelist[i]
		boxnum = curboxpath.split('_')[-3]
		labelnum = curfilepath.split('_')[-2]
		if boxnum == labelnum:
			box = np.load(curboxpath)
			print(box)
			labelImg,labelOrigin,labelShape = load_itk_image(curfilepath)
			labelImg1 = labelImg[box[0, 0]:box[0, 1],
				  box[1,0]:box[1,1],
				  box[2,0]:box[2,1]]
			# labelName = savepath + '/'+curfilepath.split('/')[-1]
			labelName = savepath + '\\' + curfilepath.split('\\')[-1]
			temp = labelName.split('.nii')
			labelName1 = temp[0]+'_label.nii'+temp[1]
			save_itk(labelImg1,labelOrigin,labelShape,labelName1)

		else:
			for curboxpath in boxfilelist:
				boxnum = curboxpath.split('_')[-3]
				if boxnum == labelnum:
					box = np.load(curboxpath)
					print(box)
					labelImg, labelOrigin, labelShape = load_itk_image(curfilepath)
					labelImg1 = labelImg[box[0, 0]:box[0, 1],
								box[1, 0]:box[1, 1],
								box[2, 0]:box[2, 1]]
					# labelName = savepath +'/'+ curfilepath.split('/')[-1]
					labelName = savepath + '\\' + curfilepath.split('\\')[-1]
					temp = labelName.split('.nii')
					labelName1 = temp[0] + '_label.nii' + temp[1]
					save_itk(labelImg1, labelOrigin, labelShape, labelName1)

					break
		i = i+1
		print(labelnum )






if __name__ == '__main__':
	# inputpath ="/home/xy/pyProjects/code/code/airway_segmentation/testdata/imagesTr";
	# inputpath = "../../airway_segmentation/airway_segmentation/testdata/imagesTr1";
	# inputpath = "../../airway_segmentation/airway_segmentation/testdata/labelsTr1";
	# inputpath = "../../airway_segmentation/airway_segmentation/testdata";
	# inputpath1 = "G:/WorkSpace/DataSets/airway seg/ATM22/TrainBatch1/imagesTr118"#/////
	#
	# assert (os.path.exists(inputpath1))
	savepath = "./result_lung682703";
	# preprocess_CT(inputpath1,savepath);

	boxpath = "./result_lung"
	inputpath2 = "G:\\WorkSpace\\DataSets\\airway seg\\ATM22\\TrainBatch2\\labelsTr\\682703"
	boxpath = 'G:\\WorkSpace\\DataSets\\airway seg\\ATM22\\segLungRegion\\ATM2\\682703'
	preprocess_CTlabel(inputpath2,boxpath, savepath)