import os import sys import glob import shutil import numpy as np import pandas as pd from tqdm import tqdm import SimpleITK as sitk from collections import defaultdict from skimage.measure import * sys.path.append('../utils') from respacing_func import * import logging logger = logging.getLogger() fh = logging.FileHandler('PrepData.log',encoding='utf-8') sh = logging.StreamHandler() formatter = logging.Formatter('%(asctime)s %(filename)s[line:%(lineno)d] %(levelname)s %(message)s') fh.setFormatter(formatter) sh.setFormatter(formatter) logger.addHandler(fh) logger.addHandler(sh) logger.setLevel(10) def ReadITKImage(path): itk_img = sitk.ReadImage(path) img_np = sitk.GetArrayFromImage(itk_img) return itk_img,img_np def WriteMask(mask,origin=None,spacing=None,path=''): if origin is None and spacing is None: itk_mask = mask else: itk_mask = sitk.GetImageFromArray(mask) itk_mask = sitk.Cast(itk_mask,sitk.sitkUInt8) itk_mask.SetSpacing(spacing) itk_mask.SetOrigin(origin) print ('Saving img to %s '%path) itkWriter = sitk.ImageFileWriter() itkWriter.SetUseCompression(True) itkWriter.SetFileName(path) itkWriter.Execute(itk_mask) def ReadCSV(path): if 'csv' in path: df = pd.read_csv(path,encoding='utf_8_sig') return df elif 'xlsx' in path: df = pd.read_excel(path,encoding='utf_8_sig') return df def LoadDFs(paths): result = [] for path in paths: if ('csv' not in path) and ('xlsx' not in path): continue df = ReadCSV(path) result.append(df) if len(result)>0: return pd.concat(result,ignore_index=True) class COVIDPre(object): def __init__(self,base_paths,df_paths,respac_df_path,area_th,bbox_df_path,uid_paths,health_uid_path,pos_uid_path,health_save_path=None,pos_save_path=None): self.img_base_path = base_paths[0] self.mask_base_path = base_paths[1] self.lung_mask_base_path = base_paths[2] self.target_spac = [1.0 for _ in range(3)] self.valtest_ratio = 0 self.test_ratio = 0.5 try: self.df = LoadDFs(df_paths) self.df_path = df_paths[0] except: self.df = None self.df_path = df_paths[0] self.respac_df_path = respac_df_path self.area_th = area_th self.bbox_df_path = bbox_df_path self.train_uid_path,self.val_uid_path,self.test_uid_path = uid_paths ##################### self.health_uid_path = health_uid_path self.pos_uid_path = pos_uid_path self.health_save_path = health_save_path self.pos_save_path = pos_save_path self.info_list = [] self.connected_component_choice = False def __call__(self): ''' 1. 将新处理的数据移到别处 ''' logger.info('step 1: shutildata') self._ShutilData() ''' 2. 针对self.df中的数据做resample ''' logger.info('step 2: resample image') self._ResampleBatch() ''' 3. 将respacaing后的数据的相关信息进行存储 ''' logger.info('step 3: Record respacing related info') self._GatherInfo() ''' 对respacing后的数据进行连通域分析(如果没有标注的mask可以不进行这一步) 如果不进行这一步，用uid产生合理的信息 ''' logger.info('step 4: Get bbox info') if self.connected_component_choice: self._ExtractConnectedComponents() else: health_uids = self._LoadUid(self.health_uid_path) pos_uids = self._LoadUid(self.pos_uid_path) self._GenerateInfoWithLungMask(uids=health_uids,base_save_path=self.health_save_path,kw='health', area=10000000000000) self._GenerateInfoWithLungMask(uids=pos_uids,base_save_path=self.pos_save_path,kw='pos',area=float('inf')) ''' 将respacing后的数据存储成npy格式 ''' logger.info('step 5: Convert To NPY') self._ConvertToNPY() def _GatherInfo(self): self.info_list = np.array(self.info_list) columns = ['uid','image_path','mask_path','lung_mask_path'] output_df = pd.DataFrame(self.info_list,columns=columns) output_df.to_csv(self.respac_df_path) self.output_df = output_df def _ResampleBatch(self): pbar = tqdm(np.arange(len(self.df))) for idx in pbar: try: self._ResampleImage(idx,self.target_spac) # print ('idx is %d, length of info list is %d'%(idx,len(self.info_list))) except Exception as err: print ('err',err) def _ResampleImage(self,idx,target_spacing): df = self.df uid,image_path,mask_path,lung_mask_path = df.uid[idx],df.image_path[idx],df.mask_path[idx],df.lung_mask_path[idx] if not os.path.exists(image_path) or not os.path.exists(mask_path) or not os.path.exists(lung_mask_path): return target_img_save_path ='%s/%s.nii.gz'%(self.img_base_path,uid) target_mask_save_path ='%s/%s.nii.gz'%(self.mask_base_path,uid) target_lung_mask_save_path ='%s/%s.nii.gz'%(self.lung_mask_base_path,uid) if os.path.exists(target_img_save_path) and os.path.exists(target_mask_save_path) and os.path.exists(target_lung_mask_save_path): self.info_list.append([uid,target_img_save_path,target_mask_save_path,target_lung_mask_save_path]) return else: image_itk,image_np = ReadITKImage(image_path) mask_itk,mask_np = ReadITKImage(mask_path) lung_mask_itk,lung_mask_np = ReadITKImage(lung_mask_path) respac_img = LinearResample(image_itk,target_spacing)[0] respac_mask = NearestResample(mask_itk,target_spacing) respac_lung_mask = NearestResample(lung_mask_itk,target_spacing) WriteMask(respac_img,path=target_img_save_path) WriteMask(respac_mask,path=target_mask_save_path) WriteMask(respac_lung_mask,path=target_lung_mask_save_path) self.info_list.append([uid,target_img_save_path,target_mask_save_path,target_lung_mask_save_path]) def _LoadUid(self,path): with open(path,'r') as f: uids = f.readlines() f.close() uids = [x.strip() for x in uids] return uids def _WriteUid(self,path,uids): with open(path,'w') as f: for uid in uids: f.write(uid+'\n') f.close() def _ExtractConnectedComponents(self): map_dict = defaultdict(list) pbar = tqdm(self.self.output_df['mask_path'].drop_duplicates().values) for mask_path in pbar: mask = sitk.GetArrayFromImage(sitk.ReadImage(mask_path)) mask = np.array(mask>0).astype(np.uint8) label_mask = label(mask) current_record = [] for region in regionprops(label_mask): area = region.area bbox = region.bbox if area>self.area_th: bbox_with_area = [int(val) for val in bbox]+[area] current_record.append(bbox_with_area) map_dict[mask_path] = current_record records = self.output_df[['uid','image_path','mask_path','lung_mask_path']].values uids = list(set(val[0] for val in records)) outputs = [] for record in records: uid,image_papth,mask_path,lung_mask_path = record bbox_list = map_dict[mask_path] for bbox in bbox_list: output_record = list(record)+list(bbox) outputs.append(output_record) output_columns = ['uid','image_path','mask_path','lung_mask_path','z_min','z_max','y_min','y_max','x_min','x_max','area'] bbox_df = pd.DataFrame(outputs,columns=output_columns) bbox_df.to_csv(self.bbox_df_path,index=False) ''' Add when to continue ''' ''' Split form ''' if os.path.exists(self.train_uid_path) and os.path.exists(self.val_uid_path) and os.path.exists(self.test_uid_path): train_uids,val_uids,test_uids = self.LoadDFs(self.train_uid_path),self.LoadDFs(self.val_uid_path),self.LoadDFs(self.test_uid_path) else: from sklearn.model_selection import train_test_split if self.valtest_ratio>0: val_flag = True train_uids,val_uids = train_test_split(uids, test_size=self.valtest_ratio, random_state=42) if self.test_ratio>0: test_flag = True val_uids,test_uids = train_test_split(val_uids, test_size=self.test_ratio, random_state=42) else: test_uids = None test_flag = False else: val_flag,test_flag = False,False val_udis,test_uids = None,None self._WriteUid(train_uids,self.train_uid_path) if val_flag: self._WriteUid(val_uids,self.val_uid_path) if test_flag: self._WriteUid(test_uids,self.test_uid_path) ''' split df based on uids & save ''' self.bbox_train = bbox_df[bbox_df['uid'].isin(train_uids)] self.bbox_train.to_csv(self.bbox_df_path.replace('.csv','_train.csv')) self.output_train = self.output_df[self.output_df['uid'].isin(train_uids)] self.output_train.to_csv(self.respac_df_path.replace('.csv','_train.csv')) logger.info('Save bbox_train to %s'%(self.bbox_df_path.replace('.csv','_train.csv'))) logger.info('Save output_train to %s'%(self.respac_df_path.replace('.csv','_train.csv'))) if val_flag: self.bbox_val = bbox_df[bbox_df['uid'].isin(val_uids)] self.bbox_val.to_csv(self.bbox_df_path.replace('.csv','_val.csv')) self.output_val = self.output_df[self.output_df['uid'].isin(val_uids)] self.output_val.to_csv(self.respac_df_path.replace('.csv','_val.csv')) logger.info('Save bbox_val to %s'%(self.bbox_df_path.replace('.csv','_val.csv'))) logger.info('Save output_val to %s'%(self.respac_df_path.replace('.csv','_val.csv'))) if test_flag: self.bbox_test = bbox_df[bbox_df['uid'].isin(test_uids)] self.bbox_test.to_csv(self.bbox_df_path.replace('.csv','_test.csv')) self.output_test = self.output_df[self.output_df['uid'].isin(test_uids)] self.output_test.to_csv(self.respac_df_path.replace('.csv','_test.csv')) logger.info('Save bbox_test to %s'%(self.bbox_df_path.replace('.csv','_test.csv'))) logger.info('Save output_test to %s'%(self.respac_df_path.replace('.csv','_test.csv'))) def _GenerateInfoWithLungMask(self,uids,base_save_path,kw,area=10000000000000): columns = [u'uid', u'image_path', u'mask_path', u'lung_mask_path', u'z_min', u'z_max', u'y_min', u'y_max', u'x_min', u'x_max', u'area'] records = [] pbar = tqdm(uids) for uid in pbar: img_path = '/fileser/CT_COVID/IMAGE/NII_1.0/%s.nii.gz'%uid lungmask_path = '/fileser/CT_COVID/LUNGMASK/NII_1.0/%s.nii.gz'%uid if not os.path.exists(img_path) or not os.path.exists(lungmask_path): continue mask = sitk.GetArrayFromImage(sitk.ReadImage(lungmask_path)) mask = np.array(mask>0).astype(np.uint8) label_mask = label(mask) if len(regionprops(label_mask))<2: continue for region in regionprops(label_mask)[:2]: bbox = region.bbox new_record = [uid,img_path,lungmask_path,lungmask_path]+list([int(val) for val in bbox])+[area] records.append(new_record) df = pd.DataFrame(records,columns=columns) df.to_csv('%s/%s_%d.csv'%(base_save_path,kw,len(uids))) def _ConvertToNPY(self): for base_path in [self.img_base_path,self.mask_base_path,self.lung_mask_base_path]: paths = sorted(glob.glob('%s/*'%base_path)) pbar = tqdm(paths) for path in pbar: target_path = path.replace('NII_1.0','1.0_NPY').replace('.nii.gz','.npy') if not os.path.exists(target_path): img = sitk.GetArrayFromImage(sitk.ReadImage(path)) print ('save path to %s'%target_path) np.save(target_path,img) def _ShutilData(self): paths = sorted(glob.glob('/fileser/CT_COVID/ALPHAMASK/*')) img_path_name = 'm_ptrRawImage.nii.gz' lungmask_path_name = 'm_seperate_lung_mask.nii.gz' target_img_base_path = '/fileser/CT_COVID/IMAGE/RAW/' target_lungmask_base_path = '/fileser/CT_COVID/LUNGMASK/RAW/' pbar = tqdm(paths) for path in pbar: lung_path = '%s/%s'%(path,lungmask_path_name) img_path = '%s/%s'%(path,img_path_name) if os.path.exists(lung_path) and os.path.exists(img_path): uid = lung_path.split('/')[-2] lungmask_des_path = '%s/%s.nii.gz'%(target_lungmask_base_path,uid) img_des_path = '%s/%s.nii.gz'%(target_img_base_path,uid) if not os.path.exists(lungmask_des_path): shutil.copy(lung_path,lungmask_des_path) if not os.path.exists(img_des_path): shutil.copy(img_path,img_des_path) uids = [x.split('/')[-1] for x in paths] print ('length of uids ',len(uids)) records = [[uid,'%s/%s.nii.gz'%(target_img_base_path,uid),'%s/%s.nii.gz'%(target_lungmask_base_path,uid),'%s/%s.nii.gz'%(target_lungmask_base_path,uid)] for uid in uids] self.df = pd.DataFrame(records,columns=['uid','image_path','mask_path','lung_mask_path']) self.df.to_csv(self.df_path,index=False) if __name__=='__main__': base_paths = [] for kw in ['IMAGE','MASK','LUNGMASK']: base_paths.append('/fileser/CT_COVID/%s/NII_1.0'%kw) df_paths = ['/fileser/zrx/COVID/label/279/279_info.csv'] respac_df_path = '/fileser/zrx/COVID/label/279/279_info_respac.csv' area_th = 50 bbox_df_path = '/fileser/zrx/COVID/label/279/279_info_respac_bbox.csv' uid_paths = [] uid_base_path = '/fileser/zrx/COVID/label/279/uids' for kw in ['train','val','test']: uid_paths.append('%s/%s_uid.txt'%(uid_base_path,kw)) health_uid_path = '/fileser/zrx/COVID/label/279/task_4069_health_uids.txt' pos_uid_path = '/fileser/zrx/COVID/label/279/task_4069_pos_uids.txt' health_save_path = '/fileser/CT_COVID/INFO/ALG_HEADER/287/NII_1.0_FORTRAIN/NEGDATA/' pos_save_path = '/fileser/CT_COVID/INFO/ALG_HEADER/287/NII_1.0_FORTRAIN/FAKEPOS/' COVID_obj = COVIDPre(base_paths,df_paths,respac_df_path,area_th,bbox_df_path,uid_paths,health_uid_path,pos_uid_path,health_save_path,pos_save_path) COVID_obj()