# -*- 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 torch.utils.data import DataLoader
from torch.nn.utils import clip_grad_norm_
# from matplotlib import pyplot as plt
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
from .utils.OnlineEval import eval_net
from .utils.OfflineEval import CalculateClsScore,CalculateAuc,CalculateClsScoreByTh
from .utils.loss_func import BCEFocalLoss,PeerLoss,MIFocalLoss,FocalLoss
from .utils.gradcam import GradCam,GuidedBackpropReLUModel
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)
class Classification3D(object):
def __init__(self,filename,mode):
self.mode = mode
self.filename = filename
self.cfg = LoadJson(self.filename)[mode]
self.basic_file_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
current_DT = datetime.datetime.now()
self.current_day = '%02d%02d'%(current_DT.month,current_DT.day)
self.convert_dict = self.cfg['convert_dict'] if 'convert_dict' in self.cfg.keys() else {}
self._clsInfoDecode()
writer_path = os.path.join(self.save_path, 'runs')
train_writer_path = os.path.join(writer_path, 'train')
val_writer_path = os.path.join(writer_path, 'val')
self.writer = SummaryWriter(writer_path) if self.mode=='training' else None
self.train_writer = SummaryWriter(train_writer_path) if self.mode=='training' else None
self.val_writer = SummaryWriter(val_writer_path) if self.mode=='training' else None
if mode=='training':
self.trainfile_conf_path = os.path.join(self.basic_file_path,self.cfg['train_file'])
self.valfile_conf_path = os.path.join(self.basic_file_path,self.cfg['val_file'])
self.trainfiles = self._DecodeDataParam(self.trainfile_conf_path)
self.valfiles = self._DecodeDataParam(self.valfile_conf_path)
self.testfiles = []
else:
self.trainfiles = []
self.valfiles = []
self.test_conf_path = os.path.join(self.basic_file_path,self.cfg['test_file'])
self.testfiles = self._DecodeDataParam(self.test_conf_path)
self.ParametersDecode()
self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
def __call__(self):
self._buildModel()
self._compileModel()
self._LoadData(self.mode)
if self.mode=='training':
self._trainModel()
else:
import time
start_time = time.time()
self._testModel()
end_time = time.time()
dur = end_time - start_time
print ('========================================== It takes %s s to make inference'%str(dur))
if self.writer:
self.writer.close()
if self.train_writer:
self.train_writer.close()
if self.val_writer:
self.val_writer.close()
def _clsInfoDecode(self):
self.cls_map_dict = eval(self.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.cfg['patch_height'],self.cfg['patch_height'],self.cfg['patch_depth']]
self.cls_name_map_dict = eval(self.cfg['cls_name_map_dict'])
model_order = self.cfg['model_order'] if 'model_order' in self.cfg.keys() else 1
if type(self.cfg['cls_label_index'])==int:
cls_related_info = self.cls_name_map_dict[self.cfg['cls_label_index']]
else:
cls_related_info = '_'.join([self.cls_name_map_dict[val] for val in self.cfg['cls_label_index']])
self.save_path = '%s/%s_%s_%dCls_%s_Patch%03d_model%s/%02d_%s/'%(self.cfg['base_path'],self.cfg['weights_pre'],cls_related_info,self.cfg['num_class'],self.current_day,self.input_shape[0],self.cfg['model_name'],model_order,self.cfg['weight_memo'])
if (not os.path.exists(self.save_path)) and self.mode == 'training':
os.makedirs(self.save_path)
if type(self.cfg['cls_label_index'])==list:
self.cfg['num_task'] = len(self.cfg['cls_label_index'])
else:
self.cfg['num_task'] = 1
def _buildModel(self):
cfg = self.model_params
cfg['num_task'] = self.cfg['num_task']
cfg['num_class'] = self.cfg['num_class']
if cfg['num_class'] == 2:
cfg['num_class'] = 1
self.net = build_model(self.cfg['model_name'],cfg)
logger.info("After building model")
self.net.apply(weight_init)
logger.info("After Weight Init")
def _compileModel(self):
self.optimizer = optim.Adam(self.net.parameters(),lr=self.cfg['learning_rate'])
self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(self.optimizer,patience=self.cfg['patience'])
if self.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 _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):
if mode == 'training':
self.train_data,self.train_info = load_data([self.trainfiles['data_paths'],self.trainfiles['info_paths']])
self.val_data,self.val_info = load_data([self.valfiles['data_paths'],self.valfiles['info_paths']])
else:
self.test_data,self.test_info = load_data([self.testfiles['data_paths'],self.testfiles['info_paths']])
def _trainModel(self):
########## Add parameters
gamma = self.cfg['gamma'] if 'gamma' in self.cfg.keys() else 1
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
nonfix_crop=self.cfg['nonfix_crop'] if 'nonfix_crop' in self.cfg.keys() else False
train_aug_choice = self.cfg['train_aug_choice'] if 'train_aug_choice' in self.cfg.keys() else False
train_dataset = SurDataSet(self.train_data,self.train_info,cls_index=self.cfg['cls_label_index'],cls_map_dict=self.cls_map_dict,
diameter_index=self.cfg['diameter_index'],num_class=self.cfg['num_class'],HU_min=self.cfg['HU_min'],
HU_max=self.cfg['HU_max'],input_shape=self.input_shape,
cls_types=self.ori_cls_types,aug=train_aug_choice,
nonfix_crop=nonfix_crop,
diameter_max_th = self.cfg['diameter_max_th'],
diameter_min_th = self.cfg['diameter_min_th'])
val_dataset = SurDataSet(self.val_data,self.val_info,cls_index=self.cfg['cls_label_index'],cls_map_dict=self.cls_map_dict,
diameter_index=self.cfg['diameter_index'],num_class=self.cfg['num_class'],HU_min=self.cfg['HU_min'],
HU_max=self.cfg['HU_max'],input_shape=self.input_shape,
cls_types=self.ori_cls_types,
nonfix_crop=nonfix_crop,
diameter_max_th = self.cfg['diameter_max_th'],
diameter_min_th = self.cfg['diameter_min_th'])
if self.cfg['focal_loss'] and self.cfg['num_task']==1:
ratios =train_dataset._GenerateRatios()
if self.cfg['num_class']<=2:
alpha = ratios[0]
print ('='*60)
print ('alpha is ',alpha)
train_focal_loss = self.cfg['train_focal_loss'] if 'train_focal_loss' in self.cfg.keys() else self.cfg['focal_loss']
loss_choice = self.cfg['loss_func'] if 'loss_func' in self.cfg.keys() else 'focal'
print ('loss choice is ',loss_choice)
if loss_choice != 'focal':
train_focal_loss = False
if train_focal_loss:
self.criterion = BCEFocalLoss(alpha=alpha,gamma=gamma)
elif loss_choice == 'peer':
self.criterion = PeerLoss(alpha=alpha,gamma=gamma)
elif loss_choice == 'MIFocalLoss':
self.criterion = MIFocalLoss(alpha=alpha,gamma=gamma)
else:
alpha = ratios
sigmoid_choice = self.cfg['sigmoid_choice'] if 'sigmoid_choice' in self.cfg.keys() else False
######### sigmoid + bce loss
self.criterion = FocalLoss(self.cfg['num_class'],alpha=ratios,gamma=gamma)
else:
alpha = None
# ratios = train_dataset._GenerateRatios()
balance_data_kws = self.cfg['balance_data'] if 'balance_data' in self.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])
train_loader = DataLoader(train_dataset,batch_size=self.cfg['batch_size'],shuffle=True)
val_loader = DataLoader(val_dataset,batch_size=self.cfg['batch_size'],shuffle=True)
lr=self.cfg['learning_rate']
batch_size = self.cfg['batch_size']
# writer = SummaryWriter(comment=f"LR_{lr}_BS_{batch_size}")
global_step = 0
##### add config
epochs = self.cfg['num_epoch']
criterion = self.criterion
for epoch in range(epochs):
self.net.train()
self.net.to(device)
self.net.cuda()
epoch_loss = 0
print ('size of train dataset is ',len(train_dataset))
with tqdm(total=int(len(train_dataset)/batch_size)+1,desc=f'Epoch{epoch+1}/{epochs}',unit='img') as pbar:
batch_count = 0
for batch in train_loader:
batch_count +=1
imgs = batch[0][0]
label = batch[1][0]
imgs = imgs.cuda()
if self.cfg['num_class']>2:
label = label.cuda()
label = label.to(device=device,dtype=torch.long)
else:
if self.cfg['num_task']>1:
label = [val.cuda() for val in label]
label = [val.to(device=device,dtype=torch.float32) for val in label]
else:
label = label.cuda()
label = label.to(device=device,dtype=torch.float32)
imgs = imgs.to(device=device,dtype=torch.float32)
prediction = self.net(imgs)
loss_values = []
if self.cfg['num_task']>1:
for x,y in zip(label,prediction):
probs = torch.sigmoid(y)
y = y.squeeze(-1)
loss = criterion([y,x])
epoch_loss += loss.item()
loss_values.append(loss.item())
else:
# print ('prediction before',prediction)
if self.cfg['num_class']>2:
prediction = prediction
else:
if type(prediction) == list:
prediction = prediction[0]
prediction = prediction.squeeze(-1)
loss = criterion([prediction,label])
epoch_loss += loss.item()
self.optimizer.zero_grad()
loss.backward()
nn.utils.clip_grad_value_(self.net.parameters(),1e-1)
self.optimizer.step()
global_step += 1
# if(global_step %(len(train_dataset)//(10*batch_size)))==0:
# for tag,value in self.net.named_parameters():
# tag = tag.replace('.','/')
# writer.add_histogram('weights/'+tag,value.data.cpu().numpy,global_step)
# writer.add_histogram('grads/'+tag,value.grad.data.cpu().numpy,global_step)
############ Add eval net func(for samples in val generator make prediction and calculating loss)
############## Stragety 1.
pbar.update(1)
if self.cfg['select_samples'] == 'partial':
loss_list = self._CalculateLossTrain(train_dataset,device)
train_dataset._updateIndices(loss_list,ratio_pos=self.cfg['ratio_pos'],num_hard=100)
elif self.cfg['select_samples'] == 'boost':
############# Stragety 2. _CalculateDiffTrain
loss_list = self._CalculateLossTrain(train_dataset,device)
initial_epoch = self.cfg['initial_epoch'] if 'initial_epoch' in self.cfg.keys() else 0
epoch_interval = self.cfg['epoch_interval'] if 'epoch_interval' in self.cfg.keys() else 1
if epoch>=initial_epoch:
if (epoch-epoch_interval)%epoch_interval == 0:
diff_list = self._CalculateDiffTrain(train_dataset,device)
increase_choice = self.cfg['increase_choice'] if 'increase_choice' in self.cfg.keys() else False
train_dataset._updateWeights(diff_list,increase=increase_choice)
train_dataset._SelectSamplesBasedOnWeights()
new_ratios = train_dataset._GenerateRatios()
if self.cfg['num_class']<2:
self.criterion = BCEFocalLoss(alpha=new_ratios[0],gamma=gamma)
# else:
# self.criterion = FocalLoss(self.cfg['num_class'],alpha=new_ratios,gamma=gamma)
print ('batch_count is',batch_count)
random.shuffle(train_dataset.indices)
val_score = eval_net(self.net, val_dataset,val_loader, device,alpha=alpha,cfg=self.cfg)*100
############ change lr based on val result
self.scheduler.step(val_score)
###########3 param_groups?
# writer.add_scalar('learning_rate',self.optimizer.param_groups[0]['lr'],global_step)
print ('epoch_loss is',epoch_loss,len(train_dataset))
logger.info("Train loss: {}".format(epoch_loss/float(len(train_dataset))*100))
logging.info("Validation loss: {}".format(val_score))
if self.cfg['save_cp']:
model_save_path_current_epoch = self.save_path + f'/CP_epoch{epoch + 1}_loss_%.2f.pth' % val_score
torch.save(self.net.state_dict(), model_save_path_current_epoch)
logging.info(f'Checkpoint{epoch + 1} saved to %s!' % model_save_path_current_epoch)
if self.train_writer:
self.train_writer.add_scalar('loss', epoch_loss/float(len(train_dataset))*100, epoch)
if self.val_writer:
self.val_writer.add_scalar('loss', val_score, epoch)
def _testModel(self):
if '*' in self.cfg['weight_path']:
model_paths = sorted(glob.glob(self.cfg['weight_path']))
model_paths = [path for path in model_paths if '.pth' in path]
else:
model_paths = [self.cfg['weight_path']]
results = []
for model_path in model_paths:
print ('='*60)
print ('model_path is ',model_path)
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
self.net.to(device=device)
self.net.load_state_dict(torch.load(model_path,map_location=device))
logging.info('Model loaded !')
########## 将模型设置为eval 固定BN/dropout的参数
self.net.eval()
nonfix_crop=self.cfg['nonfix_crop'] if 'nonfix_crop' in self.cfg.keys() else False
test_dataset = SurDataSet(self.test_data,self.test_info,cls_index=self.cfg['cls_label_index'],cls_map_dict=self.cls_map_dict,
diameter_index=self.cfg['diameter_index'],num_class=self.cfg['num_class'],HU_min=self.cfg['HU_min'],
HU_max=self.cfg['HU_max'],input_shape=self.input_shape,
cls_types=self.ori_cls_types,
nonfix_crop = nonfix_crop,
aug=False,
mode='test',
diameter_max_th = self.cfg['diameter_max_th'],
diameter_min_th = self.cfg['diameter_min_th'])
test_dataset._GenerateTrainData()
self.label_list = []
self.result_list = []
self.badcase_lists = []
for idx in range(len(test_dataset)):
current_data = test_dataset[idx]
img = current_data[0][0]
label = current_data[1][0]
img = img[np.newaxis,...]
img = torch.from_numpy(img)
img = img.to(device=device,dtype=torch.float32)
with torch.no_grad():
output = self.net(img)
if type(output) == list:
output = output[0]
if self.cfg['num_task']==1:
if self.net.num_class>2:
probs = torch.softmax(output,dim=1)
else:
probs = torch.sigmoid(output)
probs = probs.squeeze(0)
probs = probs.data.cpu().numpy()
else:
current_probs = []
for task_idx in range(self.cfg['num_task']):
if self.net.num_class>1:
probs = F.softmax(output[task_idx],dim=1)
else:
probs = torch.sigmoid(output[task_idx])
probs = probs.squeeze(0)
current_probs.append(probs.cpu().numpy()[0])
probs = current_probs
self.label_list.append(label)
self.result_list.append(probs)
print ('length of label and result',len(self.label_list),len(self.result_list))
print ('shape of result list is ',np.array(self.result_list).shape)
print ('max val of self.result_list',np.amin(self.result_list),np.amax(self.result_list))
if self.net.num_class<2:
if self.cfg['num_task']==1:
th_result = [val>0.5 for val in self.result_list]
CalculateClsScore(th_result,self.label_list)
auc_val = CalculateAuc(self.result_list,self.label_list)
results.append(auc_val)
if len(model_paths)==1:
CalculateClsScoreByTh(self.result_list,self.label_list,acc_flag=True)
else:
for task_idx in range(self.cfg['num_task']):
print ('='*60)
print ('Current task is %s'%self.cls_name_map_dict[self.cfg['cls_label_index'][task_idx]])
current_label_list = np.squeeze(np.array([val[task_idx] for val in self.label_list]))
current_result_list = np.squeeze(np.array([val[task_idx] for val in self.result_list]))
th_result = np.squeeze(np.array([val>0.5 for val in current_result_list]))
print (np.array(current_label_list).shape,np.array(current_result_list).shape,np.array(th_result).shape)
CalculateClsScore(th_result,current_label_list)
auc_val = CalculateAuc(current_result_list,current_label_list)
else:
result = np.squeeze(np.array(self.result_list))
result = np.argmax(result,axis=1)
kappa_val = CalculateClsScore(result,self.label_list)
results.append(kappa_val)
# self.badcase_lists = [case_id for case_id in range(len(test_dataset)) if self.result_list[case_id]!=self.label_list[case_id]]
if len(results)>0:
print ('results is ',results)
pos = np.argmax(results)
print ('best weights is ',model_paths[pos])
self.test_dataset = test_dataset
# def BadCaseVis(self):
# for idx in self.badcase_lists:
# current_data = test_dataset[idx]
# img = np.squeeze(current_data[0])
# label = current_data[1]
# center_slice_idx = int(current_data.shape[0]/2)
# slice_range = 2
# print ('='*60)
# print ('Current label is ',label)
# for slice_idx in range(center_slice_idx-slice_range,center_slice_idx+slice_range):
# plt.imshow(img[slice_idx],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']]
def _CalculateLossTrain(self,train_dataset,device):
loss_list = []
print ('number of data inside train dataset is ',len(train_dataset.data))
train_dataset.mode = 'test'
self.net.eval()
for idx in range(len(train_dataset.data)):
data,label,_ = train_dataset[idx]
data = data[0]
data = data[np.newaxis,...]
data = torch.from_numpy(data)
label = label[0]
label = np.array([label])
label = torch.from_numpy(label)
data = data.to(device=device,dtype=torch.float32)
if self.cfg['num_task']==1:
label = label.to(device=device,dtype=torch.float32)
else:
label = [val.to(device=device,dtype=torch.float32) for val in label]
with torch.no_grad():
prediction = self.net(data)
#prediction = prediction.squeeze(-1)
if self.cfg['num_task']==1:
prediction = prediction.squeeze(-1)
if self.cfg['num_class']>2:
label = label.to(device=device,dtype=torch.long)
loss_list.append([idx,self.criterion([prediction,label]).item()])
else:
current_loss = []
for x,y in zip(prediction,label):
x = x.squeeze(-1)
current_loss.append(self.criterion([x,y]).item())
loss_list.append([idx,current_loss])
train_dataset.mode = 'train'
self.net.train()
return loss_list
def _CalculateDiffTrain(self,train_dataset,device):
diff_list = []
print ('number of data inside train dataset is ',len(train_dataset.data))
train_dataset.mode = 'test'
self.net.eval()
for idx in range(len(train_dataset.data)):
data,label,_ = train_dataset[idx]
data = data[0]
data = data[np.newaxis,...]
data = torch.from_numpy(data)
label = label[0]
label = np.array([label])
label = torch.from_numpy(label)
data = data.to(device=device,dtype=torch.float32)
if self.cfg['num_task']==1:
label = label.to(device=device,dtype=torch.float32)
else:
label = [val.to(device=device,dtype=torch.float32) for val in label]
with torch.no_grad():
prediction = self.net(data)
prediction = prediction.squeeze(-1)
prediction = torch.sigmoid(prediction)
diff = abs(prediction-label)
diff_list.append([idx,diff])
train_dataset.mode = 'train'
self.net.train()
return diff_list