# -*- 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