def main():
# init or load model
print("init model with input shape",config["input_shape"])
model = NvNet(config=config,input_shape=config["input_shape"], seg_outChans=config["n_labels"])
parameters = model.parameters()
optimizer = optim.Adam(parameters,
lr=config["initial_learning_rate"],
weight_decay = config["L2_norm"])
start_epoch = 1
if config["VAE_enable"]:
loss_function = CombinedLoss(k1=config["loss_k1_weight"], k2=config["loss_k2_weight"])
else:
loss_function = SoftDiceLoss()
# data_generator
print("data generating")
training_data = BratsDataset(phase="train", config=config)
train_loader = torch.utils.data.DataLoader(dataset=training_data,
batch_size=config["batch_size"],
shuffle=True,
pin_memory=True)
valildation_data = BratsDataset(phase="validate", config=config)
valildation_loader = torch.utils.data.DataLoader(dataset=valildation_data,
batch_size=config["batch_size"],
shuffle=True,
pin_memory=True)
train_logger = Logger(model_name=config["model_file"],header=['epoch', 'loss', 'acc', 'lr'])
if config["cuda_devices"] is not None:
model = model.cuda()
loss_function = loss_function.cuda()
# if not config["overwrite"] and os.path.exists(config["model_file"]) or os.path.exists(config["saved_model_file"]):
# model, start_epoch, optimizer = load_old_model(model, optimizer, saved_model_path=config["saved_model_file"])
scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, 'min', factor=config["lr_decay"],patience=config["patience"])
print("training on label:{}".format(config["labels"]))
for i in range(start_epoch,config["epochs"]):
train_epoch(epoch=i,
data_loader=train_loader,
model=model,
model_name=config["model_file"],
criterion=loss_function,
optimizer=optimizer,
opt=config,
epoch_logger=train_logger)
val_loss = val_epoch(epoch=i,
data_loader=valildation_loader,
model=model,
criterion=loss_function,
opt=config,
optimizer=optimizer,
logger=train_logger)
scheduler.step(val_loss)
def predict(name_list, model):
model.eval()
config["test_patients"] = name_list
# config["tta_idx"] = 0 # 0 indices no test-time augmentation;
if not os.path.exists(config["prediction_dir"]):
os.mkdir(config["prediction_dir"])
tmp_dir = "../tmp_result_{}".format(config["checkpoint_file"][:-4])
if not os.path.exists(tmp_dir):
os.mkdir(tmp_dir)
# For testing time data augment
tta_idx_limit = 8 if tta else 1
for tta_idx in range(tta_idx_limit):
config["tta_idx"] = tta_idx
if tta:
print(
"starting evaluation of the {} mirror flip of Test-Time-Augmentation"
.format(tta_idx))
data_set = BratsDataset(phase="test", config=config)
valildation_loader = torch.utils.data.DataLoader(
dataset=data_set,
batch_size=config["batch_size"],
shuffle=False,
pin_memory=True)
predict_process = tqdm(valildation_loader)
for idx, inputs in enumerate(predict_process):
if idx > 0:
predict_process.set_description(
"processing {} picture".format(idx))
if config["cuda_devices"] is not None:
inputs = inputs.type(torch.FloatTensor)
inputs = inputs.cuda()
with torch.no_grad():
if config["VAE_enable"]:
outputs, distr = model(inputs)
else:
outputs = model(inputs)
output_array = np.array(
outputs.cpu()) # can't convert tensor in GPU directly
output_array = output_array[:, :
3, :, :, :] # (2, 7, 128, 192, 160)
print(output_array.shape)
output_array = test_time_flip_recovery(output_array,
config["tta_idx"])
# save to tmp
for i in range(config["batch_size"]):
file_idx = idx * config["batch_size"] + i
if file_idx < len(name_list):
patient_filename = name_list[file_idx]
np.save(
os.path.join(
tmp_dir,
"flip_{}_{}.npy".format(config["tta_idx"],
patient_filename)),
output_array[i])
# after all flips
if tta:
config["prediction_dir"] += "_TTA"
if config["predict_from_train_data"]:
config["prediction_dir"] += "_train"
if config["predict_from_test_data"]:
config["prediction_dir"] += "_testing"
if not os.path.exists(config["prediction_dir"]):
os.mkdir(config["prediction_dir"])
for patient_filename in name_list:
flip_arrays = []
for tta_idx in range(tta_idx_limit):
flip_array = np.load(
os.path.join(
tmp_dir, "flip_{}_{}.npy".format(config["tta_idx"],
patient_filename)))
flip_arrays.append(flip_array)
probsMap_array = np.array(flip_arrays).mean(axis=0)
preds_array = np.array(probsMap_array > 0.5,
dtype=float) # (1, 3, 128, 192, 160)
preds_array = dim_recovery(preds_array) # (1, 3, 155, 240, 240)
preds_array = preds_array.swapaxes(
-3,
-1) # convert channel first (SimpleTIK) to channel last (Nibabel)
preds_array = combine_labels_predicting(preds_array)
affine = nib.load(
os.path.join(config["test_path"], patient_filename,
patient_filename + '_t1.nii.gz')).affine
output_image = nib.Nifti1Image(preds_array, affine)
output_image.to_filename(
os.path.join(config["prediction_dir"],
patient_filename + '.nii.gz'))
propbsMap_dir = config["prediction_dir"] + "_probabilityMap"
if not os.path.exists(propbsMap_dir):
os.mkdir(propbsMap_dir)
np.save(os.path.join(propbsMap_dir, patient_filename + ".npy"),
probsMap_array)
os.system("rm -r " + tmp_dir)
eval_process.set_description("Processing Patient:%d" % (i))
# read preds
pred_name = config["validation_patients"][i] + ".nii.gz"
cur_pred_output = os.path.join(folder_path, pred_name)
sitkImage = sitk.ReadImage(cur_pred_output)
output_with_oriLabel = sitk.GetArrayFromImage(sitkImage)
output = preprocess_label(output_with_oriLabel)
acc, _ = calculate_accuracy(
torch.Tensor(output[np.newaxis, :, :, :, :]), targets)
df.loc[i, "WT"] = acc["dice_wt"].item()
df.loc[i, "TC"] = acc["dice_tc"].item()
df.loc[i, "ET"] = acc["dice_et"].item()
print(round(df["WT"].mean(), 4))
df.to_excel(excel_name, index=None)
if __name__ == "__main__":
config["test_path"] = os.path.join(config["base_path"], "data",
"MICCAI_BraTS2020_TrainingData")
mapping_file_path = os.path.join(config["test_path"], "name_mapping.csv")
name_mapping = pd.read_csv(mapping_file_path)
config["validation_patients"] = name_mapping[
"BraTS_2020_subject_ID"].tolist()
# config["validation_patients"] = ["Brats18_CBICA_AXL_1"]
config["seg_label"] = None
config["input_shape"] = None
pred_name = config["validation_patients"]
evaluation_data = BratsDataset(phase="evaluation", config=config)
evaluate(evaluation_data)
def main():
# init or load model
print("init model with input shape", config["input_shape"])
if config["attention"]:
model = AttentionVNet(config=config)
else:
model = NvNet(config=config)
parameters = model.parameters()
optimizer = optim.Adam(parameters,
lr=config["initial_learning_rate"],
weight_decay=config["L2_norm"])
start_epoch = 1
if config["VAE_enable"]:
loss_function = CombinedLoss(new_loss=config["new_SoftDiceLoss"],
k1=config["loss_k1_weight"],
k2=config["loss_k2_weight"],
alpha=config["focal_alpha"],
gamma=config["focal_gamma"],
focal_enable=config["focal_enable"])
else:
loss_function = SoftDiceLoss(new_loss=config["new_SoftDiceLoss"])
with open('valid_list_v2.txt', 'r') as f:
val_list = f.read().splitlines()
# with open('train_list.txt', 'r') as f:
with open('train_list_v2.txt', 'r') as f:
tr_list = f.read().splitlines()
config["training_patients"] = tr_list
config["validation_patients"] = val_list
# data_generator
print("data generating")
training_data = BratsDataset(phase="train", config=config)
# x = training_data[0] # for test
valildation_data = BratsDataset(phase="validate", config=config)
train_logger = Logger(
model_name=config["model_name"] + '.h5',
header=['epoch', 'loss', 'wt-dice', 'tc-dice', 'et-dice', 'lr'])
if not config["overwrite"] and config["saved_model_file"] is not None:
if not os.path.exists(config["saved_model_file"]):
raise Exception("Invalid model path!")
model, start_epoch, optimizer_resume = load_old_model(
model, optimizer, saved_model_path=config["saved_model_file"])
parameters = model.parameters()
optimizer = optim.Adam(
parameters,
lr=optimizer_resume.param_groups[0]["lr"],
weight_decay=optimizer_resume.param_groups[0]["weight_decay"])
if config["cuda_devices"] is not None:
model = model.cuda()
loss_function = loss_function.cuda()
model = nn.DataParallel(model) # multi-gpu training
for state in optimizer.state.values():
for k, v in state.items():
if isinstance(v, torch.Tensor):
state[k] = v.cuda()
# scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, 'min', factor=config["lr_decay"], patience=config["patience"])
scheduler = lr_scheduler.LambdaLR(
optimizer=optimizer,
lr_lambda=poly_lr_scheduler) # can't restore lr correctly
max_val_WT_dice = 0.
max_val_AVG_dice = 0.
for i in range(start_epoch, config["epochs"]):
train_epoch(epoch=i,
data_set=training_data,
model=model,
criterion=loss_function,
optimizer=optimizer,
opt=config,
logger=train_logger)
val_loss, WT_dice, TC_dice, ET_dice = val_epoch(
epoch=i,
data_set=valildation_data,
model=model,
criterion=loss_function,
opt=config,
optimizer=optimizer,
logger=train_logger)
scheduler.step()
# scheduler.step(val_loss)
dices = np.array([WT_dice, TC_dice, ET_dice])
AVG_dice = dices.mean()
if config["checkpoint"] and (WT_dice > max_val_WT_dice
or AVG_dice > max_val_AVG_dice
or WT_dice >= 0.912):
max_val_WT_dice = WT_dice
max_val_AVG_dice = AVG_dice
# save_dir = os.path.join(config["result_path"], config["model_file"].split("/")[-1].split(".h5")[0])
save_dir = config["result_path"]
if not os.path.exists(save_dir):
os.makedirs(save_dir)
save_states_path = os.path.join(
save_dir,
'epoch_{0}_val_loss_{1:.4f}_WTdice_{2:.4f}_AVGDice:{3:.4f}.pth'
.format(i, val_loss, WT_dice, AVG_dice))
if config["cuda_devices"] is not None:
state_dict = model.module.state_dict()
else:
state_dict = model.state_dict()
states = {
'epoch': i,
'state_dict': state_dict,
'optimizer': optimizer.state_dict(),
}
torch.save(states, save_states_path)
save_model_path = os.path.join(save_dir, "best_model.pth")
if os.path.exists(save_model_path):
os.system("rm " + save_model_path)
torch.save(model, save_model_path)
print(
"batch {0:d} finished, validation loss:{1:.4f}; WTDice:{2:.4f}; AVGDice:{3:.4f}"
.format(i, val_loss, WT_dice, AVG_dice))
def main():
# convert input images into an hdf5 file
if config["overwrite"] or not os.path.exists(config["data_file"]):
training_files, subject_ids = fetch_training_data_files(return_subject_ids=True)
write_data_to_file(training_files, config["data_file"], image_shape=config["image_shape"],
subject_ids=subject_ids)
# init or load model
print("init model with input shape",config["input_shape"])
model = NvNet(config=config)
parameters = model.parameters()
optimizer = optim.Adam(parameters,
lr=config["initial_learning_rate"],
weight_decay = config["L2_norm"])
start_epoch = 1
if config["VAE_enable"]:
loss_function = CombinedLoss(k1=config["loss_k1_weight"], k2=config["loss_k2_weight"])
else:
loss_function = SoftDiceLoss()
# data_generator
print("data generating")
training_data = BratsDataset(phase="train", config=config)
valildation_data = BratsDataset(phase="validate", config=config)
train_logger = Logger(model_name=config["model_file"],header=['epoch', 'loss', 'acc', 'lr'])
if config["cuda_devices"] is not None:
# model = nn.DataParallel(model) # multi-gpu training
model = model.cuda()
loss_function = loss_function.cuda()
if not config["overwrite"] and config["saved_model_file"] is not None:
if not os.path.exists(config["saved_model_file"]):
raise Exception("Invalid model path!")
model, start_epoch, optimizer = load_old_model(model, optimizer, saved_model_path=config["saved_model_file"])
scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, 'min', factor=config["lr_decay"],patience=config["patience"])
print("training on label:{}".format(config["labels"]))
max_val_acc = 0.
for i in range(start_epoch,config["epochs"]):
train_epoch(epoch=i,
data_set=training_data,
model=model,
criterion=loss_function,
optimizer=optimizer,
opt=config,
logger=train_logger)
val_loss, val_acc = val_epoch(epoch=i,
data_set=valildation_data,
model=model,
criterion=loss_function,
opt=config,
optimizer=optimizer,
logger=train_logger)
scheduler.step(val_loss)
if config["checkpoint"] and val_acc > max_val_acc:
max_val_acc = val_acc
save_dir = os.path.join(config["result_path"], config["model_file"].split("/")[-1].split(".h5")[0])
if not os.path.exists(save_dir):
os.makedirs(save_dir)
save_states_path = os.path.join(save_dir,'epoch_{0}_val_loss_{1:.4f}_acc_{2:.4f}.pth'.format(i, val_loss, val_acc))
states = {
'epoch': i + 1,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
}
torch.save(states, save_states_path)
save_model_path = os.path.join(save_dir, "best_model_file.pth")
if os.path.exists(save_model_path):
os.system("rm "+save_model_path)
torch.save(model, save_model_path)