def __init__(self, params, use_cuda=False):
super(MModel, self).__init__()
self.params = params
self.src_mask_delta_UNet = U_Net(params, img_ch=3, output_ch=11)
self.FG_UNet = U_Net(params, img_ch=30, output_ch=3)
# self.BG_UNet = U_Net(params, img_ch=4, output_ch=3)
self.use_cuda = use_cuda
def train():
model = U_Net(n_channels=num_channels, n_classes=num_classes).to(device)
model_graph = SummaryWriter(comment="UNet")
input_c = torch.rand(1, 3, 256, 256)
# model_graph.add_graph(model, (input_c.to(device),))
model.train()
# criterion = nn.BCELoss()
criterion = nn.NLLLoss2d()
# criterion=MulticlassDiceLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
dataset = My_Dataset(root,
num_classes,
size,
transform=img_transforms,
mask_transform=mask_transforms)
data_loaders = DataLoader(dataset,
batch_size=batch_size,
shuffle=True,
num_workers=4)
train_model(model,
criterion,
optimizer,
data_loaders,
model_graph=model_graph,
num_epochs=num_epochs)
model_graph.close()
def build_model(self):
"""Build generator and discriminator."""
if self.model_type == 'U_Net':
self.unet = U_Net(img_ch=1, output_ch=1)
elif self.model_type == 'R2U_Net':
self.unet = R2U_Net(img_ch=1, output_ch=1, t=self.t)
#init_weights(self.unet, 'normal')
self.optimizer = optim.Adam(list(self.unet.parameters()), self.lr,
(self.beta1, self.beta2))
self.unet.to(self.device)
def build_model(self):
"""Build generator and discriminator."""
if self.model_type == 'U_Net':
self.unet = U_Net(img_ch=self.img_ch, output_ch=1)
elif self.model_type == 'R2U_Net':
self.unet = R2U_Net(img_ch=self.img_ch, output_ch=1, t=self.t)
elif self.model_type == 'AttU_Net':
self.unet = AttU_Net(img_ch=self.img_ch, output_ch=1)
elif self.model_type == 'R2AttU_Net':
self.unet = R2AttU_Net(img_ch=self.img_ch, output_ch=1, t=self.t)
elif self.model_type == 'MixU_Net':
self.unet = MixU_Net(img_ch=self.img_ch, output_ch=1)
elif self.model_type == 'MixAttU_Net':
self.unet = MixAttU_Net(img_ch=self.img_ch, output_ch=1)
elif self.model_type == 'MixR2U_Net':
self.unet = MixR2U_Net(img_ch=self.img_ch, output_ch=1)
elif self.model_type == 'MixR2AttU_Net':
self.unet = MixR2AttU_Net(img_ch=self.img_ch, output_ch=1)
elif self.model_type == 'GhostU_Net':
self.unet = GhostU_Net(img_ch=self.img_ch, output_ch=1)
elif self.model_type == 'GhostU_Net1':
self.unet = GhostU_Net1(img_ch=self.img_ch, output_ch=1)
elif self.model_type == 'GhostU_Net2':
self.unet = GhostU_Net2(img_ch=self.img_ch, output_ch=1)
#pytorch_total_params = sum(p.numel() for p in self.unet.parameters() if p.requires_grad)
#print (pytorch_total_params)
#raise
self.optimizer = optim.Adam(list(self.unet.parameters()), self.lr,
[self.beta1, self.beta2])
self.unet.to(self.device)
def build_model(self):
"""Build our deep learning model."""
if self.model_type == 'U_Net':
self.unet = U_Net(img_ch=1,
output_ch=1,
first_layer_numKernel=self.first_layer_numKernel)
elif self.model_type == 'R2U_Net':
self.unet = R2U_Net(
img_ch=1,
output_ch=1,
t=self.t,
first_layer_numKernel=self.first_layer_numKernel)
elif self.model_type == 'AttU_Net':
self.unet = AttU_Net(
img_ch=1,
output_ch=1,
first_layer_numKernel=self.first_layer_numKernel)
elif self.model_type == 'R2AttU_Net':
self.unet = R2AttU_Net(
img_ch=1,
output_ch=1,
t=self.t,
first_layer_numKernel=self.first_layer_numKernel)
elif self.model_type == 'ResAttU_Net':
self.unet = ResAttU_Net(
UnetLayer=self.UnetLayer,
img_ch=1,
output_ch=1,
first_layer_numKernel=self.first_layer_numKernel)
if self.initialization != 'NA':
init_weights(self.unet, init_type=self.initialization)
self.unet.to(self.device)
def __init__(self, params, use_cuda=False):
super(MModel, self).__init__()
self.params = params
self.src_mask_delta_UNet = U_Net(params, img_ch=3, output_ch=11)
self.FG_UNet = U_Net(params, img_ch=30, output_ch=3, last_conv=False)
self.FG_tgt_Conv = nn.Conv2d(64,
3,
kernel_size=3,
stride=1,
padding=1,
padding_mode='replicate')
self.FG_mask_Conv = nn.Conv2d(64,
1,
kernel_size=3,
stride=1,
padding=1,
padding_mode='replicate')
self.BG_UNet = U_Net(params, img_ch=4, output_ch=3)
self.use_cuda = use_cuda
def get_model(model_name=None):
if model_name==None:
print('Undefined model - Model Name = None')
return None
elif model_name == 'unet':
model = U_Net(UnetLayer=5,img_ch=INPUT_CHANNELS,output_ch=OUTPUT_CHANNELS)
elif model_name == 'res_att_unet':
model = ResAttU_Net(UnetLayer=5,img_ch=INPUT_CHANNELS,output_ch=OUTPUT_CHANNELS)
return model
def build_model(self):
"""Build generator and discriminator."""
if self.model_type == 'U_Net':
self.unet = U_Net(img_ch=3, output_ch=2)
elif self.model_type == 'R2U_Net':
self.unet = R2U_Net(img_ch=3, output_ch=1, t=self.t)
elif self.model_type == 'AttU_Net':
self.unet = AttU_Net(img_ch=3, output_ch=1)
elif self.model_type == 'R2AttU_Net':
self.unet = R2AttU_Net(img_ch=3, output_ch=1, t=self.t)
self.optimizer = optim.Adam(list(self.unet.parameters()), self.lr,
[self.beta1, self.beta2])
self.unet.to(self.device)
def build_model(self):
# Load required model
if self.model_type == 'U_Net':
self.unet = U_Net(img_ch=3, output_ch=1)
elif self.model_type == 'R2U_Net':
self.unet = R2U_Net(img_ch=3, output_ch=1, t=self.t)
elif self.model_type == 'AttU_Net':
self.unet = AttU_Net(img_ch=3, output_ch=1)
elif self.model_type == 'R2AttU_Net':
self.unet = R2AttU_Net(img_ch=3, output_ch=1, t=self.t)
# Load optimizer
self.optimizer = optim.Adam(list(self.unet.parameters()), self.lr,
[self.beta1, self.beta2])
# Move model to device
self.unet.to(self.device)
def build_model(self, config):
"""Build generator and discriminator."""
if self.model_type == 'U_Net':
self.unet = U_Net(img_ch=3, output_ch=1)
elif self.model_type == 'R2U_Net':
self.unet = R2U_Net(img_ch=3, output_ch=1, t=self.t)
elif self.model_type == 'AttU_Net':
self.unet = AttU_Net(img_ch=3, output_ch=1)
elif self.model_type == 'R2AttU_Net':
self.unet = R2AttU_Net(img_ch=3, output_ch=1, t=self.t)
# # Load the pretrained Encoder
# if os.path.isfile(config.pretrained):
# self.unet.load_state_dict(torch.load(config.pretrained))
# print('%s is Successfully Loaded from %s'%(self.model_type,config.pretrained))
self.optimizer = optim.Adam(list(self.unet.parameters()), self.lr,
[self.beta1, self.beta2])
self.unet.to(self.device)
def build_model(self):
"""Build generator and discriminator."""
if self.model_type == 'U_Net':
self.unet = U_Net(UnetLayer=self.UnetLayer,
img_ch=self.img_ch,
output_ch=self.output_ch,
first_layer_numKernel=self.first_layer_numKernel)
elif self.model_type == 'R2U_Net':
self.unet = R2U_Net(
img_ch=self.img_ch,
output_ch=self.output_ch,
t=self.t,
first_layer_numKernel=self.first_layer_numKernel)
elif self.model_type == 'AttU_Net':
self.unet = AttU_Net(
img_ch=self.img_ch,
output_ch=self.output_ch,
first_layer_numKernel=self.first_layer_numKernel)
elif self.model_type == 'R2AttU_Net':
self.unet = R2AttU_Net(
img_ch=self.img_ch,
output_ch=self.output_ch,
t=self.t,
first_layer_numKernel=self.first_layer_numKernel)
elif self.model_type == 'ResAttU_Net':
self.unet = ResAttU_Net(
UnetLayer=self.UnetLayer,
img_ch=self.img_ch,
output_ch=self.output_ch,
first_layer_numKernel=self.first_layer_numKernel)
if self.optimizer_choice == 'Adam':
self.optimizer = optim.Adam(list(self.unet.parameters()),
self.initial_lr,
[self.beta1, self.beta2])
elif self.optimizer_choice == 'SGD':
self.optimizer = optim.SGD(list(self.unet.parameters()),
self.initial_lr, self.momentum)
else:
pass
self.unet.to(self.device)
def build_model(self):
"""Build generator and discriminator."""
if self.model_type == 'U_Net':
self.unet = U_Net(img_ch=3, output_ch=3)
elif self.model_type == 'R2U_Net':
print("------> using R2U <--------")
self.unet = R2U_Net(img_ch=3, output_ch=3, t=self.t)
elif self.model_type == 'AttU_Net':
print("------> using AttU <--------")
self.unet = AttU_Net(img_ch=3, output_ch=3)
elif self.model_type == 'R2AttU_Net':
print("------> using R2-AttU <--------")
self.unet = R2AttU_Net(img_ch=3, output_ch=3, t=self.t)
elif self.model_type == 'ABU_Net':
print("------> using ABU_Net <--------")
self.unet = U_Net_AB(img_ch=3, output_ch=1)
elif self.model_type == 'Multi_Task':
print("------> using Multi_Task Learning <--------")
model = torch.hub.load('pytorch/vision',
'mobilenet_v2',
pretrained=True)
model_infeatures_final_layer = model.classifier[1].in_features
model.classifier = torch.nn.Sequential(
*list(model.classifier.children())[:-1])
for param in model.parameters():
param.requires_grad = True
for param in model.features[18].parameters():
param.requires_grad = True
for param in model.classifier.parameters():
param.requires_grad = True
model_trained_mobilenet = model
print("All trainable parameters of model are")
for name, param in model_trained_mobilenet.named_parameters():
if param.requires_grad:
print(name, param.shape)
self.unet = multi_task_model_classification(
model_trained_mobilenet)
self.optimizer = optim.AdamW(list(self.unet.parameters()), self.lr,
[self.beta1, self.beta2])
self.unet.to(self.device)
def build_model(self):
"""Build generator and discriminator."""
if self.model_type == 'U_Net':
self.unet = U_Net(img_ch=1, output_ch=1)
elif self.model_type == 'R2U_Net':
self.unet = R2U_Net(img_ch=1, output_ch=1, t=self.t)
elif self.model_type == 'AttU_Net':
self.unet = AttU_Net(img_ch=1, output_ch=1)
elif self.model_type == 'R2AttU_Net':
self.unet = R2AttU_Net(img_ch=1, output_ch=1, t=self.t)
if self.optimizer_choice == 'Adam':
self.optimizer = optim.Adam(list(self.unet.parameters()),
self.initial_lr,
[self.beta1, self.beta2])
elif self.optimizer_choice == 'SGD':
self.optimizer = optim.SGD(list(self.unet.parameters()),
self.initial_lr, self.momentum)
else:
pass
self.unet.to(self.device)
class Solver(object):
def __init__(self, config, train_loader, valid_loader, test_loader):
# data loader
self.train_loader = train_loader
self.valid_loader = valid_loader
self.test_loader = test_loader
# Models
self.unet = None
self.optimizer = None
self.img_ch = config['img_ch']
self.output_ch = config['output_ch']
self.criterion = torch.nn.BCELoss() # binary cross entropy loss
# Hyper-parameters
self.lr = config['lr']
self.beta1 = config['beta1'] # momentum1 in Adam
self.beta2 = config['beta2'] # momentum2 in Adam
# Training settings
self.num_epochs = config['num_epochs']
self.num_epochs_decay = config['num_epoches_decay']
self.batch_size = config['batch_size']
# Path
self.model_path = config['model_path']
self.result_path = config['result_path']
self.device = torch.device(
'cuda' if torch.cuda.is_available() else 'cpu')
self.model_type = config['model_type']
self.t = config['t']
self.unet_path = os.path.join(
self.model_path, '%s-%d-%.4f-%d.pkl' %
(self.model_type, self.num_epochs, self.lr, self.num_epochs_decay))
self.best_epoch = 0
self.build_model()
def build_model(self):
"""Build generator and discriminator."""
if self.model_type == 'U_Net':
self.unet = U_Net(img_ch=1, output_ch=1)
elif self.model_type == 'R2U_Net':
self.unet = R2U_Net(img_ch=1, output_ch=1, t=self.t)
#init_weights(self.unet, 'normal')
self.optimizer = optim.Adam(list(self.unet.parameters()), self.lr,
(self.beta1, self.beta2))
self.unet.to(self.device)
def train(self):
"""Print out the network information."""
num_params = 0
for p in self.unet.parameters():
num_params += p.numel(
) # accumulate the number of mmodel parameters
print("The number of parameters: {}".format(num_params))
# ====================================== Training ===========================================#
# network train
if os.path.isfile(self.unet_path):
# Load the pretrained Encoder
self.unet.load_state_dict(torch.load(self.unet_path))
print('%s is Successfully Loaded from %s' %
(self.model_type, self.unet_path))
else:
lr = self.lr
best_unet_score = 0.0
best_epoch = 0
for epoch in range(self.num_epochs):
self.unet.train(True)
epoch_loss = 0
acc = 0. # Accuracy
SE = 0. # Sensitivity (Recall)
SP = 0. # Specificity
PC = 0. # Precision
F1 = 0. # F1 Score
JS = 0. # Jaccard Similarity
DC = 0. # Dice Coefficient
length = 0
for i, (images, GT) in enumerate(self.train_loader):
images, GT = images.to(self.device), GT.to(self.device)
# forward result
SR = self.unet(images)
SR_probs = torch.sigmoid(SR)
SR_flat = SR_probs.view(SR_probs.size(0),
-1) # size(0) is batch_size
GT_flat = GT.view(GT.size(0), -1)
loss = self.criterion(SR_flat, GT_flat)
epoch_loss += loss.item()
# Backprop + optimize
self.unet.zero_grad()
loss.backward()
self.optimizer.step()
acc += get_accuracy(SR, GT)
SE += get_sensitivity(SR, GT)
SP += get_specificity(SR, GT)
PC += get_precision(SR, GT)
F1 += get_F1(SR, GT)
JS += get_JS(SR, GT)
DC += get_DC(SR, GT)
length = length + 1
acc = acc / length
SE = SE / length
SP = SP / length
PC = PC / length
F1 = F1 / length
JS = JS / length
DC = DC / length
# Print the log info
print(
'Epoch [%d/%d], Loss: %.4f, \n[Training] Acc: %.4f, SE: %.4f, SP: %.4f, PC: %.4f,'
' F1: %.4f, JS: %.4f, DC: %.4f' %
(epoch + 1, self.num_epochs, epoch_loss, acc, SE, SP, PC,
F1, JS, DC))
train_accuracy.append(acc)
# Decay learning rate
if (epoch + 1) > (self.num_epochs - self.num_epochs_decay):
lr -= (self.lr / float(self.num_epochs_decay))
for param_group in self.optimizer.param_groups:
param_group['lr'] = lr
print('Decay learning rate to lr: {}.'.format(lr))
# ===================================== Validation ====================================#
self.unet.train(False)
self.unet.eval()
acc = 0. # Accuracy
SE = 0. # Sensitivity (Recall)
SP = 0. # Specificity
PC = 0. # Precision
F1 = 0. # F1 Score
JS = 0. # Jaccard Similarity
DC = 0. # Dice Coefficient
length = 0
for i, (images, GT) in enumerate(self.valid_loader):
images, GT = images.to(self.device), GT.to(self.device)
SR = torch.sigmoid(self.unet(images))
acc += get_accuracy(SR, GT)
SE += get_sensitivity(SR, GT)
SP += get_specificity(SR, GT)
PC += get_precision(SR, GT)
F1 += get_F1(SR, GT)
JS += get_JS(SR, GT)
DC += get_DC(SR, GT)
length = length + 1
acc = acc / length
SE = SE / length
SP = SP / length
PC = PC / length
F1 = F1 / length
JS = JS / length
DC = DC / length
unet_score = JS + DC
print('[Validation] Acc: %.4f, SE: %.4f, SP: %.4f, PC: %.4f, '
'F1: %.4f, JS: %.4f, DC: %.4f' %
(acc, SE, SP, PC, F1, JS, DC))
validation_accuracy.append(acc)
if unet_score > best_unet_score:
best_unet_score = unet_score
self.best_epoch = epoch
best_unet = self.unet.state_dict(
) # contain best parameters for each layer
print('Best %s model score : %.4f' %
(self.model_type, best_unet_score))
torch.save(best_unet, self.unet_path)
def test(self):
self.unet.load_state_dict(torch.load(self.unet_path))
self.unet.eval()
acc = 0. # Accuracy
SE = 0. # Sensitivity (Recall)
SP = 0. # Specificity
PC = 0. # Precision
F1 = 0. # F1 Score
JS = 0. # Jaccard Similarity
DC = 0. # Dice Coefficient
length = 0
result = []
for i, (images, GT) in enumerate(self.test_loader):
images = images.to(self.device)
GT = GT.to(self.device)
SR = torch.sigmoid(self.unet(images))
acc += get_accuracy(SR, GT)
SE += get_sensitivity(SR, GT)
SP += get_specificity(SR, GT)
PC += get_precision(SR, GT)
F1 += get_F1(SR, GT)
JS += get_JS(SR, GT)
DC += get_DC(SR, GT)
length = length + 1
SR = SR.to('cpu')
SR = SR.detach().numpy()
result.extend(SR)
acc = acc / length
SE = SE / length
SP = SP / length
PC = PC / length
F1 = F1 / length
JS = JS / length
DC = DC / length
unet_score = JS + DC
reconstruct_image(self, np.array(result))
f = open(os.path.join(self.result_path, 'result.csv'),
'a',
encoding='utf-8',
newline='')
wr = csv.writer(f)
wr.writerow([
self.model_type,
acc,
SE,
SP,
PC,
F1,
JS,
DC,
self.lr,
self.best_epoch,
self.num_epochs,
self.num_epochs_decay,
])
f.close()
# noinspection PyArgumentList
train, val, test = random_split(dataset, [228, 32, 66])
train_loader = DataLoader(dataset=train, batch_size=batch_size, shuffle=True, num_workers=4)
val_loader = DataLoader(dataset=val, batch_size=batch_size // 2, shuffle=True, num_workers=4)
test_loader = DataLoader(dataset=test, batch_size=batch_size // 4, shuffle=True, num_workers=4)
from nissl_dataset import Nissl_mask_dataset
from network import U_Net
from network import ResAttU_Net
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)
modelunet = U_Net(UnetLayer=5, img_ch=3, output_ch=4).to(device)
modelresunet = ResAttU_Net(UnetLayer=5,img_ch=3,output_ch=4).to(device)
modelunet.load_state_dict(torch.load('/gdrive/MyDrive/models/unet'), strict=False)
# output = modelunet(image.to(device))
modelresunet.load_state_dict(torch.load('/gdrive/MyDrive/models/resunet'), strict=False)
def gt_to_colorimg(masks):
#colors = np.asarray([(201, 58, 64), (242, 207, 1), (0, 152, 75), (101, 172, 228)])#,(56, 34, 132), (160, 194, 56)])
colors = np.asarray([(0,0,0), (255,0,0), (0,255,0), (0,0,255)])
colorimg = np.ones((masks.shape[1], masks.shape[2], 3), dtype=np.float32) * 255
channels, height, width = masks.shape
def train(cycle_num,
dirs,
path_to_net,
plotter,
batch_size=12,
test_split=0.3,
random_state=666,
epochs=100,
learning_rate=0.0001,
momentum=0.9,
num_folds=5,
num_slices=155,
n_classes=4):
"""
Applies training on the network
Args:
cycle_num (int): number of cycle in n-fold (num_folds) cross validation
dirs (string): path to dataset subject directories
path_to_net (string): path to directory where to save network
plotter (callable): visdom plotter
batch_size - default (int): batch size
test_split - default (float): percentage of test split
random_state - default (int): seed for k-fold cross validation
epochs - default (int): number of epochs
learning_rate - default (float): learning rate
momentum - default (float): momentum
num_folds - default (int): number of folds in cross validation
num_slices - default (int): number of slices per volume
n_classes - default (int): number of classes (regions)
"""
print('Setting started', flush=True)
# Creating data indices
# arange len of list of subject dirs
indices = np.arange(len(glob.glob(dirs + '*')))
test_indices, trainset_indices = get_test_indices(indices, test_split)
# kfold index generator
for cv_num, (train_indices, val_indices) in enumerate(
get_train_cv_indices(trainset_indices, num_folds, random_state)):
# splitted the 5-fold CV in 5 jobs
if cv_num != int(cycle_num):
continue
net = U_Net()
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
num_GPU = torch.cuda.device_count()
if num_GPU > 1:
print('Let us use {} GPUs!'.format(num_GPU), flush=True)
net = nn.DataParallel(net)
net.to(device)
criterion = nn.CrossEntropyLoss()
if cycle_num % 2 == 0:
optimizer = optim.SGD(net.parameters(),
lr=learning_rate,
momentum=momentum)
else:
optimizer = optim.Adam(net.parameters(), lr=learning_rate)
scheduler = ReduceLROnPlateau(optimizer, threshold=1e-6, patience=0)
print('cv cycle number: ', cycle_num, flush=True)
start = time.time()
print('Start Train and Val loading', flush=True)
MRIDataset_train = dataset.MRIDataset(dirs, train_indices)
MRIDataset_val = dataset.MRIDataset(dirs, val_indices)
datalengths = {
'train': len(MRIDataset_train),
'val': len(MRIDataset_val)
}
dataloaders = {
'train': get_dataloader(MRIDataset_train, batch_size, num_GPU),
'val': get_dataloader(MRIDataset_val, batch_size, num_GPU)
}
print('Train and Val loading took: ', time.time() - start, flush=True)
# make loss and acc history for train and val separatly
# Setup Metrics
running_metrics_val = runningScore(n_classes)
running_metrics_train = runningScore(n_classes)
val_loss_meter = averageMeter()
train_loss_meter = averageMeter()
itr = 0
iou_best = 0.
for epoch in tqdm(range(epochs), desc='Epochs'):
print('Epoch: ', epoch + 1, flush=True)
phase = 'train'
print('Phase: ', phase, flush=True)
start = time.time()
# Set model to training mode
net.train()
# Iterate over data.
for i, data in tqdm(enumerate(dataloaders[phase]),
desc='Data Iteration ' + phase):
if (i + 1) % 100 == 0:
print('Number of Iteration [{}/{}]'.format(
i + 1, int(datalengths[phase] / batch_size)),
flush=True)
# get the inputs
inputs = data['mri_data'].to(device)
GT = data['seg'].to(device)
subject_slice_path = data['subject_slice_path']
# Clear all accumulated gradients
optimizer.zero_grad()
# Predict classes using inputs from the train set
SR = net(inputs)
# Compute the loss based on the predictions and
# actual segmentation
loss = criterion(SR, GT)
# Backpropagate the loss
loss.backward()
# Adjust parameters according to the computed
# gradients
# -- weight update
optimizer.step()
# Trake and plot metrics and loss, and save network
predictions = SR.data.max(1)[1].cpu().numpy()
GT_cpu = GT.data.cpu().numpy()
running_metrics_train.update(GT_cpu, predictions)
train_loss_meter.update(loss.item(), n=1)
if (i + 1) % 100 == 0:
itr += 1
score, class_iou = running_metrics_train.get_scores()
for k, v in score.items():
plotter.plot(k, 'itr', phase, k, itr, v)
for k, v in class_iou.items():
print('Class {} IoU: {}'.format(k, v), flush=True)
plotter.plot(
str(k) + ' Class IoU', 'itr', phase,
str(k) + ' Class IoU', itr, v)
print('Loss Train', train_loss_meter.avg, flush=True)
plotter.plot('Loss', 'itr', phase, 'Loss Train', itr,
train_loss_meter.avg)
print('Phase {} took {} s for whole {}set!'.format(
phase,
time.time() - start, phase),
flush=True)
# Validation Phase
phase = 'val'
print('Phase: ', phase, flush=True)
start = time.time()
# Set model to evaluation mode
net.eval()
start = time.time()
with torch.no_grad():
# Iterate over data.
for i, data in tqdm(enumerate(dataloaders[phase]),
desc='Data Iteration ' + phase):
if (i + 1) % 100 == 0:
print('Number of Iteration [{}/{}]'.format(
i + 1, int(datalengths[phase] / batch_size)),
flush=True)
# get the inputs
inputs = data['mri_data'].to(device)
GT = data['seg'].to(device)
subject_slice_path = data['subject_slice_path']
# Clear all accumulated gradients
optimizer.zero_grad()
# Predict classes using inputs from the train set
SR = net(inputs)
# Compute the loss based on the predictions and
# actual segmentation
loss = criterion(SR, GT)
# Trake and plot metrics and loss
predictions = SR.data.max(1)[1].cpu().numpy()
GT_cpu = GT.data.cpu().numpy()
running_metrics_val.update(GT_cpu, predictions)
val_loss_meter.update(loss.item(), n=1)
if (i + 1) % 100 == 0:
itr += 1
score, class_iou = running_metrics_val.get_scores()
for k, v in score.items():
plotter.plot(k, 'itr', phase, k, itr, v)
for k, v in class_iou.items():
print('Class {} IoU: {}'.format(k, v), flush=True)
plotter.plot(
str(k) + ' Class IoU', 'itr', phase,
str(k) + ' Class IoU', itr, v)
print('Loss Val', val_loss_meter.avg, flush=True)
plotter.plot('Loss ', 'itr', phase, 'Loss Val', itr,
val_loss_meter.avg)
if (epoch + 1) % 10 == 0:
if score['Mean IoU'] > iou_best:
save_net(path_to_net, batch_size, epoch, cycle_num,
train_indices, val_indices, test_indices, net,
optimizer)
iou_best = score['Mean IoU']
save_output(epoch, path_to_net, subject_slice_path,
SR.data.cpu().numpy(), GT_cpu)
print('Phase {} took {} s for whole {}set!'.format(
phase,
time.time() - start, phase),
flush=True)
# Call the learning rate adjustment function after every epoch
scheduler.step(val_loss_meter.avg)
# save network after training
save_net(path_to_net,
batch_size,
epochs,
cycle_num,
train_indices,
val_indices,
test_indices,
net,
optimizer,
iter_num=None)
from utils import COLOR_DICT
from utils import dense_crf
from utils import intersectionAndUnion
from PIL import Image
os.environ["CUDA_VISIBLE_DEVICES"] = "3"
# 是否使用cuda
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
num_classes = 2
num_channels = 3
batch_size = 4
size = (256, 256)
root = "data/membrane/test"
img_file = search_file(root, [".png"])
# print(img_file)
if __name__ == "__main__":
model = U_Net(num_channels, num_classes).to(device)
model.load_state_dict(torch.load('UNet_weights_bilinear_weight.pth'))
model.eval()
with torch.no_grad():
for i in range(1):
print(img_file[i])
input = cv2.imread(img_file[i], cv2.IMREAD_COLOR)
input = cv2.resize(input, size)
original_img = input
print(
os.path.join(
"data/membrane/result1",
os.path.splitext(os.path.basename(img_file[i]))[0] +
"_predict.png"), )
label = cv2.imread(
os.path.join(