def eval_net_dice(net, criterion, phase, network_switch, dataset, preview=True, gpu=False, visualize_batch=None, epoch=None, slice=20, root_path='no_root_path'):
"""Evaluation without the densecrf with the dice coefficient"""
net.eval()
tot = 0
running_loss = 0
for i, b in enumerate(dataset):
img = b['image'][:, 2:3].float()
true_mask = b['mask'][:].float()
if gpu:
img = img.cuda()
true_mask = true_mask.cuda()
mask_pred = net(img, phase, network_switch)[0]
# mask_pred = (mask_pred > 0.5).float()
tot += dice_coeff(mask_pred, true_mask).item()
# test_image = np.transpose((true_mask.cpu().detach().numpy()[0]).astype(float), [1, 2, 0])
# result_image = np.transpose((mask_pred.cpu().detach().numpy()[0][0]).astype(float), [1, 2, 0])
#
# tot += evaluation.do(sitk.GetImageFromArray(test_image), sitk.GetImageFromArray(result_image))[0]
loss = criterion[0](mask_pred.float(), true_mask.float())
running_loss += loss.item() * img.size(0)
if preview:
if visualize_batch is not None:
if i == int(visualize_batch):
outputs_vis = mask_pred.cpu().detach().numpy()
inputs_vis = img.cpu().detach().numpy()
labels_vis = true_mask.cpu().detach().numpy()
fig = visualize_Seg(inputs_vis[0][0], labels_vis[0], outputs_vis[0][0], figsize=(6, 6), epoch=epoch, slice=slice)
cm.mkdir(root_path + 'preview/val/Labeled')
plt.savefig(root_path + 'preview/val/Labeled/' + 'epoch_%s.jpg' % epoch)
# plt.show(block=False)
# plt.pause(1.0)
plt.close(fig)
if i == 0:
tot = tot
running_loss = running_loss
else:
tot = tot / (i+1)
running_loss = running_loss / (i+1)
return tot, running_loss
def eval_net_mse(net, criterion, phase, network_switch, dataset, preview=True, gpu=False, visualize_batch=None, epoch=None, slice=20, root_path='no_root_path'):
"""Evaluation without the densecrf with the dice coefficient"""
net.eval()
running_loss = 0
for i, b in enumerate(dataset):
img = b['image'][:, 2:3].float()
true_mask = b['image'][:, 2:3].float()
if gpu:
img = img.cuda()
true_mask = true_mask.cuda()
mask_predL, mask_pred = net(img, phase, network_switch)
# mask_pred = (mask_pred > 0.5).float()
mask_pred = mask_pred.float()
outputsU_back = mask_pred[:, 0:1]
outputsU_1 = mask_pred[:, 1:2]
labels_back = (1.0 - mask_predL) * true_mask.float()
labels_1 = mask_predL * true_mask.float()
loss = criterion[1](outputsU_1.float(), labels_1.float())
running_loss += loss.item() * img.size(0)
if preview:
if visualize_batch is not None:
if i == int(visualize_batch):
outputsU_back_vis = outputsU_back.cpu().detach().numpy()
outputsU_1_vis = outputsU_1.cpu().detach().numpy()
inputs_vis = img.cpu().detach().numpy()
labels_back_vis = labels_back.cpu().detach().numpy()
labels_1_vis = labels_1.cpu().detach().numpy()
fig = visualize_Rec(inputs_vis[0][0], labels_back_vis[0, 0], labels_1_vis[0, 0],
outputsU_back_vis[0, 0], outputsU_1_vis[0, 0], figsize=(6, 6), epoch=epoch)
cm.mkdir(root_path + 'preview/val/Unlabeled')
plt.savefig(root_path + 'preview/val/Unlabeled/' + 'epoch_%s.jpg' % epoch)
# plt.show(block=False)
# plt.pause(1.0)
plt.close(fig)
if i == 0:
running_loss = running_loss
else:
running_loss = running_loss / (i+1)
return running_loss, running_loss
def network_training_ssl_epoch(Test_only, job, data_seed, data_split, device,
data_sizes, modelDataLoader, num_epoch,
folder_name, TSNE):
val_dice = 0
test_results = 0
device = device
dataset_sizes = data_sizes
print('-' * 64)
print('Training start')
basic_path = folder_name + str(job) + '/' + str(data_split)
if job == 'CNN_baseline':
switch = {
'trainL_encoder': True,
'trainL_decoder_seg': True,
'trainL_decoder_rec': False,
'trainU_encoder': False,
'trainU_decoder_seg': False,
'trainU_decoder_rec': False
}
root_path = basic_path + '/seed' + str(data_seed) + '/'
cm.mkdir(root_path + 'model')
cm.mkdir(root_path + 'preview')
cm.mkdir(root_path + 'preview/train/Labeled')
cm.mkdir(root_path + 'preview/train/Unlabeled')
base_features = 16
model = ssl_3d_sep.MSSL_norm(1, 1, base_features).to(device)
# model = ssl_3d_sep.semiSupervised3D_sep(1, 1, base_features).to(device)
Pretrain = False
# pretrain
if Pretrain:
model.load_state_dict(torch.load(root_path + 'model/val_unet.pth'))
if not Test_only:
criterionDICE = DiceCoefficientLF(device)
criterionMSE = nn.MSELoss()
criterion = (criterionDICE, criterionMSE)
# optimizer_ft = optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=0.0000)
optimizer_ft = (optim.Adam(model.parameters(), lr=1e-2),
optim.Adam(model.parameters(), lr=1e-3))
exp_lr_scheduler = (lr_scheduler.StepLR(optimizer_ft[0],
step_size=500,
gamma=0.5),
lr_scheduler.StepLR(optimizer_ft[1],
step_size=500,
gamma=0.5))
# save training information
train_info = (
'job: {}\n\ndata random seed: {}\n\ndata_split: {}\n\ndataset sizes: {}\n\nmodel: {}\n\nbase features: {}\n\nnetwork_switch: {}\n\nloss function: {}\n\n'
'optimizer: {}\n\nlr scheduler: {}\n\n'.format(
job, data_seed, data_split, dataset_sizes, type(model),
base_features, switch, criterion, optimizer_ft,
exp_lr_scheduler))
cm.history_log(root_path + 'info.txt', train_info, 'w')
print('data random seed: ', data_seed)
print('device: ', device)
print('dataset sizes: ', dataset_sizes)
print('-' * 64)
model, val_dice = train_model(model,
modelDataLoader,
device,
root_path,
switch,
criterion,
optimizer_ft,
exp_lr_scheduler,
num_epochs=num_epoch,
unet=True)
# Testing model
test_results = test_net_dice(root_path,
basic_path,
model,
switch,
modelDataLoader['test'],
TSNE,
gpu=True)
print('CNN_baseline training finished')
else:
# Testing model
test_results = test_net_dice(root_path,
basic_path,
model,
switch,
modelDataLoader['test'],
TSNE,
gpu=True)
print('CNN_baseline training finished')
elif job == 'MSSL_pretrain_Decoder':
# # Autoencoder:
switch = {
'trainL_encoder': False,
'trainL_decoder_seg': False,
'trainL_decoder_rec': False,
'trainU_encoder': True,
'trainU_decoder_seg': False,
'trainU_decoder_rec': True
}
root_path = basic_path + '/seed' + str(data_seed) + '/'
cm.mkdir(root_path + 'model')
cm.mkdir(root_path + 'preview')
cm.mkdir(root_path + 'preview/train/Labeled')
cm.mkdir(root_path + 'preview/train/Unlabeled')
base_features = 16
model = ssl_3d_sep.MSSL_norm(1, 1, base_features).to(device)
if not Test_only:
criterionDICE = DiceCoefficientLF(device)
criterionMSE = nn.MSELoss()
criterion = (criterionDICE, criterionMSE)
# optimizer_ft = optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=0.0000)
optimizer_ft = (optim.Adam(model.parameters(), lr=1e-2),
optim.Adam(model.parameters(), lr=1e-3))
exp_lr_scheduler = (lr_scheduler.StepLR(optimizer_ft[0],
step_size=500,
gamma=0.5),
lr_scheduler.StepLR(optimizer_ft[1],
step_size=500,
gamma=0.5))
# save training information
train_info = (
'job: {}\n\ndata random seed: {}\n\ndata_split: {}\n\ndataset sizes: {}\n\nmodel: {}\n\nbase features: {}\n\nnetwork_switch: {}\n\nloss function: {}\n\n'
'optimizer: {}\n\nlr scheduler: {}\n\n'.format(
job, data_seed, data_split, dataset_sizes, type(model),
base_features, switch, criterion, optimizer_ft,
exp_lr_scheduler))
cm.history_log(root_path + 'info.txt', train_info, 'w')
print('data random seed: ', data_seed)
print('device: ', device)
print('dataset sizes: ', dataset_sizes)
print('-' * 64)
model, val_dice = train_model(model,
modelDataLoader,
device,
root_path,
switch,
criterion,
optimizer_ft,
exp_lr_scheduler,
num_epochs=num_epoch)
# Save model and results
torch.save(model.state_dict(), root_path + 'model/best_unet.pth')
print('Autoencoder pretraining finished')
# Decoder_Seg:
switch = {
'trainL_encoder': False,
'trainL_decoder_seg': True,
'trainL_decoder_rec': False,
'trainU_encoder': False,
'trainU_decoder_seg': False,
'trainU_decoder_rec': False
}
base_features = 16
model = ssl_3d_sep.MSSL_norm(1, 1, base_features).to(device)
# pretrain
model.load_state_dict(torch.load(root_path + 'model/val_unet.pth'))
# reset root path
root_path = basic_path + '/seed' + str(
data_seed) + '/' + 'Decoder_Seg/'
cm.mkdir(root_path + 'model')
cm.mkdir(root_path + 'preview')
cm.mkdir(root_path + 'preview/train/Labeled')
cm.mkdir(root_path + 'preview/train/Unlabeled')
if not Test_only:
criterionDICE = DiceCoefficientLF(device)
criterionMSE = nn.MSELoss()
criterion = (criterionDICE, criterionMSE)
# optimizer_ft = optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=0.0000)
optimizer_ft = (optim.Adam(model.parameters(), lr=1e-2),
optim.Adam(model.parameters(), lr=1e-3))
exp_lr_scheduler = (lr_scheduler.StepLR(optimizer_ft[0],
step_size=500,
gamma=0.5),
lr_scheduler.StepLR(optimizer_ft[1],
step_size=500,
gamma=0.5))
# save training information
train_info = (
'job: {}\n\ndata random seed: {}\n\ndata_split: {}\n\ndataset sizes: {}\n\nmodel: {}\n\nbase features: {}\n\nnetwork_switch: {}\n\nloss function: {}\n\n'
'optimizer: {}\n\nlr scheduler: {}\n\n'.format(
job, data_seed, data_split, dataset_sizes, type(model),
base_features, switch, criterion, optimizer_ft,
exp_lr_scheduler))
cm.history_log(root_path + 'info.txt', train_info, 'w')
print('data random seed: ', data_seed)
print('device: ', device)
print('dataset sizes: ', dataset_sizes)
print('-' * 64)
model, val_dice = train_model(model,
modelDataLoader,
device,
root_path,
switch,
criterion,
optimizer_ft,
exp_lr_scheduler,
num_epochs=num_epoch)
# Testing model
test_results = test_net_dice(root_path,
basic_path,
model,
switch,
modelDataLoader['test'],
TSNE,
gpu=True)
print('Decoder_Seg training finished')
else:
# Testing model
test_results = test_net_dice(root_path,
basic_path,
model,
switch,
modelDataLoader['test'],
TSNE,
gpu=True)
print('Decoder_Seg testing finished')
elif job == 'MSSL_pretrain_CNN':
# Autoencoder:
switch = {
'trainL_encoder': False,
'trainL_decoder_seg': False,
'trainL_decoder_rec': False,
'trainU_encoder': True,
'trainU_decoder_seg': False,
'trainU_decoder_rec': True
}
root_path = basic_path + '/seed' + str(data_seed) + '/'
cm.mkdir(root_path + 'model')
cm.mkdir(root_path + 'preview')
cm.mkdir(root_path + 'preview/train/Labeled')
cm.mkdir(root_path + 'preview/train/Unlabeled')
base_features = 16
model = ssl_3d_sep.MSSL_norm(1, 1, base_features).to(device)
if not Test_only:
criterionDICE = DiceCoefficientLF(device)
criterionMSE = nn.MSELoss()
criterion = (criterionDICE, criterionMSE)
# optimizer_ft = optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=0.0000)
optimizer_ft = (optim.Adam(model.parameters(), lr=1e-2),
optim.Adam(model.parameters(), lr=1e-3))
exp_lr_scheduler = (lr_scheduler.StepLR(optimizer_ft[0],
step_size=500,
gamma=0.5),
lr_scheduler.StepLR(optimizer_ft[1],
step_size=500,
gamma=0.5))
# save training information
train_info = (
'job: {}\n\ndata random seed: {}\n\ndata_split: {}\n\ndataset sizes: {}\n\nmodel: {}\n\nbase features: {}\n\nnetwork_switch: {}\n\nloss function: {}\n\n'
'optimizer: {}\n\nlr scheduler: {}\n\n'.format(
job, data_seed, data_split, dataset_sizes, type(model),
base_features, switch, criterion, optimizer_ft,
exp_lr_scheduler))
cm.history_log(root_path + 'info.txt', train_info, 'w')
print('data random seed: ', data_seed)
print('device: ', device)
print('dataset sizes: ', dataset_sizes)
print('-' * 64)
model, val_dice = train_model(model,
modelDataLoader,
device,
root_path,
switch,
criterion,
optimizer_ft,
exp_lr_scheduler,
num_epochs=num_epoch)
# Save model and results
torch.save(model.state_dict(), root_path + 'model/best_unet.pth')
print('Autoencoder pretraining finished')
# Decoder_Seg:
switch = {
'trainL_encoder': True,
'trainL_decoder_seg': True,
'trainL_decoder_rec': False,
'trainU_encoder': False,
'trainU_decoder_seg': False,
'trainU_decoder_rec': False
}
base_features = 16
model = ssl_3d_sep.MSSL_norm(1, 1, base_features).to(device)
# pretrain
model.load_state_dict(torch.load(root_path + 'model/best_unet.pth'))
# reset root path
root_path = basic_path + '/seed' + str(data_seed) + '/' + 'UNet/'
cm.mkdir(root_path + 'model')
cm.mkdir(root_path + 'preview')
cm.mkdir(root_path + 'preview/train/Labeled')
cm.mkdir(root_path + 'preview/train/Unlabeled')
if not Test_only:
criterionDICE = DiceCoefficientLF(device)
criterionMSE = nn.MSELoss()
criterion = (criterionDICE, criterionMSE)
# optimizer_ft = optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=0.0000)
optimizer_ft = (optim.Adam(model.parameters(), lr=1e-2),
optim.Adam(model.parameters(), lr=1e-3))
exp_lr_scheduler = (lr_scheduler.StepLR(optimizer_ft[0],
step_size=500,
gamma=0.5),
lr_scheduler.StepLR(optimizer_ft[1],
step_size=500,
gamma=0.5))
# save training information
train_info = (
'job: {}\n\ndata random seed: {}\n\ndata_split: {}\n\ndataset sizes: {}\n\nmodel: {}\n\nbase features: {}\n\nnetwork_switch: {}\n\nloss function: {}\n\n'
'optimizer: {}\n\nlr scheduler: {}\n\n'.format(
job, data_seed, data_split, dataset_sizes, type(model),
base_features, switch, criterion, optimizer_ft,
exp_lr_scheduler))
cm.history_log(root_path + 'info.txt', train_info, 'w')
print('data random seed: ', data_seed)
print('device: ', device)
print('dataset sizes: ', dataset_sizes)
print('-' * 64)
model, val_dice = train_model(model,
modelDataLoader,
device,
root_path,
switch,
criterion,
optimizer_ft,
exp_lr_scheduler,
num_epochs=num_epoch)
# Testing model
test_results = test_net_dice(root_path,
basic_path,
model,
switch,
modelDataLoader['test'],
TSNE,
gpu=True)
print('Pretrain UNet testing finished')
else:
# Testing model
test_results = test_net_dice(root_path,
basic_path,
model,
switch,
modelDataLoader['test'],
TSNE,
gpu=True)
print('Pretrain UNet testing finished')
elif job == 'MSSL_alter_rec':
switch = {
'trainL_encoder': True,
'trainL_decoder_seg': True,
'trainL_decoder_rec': False,
'trainU_encoder': True,
'trainU_decoder_seg': False,
'trainU_decoder_rec': True
}
root_path = basic_path + '/seed' + str(data_seed) + '/'
cm.mkdir(root_path + 'model')
cm.mkdir(root_path + 'preview')
cm.mkdir(root_path + 'preview/train/Labeled')
cm.mkdir(root_path + 'preview/train/Unlabeled')
base_features = 16
model = ssl_3d_sep.MSSL_norm(1, 1, base_features).to(device)
Pretrain = False
# pretrain
if Pretrain:
model.load_state_dict(torch.load(root_path + 'model/val_unet.pth'))
if not Test_only:
criterionDICE = DiceCoefficientLF(device)
criterionMSE = nn.MSELoss()
criterion = (criterionDICE, criterionMSE)
# optimizer_ft = optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=0.0000)
optimizer_ft = (optim.Adam(model.parameters(), lr=1e-2),
optim.Adam(model.parameters(), lr=1e-3))
exp_lr_scheduler = (lr_scheduler.StepLR(optimizer_ft[0],
step_size=500,
gamma=0.5),
lr_scheduler.StepLR(optimizer_ft[1],
step_size=500,
gamma=0.5))
# save training information
train_info = (
'job: {}\n\ndata random seed: {}\n\ndata_split: {}\n\ndataset sizes: {}\n\nmodel: {}\n\nbase features: {}\n\nnetwork_switch: {}\n\nloss function: {}\n\n'
'optimizer: {}\n\nlr scheduler: {}\n\n'.format(
job, data_seed, data_split, dataset_sizes, type(model),
base_features, switch, criterion, optimizer_ft,
exp_lr_scheduler))
cm.history_log(root_path + 'info.txt', train_info, 'w')
print('data random seed: ', data_seed)
print('device: ', device)
print('dataset sizes: ', dataset_sizes)
print('-' * 64)
model, val_dice = train_model(model,
modelDataLoader,
device,
root_path,
switch,
criterion,
optimizer_ft,
exp_lr_scheduler,
num_epochs=num_epoch,
num_optimizer='one')
# Testing model
test_results = test_net_dice(root_path,
basic_path,
model,
switch,
modelDataLoader['test'],
TSNE,
gpu=True)
print('MSSL_alter_rec training finished')
else:
# Testing model
test_results = test_net_dice(root_path,
basic_path,
model,
switch,
modelDataLoader['test'],
TSNE,
gpu=True)
print('MSSL_alter_rec training finished')
elif job == 'MSSL_jointly':
switch = {
'trainL_encoder': True,
'trainL_decoder_seg': True,
'trainL_decoder_rec': True,
'trainU_encoder': True,
'trainU_decoder_seg': False,
'trainU_decoder_rec': True
}
root_path = basic_path + '/seed' + str(data_seed) + '/'
cm.mkdir(root_path + 'model')
cm.mkdir(root_path + 'preview')
cm.mkdir(root_path + 'preview/train/Labeled')
cm.mkdir(root_path + 'preview/train/Unlabeled')
base_features = 16
model = ssl_3d_sep.MSSL_norm(1, 1, base_features).to(device)
Pretrain = False
# pretrain
if Pretrain:
model.load_state_dict(torch.load(root_path + 'model/val_unet.pth'))
if not Test_only:
criterionDICE = DiceCoefficientLF(device)
criterionMSE = nn.MSELoss()
criterion = (criterionDICE, criterionMSE)
# optimizer_ft = optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=0.0000)
optimizer_ft = (optim.Adam(model.parameters(), lr=1e-2),
optim.Adam(model.parameters(), lr=1e-3))
exp_lr_scheduler = (lr_scheduler.StepLR(optimizer_ft[0],
step_size=500,
gamma=0.5),
lr_scheduler.StepLR(optimizer_ft[1],
step_size=500,
gamma=0.5))
# save training information
train_info = (
'job: {}\n\ndata random seed: {}\n\ndata_split: {}\n\ndataset sizes: {}\n\nmodel: {}\n\nbase features: {}\n\nnetwork_switch: {}\n\nloss function: {}\n\n'
'optimizer: {}\n\nlr scheduler: {}\n\n'.format(
job, data_seed, data_split, dataset_sizes, type(model),
base_features, switch, criterion, optimizer_ft,
exp_lr_scheduler))
cm.history_log(root_path + 'info.txt', train_info, 'w')
print('data random seed: ', data_seed)
print('device: ', device)
print('dataset sizes: ', dataset_sizes)
print('-' * 64)
model, val_dice = train_model(model,
modelDataLoader,
device,
root_path,
switch,
criterion,
optimizer_ft,
exp_lr_scheduler,
num_epochs=num_epoch,
jointly=True,
self=False)
# Testing model
test_results = test_net_dice(root_path,
basic_path,
model,
switch,
modelDataLoader['test'],
TSNE,
gpu=True)
print('MSSL jointly training finished')
else:
# Testing model
test_results = test_net_dice(root_path,
basic_path,
model,
switch,
modelDataLoader['test'],
TSNE,
gpu=True)
print('MSSL jointly testing finished')
return val_dice, test_results
# FLAIRslice = FLAIRVol[i + start_slice]
# maskslice = maskVol[i + start_slice]
# ax = plt.subplot(2, 3, i + 1)
# plt.tight_layout(rect=[0, 0.03, 1, 0.95])
# ax.set_title('Slice {}'.format(i + start_slice))
# ax.axis('off')
# plt.imshow(FLAIRslice, cmap='gray', alpha=transparent1)
# plt.imshow(maskslice, cmap=my_cmap, alpha=transparent2)
# plt.pause(0.001)
# plt.show()
# plt.close()
imageVol = np.concatenate(
(np.expand_dims(T1Vol, axis=0), np.expand_dims(T2Vol, axis=0),
np.expand_dims(FLAIRVol, axis=0), np.expand_dims(T1cVol, axis=0)),
axis=0)
mkd.mkdir('data')
mkd.mkdir('data/BraTS2018/HGG')
np.save(
'data/BraTS2018/HGG/img_%s.npy' %
(str(Patient_dir[nb_file].split('Brats18_')[-1])), imageVol)
np.save(
'data/BraTS2018/HGG/mask_%s.npy' %
(str(Patient_dir[nb_file].split('Brats18_')[-1])), maskVol)
print('BraTS2018/HGG Image process {}/{} finished'.format(
nb_file, len(Patient_dir)))
print('finished')
def network_training_ssl_epoch(Test_only, job, data_seed, data_split, device,
data_sizes, modelDataLoader, num_epoch,
folder_name, TSNE):
val_dice = 0
test_results = 0
device = device
dataset_sizes = data_sizes
print('-' * 64)
print('Training start')
basic_path = folder_name + str(job) + '/' + str(data_split)[:]
#################################################
if job == 'MASSL_alter':
switch = {
'trainL_encoder': True,
'trainL_decoder_seg': True,
'trainL_decoder_rec': False,
'trainU_encoder': True,
'trainU_decoder_seg': False,
'trainU_decoder_rec': True
}
root_path = basic_path + '/seed' + str(data_seed) + '/'
cm.mkdir(root_path + 'model')
cm.mkdir(root_path + 'preview')
cm.mkdir(root_path + 'preview/train/Labeled')
cm.mkdir(root_path + 'preview/train/Unlabeled')
base_features = 16
model = ssl_3d_attention.MASSL_norm(1, 1, base_features).to(device)
use_existing = False
if use_existing:
model.load_state_dict(torch.load(root_path + 'model/val_unet.pth'))
if not Test_only:
criterionDICE = DiceCoefficientLF(device)
criterionMSE = nn.MSELoss()
criterion = (criterionDICE, criterionMSE)
optimizer_ft = (optim.Adam(model.parameters(), lr=1e-2),
optim.Adam(model.parameters(), lr=1e-3))
exp_lr_scheduler = (lr_scheduler.StepLR(optimizer_ft[0],
step_size=500,
gamma=0.5),
lr_scheduler.StepLR(optimizer_ft[1],
step_size=500,
gamma=0.5))
# save training information
train_info = (
'job: {}\n\ndata random seed: {}\n\ndata_split: {}\n\ndataset sizes: {}\n\nmodel: {}\n\n'
'base features: {}\n\nnetwork_switch: {}\n\nloss function: {}\n\n'
'optimizer: {}\n\nlr scheduler: {}\n\n'.format(
job, data_seed, data_split, dataset_sizes, type(model),
base_features, switch, criterion, optimizer_ft,
exp_lr_scheduler))
cm.history_log(root_path + 'info.txt', train_info, 'w')
print('data random seed: ', data_seed)
print('device: ', device)
print('dataset sizes: ', dataset_sizes)
print('-' * 64)
model, val_dice = train_model(model,
modelDataLoader,
device,
root_path,
switch,
criterion,
optimizer_ft,
exp_lr_scheduler,
num_epochs=num_epoch,
loss_weighted=True)
# Testing model
test_results = test_net_dice(root_path,
basic_path,
model,
switch,
modelDataLoader['test'],
TSNE,
gpu=True)
print('MASSL_alter finished')
else:
# Testing model
test_results = test_net_dice(root_path,
basic_path,
model,
switch,
modelDataLoader['test'],
TSNE,
gpu=True)
print('MASSL_alter finished')
###############################################
elif job == 'MASSL_joint':
switch = {
'trainL_encoder': True,
'trainL_decoder_seg': True,
'trainL_decoder_rec': True,
'trainU_encoder': True,
'trainU_decoder_seg': False,
'trainU_decoder_rec': True
}
root_path = basic_path + '/seed' + str(data_seed) + '/'
cm.mkdir(root_path + 'model')
cm.mkdir(root_path + 'preview')
cm.mkdir(root_path + 'preview/train/Labeled')
cm.mkdir(root_path + 'preview/train/Unlabeled')
base_features = 16
model = ssl_3d_attention.MASSL_norm(1, 1, base_features).to(device)
if not Test_only:
criterionDICE = DiceCoefficientLF(device)
criterionMSE = nn.MSELoss()
criterion = (criterionDICE, criterionMSE)
optimizer_ft = (optim.Adam(model.parameters(), lr=1e-2),
optim.Adam(model.parameters(), lr=1e-3))
exp_lr_scheduler = (lr_scheduler.StepLR(optimizer_ft[0],
step_size=500,
gamma=0.5),
lr_scheduler.StepLR(optimizer_ft[1],
step_size=500,
gamma=0.5))
# save training information
train_info = (
'job: {}\n\ndata random seed: {}\n\ndata_split: {}\n\ndataset sizes: {}\n\nmodel: {}\n\n'
'base features: {}\n\nnetwork_switch: {}\n\nloss function: {}\n\n'
'optimizer: {}\n\nlr scheduler: {}\n\n'.format(
job, data_seed, data_split, dataset_sizes, type(model),
base_features, switch, criterion, optimizer_ft,
exp_lr_scheduler))
cm.history_log(root_path + 'info.txt', train_info, 'w')
print('data random seed: ', data_seed)
print('device: ', device)
print('dataset sizes: ', dataset_sizes)
print('-' * 64)
model, val_dice = train_model(model,
modelDataLoader,
device,
root_path,
switch,
criterion,
optimizer_ft,
exp_lr_scheduler,
num_epochs=num_epoch,
loss_weighted=True,
jointly=True)
# Testing model
test_results = test_net_dice(root_path,
basic_path,
model,
switch,
modelDataLoader['test'],
TSNE,
gpu=True)
print('MASSL_outside_loss_jointly finished')
else:
# Testing model
test_results = test_net_dice(root_path,
basic_path,
model,
switch,
modelDataLoader['test'],
TSNE,
gpu=True)
print('MASSL_outside_loss_jointly finished')
return val_dice, test_results
maskVol = np.where(maskVol == 6, 5, maskVol)
maskVol = np.where(maskVol == 7, 6, maskVol)
maskVol = np.where(maskVol == 8, 7, maskVol)
maskVol = np.where(maskVol == 10, 1, maskVol)
maskVol = np.where(maskVol == 9, 4, maskVol)
# print(np.unique(maskVol))
# Padding and cut image
maskVol = maskVol[cut_slice:, crop_h1:(maskVol.shape[1] - crop_h2),
crop_w1:(maskVol.shape[2] - crop_w2)]
# imageVol = np.concatenate((np.expand_dims(T1Vol, axis=0), np.expand_dims(T2Vol, axis=0), np.expand_dims(FLAIRVol, axis=0), np.expand_dims(T1cVol, axis=0)), axis=0)
imageVol = np.expand_dims(T1Vol, axis=0)
print(str(Patient_dir[nb_file].split('/')[0]))
mkdir('/scratch_net/pengyou/himeva/data/MRBrains18')
np.save(
'/scratch_net/pengyou/himeva/data/MRBrains18' + '/img_%s.npy' %
(str(Patient_dir[nb_file].split('/training/')[-1][:3])), imageVol)
np.save(
'/scratch_net/pengyou/himeva/data/MRBrains18' + '/mask_%s.npy' %
(str(Patient_dir[nb_file].split('/training/')[-1][:3])), maskVol)
print('MRBrains18 Image process {}/{} finished'.format(
nb_file, len(Patient_dir)))
print('finished')
mean_all /= 210
std_all /= 210