{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# 整体训练示例" ] }, { "cell_type": "code", "execution_count": 10, "metadata": {}, "outputs": [], "source": [ "import os\n", "import sys\n", "import six\n", "import json\n", "import shutil\n", "import datetime\n", "import numpy as np\n", "import scipy.ndimage as nd\n", "\n", "sys.path.append('./models/backbones/')\n", "sys.path.append('./models/baseLayers/')\n", "sys.path.append('./models/decoders/')\n", "sys.path.append('./models/')\n", "sys.path.append('./utils/')\n", "\n", "from data_generator import UNet_3D_Generator\n", "from ImagePreprocess import Image_preprocess,DataNormlization\n", "from parameter_utils import RectifyParameters\n", "from ResultEval import dice_coefficient_post\n", "from data_aug import *\n", "from RWconfig import LoadJson,GenerateParameters,GenerateParameters,WriteJson\n", "from ReadData import load_data\n", "from ResultPostProcess import post_process\n", "\n", "from modelBuild import build_model\n", "from model_utils import get_block\n", "from Segmentation_loss import Tversky_loss,loss_dice_coefficient_error,NoduleSegAuxLoss\n", "from train_utils import get_callbacks\n", "\n", "\n", "\n", "from keras import optimizers\n", "\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# 定义对应类实现相应功能，本类包括\n", "1. init\n", "2. LoadGroupFiles:因为数据存储成多个npy文件，多组data_path/mask_path/info_path写在train/val/test files json文件中。本函数用于将所有data_path,mask_path,info_path组合成数组\n", "3. BuildSegModel：搭建分割模型（train/inference都需要） 调用了GenerateParameters、RectifyParameters产生模型参数\n", "4. GenerateTrainParameters：产生训练相关的参数。（step/generator/callbacks等）\n", "5. CompileSegModel:compileModel\n", "6. Train：调用上面的函数，开始训练\n", "7. Inference:搭建模型、inference数据预处理、预测、后处理、产生指标（目前写定了用dice）\n", "8. LoadInferenceData：inference阶段载入数据\n", "9. LoadPretrainModel：根据config文件中定义的pretrain模型路径，载入模型。如果路径不存在，则提示" ] }, { "cell_type": "code", "execution_count": 13, "metadata": {}, "outputs": [], "source": [ "class Segmentation3D(object):\n", " def __init__(self,filename):\n", " self.filename = filename\n", " current_config = LoadJson(filename)\n", " key = current_config.keys()[0]\n", " parameters = current_config[key]\n", " self.parameters = parameters\n", " ### Generate parameters\n", " \n", " '''\n", " General parameters\n", " '''\n", " \n", " self.num_slice = parameters['num_slice']\n", " self.voxel_size = parameters['voxel_size']\n", " self.n_channels = parameters['n_channels']\n", " self.input_shape = [self.num_slice,self.voxel_size,self.voxel_size,self.n_channels]\n", " self.num_classes = parameters['num_classes']\n", " self.final_ksize = parameters['final_ksize']\n", " self.num_units = parameters['num_units']\n", " self.model_parameters_txt = os.path.join('./files',parameters['model_parameters_txt']) \n", " self.model_map_parameters_txt = os.path.join('./files',parameters['model_map_parameters_txt']) \n", " \n", " self.encoder = parameters['encoder']\n", " self.decoder = parameters['decoder']\n", " self.normalizer = parameters['normalizer']\n", " self.base_path = parameters['base_path']\n", " self.pre_model_file = parameters['pre_model_file']\n", " self.model = None\n", " \n", " \n", " if 'train' in key.lower():\n", " '''\n", " parameters for training only\n", " '''\n", " lr = parameters['initial_learning_rate']\n", "\n", " self.train_file = os.path.join('./files',parameters['train_file']) \n", " self.val_file = os.path.join('./files',parameters['val_txt'])\n", " self.load_pre_trained = parameters['load_pre_trained']\n", "\n", " self.initial_learning_rate = parameters['initial_learning_rate']\n", " self.learning_rate_drop = parameters['learning_rate_drop']\n", " self.learning_rate_patience = parameters['learning_rate_patience']\n", " self.early_stop = parameters['early_stop']\n", " self.base_path = parameters['base_path']\n", " self.n_epochs = parameters['n_epochs']\n", " self.aug = parameters['aug']\n", " self.prefix = parameters['prefix']\n", " self.train_batch_size = parameters['train_batch_size']\n", " self.val_batch_size = parameters['val_batch_size']\n", " self.shuffle = parameters['shuffle']\n", " \n", " self.trainfiles = self.LoadGroupFiles(self.train_file)\n", " self.valfiles = self.LoadGroupFiles(self.val_file)\n", " \n", " else:\n", " self.test_file = os.path.join('./files',parameters['test_file'])\n", " self.testfiles = self.LoadGroupFiles(self.test_file)\n", " \n", " \n", " def LoadGroupFiles(self,filename):\n", " result = LoadJson(filename)\n", " data_paths = [record['data_path'] for record in result]\n", " mask_paths = [record['mask_path'] for record in result]\n", " info_paths = [record['info_path'] for record in result]\n", " return {\n", " 'data_paths':data_paths,\n", " 'mask_paths':mask_paths,\n", " 'info_paths':info_paths\n", " }\n", " \n", " def BuildSegModel(self):\n", " '''\n", " Generate encoder/decoder parameters with model_parameters_txt and model name\n", " '''\n", "\n", " \n", " self.encoder_parameter_list = GenerateParameters(self.model_parameters_txt,self.model_map_parameters_txt,model=self.encoder)\n", " self.decoder_parameter_list = GenerateParameters(self.model_parameters_txt,self.model_map_parameters_txt,model=self.decoder)\n", "\n", "\n", " '''\n", " Rectify parameters \n", " '''\n", " self.encoder_parameter_list = RectifyParameters(self.encoder_parameter_list,self.parameters)\n", " self.decoder_parameter_list = RectifyParameters(self.decoder_parameter_list,self.parameters)\n", " \n", "\n", " '''\n", " Build encoder and decoder\n", " '''\n", "\n", " self.encoder_result = build_model(self.encoder,input_shape=self.input_shape,parameter_list=[self.encoder_parameter_list])\n", " self.model = build_model(self.decoder,input_shape=self.input_shape,parameter_list=[self.encoder_result,self.decoder_parameter_list])\n", "\n", "# self.model.summary()\n", " \n", " def CompileSegModel(self):\n", " \n", "\n", " adam=optimizers.Adam(lr=self.initial_learning_rate,beta_1=0.9,beta_2=0.999,epsilon=1e-08)\n", " \n", " ## todo: change loss to parameters in config file\n", " seg_loss = loss_dice_coefficient_error\n", " cls_loss = NoduleSegAuxLoss\n", "\n", " final_loss = 0\n", "\n", " \n", " if self.decoder_parameter_list['deep_supervision']:\n", " final_loss = [seg_loss for _ in range(3)]\n", " elif self.decoder_parameter_list['ACF_choice']:\n", " final_loss = [seg_loss for _ in range(2)]\n", " else:\n", " final_loss = [seg_loss]\n", " \n", " if self.decoder_parameter_list['aux_task']:\n", " final_loss.append(cls_loss)\n", "\n", " if len(final_loss)==1:\n", " final_loss = final_loss[0]\n", "\n", " self.model.compile(optimizer= adam, loss=final_loss, metrics=['accuracy'])\n", " \n", " def GenerateTrainParameters(self):\n", " self.train_step_per_epoch = int(np.sum([np.load(path).shape[0] for path in self.trainfiles['info_paths']] )/self.train_batch_size)\n", " self.valid_step_per_epoch = int(np.sum([np.load(path).shape[0] for path in self.valfiles['info_paths']] )/self.val_batch_size )\n", " \n", " \n", "# self.train_step_per_epoch = 2\n", "# self.valid_step_per_epoch = 2\n", " \n", " self.train_generator = UNet_3D_Generator(self.trainfiles['data_paths'],self.trainfiles['mask_paths'],\n", " self.trainfiles['info_paths'],input_size=self.input_shape,\n", " batch_size = self.train_batch_size,aug=self.aug,\n", " HU_min = self.normalizer[0],HU_max=self.normalizer[1],crop_choice=True,\n", " config_dict=self.decoder_parameter_list)\n", " \n", " self.val_generator = UNet_3D_Generator(self.valfiles['data_paths'],self.valfiles['mask_paths'],\n", " self.valfiles['info_paths'],input_size=self.input_shape,\n", " batch_size = self.val_batch_size,aug=False,\n", " HU_min = self.normalizer[0],HU_max=self.normalizer[1],crop_choice=True,\n", " config_dict=self.decoder_parameter_list)\n", " \n", " current_DT = datetime.datetime.now()\n", " current_day = '%02d%02d' % (current_DT.month,current_DT.day)\n", "\n", " self.current_base_path = '%s/%s_%s/'%(self.base_path,self.prefix,current_day)\n", " if not os.path.exists(self.current_base_path):\n", " os.mkdir(self.current_base_path)\n", " \n", " \n", " for target_filename in [self.filename,self.model_parameters_txt,self.model_map_parameters_txt,self.train_file,self.val_file]:\n", " shutil.copyfile(self.filename,os.path.join(self.current_base_path,os.path.basename(target_filename)))\n", " \n", " \n", " self.filepath = str('%s/Train-{epoch:02d}-{val_loss:.5f}.hdf5'%(self.current_base_path))\n", " self.filepath = str(self.filepath)\n", " print ('self.filepath is ',type(self.filepath))\n", " self.callbacks_list = get_callbacks(self.filepath,initial_learning_rate = self.initial_learning_rate,\n", " learning_rate_drop = self.learning_rate_drop,learning_rate_epochs = None,\n", " learning_rate_patience = self.learning_rate_patience,\n", " early_stopping_patience = self.early_stop)\n", " \n", " \n", " def Train(self):\n", " \n", " self.BuildSegModel()\n", " print ('After building model')\n", " self.CompileSegModel()\n", " print ('After compiling model')\n", " self.LoadPretrainModel()\n", " print ('After loading pretrain model')\n", " self.GenerateTrainParameters()\n", " \n", " self.history_callbacks = self.model.fit_generator(self.train_generator,steps_per_epoch=self.train_step_per_epoch,epochs=self.n_epochs,\n", " callbacks=self.callbacks_list,shuffle=self.shuffle,\n", " validation_data=self.val_generator,validation_steps=self.valid_step_per_epoch)\n", " \n", " WriteJson(os.path.join(self.current_base_path,'history.json'),self.history_callbacks.history)\n", " \n", " \n", " def Inference(self):\n", " \n", " self.BuildSegModel()\n", " self.LoadPretrainModel()\n", " \n", " print ('Before Loading test data')\n", " self.LoadInferenceData()\n", " print ('After loading test data')\n", " \n", " deltas = [0,0,0]\n", " self.test_data = np.squeeze(np.asarray([Image_preprocess(single_data,deltas,ratio=1) for single_data in self.test_data]))\n", " self.test_mask = np.squeeze(np.asarray([Image_preprocess(single_mask,deltas,ratio=1) for single_mask in self.test_mask]))\n", "\n", " test_data_cp = DataNormlization(self.test_data,self.test_info,HU_min=self.normalizer[0],HU_max=self.normalizer[1],mode='HU')\n", " self.test_data = test_data_cp[...,np.newaxis]\n", " self.test_mask = self.test_mask[...,np.newaxis]\n", " \n", " predict_result = self.model.predict(self.test_data,batch_size=1)\n", " \n", " print ('After making prediction')\n", " inference_result_dict = {\n", " 'model_weights':self.pre_model_file,\n", " 'test_data':self.testfiles,\n", " 'result':{}\n", " \n", " }\n", " \n", " print ('length of predict_result is',len(predict_result))\n", " for branch_idx in range(min(3,len(predict_result))):\n", " print ('branch_idx is ',branch_idx)\n", " final_result = predict_result[branch_idx]>0.05\n", " postprocess_result = np.asarray([np.squeeze(post_process(mask_image,spacing_z=1.0,inp_shape=self.input_shape,debug=False)[2]) for mask_image in final_result])\n", " result_dice_post = [dice_coefficient_post(y_true,y_pred) for y_true,y_pred in zip(np.squeeze(self.test_mask),np.squeeze(postprocess_result))]\n", " result_dice_infer = [dice_coefficient_post(y_true,y_pred) for y_true,y_pred in zip(np.squeeze(self.test_mask),np.squeeze(final_result))]\n", " inference_result_dict['result']['branch_%02d_post_dice'] = [np.mean(result_dice_post),np.std(result_dice_post)]\n", " inference_result_dict['result']['branch_%02d_infer_dice'] = [np.mean(result_dice_infer),np.std(result_dice_infer)]\n", "\n", " print ('dice result(label and mask after threshold): branch %02d is %s'%(branch_idx,str(inference_result_dict['result']['branch_%02d_infer_dice'])))\n", " print ('dice result(label and mask after postprocess): branch %02d is %s'%(branch_idx,str(inference_result_dict['result']['branch_%02d_infer_dice'])))\n", " \n", " print ('Before writing json file to %s'%os.path.join(os.path.dirname(self.pre_model_file),'inference_result.json'))\n", " WriteJson(os.path.join(os.path.dirname(self.pre_model_file),'inference_result.json'),inference_result_dict)\n", " print ('After writing json file')\n", " \n", " \n", " def LoadInferenceData(self):\n", " self.test_data_list = load_data([self.testfiles['data_paths'],self.testfiles['mask_paths'],\n", " self.testfiles['info_paths']])\n", " self.test_data,self.test_mask,self.test_info = self.test_data_list\n", " \n", " \n", " \n", " def LoadPretrainModel(self):\n", " if os.path.exists(self.pre_model_file) and self.model:\n", " self.model.load_weights(self.pre_model_file)\n", " elif not os.path.exists(self.pre_model_file):\n", " print ('Pretrain model path does not exist')\n", " else:\n", " print ('model has not been defined')" ] }, { "cell_type": "code", "execution_count": 16, "metadata": {}, "outputs": [], "source": [ "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"6\"" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 参数\n", "模型相关基本参数定义在/files/train_config.json中。其他相关参数文件也定义在该文件中。\n", "可以通过修改link的文件/模型名字/其他参数，调整模型、训练过程" ] }, { "cell_type": "code", "execution_count": 14, "metadata": {}, "outputs": [], "source": [ "train_file_path = './files/train_config.json'" ] }, { "cell_type": "code", "execution_count": 15, "metadata": {}, "outputs": [], "source": [ "train_obj = Segmentation3D(train_file_path)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "Python 2", "language": "python", "name": "python2" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 2 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython2", "version": "2.7.15" } }, "nbformat": 4, "nbformat_minor": 2 }