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 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_NII(object): def __init__(self, in_path, out_path, raw_dict_dir): ''' :param in_path: path of dicom :param out_path: path of nii.gz ''' self.in_path = in_path self.out_path = out_path self.raw_dict_dir = raw_dict_dir self.target_spacing = 0.5 self.itk_resample = ItkResample() if not os.path.exists(self.out_path): os.makedirs(self.out_path) def __call__(self): # raw_dicts = {} uids = os.listdir(self.in_path) ## cpu_count() 112 # pool = Pool(int(cpu_count() * 0.7)) for uid in uids: if not os.path.exists(os.path.join(self.out_path, uid + '.nii.gz')): try: self._single_transform(uid, self.raw_dict_dir) # pool.apply_async(self._single_transform, (uid, self.raw_dict_dir)) # raw_info = raw_info.get() # raw_dicts[uid] = raw_info except Exception as err: traceback.print_exc() print('Transform dicom to nii.gz throws exception %s, with series uid %s!' % (err, uid)) # pool.close() # pool.join() # if not os.path.exists(os.path.join(os.path.dirname(self.out_path), 'raw_dicts.json')): # with open(os.path.join(os.path.dirname(self.out_path), 'raw_dicts.json'), 'w') as file: # json.dump(raw_dicts, file, indent=4) # else: # raw_dicts = json.load(open(os.path.join(os.path.dirname(self.out_path), 'raw_dicts.json'), 'r')) # return raw_dicts def _single_transform(self, uid, raw_dict_dir): # dcm -> nii.gz，并斤进行排序 print('Processing series uid %s' % uid) nii_file = os.path.join(self.out_path, uid + '.nii.gz') dcm_folder = os.path.join(self.in_path, uid) # convert dicom to nii 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) try: sitk.WriteImage(itk_image_resample, nii_file) except Exception as err: print('!!!!! Write %s throws exception %s.' % (uid, err)) return info_dict = dict() info_dict['raw_origin'] = list(np.array(itk_image_resample.GetOrigin())) info_dict['raw_spacing'] = list(np.array(itk_image_resample.GetSpacing())) info_dict['raw_size'] = list(np.array(itk_image_resample.GetSize()).astype(float)) raw_dict = {uid: info_dict} with open(os.path.join(raw_dict_dir, uid+'.json'), 'w+') as file: json.dump(raw_dict, file, indent=4) # info_dict = dict() # info_dict['raw_origin'] = list(np.array(itk_image_resample.GetOrigin())) # info_dict['raw_spacing'] = list(np.array(itk_image_resample.GetSpacing())) # info_dict['raw_size'] = list(np.array(itk_image_resample.GetSize()).astype(float)) # return info_dict