import os import sys import random import numpy as np import pandas as pd import scipy.ndimage as nd from ReadData import load_data,load_single_data from data_aug import data_augmentation,random_add from tensorflow.keras.utils import to_categorical from ImagePreprocess import Data_preprocess,CalculatePatchSizeWithDiameter # densityTypeConvertMap = { # 1:1, # 2:2, # 3:3, # 4:4, # 5:1, # 6:4 # } def UNet_3D_Generator(data_paths,mask_paths,info_paths,**kwargs): ''' Define generator for Unet @param:data_paths,mask_paths,info_paths. Pahts of data/mask/info, they should have same shape @param:clip_choice. Whether to do clip on intensity or not @param:HU_min,HU_max. If clip_choice is True then this range will be used to do clip @param:aug choice @param:batch_size, default value is 32 ''' clip_choice = kwargs.get('clip_choice',True) HU_min,HU_max = kwargs.get('HU_min'),kwargs.get('HU_max') aug = kwargs.get('aug',True) nonfix_crop = kwargs.get('nonfix_crop',False) stride_ratio = kwargs.get('stride_ratio',8) densityTypeIndex = kwargs.get('densityTypeIndex',-1) batch_size = kwargs.get('batch_size',4) input_size = kwargs.get('input_size') diameter_enlarge_ratio = kwargs.get('diameter_enlarge_ratio',1.5) direction_z_ratio = kwargs.get('direction_z_ratio',1.0) min_patch_size = stride_ratio config_dict = kwargs.get('config_dict') deep_supervision = config_dict['deep_supervision'] aux_task = config_dict['aux_task'] num_classes = config_dict['num_classes'] densityTypeConvertMap = config_dict['densityTypeConvertMap'] print ("densityTypeIndex is",densityTypeIndex) batch_image,batch_label,batch_densityType = 0,0,0 result_list = load_data([data_paths,mask_paths,info_paths],aux_task=aux_task) data,mask,info = result_list print ('After loading data',data.shape,mask.shape,info.shape) test_idx = 0 # while test_idx<2: # test_idx +=1 while True: sample_indices = np.arange(data.shape[0]) random.shuffle(sample_indices) for inner_idx in range(0,len(sample_indices)): current_sample_idx = sample_indices[inner_idx] current_image,current_mask = data[current_sample_idx],mask[current_sample_idx] current_info = info[current_sample_idx] if aux_task: densityType = int(float(current_info[densityTypeIndex])) if densityType==5: densityType = 1 elif densityType == 6: densityType = 4 if densityType>=7: continue densityType -=1 current_densityType = to_categorical(densityType,num_classes=4) current_densityType = np.reshape(np.array(current_densityType),(1,4)) else: current_densityType = 0 # print ('shape of current_densityType',current_densityType.shape) #### Crop patch if nonfix_crop: patch_shape = CalculatePatchSizeWithDiameter(current_mask,stride_ratio=stride_ratio,diameter_enlarge_ratio=diameter_enlarge_ratio,direction_z_ratio=direction_z_ratio) else: patch_shape = input_size final_image,final_mask = Data_preprocess(current_image,current_mask,aug=aug,input_shape=patch_shape,shift_ratio=1.0) if aug: if np.random.choice([0,1]): final_image = np.swapaxes(final_image,2,3) final_mask = np.swapaxes(final_mask,2,3) if np.random.choice([0,1]): theta = np.random.uniform(-5, 5) cval = HU_min if HU_min else np.amin(final_image) final_image = nd.interpolation.rotate(final_image,theta,axes=(2,3),reshape=False,cval=cval) final_mask = nd.interpolation.rotate(final_mask,theta,axes=(2,3),reshape=False) final_mask = final_mask>0 final_image, final_mask = data_augmentation(np.squeeze(final_image), np.squeeze(final_mask)) if clip_choice: final_image = np.clip(final_image,HU_min,HU_max) final_image = (final_image - HU_min)/float(HU_max - HU_min) if aug: final_image = random_add(final_image) final_image = np.squeeze(final_image)[np.newaxis,...,np.newaxis] final_mask = np.squeeze(final_mask)[np.newaxis,...,np.newaxis] if type(batch_image) == int: batch_image = np.asarray(final_image) batch_mask = np.asarray(final_mask) batch_densityType = np.array([current_densityType]) else: batch_image = np.concatenate([batch_image,final_image],axis=0) batch_mask = np.concatenate([batch_mask,final_mask],axis=0) batch_densityType = np.concatenate([batch_densityType,[current_densityType]],axis=0) if batch_mask.shape[0] == batch_size: # batch_densityType = np.squeeze(batch_densityType) # print ('batch_densityType shape',batch_densityType.shape) if deep_supervision: batch_mask = [batch_mask for _ in range(3)] if aux_task: batch_densityType = np.reshape(batch_densityType,(batch_size,4)) if deep_supervision: batch_mask.append(batch_densityType) else: batch_mask = [batch_mask,batch_densityType] yield batch_image,batch_mask batch_image,batch_label,batch_densityType = 0,0,0 def DataStat(info_paths,diameter_ths,**kwargs): infos = load_single_data(info_paths,aux_task=True) densityTypeIndex = kwargs.get('densityTypeIndex',-4) from collections import defaultdict diameter_index = kwargs.get('diameter_index',4) uids = set([x[0] for x in infos]) print ('Number of records in total is %d in %d series'%(len(infos),len(uids))) density_map_dict = { 1:'solid', 2:'mix', 3:'ggo', 4:'calc', 5:'solid', 6:'clac', 7:'p_patch' } keys = ['solid','mix','ggo','calc','p_patch'] for diam_th_vals in diameter_ths: diameter_min,diameter_max = diam_th_vals density_stat_map = defaultdict(int) count = 0 for record in infos: diameter = float(record[diameter_index]) densityType = int(float(record[densityTypeIndex])) if densityType not in density_map_dict.keys(): continue if diameter>=diameter_min and diameter