import json import os import sys import traceback from multiprocessing import Pool, cpu_count import SimpleITK as sitk from pydicom import dicomio import lmdb import numpy as np from datetime import datetime from resample import ItkResample # from resample import ItkResample BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) # 添加路径 sys.path.append(BASE_DIR) class MyEncoder(json.JSONEncoder): def default(self, obj): if isinstance(obj, np.integer): return int(obj) elif isinstance(obj, np.floating): return float(obj) elif isinstance(obj, np.ndarray): return obj.tolist() elif isinstance(obj, datetime): return obj.__str__() else: return super(MyEncoder, self).default(obj) def load_ct_from_dicom(dcm_path, sort_by_distance=True): class DcmInfo(object): def __init__(self, dcm_path, series_instance_uid, acquisition_number, sop_instance_uid, instance_number, image_orientation_patient, image_position_patient): super(DcmInfo, self).__init__() self.dcm_path = dcm_path self.series_instance_uid = series_instance_uid self.acquisition_number = acquisition_number self.sop_instance_uid = sop_instance_uid self.instance_number = instance_number self.image_orientation_patient = image_orientation_patient self.image_position_patient = image_position_patient self.slice_distance = self._cal_distance() def _cal_distance(self): normal = [self.image_orientation_patient[1] * self.image_orientation_patient[5] - self.image_orientation_patient[2] * self.image_orientation_patient[4], self.image_orientation_patient[2] * self.image_orientation_patient[3] - self.image_orientation_patient[0] * self.image_orientation_patient[5], self.image_orientation_patient[0] * self.image_orientation_patient[4] - self.image_orientation_patient[1] * self.image_orientation_patient[3]] distance = 0 for i in range(3): distance += normal[i] * self.image_position_patient[i] return distance def is_sop_instance_uid_exist(dcm_info, dcm_infos): for item in dcm_infos: if dcm_info.sop_instance_uid == item.sop_instance_uid: return True return False def get_dcm_path(dcm_info): return dcm_info.dcm_path reader = sitk.ImageSeriesReader() if sort_by_distance: dcm_infos = [] files = os.listdir(dcm_path) for file in files: file_path = os.path.join(dcm_path, file) dcm = dicomio.read_file(file_path, force=True) _series_instance_uid = dcm.SeriesInstanceUID _sop_instance_uid = dcm.SOPInstanceUID _instance_number = dcm.InstanceNumber _acquisition_number = dcm.AcquisitionNumber _image_orientation_patient = dcm.ImageOrientationPatient _image_position_patient = dcm.ImagePositionPatient dcm_info = DcmInfo(file_path, _series_instance_uid, _acquisition_number, _sop_instance_uid, _instance_number, _image_orientation_patient, _image_position_patient) if is_sop_instance_uid_exist(dcm_info, dcm_infos): continue dcm_infos.append(dcm_info) dcm_infos.sort(key=lambda x: x.slice_distance) dcm_series = list(map(get_dcm_path, dcm_infos)) else: dcm_series = reader.GetGDCMSeriesFileNames(dcm_path) reader.SetFileNames(dcm_series) sitk_image = reader.Execute() return sitk_image class Trans_Dicom_2_NPY(object): def __init__(self, in_path, out_path, datainfos): ''' :param in_path: path of dicom :param out_path: path of nii.gz ''' self.in_path = in_path self.out_path = out_path self.datainfos = datainfos self.itk_resample = ItkResample() self.target_spacing = 1 self.suid_list = [] for info in datainfos: self.suid_list.append(info[-2][0].split('/')[-2]) # print('self.suid_list', self.suid_list) # print(self.suid_list) # self.target_spacing = 1 # self.itk_resample = ItkResample() if not os.path.exists(self.out_path): os.makedirs(self.out_path) def __call__(self): uids = os.listdir(self.in_path) print(f'num cpus: {cpu_count()}') # pool = Pool(int(cpu_count() * 0.7)) for uid in uids: if not os.path.exists(os.path.join(self.out_path, uid + '_data.npy')) or not os.path.exists(os.path.join(self.out_path, uid + '_info.npy')): # pool.apply_async(self._single_transform, (uid, )) # np.save(os.path.join(self.out_path, uid + '_data.npy'), whole_data) # np.save(os.path.join(self.out_path, uid + '_info.npy'), whole_info) self._single_transform(uid) # pool.close() # pool.join() def _single_transform(self, uid): # dcm -> npy print('Processing series uid %s' % uid) dcm_folder = os.path.join(self.in_path, uid) # get itkimage try: itk_image = load_ct_from_dicom(dcm_folder) except Exception as err: print('!!!!! Read %s throws exception %s.' % (uid, err)) return # resample spacing print('Resample uid %s to spacing %.1f.' % (uid, self.target_spacing)) itk_image_resample = self.itk_resample.resample_to_spacing( itk_image=itk_image, target_spacing=(tuple([self.target_spacing] * 3)), interpolator=sitk.sitkBSpline) itk_image_resample.SetOrigin(itk_image.GetOrigin()) itk_image_resample.SetSpacing([self.target_spacing] * 3) # 将itk数据准换位npy格式，并且进行裁切 origin = itk_image_resample.GetOrigin() spacing = itk_image_resample.GetSpacing() array_image = sitk.GetArrayFromImage(itk_image_resample) bz, by, bx = array_image.shape # print(uid) # print(self.suid_list) work_index = self.suid_list.index(uid) info = self.datainfos[work_index][-1] mask_path = None for i in info: if 'nii.gz' in i: mask_path = i break if mask_path is None: return # label_infos = json.load(open(self.datainfos[work_index][-1][0], 'r'))['annotationSessions'] print(f'mask path {mask_path}') mask = sitk.ReadImage(mask_path) itk_mask_resample = self.itk_resample.resample_to_spacing( itk_image=mask, target_spacing=(tuple([self.target_spacing] * 3)), interpolator=sitk.sitkBSpline) itk_mask_resample.SetOrigin(itk_image.GetOrigin()) itk_mask_resample.SetSpacing([self.target_spacing] * 3) spacing_mask = itk_mask_resample.GetSpacing() array_mask = sitk.GetArrayFromImage(itk_mask_resample) print(f'img shape: {array_image.shape}, spacing {spacing}') print(f'mask shape: {array_mask.shape}, mask spacing {spacing_mask}') ## 每个mask id切个图 mask_ids = np.unique(array_mask) cut_size = 32 shape_min = min(array_image.shape) cut_size = min(cut_size, shape_min) save_dir = os.path.join(self.out_path, uid) os.makedirs(save_dir, exist_ok=True) for i in mask_ids: try: mask_i = np.where(array_mask==i, 1, 0) zs, ys, xs = np.where(mask_i==1) z_min, y_min, x_min = min(zs), min(ys), min(xs) z_max, y_max, x_max = max(zs), max(ys), max(xs) print(f'cut coords: {z_min}, {z_max}, {y_min}, {y_max}, {x_min}, {x_max}') # zc, yc, xc = int((z_min+z_max)/2), int((y_min+y_max)/2), int((x_min+x_max)/2) zc, yc, xc = int(np.median(zs)), int(np.median(ys)), int(np.median(xs)) z_min_cut, y_min_cut, x_min_cut = int(zc-cut_size/2), int(yc-cut_size/2), int(xc-cut_size/2) z_max_cut, y_max_cut, x_max_cut = int(zc+cut_size/2), int(yc+cut_size/2), int(xc+cut_size/2) if z_min_cut < 0: z_min_cut = 0 z_max_cut = cut_size if y_min_cut < 0: y_min_cut = 0 y_max_cut = cut_size if x_min_cut < 0: x_min_cut = 0 x_max_cut = cut_size if z_max_cut > array_image.shape[0]: z_max_cut = array_image.shape[0] z_min_cut = array_image.shape[0]-cut_size if y_max_cut > array_image.shape[1]: y_max_cut = array_image.shape[1] y_min_cut = array_image.shape[1]-cut_size if x_max_cut > array_image.shape[2]: x_max_cut = array_image.shape[2] x_min_cut = array_image.shape[2]-cut_size mask_i = mask_i[z_min_cut: z_max_cut, y_min_cut: y_max_cut, x_min_cut: x_max_cut] img_i = array_image[z_min_cut: z_max_cut, y_min_cut: y_max_cut, x_min_cut: x_max_cut] print(f'img_i shape: {img_i.shape}, mask_i shape: {mask_i.shape}') save_mask_path = os.path.join(save_dir, str(i) + '_mask.npy') save_img_path = os.path.join(save_dir, str(i) + '_data.npy') save_info_path = os.path.join(save_dir, str(i) + '_info.npy') np.save(save_mask_path, mask_i) np.save(save_img_path, img_i) np.save(save_info_path, np.array(spacing)) except: print(f'mask id {i} cutting failed')