# -*- coding:utf-8 -*- import os import sys import six import glob import json import random import shutil import datetime import numpy as np import pandas as pd import scipy.ndimage as nd from tqdm import tqdm import datetime import logging import torch import torch.nn as nn from torch import optim from matplotlib import pyplot as plt from torch.utils.data import DataLoader from torch.nn.utils import clip_grad_norm_ import copy from tensorboardX import SummaryWriter sys.path.append('./utils') sys.path.append('./model') from .utils.RWconfig import LoadJson from .utils.ReadData import load_data, load_single_data from .model.modelBuild import build_model from .model.BasicModules import weight_init from .utils.DataGen import SurDataSet, InferDataSet, TestDataSet from .utils.OnlineEval import eval_net, eval_netSE, eval_netFSE from .utils.OfflineEval import CalculateClsScore, CalculateAuc, CalculateClsScoreByTh from .utils.loss_func import * from .utils.gradcam import GradCam, GuidedBackpropReLUModel from .utils.ImagePro import ShowEdges, HFFilter from .utils.TripleLoss import TripletFocalLoss from .utils.TripleBin import TripletBinFocal from .utils.GradientBoost import GradientBoostFocal from k8s_utils import CMetricsWriter import torch.distributed as dist from sklearn.metrics import precision_score, recall_score import copy from pydicom import dicomio logger = logging.getLogger() fh = logging.FileHandler('Train.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) import os import numpy as np import SimpleITK as sitk import cv2 import imageio from tqdm import tqdm import sys from PIL import Image, ImageDraw, ImageFont def norm(image, hu_min=-1000.0, hu_max=600.0): image = (np.clip(image.astype(np.float32), hu_min, hu_max) - hu_min) / float(hu_max - hu_min) return image * 255. def img2gif(tmp_path, raw_path): target_path = os.path.join(os.path.dirname(tmp_path), 'save_images') if not os.path.exists(target_path): os.makedirs(target_path) tmp_files = os.listdir(tmp_path) for file in tmp_files: print('deal with {}'.format(file)) npy_info = np.load(os.path.join(tmp_path, file, 'all_rpns.npy')) raw_img = sitk.ReadImage(os.path.join(raw_path, file + '.nii.gz')) raw_img = sitk.GetArrayFromImage(raw_img) all_z = [int(i) for i in list(npy_info[:,2])] min_z, max_z = min(all_z)-3, max(all_z)+3 used_slices = [] for z in all_z: temp = [] for i in range(z-3, z+3): temp.append(i) used_slices.append(temp) all_images = [] for it_z in range(min_z, max_z): u_idx = [] for key, value in enumerate(used_slices): if it_z in value: u_idx.append(key) u_idx = sorted(u_idx) print('have {} box'.format(len(u_idx))) if len(u_idx) == 0: all_images.append(cv2.cvtColor(norm(raw_img[it_z,:,:]), cv2.COLOR_GRAY2BGR)) else: worked_slice = norm(raw_img[it_z,:,:]) worked_slice = cv2.cvtColor(worked_slice, cv2.COLOR_GRAY2BGR) for idx in u_idx: z,y,x,dim_z,dim_y,dim_x = npy_info[idx][1:] start_y = max(0, y - dim_y / 2) end_y = min(worked_slice.shape[0], y + dim_y / 2) start_z = max(0, z - dim_z / 2) end_z = min(worked_slice.shape[1], z + dim_z / 2) cv2.rectangle(worked_slice, (int(start_y), int(start_z)), (int(end_y), int(end_z)), (0,0,255), 2) all_images.append(worked_slice) for idx, img in enumerate(all_images): save_path = os.path.join(target_path, file) if not os.path.exists(save_path): os.makedirs(save_path) cv2.imwrite(os.path.join(save_path, '{:03d}.png'.format(idx)), img) output_gif_path = os.path.join(os.path.dirname(target_path), 'save_gifs') if not os.path.exists(output_gif_path): os.makedirs(output_gif_path) for file in os.listdir(target_path): files = [] for f in os.listdir(os.path.join(target_path, file)): files.append(f) files.sort(key=lambda x: x[:-4]) frames = [] for i in tqdm(range(len(files))): im = imageio.imread(os.path.join(target_path, file) + '/' + files[i]) frames.append(im) fps = 24.0 imageio.mimsave(os.path.join(output_gif_path, file+'.gif'), frames, 'GIF', duration=1/fps) 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 Metric(object): def __init__(self, name): self.name = name self.sum = torch.tensor(0.) self.n = torch.tensor(0.) def update(self, val): val = val.clone().detach() output_tensors = [val.clone().to(val) for _ in range(torch.distributed.get_world_size())] dist.all_gather(output_tensors, val) self.sum += torch.sum(torch.stack(output_tensors)).cpu().detach() self.n += 1 @property def avg(self): return self.sum / self.n def accuracy(outputs, labels): _, preds = torch.max(outputs, dim=1) # get the accuracy of number preds correctly # if math.isnan(torch.sum(preds == labels).item() / len(preds)): # print('preds:', preds) # print('check: ', preds == labels) return torch.tensor(torch.sum(preds == labels).item() / len(preds)) def precision(outputs, labels): _, preds = torch.max(outputs, dim=1) # get the accuracy of number preds correctly return torch.tensor(precision_score(labels.cpu().numpy(), preds.cpu().numpy(), average='macro')) def recall(outputs, labels): _, preds = torch.max(outputs, dim=1) # get the accuracy of number preds correctly return torch.tensor(recall_score(labels.cpu().numpy(), preds.cpu().numpy(), average='macro')) class Classification3D(object): def __init__(self, cfg, mode, log_msg): self.cfg = cfg self.mode = mode self.log_msg = log_msg if mode == 'training': self.mode_cfg = cfg.training else: logger.info('load testing parameters...') self.mode_cfg = cfg.testing if mode == 'training': self.metrics_writer = CMetricsWriter( filename=cfg.training['train_metrics'], headernames=['epoch', 'accuracy', 'f1_score', 'precision', 'recall', 'lr', 'lossTotal', 'performanceAccuracy'] ) self.basic_file_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) self.n_input_channel = 1 self._clsInfoDecode() if mode == 'training': # self.traindata_npy = self.mode_cfg['train_data_path'] # self.traininfo_npy = self.mode_cfg['train_info_path'] # self.valdata_npy = self.mode_cfg['val_data_path'] # self.valinfo_npy = self.mode_cfg['val_info_path'] # self.testdata_npy = [] # self.testinfo_npy = [] self.traindata = self.mode_cfg['train_data_path'] self.valdata = self.mode_cfg['val_data_path'] self.testdata = [] print(self.traindata, self.valdata) else: # self.traindata_npy = [] # self.traininfo_npy = [] # self.valdata_npy = [] # self.valinfo_npy = [] # self.testdata_npy = self.mode_cfg['test_data_path'] # self.testinfo_npy = self.mode_cfg['test_info_path'] self.traindata = [] self.valdata = [] self.testdata = self.mode_cfg['test_data_path'] self.ParametersDecode() self.mode_cfg['model_mode'] = 'basic' # self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') def __call__(self): ''' 训练/infer流程 1. buildModel&compileModel 2. 载入数据 3. 根据选择模式，进行训练or inference ''' # self._LoadData(self.mode) if self.mode == 'training': self._buildModel() self._compileModel() self._trainModel() elif self.mode == 'infer': self._buildModel() self._inferModel() else: self._buildModel() self._testModel() def _clsInfoDecode(self): self.cls_map_dict = eval(self.mode_cfg['cls_map_dict']) self.ori_cls_types = set([self.cls_map_dict[key] for key in self.cls_map_dict.keys()]) self.input_shape = [self.mode_cfg['patch_height'], self.mode_cfg['patch_height'], self.mode_cfg['patch_depth']] self.cls_name_map_dict = eval(self.mode_cfg['cls_name_map_dict']) # self.save_path = self.mode_cfg['save_path'] # if (not os.path.exists(self.save_path)) and self.mode == 'training': # os.makedirs(self.save_path) if type(self.mode_cfg['cls_label_index']) == list: self.mode_cfg['num_task'] = len(self.mode_cfg['cls_label_index']) else: self.mode_cfg['num_task'] = 1 def _buildModel(self): ''' buildModel 流程 1. 读取模型参数并修改部分 1. num_task：有几个要分类的任务 2. num_class:有几类（目前多个要分类的任务类别数目必须相同） 3. arcface_norm_choice：arcface是否进行归一化 4. easy_margin: arcface相关参数 如果设置为True,只要cos_theta大于0，就设置为乘性压缩的角度 否则，要cos_theta大于th，设置为乘性压缩角度，cos_theta小于th的情况，设置为减性压缩角度 5. freeze_blocks:不训练的block ''' cfg = self.model_params cfg['num_task'] = self.mode_cfg['num_task'] cfg['num_class'] = self.mode_cfg['num_class'] cfg['arcface_norm_choice'] = self.mode_cfg.get('norm_choice', False) cfg['easy_margin'] = self.mode_cfg.get('easy_margin', True) cfg['freeze_blocks'] = eval(self.mode_cfg.get('freeze_blocks', "[]")) if cfg['num_class'] == 2: cfg['num_class'] = 1 self.n_input_channel = cfg['n_input_channel'] if 'n_input_channel' in cfg.keys() else self.n_input_channel self.net = build_model(self.mode_cfg['model_name'], cfg) # print('model:',self.net) self.net_init = copy.deepcopy(self.net) if self.cfg.pretrain_msg: print('pretrain weight:', self.cfg.pretrain_msg) # if self.cfg.pretrain_msg.split('.')[-1] == 'pt': try: logger.info("load model from %s" % self.cfg.pretrain_msg) need_keys = [key for key,_ in self.net.named_parameters()] ## print('get need_keys') logger.info('get need_keys') checkpoints = torch.load(self.cfg.pretrain_msg, map_location=torch.device('cpu')) print('get checkpoints') # print(checkpoints.keys()) logger.info('get checkpoints') if self.cfg.mode == 'testing': new_checkpoints = {key[7:]:value for key, value in checkpoints.items() if key[7:] in need_keys} else: new_checkpoints = {key:value for key, value in checkpoints.items() if key in need_keys} print('Start load ...') # print(list(new_checkpoints.keys())[0], new_checkpoints[list(new_checkpoints.keys())[0]]) logger.info('Start load ...') self.net_init.load_state_dict(new_checkpoints) print('Have already load pretrain_msg !') logger.info('Have already load pretrain_msg !') except: self.net_init = torch.jit.load(self.cfg.pretrain_msg) print('load pt!!!!!') # for a,b in self.net.named_parameters(): # break # logger.info(b) # print(b) ## DDP logger.info(self.mode) if self.mode == 'training': print('start') logger.info('start') self.net = self.net.cuda() logger.info('cuda') logger.info(self.cfg.local_rank) logger.info(torch.cuda.device_count()) self.net = torch.nn.parallel.DistributedDataParallel(self.net, device_ids=[int(self.cfg.local_rank)]) # broadcast_buffers=False, # find_unused_parameters=True) logger.info('have ddp') else: self.net = self.net_init def _compileModel(self): self.optimizer = optim.Adam(filter(lambda p: p.requires_grad, self.net.parameters()), lr=self.cfg.lr) self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, patience=self.mode_cfg['patience'], factor=0.5, verbose=True) if self.mode_cfg['num_class'] > 2: self.criterion = nn.CrossEntropyLoss() else: self.criterion = nn.BCEWithLogitsLoss() # conf_paths = [self.filename, self.model_param_name] # if self.mode == 'training': # conf_paths += [self.trainfile_conf_path, self.valfile_conf_path] # for conf_path in conf_paths: # filename = conf_path.split('/')[-1] # target_save_path = os.path.join(self.save_path, filename) # shutil.copy(conf_path, target_save_path) # else: # conf_paths += [self.test_conf_path] def _DecodeGeneratorParam(self): high_freq_filter = self.mode_cfg['high_freq_filter'] if 'high_freq_filter' in self.mode_cfg.keys() else False rescale_intensity = self.mode_cfg['rescale_intensity'] if 'rescale_intensity' in self.mode_cfg.keys() else False nonfix_crop = self.mode_cfg['nonfix_crop'] if 'nonfix_crop' in self.mode_cfg.keys() else False aug = True if self.mode == 'training' else False model_mode = self.mode_cfg['model_mode'] if 'model_mode' in self.mode_cfg.keys() else 'basic' reg_map_dict = eval(self.mode_cfg['reg_map_dict']) if 'reg_map_dict' in self.mode_cfg.keys() else None self.regression_choice = self.mode_cfg['regression_choice'] if 'regression_choice' in self.mode_cfg.keys() else False self.generator_parameters = { 'cls_index': self.mode_cfg['cls_label_index'], 'cls_map_dict': self.cls_map_dict, 'aug': aug, 'nonfix_crop': nonfix_crop, 'diameter_index': self.mode_cfg['diameter_index'], 'diameter_max_th': self.mode_cfg['diameter_max_th'], 'diameter_min_th': self.mode_cfg['diameter_min_th'], 'num_class': self.mode_cfg['num_class'], 'HU_min': self.mode_cfg['HU_min'], 'HU_max': self.mode_cfg['HU_max'], 'cls_types': self.ori_cls_types, 'n_input_channel': self.n_input_channel, 'high_freq_filter': high_freq_filter, 'rescale_intensity': rescale_intensity, "input_shape": self.input_shape, "model_mode": model_mode, "regression_choice": self.regression_choice, "reg_map_dict": reg_map_dict, 'cls_balance': self.mode_cfg.get('cls_balance', False) } def _DecodeDataParam(self, filename): paths = LoadJson(filename) data_paths = [record['data_path'] for record in paths] info_paths = [record['info_path'] for record in paths] return {'data_paths': data_paths, 'info_paths': info_paths} # def _LoadData(self, mode): # ''' # 载入npy格式的数据 # ''' # if mode == 'training': # self.train_data, self.train_info = load_data([self.traindata_npy, self.traininfo_npy]) # self.val_data, self.val_info = load_data([self.valdata_npy, self.valinfo_npy]) # else: # print('testdata_npy and info: ',self.testdata_npy, self.testinfo_npy) # self.test_data, self.test_info = load_data([self.testdata_npy, self.testinfo_npy]) def _GetLoss(self, prob_diff, gamma): model_mode = self.mode_cfg['model_mode'] if 'model_mode' in self.mode_cfg.keys() else 'basic' loss_choice = self.mode_cfg['loss_func'] if 'loss_func' in self.mode_cfg.keys() else 'focal' lambda_val = self.mode_cfg['lambda_val'] if 'lambda_val' in self.mode_cfg.keys() else 1 margin_val = self.mode_cfg['margin_val'] if 'margin_val' in self.mode_cfg.keys() else 0.0 print('model_mode is ', model_mode) ''' gradientBoost的参数，是否使用gradient_boost/选择保留几个类别 ''' gradientBoost_choice = self.mode_cfg['gradientBoost'] if 'gradientBoost' in self.mode_cfg.keys() else False gradientBoost_k_val = self.mode_cfg['gradientBoost_k'] if 'gradientBoost_k' in self.mode_cfg.keys() else 1 ''' 基于focal loss基本的weights(1/该类样本量->归一化)，是否要对某一类的weights进行修改，如果否就设置为None，否则设置为对应的loss ''' alpha_multi_ratio = eval(self.mode_cfg['alpha_multi_ratio']) if 'alpha_multi_ratio' in self.mode_cfg.keys() else None ''' 得到regression_task的loss ''' self.regression_loss = torch.nn.MSELoss() if self.regression_choice else None self.regression_loss = torch.nn.SmoothL1Loss() if self.mode_cfg.get('regression_loss', 'MSE') == 'SmoothL1Loss' else self.regression_loss self.convert_softmax = True print("self.mode_cfg['num_task']: ", self.mode_cfg['num_task']) if self.mode_cfg['focal_loss'] and self.mode_cfg['num_task'] == 1: # ratios = self.train_dataset._GenerateRatios() ratios = [0.25, 0.25, 0.25, 0.25] if self.mode_cfg['num_class'] <= 2: alpha = ratios[0] if model_mode == 'deux': self.criterion = SeLossWithBCEFocalLoss(alpha=alpha, gamma=gamma, lambda_val=lambda_val, prob_diff=prob_diff) elif model_mode == 'triple': self.criterion = TripletBinFocal(self.cfg['num_class'], alpha=alpha, gamma=gamma, lambda_val=lambda_val, margin=margin_val) else: self.criterion = BCEFocalLoss(alpha=alpha, gamma=gamma) else: ''' 多类情况下的loss ''' alpha = ratios if model_mode == 'deux': # ratios = [0.25, 0.25, 0.25, 0.25] print('parameters-alpha: ', ratios) self.criterion = SE_FocalLoss(self.mode_cfg['num_class'], alpha=ratios, gamma=gamma, lambda_val=lambda_val, gradient_boost=gradientBoost_choice, alpha_multi_ratio=alpha_multi_ratio, convert_softmax=self.convert_softmax, k=gradientBoost_k_val, margin=margin_val, new_se_loss=self.mode_cfg.get('new_se_loss', False)) print('self.alpha: ', self.criterion.alpha) elif model_mode == 'triple': self.criterion = TripletFocalLoss(self.mode_cfg['num_class'], alpha=ratios, gamma=gamma, lambda_val=lambda_val, margin=margin_val) else: if gradientBoost_choice: self.criterion = GradientBoostFocal(self.mode_cfg['num_class'], alpha=ratios, gamma=gamma) else: # ratios = [0.25, 0.25, 0.25, 0.25] print('parameters-alpha: ', ratios) self.criterion = FocalLoss(self.mode_cfg['num_class'], alpha=ratios, gamma=gamma, alpha_multi_ratio=alpha_multi_ratio, convert_softmax=self.convert_softmax) print('self.alpha: ', self.criterion.alpha) else: alpha = None # return alpha def _trainOnBatch(self, epoch, batch, label_dtype): imgs, labels, label_diff = batch[0] label_diff = [single_label_diff.type(label_dtype).cuda() for single_label_diff in label_diff] imgs = [single_set_img.type(torch.float32).cuda() for single_set_img in imgs] labels = [single_label.type(label_dtype).cuda() for single_label in labels] prediction = self.net(imgs, labels=labels) # 计算训练输出 whole_pred = torch.cat(prediction, dim=0) whole_labels = torch.cat(labels, dim=0) acc = accuracy(whole_pred, whole_labels) out_precision = precision(whole_pred, whole_labels) out_recall = recall(whole_pred, whole_labels) out_f1 = 2*out_precision*out_recall / (out_precision + out_recall) ''' 网络预测的前len(labels)是logits。len(labels)->2*len(labels)部分如果存在就是features ''' logits = [pred.squeeze(-1) for pred in prediction[:len(labels)]] features = [pred.squeeze(-1) for pred in prediction[len(labels):len(labels) * 2]] cons = prediction[-1].squeeze(-1) loss_input = [logits] + [labels] + [label_diff] loss = self.criterion(loss_input) # self.metrics_writer.append_one_line( # [epoch, acc.item(), out_f1.item(), out_precision.item(), out_recall.item(), self.optimizer.param_groups[0]['lr'], loss.item(), acc.item()] # ) return loss, acc, out_f1, out_precision, out_recall def _trainOnBatchWithReg(self, epoch, batch, label_dtype, device): regression_ratio = self.mode_cfg.get('regression_ratio', 0.05) imgs, labels, label_diff = batch[0] label_diff = [single_label_diff.cuda().to(device=device, dtype=label_dtype) for single_label_diff in label_diff] imgs = [single_set_img.cuda().to(device=device, dtype=torch.float32) for single_set_img in imgs] cls_labels = [single_label[0].cuda().to(device=device, dtype=label_dtype) for single_label in labels] reg_labels = [single_label[1].cuda().to(device=device, dtype=torch.float32) for single_label in labels] prediction = self.net(imgs, labels=cls_labels) logits = [pred[0].squeeze(-1) for pred in prediction[:len(labels)]] reg_outputs = [pred[1].squeeze(-1) for pred in prediction[:len(labels)]] loss_input = [logits] + [cls_labels] + [label_diff] loss = self.criterion(loss_input) reg_loss = 0 for sample_idx in range(len(reg_outputs)): reg_loss_val = self.regression_loss(reg_outputs[sample_idx], reg_labels[sample_idx]) if sample_idx == 0: reg_loss = reg_loss_val else: reg_loss += reg_loss_val loss = regression_ratio * reg_loss + loss # 计算训练输出 whole_pred = torch.cat(prediction, dim=0) whole_labels = torch.cat(labels, dim=0) acc = accuracy(whole_pred, whole_labels) out_precision = precision(whole_pred, whole_labels) out_recall = recall(whole_pred, whole_labels) out_f1 = 2*out_precision*out_recall / (out_precision + out_recall) # self.metrics_writer.append_one_line( # [epoch, acc, out_f1, out_precision, out_recall, self.optimizer.param_groups[0]['lr'], loss.item()] # ) return def _trainModel(self): ########## Add parameters ''' 设置训练的device ''' # device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') ''' 如果pretrain_weight_path不为空，则load进来 ''' # if self.cfg.pretrain_msg and self.cfg.rank == 0: # logger.info("load model from %s" % self.cfg.pretrain_msg) # need_keys = [key for key,_ in self.net.named_parameters()] # checkpoints = torch.load(self.cfg.pretrain_msg) # new_checkpoints = {key:value for key, value in checkpoints.items() if key in need_keys} # self.net.load_state_dict(new_checkpoints) # for name, value in self.net.named_parameters(): # print(name, value.shape) ''' 读取关于loss和数据处理的参数 nonfix_crop：输入的大小是否一致 train_aug_choice:train的数据是否进行augmentation gamma:focal loss参数 feature_sim： prob_diff: ''' nonfix_crop = self.mode_cfg['nonfix_crop'] if 'nonfix_crop' in self.mode_cfg.keys() else False train_aug_choice = self.mode_cfg['train_aug_choice'] if 'train_aug_choice' in self.mode_cfg.keys() else False # print('type: ', type(self.cfg)) # print('cfg: ', self.cfg) gamma = self.cfg.lr_gamma if hasattr(self.cfg, 'lr_gamma') else 1 prob_diff = self.mode_cfg['prob_diff'] if 'prob_diff' in self.mode_cfg.keys() else False ############ generate params for generator self._DecodeGeneratorParam() # train_dataset = SurDataSet(self.train_data, self.train_info, **self.generator_parameters) train_dataset = SurDataSet(self.traindata, **self.generator_parameters) self.generator_parameters['aug'] = False # val_dataset = SurDataSet(self.val_data, self.val_info, **self.generator_parameters) val_dataset = SurDataSet(self.valdata, **self.generator_parameters) balance_data_kws = self.mode_cfg['balance_data'] if 'balance_data' in self.mode_cfg.keys() else [0, 0] modes = ['val' if val == 0 else 'train' for val in balance_data_kws] # train_dataset._GenerateTrainData(mode=modes[0]) # val_dataset._GenerateTrainData(mode=modes[1]) self.train_dataset = train_dataset ''' 根据参数情况得到训练的loss ''' self._GetLoss(prob_diff, gamma) train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset) val_sampler = torch.utils.data.distributed.DistributedSampler(val_dataset) train_loader = DataLoader(train_dataset, batch_size=self.cfg.batch_size, sampler=train_sampler) val_loader = DataLoader(val_dataset, batch_size=self.cfg.batch_size, sampler=val_sampler) lr = self.cfg.lr batch_size = self.cfg.batch_size global_step = 0 epochs = self.cfg.epochs criterion = self.criterion val_criterion = self.criterion label_dtype = torch.long if self.mode_cfg['num_class'] > 2 else torch.float32 n_train = int(round(len(train_dataset) / self.cfg.batch_size)) best_val_score = 1000000 for epoch in range(epochs): train_sampler.set_epoch(epoch) self.net.train() # self.net.to(device) self.net.cuda() epoch_loss = 0 print('size of train dataset is ', len(train_dataset)) if torch.distributed.get_rank() == 0: epoch_acc = 0 epoch_f1 = 0 epoch_precision = 0 epoch_recall = 0 epoch_num = 0 with tqdm(total=n_train, desc=f'Epoch{epoch + 1}/{epochs}', unit='img') as pbar: batch_count = 0 for batch in zip(train_loader): batch_count += 1 loss_values = [] # if self.regression_choice: # loss = self._trainOnBatchWithReg(epoch, batch, label_dtype, device) # else: loss, acc, out_f1, out_precision, out_recall = self._trainOnBatch(epoch, batch, label_dtype) # logger.info('torch.distributed.get_rank(): ', torch.distributed.get_rank()) if torch.distributed.get_rank() == 0: # logger.info('deal with epoch_num ...') epoch_acc += acc.item() epoch_f1 += out_f1.item() epoch_precision += out_precision.item() epoch_recall += out_recall.item() epoch_num += 1 self.optimizer.zero_grad() loss.backward() epoch_loss += loss.item() nn.utils.clip_grad_value_(self.net.parameters(), 1e-1) self.optimizer.step() global_step += 1 pbar.update() if torch.distributed.get_rank() == 0: print('self.metrics_writer.append_one_line -- loss: ', loss) self.metrics_writer.append_one_line( [epoch, epoch_acc / epoch_num, epoch_f1 / epoch_num, epoch_precision / epoch_num, epoch_recall / epoch_num, self.optimizer.param_groups[0]['lr'], loss.item(), epoch_acc / epoch_num] ) # random.shuffle(train_dataset.indices) print('val_criterion: ', val_criterion) val_score = eval_netSE(self.net, val_dataset, val_loader, base_cfg=self.cfg, cfg=self.mode_cfg, criterion=val_criterion, reg_loss=self.regression_loss) ############ change lr based on val result self.scheduler.step(val_score) logger.info("Train loss: {}".format(epoch_loss / n_train)) logging.info("Validation loss: {}".format(val_score)) if val_score <= best_val_score and self.cfg.rank == 0: torch.save(self.net.state_dict(), self.mode_cfg['save_path'] + '/model.pth') best_val_score = val_score with open(self.mode_cfg['train_result'], 'w+') as file: json.dump( { "successFlag":"TRAINING", "bestModelEpoch": epoch }, file, indent=4) # 更新eval中的performance.json with open(self.mode_cfg['eval_pjson'], 'w+') as file: json.dump( { "loss": val_score, }, file, indent=4) # 更新eval中的performance.md with open(self.mode_cfg['eval_pmd'], 'w+') as file: file.write('# overall performance \n') file.write('| loss | \n') file.write('| -------- | \n') file.write(fr'| {val_score} | \n') if self.cfg.rank == 0: json_info = json.load(open(self.mode_cfg['train_result'], 'r')) json_info['successFlag'] = 'SUCCESS' with open(self.mode_cfg['train_result'], 'w') as file: json.dump(json_info, file, indent=4) shutil.copy(self.cfg.pretrain_msg, self.mode_cfg['save_path'] + '/model.pth') print('save best ok') def _testModel_ori(self): # device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') # model_path = self.mode_cfg['weight_path'] # self.net.load_state_dict(torch.load(model_path, map_location=device)) # logger.info("load model from %s" % model_path) # self.net.to(device=device) assert torch.cuda.is_available() self.net = self.net.cuda() self.net.eval() gamma = self.cfg.lr_gamma if hasattr(self.cfg, 'lr_gamma') else 1 prob_diff = self.mode_cfg['prob_diff'] if 'prob_diff' in self.mode_cfg.keys() else False ############ generate params for generator self._DecodeGeneratorParam() self.generator_parameters['aug'] = False val_dataset = SurDataSet(self.test_data, self.test_info, **self.generator_parameters) # print('val_dataset: ',len(val_dataset)) val_dataset._GenerateTrainData() self._GetLoss(prob_diff, gamma) val_loader = DataLoader(val_dataset, batch_size=self.mode_cfg['batch_size'], shuffle=False) val_criterion = self.criterion val_score = eval_netSE(self.net, val_dataset, val_loader, base_cfg=self.cfg, cfg=self.mode_cfg, criterion=val_criterion, reg_loss=self.regression_loss) # 更新eval中的performance.json with open(self.mode_cfg['eval_pjson'], 'w+') as file: json.dump( { "precision": val_score, }, file, indent=4) # 更新eval中的performance.md with open(self.mode_cfg['eval_pmd'], 'w+') as file: file.write('# overall performance \n') file.write('| precision | \n') file.write('| -------- | \n') file.write(fr'| {val_score} | \n') with open(self.mode_cfg['eval_result'], 'w+') as file: json.dump( { "successFlag": "SUCCESS", },file, indent=4) def _testModel(self): device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') # model_path = self.mode_cfg['weight_path'] # self.net.load_state_dict(torch.load(model_path, map_location=device)) # logger.info("load model from %s" % model_path) # self.net.to(device=device) print('all cfg:') print(self.cfg) print('model cfg') print(self.mode_cfg) assert torch.cuda.is_available() self.net = torch.jit.load(self.cfg.pretrain_msg) self.net = self.net.cuda() self.net.eval() gamma = self.cfg.lr_gamma if hasattr(self.cfg, 'lr_gamma') else 1 prob_diff = self.mode_cfg['prob_diff'] if 'prob_diff' in self.mode_cfg.keys() else False ############ generate params for generator self._DecodeGeneratorParam() self.generator_parameters['aug'] = False val_dataset = TestDataSet(self.testdata, **self.generator_parameters) # val_loader = DataLoader(val_dataset, batch_size=self.mode_cfg['batch_size'], shuffle=False) # batch_size = self.mode_cfg['batch_size'] # print(f'batch size: {batch_size}') labels = [] preds = [] limits = [] # for data, label, is_limit in val_loader: for data, label, is_limit in val_dataset: # data = data.float().cuda() data = torch.from_numpy(data[None, ...]).to(device=device, dtype=torch.float32) pred = self.net(data) # print(pred) labels.append(label) pred = torch.softmax(pred[0], -1) pred = torch.max(pred, -1)[1] # pred = torch.argmax(pred[0]) # print('pred', pred, pred.shape) preds.append(pred) # limits += is_limit.reshape(-1).tolist() limits.append(is_limit) # print(f'data shape: {data.shape}, label shape: {label.shape}, pred shape: {pred.shape}, is limit shape: {is_limit.shape}') # labels = torch.cat(labels, 0).cpu().numpy() labels = np.array(labels) preds = torch.cat(preds, 0).cpu().numpy() # print(labels.shape, preds.shape) # preds = torch.cat(preds, 0).cpu().numpy() # print(f'labels shape: {labels.shape}, preds shape: {preds.shape}, limits shape: {len(limits)}') # print(f'labels: {labels}') # print(f'preds: {preds}') preds_new = [] for pred_cls_ori, limit in zip(preds, limits): if pred_cls_ori == 1: pred_cls = 2 elif pred_cls_ori == 2: pred_cls = 1 else: pred_cls = pred_cls_ori # if pred_cls_ori == 3: # if limit: # pred_cls = 5 # else: # if limit: # pred_cls = 4 preds_new.append(pred_cls) preds_new = np.array(preds_new) acc = np.sum(preds_new == labels) / preds_new.shape[0] num_gts = preds_new.shape[0] num_tps = np.sum(preds_new == labels) print('gt info:----------------') elems, elems_count = np.unique(labels, return_counts=True) print('gts:', elems) print('gt counts:', elems_count) print('pred info:----------------') elems, elems_count = np.unique(preds_new, return_counts=True) print('preds:', elems) print('pred counts:', elems_count) print('=================') print('num gts:', num_gts) print('num tps:', num_tps) # 更新eval中的performance.json with open(self.mode_cfg['eval_pjson'], 'w+') as file: json.dump( { "accuracy": acc, }, file, indent=4) # 更新eval中的performance.md with open(self.mode_cfg['eval_pmd'], 'w+') as file: file.write('# overall performance \n') file.write('| accuracy |\n') file.write('| -------- |\n') file.write(fr'| {acc} |\n') with open(self.mode_cfg['eval_result'], 'w+') as file: json.dump( { "successFlag": "SUCCESS", },file, indent=4) def make_infer_data_info(self, info_path, data_path): info_dict = {} with open(info_path, 'r+') as f: infer_data_info = json.load(f) for i in infer_data_info['dataList']: uid = i['rawDataUrls'][0].split('/')[-1].split('.zip')[0] if uid not in info_dict: info_dict[uid] = {'dcm_path': os.path.join(data_path, uid)} for j in i['annotations']: ann_path = j['annotationUrls'][0] with open(ann_path, 'r+') as f: ann = json.load(f) for index, k in enumerate(ann['annotationSessions'][0]['annotationSet']): info_dict[uid][index] = k['coordinates'] return info_dict def _inferModel(self): class_map = {0: '肺内磨玻璃结节', 1:'肺内混合结节', 2: '肺内实性结节', 3: '肺内钙化结节', 4: '胸膜结节或斑块', 5: '胸膜钙化结节'} class_map2 = {'肺内磨玻璃结节': 'mbl', '肺内混合结节': 'hh', '肺内实性结节': 'sx', '肺内钙化结节': 'gh', '胸膜结节或斑块': 'xm', '胸膜钙化结节': 'xmgh'} print(class_map) assert torch.cuda.is_available() self.net = self.net.cuda() self.net.eval() gamma = self.cfg.lr_gamma if hasattr(self.cfg, 'lr_gamma') else 1 prob_diff = self.mode_cfg['prob_diff'] if 'prob_diff' in self.mode_cfg.keys() else False ############ generate params for generator self._DecodeGeneratorParam() self.generator_parameters['aug'] = False val_dataset = InferDataSet(self.testdata, **self.generator_parameters) batch_size = self.mode_cfg['batch_size'] print(f'batch size: {batch_size}') # preds = [] info_dict = self.make_infer_data_info(self.cfg.dataset_info_path, self.cfg.dcm_data_path) print(f'info dict: {info_dict}') pred_dict = {} font = cv2.FONT_HERSHEY_SIMPLEX font_scale = 0.2 font_color = (0,0,255) font_thickness = 1 font_size = 10 # text_position = (0,0) text = '肺结节' (text_w, text_h), _ = cv2.getTextSize(text, font, font_scale, font_thickness) adjusted_position = (20-text_w, 20+text_h) font_path = '/shared/msyhl.ttc' for data, info in val_dataset: uid = info[4] is_limit = info[5] if uid not in pred_dict: pred_dict[uid] = {} index = info[1] # print(data.shape) data = torch.from_numpy(data[None, ...]).float().cuda() pred = self.net(data) pred = torch.softmax(pred[0], -1) pred = torch.max(pred, -1)[1] pred = pred.cpu().numpy()[0] if pred == 3: if is_limit: pred = 5 else: if is_limit: pred = 4 pred_cls = class_map[pred] pred_dict[uid][index] = pred_cls # print('pred', pred, pred.shape) # preds.append(pred) # print(f'data shape: {data.shape}, pred shape: {pred.shape}') # preds = torch.cat(preds, 0).cpu().numpy() # for i in preds: # pred_cls = class_map[i] # print(f'pred class: {pred_cls}') print(f'pred dict: {pred_dict}') target_path = self.cfg.save_img_dir for uid, pred in pred_dict.items(): dcm_path = info_dict[uid]['dcm_path'] itk_image = load_ct_from_dicom(dcm_path) origin = itk_image.GetOrigin() spacing = itk_image.GetSpacing() image = sitk.GetArrayFromImage(itk_image) x_origin, y_origin, z_origin = origin x_spacing, y_spacing, z_spacing = spacing for index, pred_cls in pred.items(): gt_info = info_dict[uid][index] x1, y1, z1, x2, y2, z2 = gt_info[0]['x'], gt_info[0]['y'], gt_info[0]['z'], gt_info[1]['x'], gt_info[1]['y'], gt_info[1]['z'] x1 = int((x1 - x_origin)/x_spacing) y1 = int((y1 - y_origin)/y_spacing) z1 = int((z1 - z_origin)/z_spacing) x2 = int((x2 - x_origin)/x_spacing) y2 = int((y2 - y_origin)/y_spacing) z2 = int((z2 - z_origin)/z_spacing) all_images = [] z_min = max(0, z1-4) z_max = min(image.shape[0]-1, z2+4) for z in range(z_min, z_max, 1): img_z = image.copy()[z, :, :] img_z = norm(img_z) img_z = cv2.cvtColor(img_z, cv2.COLOR_GRAY2BGR) # cv2.rectangle(img_z, (int(y1), int(x1)), (int(y2), int(x2)), (0,0,255), 2) if z>=z1 and z<=z2: cv2.rectangle(img_z, (int(x1), int(y1)), (int(x2), int(y2)), (0,0,255), 2) cut_x_min = max(0, x1-64) cut_y_min = max(0, y1-64) cut_x_max = min(image.shape[2]-1, x2+64) cut_y_max = min(image.shape[1]-1, y2+64) img_z = img_z[cut_y_min: cut_y_max, cut_x_min: cut_x_max, :] # img_z = img_z[cut_x_min: cut_x_max, cut_y_min: cut_y_max, :] img_pil = Image.fromarray(cv2.cvtColor(img_z, cv2.COLOR_BGR2RGB).astype(np.uint8)) draw = ImageDraw.Draw(img_pil) font = ImageFont.truetype(font_path, font_size) draw.text((0, 0), pred_cls, font=font, fill=(255,0,0)) img_z = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR) # cv2.putText(img_z, text, adjusted_position, font, font_scale, font_color, font_thickness) all_images.append(img_z) for idx, img in enumerate(all_images): save_path = os.path.join(target_path, uid+'_'+str(index)+'_'+class_map2[pred_cls]) if not os.path.exists(save_path): os.makedirs(save_path) # print(img) cv2.imwrite(os.path.join(save_path, '{:03d}.png'.format(idx)), img) output_gif_path = self.cfg.save_gif_dir if not os.path.exists(output_gif_path): os.makedirs(output_gif_path) for file in os.listdir(target_path): files = [] for f in os.listdir(os.path.join(target_path, file)): files.append(f) files.sort(key=lambda x: x[:-4]) frames = [] for i in tqdm(range(len(files))): im = imageio.imread(os.path.join(target_path, file) + '/' + files[i]) frames.append(im) fps = 24.0 imageio.mimsave(os.path.join(output_gif_path, file+'.gif'), frames, 'GIF', duration=1/fps) def _GetTargetClassProb(self, labels, probs): target_class_probs = [prob[0][label] for prob, label in zip(probs, labels)] return target_class_probs def ProbVis(self, labels, probs): values = np.unique(labels) from matplotlib import pyplot as plt for val in values: print('=' * 60) print('Current label is %d' % val) indices = [idx for idx in range(len(labels)) if labels[idx] == val] current_probs = np.take(probs, indices, axis=0) print('Number of samples are %d' % len(current_probs)) plt.hist(current_probs, range=(0, 1)) plt.show() plt.show() def BadCaseVis(self, num_incies, show_edges=False, gaussian_choice=False, sigma=1): for idx in self.badcase_lists[:num_incies]: current_data = self.test_dataset[idx] img = np.squeeze(current_data[0]) label = current_data[1] center_slice_idx = int(img.shape[0] / 2) slice_range = 1 print('=' * 60) print('Current label is ', label, ' result is ', self.result[idx]) print('prob is', self.result_list[idx]) if not show_edges: _, axs = plt.subplots(1, 2 * slice_range + 1) else: _, axs = plt.subplots(2, 2 * slice_range + 1) for slice_idx in range(center_slice_idx - slice_range, center_slice_idx + slice_range + 1): if not show_edges: axs[slice_idx - (center_slice_idx - slice_range)].imshow(img[slice_idx], cmap='bone') else: edges = ShowEdges(img[slice_idx], gaussian_choice=gaussian_choice, sigma=sigma) # edges = HFFilter(img[slice_idx]) axs[0][slice_idx - (center_slice_idx - slice_range)].imshow(img[slice_idx], cmap='bone') axs[1][slice_idx - (center_slice_idx - slice_range)].imshow(edges, cmap='bone') plt.show() def ParametersDecode(self): ''' Generate model parameters ''' # self.model_param_name = os.path.join(self.basic_file_path, self.cfg['model_params']) # self.model_params = LoadJson(self.model_param_name)[self.cfg['model_name']] if self.mode_cfg['model_name'] == 'VGG3D': self.model_params = self.cfg.VGG3D elif self.mode_cfg['model_name'] == 'ResNet': self.model_params = self.cfg.ResNet elif self.mode_cfg['model_name'] == 'Se_ResNet': self.model_params = self.cfg.Se_ResNet elif self.mode_cfg['model_name'] == 'Darknet': self.model_params = self.cfg.Darknet elif self.mode_cfg['model_name'] == 'FSe_ResNet': self.model_params = self.cfg.FSe_ResNet