def create_loss_ops(self, logits_pred, weight, y_true):
# soft means non-thresholded, so takes values in between [0,1]
y_pred_soft, y_pred, _ = prediction(logits_pred)
y_pred_soft = tf.expand_dims(y_pred_soft, -1)
loss_d = losses.dice_loss(y_pred_soft, y_true, weight)
loss_j = losses.iou_loss(y_pred_soft, y_true, weight)
loss_p = losses.pixel_wise_loss(logits_pred, y_true, weight)
if self.cfg.loss == 'dice':
loss = loss_d
elif self.cfg.loss == 'jacc':
loss = loss_j
elif self.cfg.loss == 'sce':
loss = loss_p
else:
raise ValueError
# Add summaries.
tf.summary.scalar("step_losses/dice_loss", loss_d)
tf.summary.scalar("step_losses/jacc_loss", loss_j)
tf.summary.scalar("step_losses/cross_entropy_loss", loss_p)
return [loss, loss_d, loss_p, loss_j]
def main(argv=None):
os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu_list
# check if checkpoint path exists
if not os.path.exists(FLAGS.checkpoint_path):
os.mkdir(FLAGS.checkpoint_path)
else:
#if not FLAGS.restore:
# shutil.rmtree(FLAGS.checkpoint_path)
# os.mkdir(FLAGS.checkpoint_path)
shutil.rmtree(FLAGS.checkpoint_path)
os.mkdir(FLAGS.checkpoint_path)
train_data_generator = data_processor.generator(FLAGS)
train_samples_count = data_processor.count_samples(FLAGS)
val_data = data_processor.load_data(FLAGS)
if len(gpus) <= 1:
print('Training with 1 GPU')
if FLAGS.drn:
east = EAST_DRN_model(input_size=FLAGS.input_size)
else:
east = EAST_model(FLAGS.input_size)
parallel_model = east.model
else:
print('Training with %d GPUs' % len(gpus))
with tf.device("/cpu:0"):
east = EAST_model(FLAGS.input_size)
if FLAGS.restore_model is not '':
east.model.load_weights(FLAGS.restore_model)
parallel_model = multi_gpu_model(east.model, gpus=len(gpus))
score_map_loss_weight = K.variable(0.01, name='score_map_loss_weight')
small_text_weight = K.variable(0., name='small_text_weight')
lr_scheduler = LearningRateScheduler(lr_decay)
ckpt = CustomModelCheckpoint(model=east.model, path=FLAGS.checkpoint_path + '/model-{epoch:02d}.h5', period=FLAGS.save_checkpoint_epochs, save_weights_only=True)
tb = CustomTensorBoard(log_dir=FLAGS.checkpoint_path + '/train', score_map_loss_weight=score_map_loss_weight, small_text_weight=small_text_weight, data_generator=train_data_generator, write_graph=True)
small_text_weight_callback = SmallTextWeight(small_text_weight)
validation_evaluator = ValidationEvaluator(val_data, validation_log_dir=FLAGS.checkpoint_path + '/val')
callbacks = [lr_scheduler, ckpt, tb, small_text_weight_callback, validation_evaluator]
opt = AdamW(FLAGS.init_learning_rate)
parallel_model.compile(loss=[dice_loss(east.overly_small_text_region_training_mask, east.text_region_boundary_training_mask, score_map_loss_weight, small_text_weight),
rbox_loss(east.overly_small_text_region_training_mask, east.text_region_boundary_training_mask, small_text_weight, east.target_score_map)],
loss_weights=[1., 1.],
optimizer=opt)
east.model.summary()
model_json = east.model.to_json()
with open(FLAGS.checkpoint_path + '/model.json', 'w') as json_file:
json_file.write(model_json)
history = parallel_model.fit_generator(train_data_generator, epochs=FLAGS.max_epochs, steps_per_epoch=train_samples_count/FLAGS.batch_size, workers=FLAGS.nb_workers, use_multiprocessing=True, callbacks=callbacks, verbose=1)
def main(argv=None):
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
# os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu_list
# check if checkpoint path exists
if not os.path.exists(FLAGS.checkpoint_path):
os.mkdir(FLAGS.checkpoint_path)
else:
#if not FLAGS.restore:
# shutil.rmtree(FLAGS.checkpoint_path)
# os.mkdir(FLAGS.checkpoint_path)
shutil.rmtree(FLAGS.checkpoint_path)
os.mkdir(FLAGS.checkpoint_path)
# train_data_generator = data_processor.generator(FLAGS)
# train_samples_count = data_processor.count_samples(FLAGS)
train_data_generator = data_processor.DataGenerator(FLAGS)
# train_data_generator = data_processor.DataGenerator(FLAGS)
# val_data = data_processor.load_data(FLAGS)
val_data = data_processor.val_generator(FLAGS)
val_samples_count = data_processor.count_val_samples(FLAGS)
east = EAST_model(FLAGS.input_size)
if FLAGS.restore_model != '':
east.model.load_weights(FLAGS.restore_model)
score_map_loss_weight = K.variable(0.01, name='score_map_loss_weight')
small_text_weight = K.variable(0., name='small_text_weight')
lr_scheduler = LearningRateScheduler(lr_decay)
ckpt = CustomModelCheckpoint(model=east.model, path=FLAGS.checkpoint_path + '/model-{epoch:02d}.h5', period=FLAGS.save_checkpoint_epochs, save_weights_only=False)
small_text_weight_callback = SmallTextWeight(small_text_weight)
validation_evaluator = ValidationEvaluator(val_data, val_samples_count)
callbacks = [lr_scheduler, ckpt, small_text_weight_callback, validation_evaluator]
opt = tfa.optimizers.AdamW(weight_decay=1e-4, learning_rate=FLAGS.init_learning_rate, name='AdamW')
east.model.compile(
loss=[dice_loss(east.overly_small_text_region_training_mask, east.text_region_boundary_training_mask, score_map_loss_weight, small_text_weight),
rbox_loss(east.overly_small_text_region_training_mask, east.text_region_boundary_training_mask, small_text_weight, east.target_score_map)],
loss_weights=[1., 1.],
optimizer=opt,
# metrics=[['accuracy'], ['accuracy']]
)
east.model.summary()
model_json = east.model.to_json()
with open(FLAGS.checkpoint_path + '/model.json', 'w') as json_file:
json_file.write(model_json)
# history = east.model.fit(train_data_generator, epochs=FLAGS.max_epochs, steps_per_epoch=train_samples_count/FLAGS.batch_size, workers=FLAGS.nb_workers, use_multiprocessing=True, max_queue_size=10, callbacks=callbacks, verbose=1)
history = east.model.fit(train_data_generator, epochs=FLAGS.max_epochs, workers=FLAGS.nb_workers, use_multiprocessing=True, max_queue_size=10, callbacks=callbacks, verbose=1)
east.model.save('east_retrained.h5')
def train(train_loader, model, criterion, optimizer, epoch, args):
batch_time = AverageMeter('Time', ':6.3f')
data_time = AverageMeter('Data', ':6.3f')
losses = AverageMeter('Loss', ':.4e')
top1 = AverageMeter('Acc@1', ':6.2f')
dice = AverageMeter('Dice', ':6.2f')
progress = ProgressMeter(len(train_loader),
[batch_time, data_time, losses, top1, dice],
prefix="Epoch: [{}]".format(epoch))
# switch to train mode
if args.freeze_bn:
model.eval()
else:
model.train()
end = time.time()
for i, (images, target) in enumerate(train_loader):
if args.debug and i > 10: break
# measure data loading time
data_time.update(time.time() - end)
if args.gpu is not None:
images = images.cuda(args.gpu, non_blocking=True)
target = target.cuda(args.gpu, non_blocking=True)
# compute output
output_coarse, output = model(images)
loss = 0.6 * criterion(output, target) + 0.7 * criterion(
output_coarse, target)
# measure accuracy and record loss
acc1 = accuracy(output.permute(0, 2, 3,
1).reshape(-1, args.num_classes),
target.view(-1),
topk=(1, ))
losses.update(loss.item(),
output.size(0) * output.size(2) * output.size(3))
top1.update(acc1[0][0].item(),
output.size(0) * output.size(2) * output.size(3))
dice.update(1 - dice_loss(output, target).item(), output.size(0))
# compute gradient and do SGD step
optimizer.zero_grad()
loss.backward()
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % args.print_freq == 0:
progress.display(i)
writer.add_scalar('acc1/acc1_train', top1.avg, epoch)
writer.add_scalar('dice/dice_train', dice.avg, epoch)
def main():
train_data_generator = DataGenerator(input_size=FLAGS.input_size, batch_size=FLAGS.batch_size,
data_path=FLAGS.training_data_path, FLAGS=FLAGS, is_train=True)
train_samples_count = len(train_data_generator.image_paths)
validation_data_generator = DataGenerator(input_size=FLAGS.input_size, batch_size=FLAGS.batch_size,
data_path=FLAGS.validation_data_path, FLAGS=FLAGS, is_train=False)
east = EastModel(FLAGS.input_size)
if FLAGS.pretrained_weights_path != '':
print(f'Loading pre-trained model at {FLAGS.pretrained_weights_path}')
east.model.load_weights(FLAGS.pretrained_weights_path)
score_map_loss_weight = K.variable(0.01, name='score_map_loss_weight')
small_text_weight = K.variable(0., name='small_text_weight')
opt = AdamW(FLAGS.init_learning_rate)
east.model.compile(
loss=[
dice_loss(east.overly_small_text_region_training_mask, east.text_region_boundary_training_mask,
score_map_loss_weight, small_text_weight),
rbox_loss(east.overly_small_text_region_training_mask, east.text_region_boundary_training_mask,
small_text_weight, east.target_score_map)
],
loss_weights=[1., 1.],
optimizer=opt,
)
tb_callback = tensorboard_callback()
cp_callback = checkpoint_callback()
with open(os.path.join(FLAGS.checkpoint_path, 'model.json'), 'w') as json_file:
json_file.write(east.model.to_json())
east.model.fit_generator(
generator=train_data_generator,
epochs=FLAGS.max_epochs,
steps_per_epoch=train_samples_count // FLAGS.batch_size,
validation_data=validation_data_generator,
callbacks=[cp_callback, tb_callback],
workers=FLAGS.nb_workers,
use_multiprocessing=True,
max_queue_size=10,
verbose=1,
)
def train_step(model,
x,
optimizer,
overly_small_text_region_training_mask,
text_region_boundary_training_mask,
small_text_weight,
target_score_map,
target_geo_maps,
loss_weight):
with tf.GradientTape() as tape:
score_y_pred, geo_y_pred = model(x)
_dice_loss = dice_loss(overly_small_text_region_training_mask, text_region_boundary_training_mask, loss_weight, small_text_weight, target_score_map, score_y_pred)
_rbox_loss = rbox_loss(overly_small_text_region_training_mask, text_region_boundary_training_mask, small_text_weight, target_score_map, target_geo_maps, geo_y_pred)
loss = _dice_loss + _rbox_loss
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
def compute_loss(input, target, axis=(2, 3)):
target = one_hot(target.squeeze(1))
input, target = input[:, 1:], target[:, 1:]
ce = sigmoid_cross_entropy(input=input, target=target).mean(axis).mean(1)
# focal = softmax_focal_loss(input=input, target=target, axis=1, keepdim=True).mean(axis).mean(1)
dice = dice_loss(input=input.sigmoid(), target=target, axis=axis).mean(1)
loss = [
ce,
# focal,
dice,
]
assert all(l.size() == loss[0].size() for l in loss)
loss = sum(loss) / len(loss)
return loss
def infer(epoch):
sigmoid = nn.Sigmoid()
val_loss = 0
dice = 0
model.eval()
print('===> Evaluating Model')
with torch.no_grad():
for iteration, batch in enumerate(infering_data_loader, 1):
input, target = Variable(batch[0]), Variable(batch[1])
if cuda:
input = input.cuda()
target = target.cuda()
pred = model(input)
loss = crit(pred.squeeze(1), target.float()).item()
val_loss += loss
# Turn the prediction into probabilities
pred = sigmoid(pred)
mask = (pred[:, 0, :, :, :] > 0.5).float()
dice += dice_loss(mask, target.float()).item()
if iteration == 1:
non_zero = [i for i, x in enumerate(target, 0) if x.sum() != 0]
if not non_zero:
writer.add_image('Test/Input', input[0, 0:3, 32, :, :].squeeze(), epoch)
writer.add_image('Test/Prediction', pred[0, 0, 32, :, :].repeat(3, 1, 1), epoch)
writer.add_image('Test/Segmentation', mask[0, 32, :, :].repeat(3, 1, 1), epoch)
writer.add_image('Test/Ground Truth', target[0, 32, :, :].repeat(3, 1, 1), epoch)
else:
index = non_zero[0]
writer.add_image('Test/Input', input[index, 0:3, 32, :, :].squeeze(), epoch)
writer.add_image('Test/Prediction', pred[index, 0, 32, :, :].repeat(3, 1, 1), epoch)
writer.add_image('Test/Segmentation', mask[index, 32, :, :].repeat(3, 1, 1), epoch)
writer.add_image('Test/Ground Truth', target[index, 32, :, :].repeat(3, 1, 1), epoch)
print("=> Done with {} / {} Batch Loss: {:.6f}".format(iteration, len(infering_data_loader), loss))
writer.add_scalar('test/Average Loss', val_loss / len(infering_data_loader), epoch)
writer.add_scalar('test/Average DICE', dice / len(infering_data_loader), epoch)
print("===> Avg. DICE: {:.6f}".format(dice / len(infering_data_loader)))
def loss(logits, labels, nlabels, loss_type, labels_as_1hot=False):
'''
Loss to be minimised by the neural network
:param logits: The output of the neural network before the softmax
:param labels: The ground truth labels in standard (i.e. not one-hot) format
:param nlabels: The number of GT labels
:param loss_type: Can be 'crossentropy'/'dice'/'crossentropy_and_dice'
:return: The segmentation
'''
if labels_as_1hot is False:
labels = tf.one_hot(labels, depth=nlabels, axis=-1)
if loss_type == 'crossentropy':
segmentation_loss = losses.pixel_wise_cross_entropy_loss(
logits, labels)
elif loss_type == 'dice':
segmentation_loss = losses.dice_loss(logits, labels)
else:
raise ValueError('Unknown loss: %s' % loss_type)
return segmentation_loss
def main(args):
torch.backends.cudnn.benchmark = True
seed_all(args.seed)
d = Dataset(train_set_size=args.train_set_sz,
num_cls=args.num_cls,
remove_nan_center=False)
train = d.train_set
valid = d.test_set
num_cls = args.num_cls + 1 # +1 for background
net = UNet(in_dim=1, out_dim=num_cls).cuda()
best_net = UNet(in_dim=1, out_dim=num_cls)
best_val_dice = -np.inf
best_cls_val_dices = None
optimizer = torch.optim.Adam(params=net.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
scheduler_warmup = GradualWarmupScheduler(optimizer,
multiplier=10,
total_epoch=50,
after_scheduler=None)
if not os.path.exists(args.log_dir):
os.makedirs(args.log_dir, exist_ok=True)
writer = tensorboardX.SummaryWriter(log_dir=args.log_dir)
step = 1
for epoch in range(1, args.n_epochs + 1):
for iteration in range(
1,
int(np.ceil(train.dataset_sz() / args.batch_sz)) + 1):
net.train()
imgs, masks, one_hot_masks, centers, _, _, _, _ = train.next_batch(
args.batch_sz)
imgs = make_batch_input(imgs)
imgs = torch.cuda.FloatTensor(imgs)
one_hot_masks = torch.cuda.FloatTensor(one_hot_masks)
pred_logit = net(imgs)
pred_softmax = F.softmax(pred_logit, dim=1)
if args.use_ce:
ce = torch.nn.CrossEntropyLoss()
loss = ce(pred_logit, torch.cuda.LongTensor(masks))
else:
loss = dice_loss(pred_softmax,
one_hot_masks,
keep_background=False).mean()
scheduler_warmup.step()
optimizer.zero_grad()
loss.backward()
optimizer.step()
if step % args.log_freq == 0:
print(
f"step={step}\tepoch={epoch}\titer={iteration}\tloss={loss.data.cpu().numpy()}"
)
writer.add_scalar("cnn_dice_loss",
loss.data.cpu().numpy(), step)
writer.add_scalar("lr", optimizer.param_groups[0]["lr"], step)
if step % args.train_eval_freq == 0:
train_dice, cls_train_dices = do_eval(net, train.images,
train.onehot_masks,
args.batch_sz, num_cls)
train_dice = train_dice.cpu().numpy()
cls_train_dices = cls_train_dices.cpu().numpy()
writer.add_scalar("train_dice", train_dice, step)
# lr_sched.step(1-train_dice)
for j, cls_train_dice in enumerate(cls_train_dices):
writer.add_scalar(f"train_dice/{j}", cls_train_dice, step)
print(
f"step={step}\tepoch={epoch}\titer={iteration}\ttrain_eval: train_dice={train_dice}"
)
if step % args.val_eval_freq == 0:
_pickle.dump(
net.state_dict(),
open(os.path.join(args.log_dir, 'model.pth.tar'), 'wb'))
val_dice, cls_val_dices = do_eval(net, valid.images,
valid.onehot_masks,
args.batch_sz, num_cls)
val_dice = val_dice.cpu().numpy()
cls_val_dices = cls_val_dices.cpu().numpy()
writer.add_scalar("val_dice", val_dice, step)
for j, cls_val_dice in enumerate(cls_val_dices):
writer.add_scalar(f"val_dice/{j}", cls_val_dice, step)
print(
f"step={step}\tepoch={epoch}\titer={iteration}\tvalid_dice={val_dice}"
)
if val_dice > best_val_dice:
best_val_dice = val_dice
best_cls_val_dices = cls_val_dices
best_net.load_state_dict(net.state_dict().copy())
_pickle.dump(
best_net.state_dict(),
open(os.path.join(args.log_dir, 'best_model.pth.tar'),
'wb'))
f = open(
os.path.join(args.log_dir, f"best_val_dice{step}.txt"),
'w')
f.write(str(best_val_dice) + "\n")
f.write(" ".join([
str(dice_score) for dice_score in best_cls_val_dices
]))
f.close()
print(f"better val dice detected.")
# if step % 5000 == 0:
# _pickle.dump(net.state_dict(), open(os.path.join(args.log_dir, '{}.pth.tar'.format(step)),
# 'wb'))
step += 1
return best_val_dice, best_cls_val_dices
def test_simple(self):
dice_coef: float = metrics.dice_coefficient(self.outputs, self.targets)
dice_loss: float = float(losses.dice_loss(self.outputs, self.targets))
self.assertEqual(dice_coef, 1. - dice_loss)
for epoch in range(num_epochs):
train_losses = []
train_dsc = []
val_losses = []
val_dsc = []
steps = 0
for batch_idx, (images, masks) in enumerate(train_loader):
images = images.to(device)
masks = masks.to(device)
model.train()
opt.zero_grad()
preds = model(images)
loss = losses.dice_loss(preds, masks)
loss.backward()
#opt.step()
train_losses.append(loss.item())
train_dsc.append(losses.dice_score(preds, masks).item())
else:
val_loss = 0
val_acc = 0
model.eval()
with torch.no_grad():
for inputs, masks in val_loader:
inputs, masks = inputs.to(device), masks.to(device)
def test_functional(self):
loss_class = self.criterion(self.outputs, self.targets)
loss_func = losses.dice_loss(self.outputs, self.targets)
self.assertTrue(float(loss_class), float(loss_func))
def main():
if args.restart_training == 'true':
if use_multiinput_architecture is False:
if modeltype == 'unet':
model = UNet(n_classes=n_classes,
padding=True,
depth=model_depth,
wf=wf,
up_mode='upconv',
batch_norm=True,
residual=False).double().to(device)
elif modeltype == 'resunet':
model = UNet(n_classes=n_classes,
padding=True,
depth=model_depth,
wf=wf,
up_mode='upconv',
batch_norm=True,
residual=True).double().to(device)
elif use_multiinput_architecture is True:
if modeltype == 'unet':
model = Attention_UNet(
n_classes=n_classes,
padding=True,
up_mode='upconv',
batch_norm=True,
residual=False,
wf=wf,
use_attention=use_attention).double().to(device)
elif modeltype == 'resunet':
model = Attention_UNet(
n_classes=n_classes,
padding=True,
up_mode='upconv',
batch_norm=True,
residual=True,
wf=wf,
use_attention=use_attention).double().to(device)
else:
if use_multiinput_architecture is False:
if modeltype == 'unet':
model = UNet(n_classes=n_classes,
padding=True,
depth=model_depth,
wf=wf,
up_mode='upconv',
batch_norm=True,
residual=False).double().to(device)
elif modeltype == 'resunet':
model = UNet(n_classes=n_classes,
padding=True,
depth=model_depth,
wf=wf,
up_mode='upconv',
batch_norm=True,
residual=True).double().to(device)
# checkpoint = torch.load(args.model_path, map_location=lambda storage, loc: storage)
# pretrained_dict = checkpoint['model_state_dict']
# model_dict = model.state_dict()
# # 1. filter out unnecessary keys
# pretrained_dict = {k: v for k, v in pretrained_dict.items() if k not in ['last.weight', 'last.bias']}
# # 2. overwrite entries in the existing state dict
# model_dict.update(pretrained_dict)
# # 3. load the new state dict
# model.load_state_dict(model_dict)
elif use_multiinput_architecture is True:
if modeltype == 'unet':
model = Attention_UNet(
n_classes=n_classes,
padding=True,
up_mode='upconv',
batch_norm=True,
residual=False,
wf=wf,
use_attention=use_attention).double().to(device)
elif modeltype == 'resunet':
model = Attention_UNet(
n_classes=n_classes,
padding=True,
up_mode='upconv',
batch_norm=True,
residual=True,
wf=wf,
use_attention=use_attention).double().to(device)
# checkpoint = torch.load(args.model_path, map_location=lambda storage, loc: storage)
# model.load_state_dict(checkpoint['model_state_dict'])
checkpoint = torch.load(args.model_path,
map_location=lambda storage, loc: storage)
model.load_state_dict(checkpoint['model_state_dict'])
train_loader = dataloader_cxr.DataLoader(data_path,
dataloader_type='train',
batchsize=batch_size,
device=device,
image_resolution=image_resolution)
print('trainloader loaded')
valid_loader = dataloader_cxr.DataLoader(data_path,
dataloader_type='valid',
batchsize=batch_size,
device=device,
image_resolution=image_resolution)
print('validloader loaded')
loss_list_train_epoch = [None]
dice_score_list_train_epoch = [None]
epoch_data_list = [None]
loss_list_validation = [None]
loss_list_validation_index = [None]
dice_score_list_validation = [None]
dice_score_list_validation_0 = [None]
dice_score_list_validation_1 = [None]
epoch_old = 0
if load_old_lists == True:
if args.restart_training == 'false':
epoch_old = checkpoint['epochs']
if checkpoint['train_loss_list_epoch'][-1] == None:
dice_score_list_train_epoch = [None]
loss_list_train_epoch = [None]
epoch_data_list = [None]
else:
dice_score_list_train_epoch = checkpoint[
'train_dice_score_list_epoch']
loss_list_train_epoch = checkpoint['train_loss_list_epoch']
epoch_data_list = checkpoint['train_loss_index_epoch']
if checkpoint['valid_loss_list'][-1] == None:
loss_list_validation = [None]
loss_list_validation_index = [None]
dice_score_list_validation = [None]
dice_score_list_validation_0 = [None]
dice_score_list_validation_1 = [None]
else:
loss_list_validation = checkpoint['valid_loss_list']
loss_list_validation_index = checkpoint['valid_loss_index']
dice_score_list_validation = checkpoint[
'valid_dice_score_list']
dice_score_list_validation_0 = checkpoint[
'valid_dice_score_list_0']
dice_score_list_validation_1 = checkpoint[
'valid_dice_score_list_1']
best_model_accuracy = np.max(dice_score_list_validation[1:])
if len(train_loader.data_list) % batch_size == 0:
total_idx_train = len(train_loader.data_list) // batch_size
else:
total_idx_train = len(train_loader.data_list) // batch_size + 1
if len(valid_loader.data_list) % batch_size == 0:
total_idx_valid = len(valid_loader.data_list) // batch_size
else:
total_idx_valid = len(valid_loader.data_list) // batch_size + 1
if epoch_old != 0:
power_factor = epoch_old // scheduler_step_size
LR_ = LR * (scheduler_gamma**power_factor)
else:
LR_ = LR
LR_ = LR
optimizer = optim.Adam(model.parameters(), lr=LR_)
# optimizer = optim.SGD(model.parameters(), lr=LR_, momentum=0.9)
scheduler = StepLR(optimizer,
step_size=scheduler_step_size,
gamma=scheduler_gamma)
for epoch in range(epoch_old, train_epoch):
if (epoch + 1) % 10 == 0:
scheduler.step()
epoch_loss = 0.0
epoch_dice_score = 0.0
train_count = 0
model.train()
for idx in range(total_idx_train):
optimizer.zero_grad()
batch_images_input, batch_label_input = train_loader[idx]
output = model(batch_images_input)
if use_multiinput_architecture is False:
loss = losses.dice_loss(
output,
batch_label_input,
weights=torch.Tensor([gamma0, gamma1]).double().to(device))
elif use_multiinput_architecture is True:
loss = losses.dice_loss_deep_supervised(
output,
batch_label_input,
weights=torch.Tensor([gamma0, gamma1]).double().to(device))
loss.backward()
optimizer.step()
if use_multiinput_architecture is False:
score = losses.dice_score(output, batch_label_input)
else:
score = losses.dice_score(output[-1], batch_label_input)
epoch_dice_score += (score.sum().item() /
score.size(0)) * batch_images_input.shape[0]
epoch_loss += loss.item() * batch_images_input.shape[0]
train_count += batch_images_input.shape[0]
loss_list_train_epoch.append(epoch_loss / train_count)
epoch_data_list.append(epoch + 1)
dice_score_list_train_epoch.append(epoch_dice_score / train_count)
print(
'Epoch %d Training Loss: %.3f Dice Score: %.3f' %
(epoch + 1, loss_list_train_epoch[-1],
dice_score_list_train_epoch[-1]), ' Time:',
datetime.datetime.now())
plt.plot(epoch_data_list[1:],
loss_list_train_epoch[1:],
label="Training",
color='red',
marker='o',
markerfacecolor='yellow',
markersize=5)
plt.xlabel('Epoch')
plt.ylabel('Training Loss')
plt.savefig(plots_dir + '/train_loss_plot.png')
plt.clf()
plt.plot(epoch_data_list[1:],
dice_score_list_train_epoch[1:],
label="Training",
color='red',
marker='o',
markerfacecolor='yellow',
markersize=5)
plt.xlabel('Epoch')
plt.ylabel('Training Dice Score')
plt.savefig(plots_dir + '/train_dice_score_plot.png')
plt.clf()
training_pickle = open(plots_pickle_dir + "/loss_list_train.npy", 'wb')
pickle.dump(loss_list_train_epoch, training_pickle)
training_pickle.close()
training_pickle = open(plots_pickle_dir + "/epoch_list_train.npy",
'wb')
pickle.dump(epoch_data_list, training_pickle)
training_pickle.close()
training_pickle = open(
plots_pickle_dir + "/dice_score_list_train_epoch.npy", 'wb')
pickle.dump(dice_score_list_train_epoch, training_pickle)
training_pickle.close()
if (epoch + 1) % save_every == 0:
print('Saving model at %d epoch' % (epoch + 1), ' Time:',
datetime.datetime.now()
) # save every save_every mini_batch of data
torch.save(
{
'epochs': epoch + 1,
'batchsize': batch_size,
'train_loss_list_epoch': loss_list_train_epoch,
'train_dice_score_list_epoch': dice_score_list_train_epoch,
'train_loss_index_epoch': epoch_data_list,
'valid_loss_list': loss_list_validation,
'valid_dice_score_list': dice_score_list_validation,
'valid_dice_score_list_0': dice_score_list_validation_0,
'valid_dice_score_list_1': dice_score_list_validation_1,
'valid_loss_index': loss_list_validation_index,
'model_state_dict': model.state_dict(),
}, model_checkpoint_dir + '/model_%d.pth' % (epoch + 1))
if (epoch + 1) % valid_every == 0:
model.eval()
optimizer.zero_grad()
valid_count = 0
total_loss_valid = 0.0
valid_dice_score = 0.0
valid_dice_score_0 = 0.0
valid_dice_score_1 = 0.0
for idx in range(total_idx_valid):
with torch.no_grad():
batch_images_input, batch_label_input = valid_loader[idx]
output = model(batch_images_input)
if use_multiinput_architecture is False:
loss = losses.dice_loss(output, batch_label_input)
else:
loss = losses.dice_loss(output[-1], batch_label_input)
total_loss_valid += loss.item(
) * batch_images_input.shape[0]
valid_count += batch_images_input.shape[0]
if use_multiinput_architecture is False:
score = losses.dice_score(output, batch_label_input)
else:
score = losses.dice_score(output[-1],
batch_label_input)
valid_dice_score += (score.sum().item() / score.size(0)
) * batch_images_input.shape[0]
valid_dice_score_0 += score[0].item(
) * batch_images_input.shape[0]
valid_dice_score_1 += score[1].item(
) * batch_images_input.shape[0]
loss_list_validation.append(total_loss_valid / valid_count)
dice_score_list_validation.append(valid_dice_score / valid_count)
dice_score_list_validation_0.append(valid_dice_score_0 /
valid_count)
dice_score_list_validation_1.append(valid_dice_score_1 /
valid_count)
loss_list_validation_index.append(epoch + 1)
print(
'Epoch %d Valid Loss: %.3f' %
(epoch + 1, loss_list_validation[-1]), ' Time:',
datetime.datetime.now())
print('Valid Dice Score: ', dice_score_list_validation[-1],
' Valid Dice Score 0: ', dice_score_list_validation_0[-1],
' Valid Dice Score 1: ', dice_score_list_validation_1[-1])
plt.plot(loss_list_validation_index[1:],
loss_list_validation[1:],
label="Validation",
color='red',
marker='o',
markerfacecolor='yellow',
markersize=5)
plt.xlabel('Epoch')
plt.ylabel('Validation Loss')
plt.savefig(plots_dir + '/valid_loss_plot.png')
plt.clf()
plt.plot(loss_list_validation_index[1:],
dice_score_list_validation[1:],
label="Validation",
color='red',
marker='o',
markerfacecolor='yellow',
markersize=5)
plt.xlabel('Epoch')
plt.ylabel('Validation Dice Score')
plt.savefig(plots_dir + '/valid_dice_score_plot.png')
plt.clf()
plt.plot(loss_list_validation_index[1:],
dice_score_list_validation_0[1:],
label="Validation",
color='red',
marker='o',
markerfacecolor='yellow',
markersize=5)
plt.xlabel('Epoch')
plt.ylabel('Validation Dice Score')
plt.savefig(plots_dir + '/valid_dice_score_0_plot.png')
plt.clf()
plt.plot(loss_list_validation_index[1:],
dice_score_list_validation_1[1:],
label="Validation",
color='red',
marker='o',
markerfacecolor='yellow',
markersize=5)
plt.xlabel('Epoch')
plt.ylabel('Validation Dice Score')
plt.savefig(plots_dir + '/valid_dice_score_1_plot.png')
plt.clf()
validation_pickle = open(
plots_pickle_dir + "/loss_list_validation.npy", 'wb')
pickle.dump(loss_list_validation, validation_pickle)
validation_pickle.close()
validation_pickle = open(
plots_pickle_dir + "/index_list_validation.npy", 'wb')
pickle.dump(loss_list_validation_index, validation_pickle)
validation_pickle.close()
validation_pickle = open(
plots_pickle_dir + "/dice_score_list_validation.npy", 'wb')
pickle.dump(dice_score_list_validation, validation_pickle)
validation_pickle.close()
if len(loss_list_validation) >= 3:
if dice_score_list_validation[-1] > best_model_accuracy:
best_model_accuracy = dice_score_list_validation[-1]
torch.save(
{
'epochs': epoch + 1,
'batchsize': batch_size,
'train_loss_list_epoch': loss_list_train_epoch,
'train_dice_score_list_epoch':
dice_score_list_train_epoch,
'train_loss_index_epoch': epoch_data_list,
'valid_loss_list': loss_list_validation,
'valid_dice_score_list':
dice_score_list_validation,
'valid_dice_score_list_0':
dice_score_list_validation_0,
'valid_dice_score_list_1':
dice_score_list_validation_1,
'valid_loss_index': loss_list_validation_index,
'model_state_dict': model.state_dict(),
}, model_checkpoint_dir + '/model_best.pth')
else:
best_model_accuracy = dice_score_list_validation[-1]
torch.save(
{
'epochs': epoch + 1,
'batchsize': batch_size,
'train_loss_list_epoch': loss_list_train_epoch,
'train_dice_score_list_epoch':
dice_score_list_train_epoch,
'train_loss_index_epoch': epoch_data_list,
'valid_loss_list': loss_list_validation,
'valid_dice_score_list': dice_score_list_validation,
'valid_dice_score_list_0':
dice_score_list_validation_0,
'valid_dice_score_list_1':
dice_score_list_validation_1,
'valid_loss_index': loss_list_validation_index,
'model_state_dict': model.state_dict(),
}, model_checkpoint_dir + '/model_best.pth')
num_workers=1)
### DEFINE NET ARCHITECTURE
net = model.Net(n_channels_FirstLayer=params.N_CHANNELS_FIRST_LAYER)
if torch.cuda.is_available():
net.cuda()
if params.HOST == "local":
summary(
net,
input_size=tuple(
[1] +
list(np.array(params.SINGLE_IMAGE_SIZE) // params.SUBSAMPLING)))
### DEFINING THE LOSS FUNCTION AND OPTIMIZER
optimizer = optim.SGD(net.parameters(), lr=params.LR_START, momentum=0.9)
loss_fct = lambda outputs, labels: losses.dice_loss(
outputs, labels[:, 0, :, :, :], power=params.DICE_POWER)
#loss_fct = lambda outputs, labels : losses.cross_entropy_custom(outputs, labels[:,0,:,:,:])
scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=[25, 35], gamma=0.3)
if True:
#if params.HOST == "google":
if not os.path.isdir("./last_epoch_results"):
os.makedirs("./last_epoch_results")
### TRAIN THE NETWORK
for epoch in range(
params.N_EPOCHS): # loop over the dataset multiple times
scheduler.step(epoch)
lr = get_learning_rate(optimizer)
def main(_):
# check if checkpoint path exists
if not os.path.exists(FLAGS.checkpoint_path):
os.mkdir(FLAGS.checkpoint_path)
train_data_generator = data_processor.generator(FLAGS)
train_samples_count = data_processor.count_samples(FLAGS)
print('total batches per epoch : {}'.format(train_samples_count / FLAGS.batch_size))
east = EAST_model(FLAGS.input_size)
east.model.summary()
score_map_loss_weight = tf.Variable(0.01, name='score_map_loss_weight')
small_text_weight = tf.Variable(0., name='small_text_weight')
lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
FLAGS.init_learning_rate,
decay_steps=FLAGS.lr_decay_steps,
decay_rate=FLAGS.lr_decay_rate,
staircase=True)
optimizer = tf.optimizers.Adam(lr_schedule)
# set checkpoint manager
ckpt = tf.train.Checkpoint(step=tf.Variable(0), model=east)
ckpt_manager = tf.train.CheckpointManager(ckpt,
directory=FLAGS.checkpoint_path,
max_to_keep=5)
latest_ckpt = tf.train.latest_checkpoint(FLAGS.checkpoint_path)
# restore latest checkpoint
if latest_ckpt:
ckpt.restore(latest_ckpt)
print('global_step : {}, checkpoint is restored!'.format(int(ckpt.step)))
# set tensorboard summary writer
summary_writer = tf.summary.create_file_writer(FLAGS.checkpoint_path + '/train')
while int(ckpt.step) < (FLAGS.max_steps + 1):
# load data
[input_images, overly_small_text_region_training_masks, text_region_boundary_training_masks, score_maps], \
[target_score_maps, target_geo_maps] = next(train_data_generator)
# update parameter
train_step(east,
input_images,
optimizer,
overly_small_text_region_training_masks,
text_region_boundary_training_masks,
small_text_weight,
target_score_maps,
target_geo_maps,
score_map_loss_weight
)
score_y_pred, geo_y_pred = east(input_images)
_dice_loss = dice_loss(overly_small_text_region_training_masks, text_region_boundary_training_masks, score_map_loss_weight,
small_text_weight, target_score_maps, score_y_pred)
_rbox_loss = rbox_loss(overly_small_text_region_training_masks, text_region_boundary_training_masks,
small_text_weight, target_score_maps, target_geo_maps, geo_y_pred)
loss = _dice_loss + _rbox_loss
print('Step {:06d}, dice_loss {:.4f}, rbox_loss {:.4f}, total_loss {:.4f}'.format(int(ckpt.step), _dice_loss, _rbox_loss, loss))
if ckpt.step % FLAGS.save_checkpoint_steps == 0:
# save checkpoint
ckpt_manager.save(checkpoint_number=ckpt.step)
print('global_step : {}, checkpoint is saved!'.format(int(ckpt.step)))
with summary_writer.as_default():
tf.summary.scalar('loss', loss, step=int(ckpt.step))
tf.summary.scalar('pred_score_map_loss', _dice_loss, step=int(ckpt.step))
tf.summary.scalar('pred_geo_map_loss ', _rbox_loss, step=int(ckpt.step))
tf.summary.scalar('learning_rate ', optimizer.lr(ckpt.step).numpy(), step=int(ckpt.step))
tf.summary.scalar('small_text_weight', small_text_weight, step=int(ckpt.step))
tf.summary.image("input_image", tf.cast((input_images + 1) * 127.5, tf.uint8), step=int(ckpt.step), max_outputs=3)
tf.summary.image("overly_small_text_region_training_mask", tf.cast(overly_small_text_region_training_masks * 255, tf.uint8), step=int(ckpt.step), max_outputs=3)
tf.summary.image("text_region_boundary_training_mask", tf.cast(text_region_boundary_training_masks * 255, tf.uint8), step=int(ckpt.step), max_outputs=3)
tf.summary.image("score_map_target", tf.cast(target_score_maps * 255, tf.uint8), step=int(ckpt.step), max_outputs=3)
tf.summary.image("score_map_pred", tf.cast(score_y_pred * 255, tf.uint8), step=int(ckpt.step), max_outputs=3)
for i in range(4):
tf.summary.image("geo_map_%d_target" % (i), tf.cast(tf.expand_dims(target_geo_maps[:, :, :, i], axis=3) / FLAGS.input_size * 255, tf.uint8), step=int(ckpt.step), max_outputs=3)
tf.summary.image("geo_map_%d_pred" % (i), tf.cast(tf.expand_dims(geo_y_pred[:, :, :, i], axis=3) / FLAGS.input_size * 255, tf.uint8), step=int(ckpt.step), max_outputs=3)
tf.summary.image("geo_map_4_target", tf.cast((tf.expand_dims(target_geo_maps[:, :, :, 4], axis=3) + 1) * 127.5, tf.uint8), step=int(ckpt.step), max_outputs=3)
tf.summary.image("geo_map_4_pred", tf.cast((tf.expand_dims(geo_y_pred[:, :, :, 4], axis=3) + 1) * 127.5, tf.uint8), step=int(ckpt.step), max_outputs=3)
ckpt.step.assign_add(1)
def main(argv=None):
os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu_list
# check if checkpoint path exists
if not os.path.exists(FLAGS.checkpoint_path):
os.mkdir(FLAGS.checkpoint_path)
else:
#if not FLAGS.restore:
# shutil.rmtree(FLAGS.checkpoint_path)
# os.mkdir(FLAGS.checkpoint_path)
shutil.rmtree(FLAGS.checkpoint_path)
os.mkdir(FLAGS.checkpoint_path)
if len(gpus) <= 1:
print('Training with 1 GPU')
east = EAST_model(FLAGS.input_size)
parallel_model = east.model
else:
print('Training with %d GPUs' % len(gpus))
with tf.device("/cpu:0"):
east = EAST_model(FLAGS.input_size)
if FLAGS.restore_model is not '':
east.model.load_weights(FLAGS.restore_model)
parallel_model = multi_gpu_model(east.model, gpus=len(gpus))
score_map_loss_weight = K.variable(0.01, name='score_map_loss_weight')
small_text_weight = K.variable(0., name='small_text_weight')
# opt = AdamW(FLAGS.init_learning_rate)
opt = Adam(learning_rate=0.001, beta_1=0.9, beta_2=0.999, epsilon=1e-4)
parallel_model.compile(loss=[
dice_loss(east.overly_small_text_region_training_mask,
east.text_region_boundary_training_mask,
score_map_loss_weight, small_text_weight),
rbox_loss(east.overly_small_text_region_training_mask,
east.text_region_boundary_training_mask, small_text_weight,
east.target_score_map)
],
loss_weights=[1., 1.],
optimizer=opt)
east.model.summary()
model_json = east.model.to_json()
with open(FLAGS.checkpoint_path + '/model.json', 'w') as json_file:
json_file.write(model_json)
plot_model(east.model,
to_file=FLAGS.checkpoint_path + '/east-train.png',
show_shapes=True,
show_layer_names=True,
expand_nested=True)
plot_model(east.resnet,
to_file=FLAGS.checkpoint_path + '/resnet.png',
show_shapes=True,
show_layer_names=True,
expand_nested=True)
train_data_generator = data_processor.TrainDataSequence(FLAGS)
#train_data_generator = data_processor.generator(FLAGS)
#train_data_x, train_data_y = data_processor.load_train_data(FLAGS)
#print('# of train data: %d' %(len(train_data_x[0])))
train_samples_count = data_processor.count_samples(FLAGS)
print("train_samples_count: %d" % (train_samples_count))
val_data = data_processor.load_data(FLAGS)
lr_scheduler = LearningRateScheduler(lr_decay)
ckpt = CustomModelCheckpoint(model=east.model,
path=FLAGS.checkpoint_path,
period=FLAGS.save_checkpoint_epochs,
save_weights_only=False)
# tb = CustomTensorBoard(log_dir=FLAGS.checkpoint_path + '/train', score_map_loss_weight=score_map_loss_weight, small_text_weight=small_text_weight, data_generator=train_data_generator, write_graph=True)
small_text_weight_callback = SmallTextWeight(small_text_weight)
validation_evaluator = ValidationEvaluator(
val_data, validation_log_dir=FLAGS.checkpoint_path + '/val')
callbacks = [
lr_scheduler, ckpt, small_text_weight_callback, validation_evaluator
]
#callbacks = [lr_scheduler, ckpt, small_text_weight_callback]
#history = parallel_model.fit(x=train_data_x, y=train_data_y, batch_size=FLAGS.batch_size, epochs=FLAGS.max_epochs, verbose=1, callbacks=callbacks, max_queue_size=10, workers=FLAGS.nb_workers, use_multiprocessing=True)
#history = parallel_model.fit(x=train_data_generator, epochs=FLAGS.max_epochs, steps_per_epoch=train_samples_count/FLAGS.batch_size, callbacks=callbacks, max_queue_size=10, verbose=1)
history = parallel_model.fit(x=train_data_generator,
epochs=FLAGS.max_epochs,
callbacks=callbacks,
max_queue_size=10,
workers=FLAGS.nb_workers,
use_multiprocessing=False,
verbose=1)
file_name = FLAGS.checkpoint_path + '/model-train.h5'
east.model.save(file_name, overwrite=True)
plot_model(east.model_core,
to_file=FLAGS.checkpoint_path + '/east.png',
show_shapes=True,
show_layer_names=True,
expand_nested=True)
file_name = FLAGS.checkpoint_path + '/model.h5'
east.model_core.save(file_name, overwrite=True, include_optimizer=False)
tflite_model = lite.TFLiteConverter.from_keras_model_file(
file_name,
input_arrays=["input_image"],
input_shapes={
"input_image": [1, 224, 320, 3]
}).convert()
with open(file_name + '.tflite', 'wb') as tflite_file:
tflite_file.write(tflite_model)
def main():
os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu_list
# check if checkpoint path exists
if not os.path.exists(FLAGS.checkpoint_path):
os.mkdir(FLAGS.checkpoint_path)
else:
# if not FLAGS.restore:
# shutil.rmtree(FLAGS.checkpoint_path)
# os.mkdir(FLAGS.checkpoint_path)
shutil.rmtree(FLAGS.checkpoint_path)
os.mkdir(FLAGS.checkpoint_path)
# train_data_generator = data_processor.generator(FLAGS)
train_data_generator = EastSequence(FLAGS, input_size=FLAGS.input_size)
train_samples_count = data_processor.count_samples(FLAGS)
val_data = data_processor.load_data(FLAGS)
train_graph = tf.Graph()
train_sess = tf.Session(graph=train_graph)
K.set_session(train_sess)
with train_graph.as_default():
K.set_learning_phase(0)
if len(gpus) <= 1:
print('Training with 1 GPU')
east = EASTModel(FLAGS.input_size)
parallel_model = east.model
else:
print('Training with %d GPUs' % len(gpus))
with tf.device("/cpu:0"):
east = EASTModel(FLAGS.input_size)
parallel_model = multi_gpu_model(east.model, gpus=len(gpus))
score_map_loss_weight = K.variable(0.01, name='score_map_loss_weight')
small_text_weight = K.variable(0., name='small_text_weight')
tf.contrib.quantize.create_training_graph(input_graph=train_graph,
quant_delay=250000)
train_sess.run(tf.global_variables_initializer())
lr_scheduler = LearningRateScheduler(lr_decay)
ckpt = CustomModelCheckpoint(model=east.model,
path=FLAGS.checkpoint_path +
'/model-{epoch:02d}.h5',
period=FLAGS.save_checkpoint_epochs,
save_weights_only=True)
tb = TensorBoard(log_dir=FLAGS.checkpoint_path + '/train',
batch_size=FLAGS.batch_size,
write_images=True)
small_text_weight_callback = SmallTextWeight(small_text_weight)
validation_evaluator = ValidationEvaluator(
val_data, validation_log_dir=FLAGS.checkpoint_path + '/val')
callbacks = [
lr_scheduler, ckpt, tb, small_text_weight_callback,
validation_evaluator
]
# opt = AdamW(FLAGS.init_learning_rate)
parallel_model.compile(loss=[
dice_loss(east.overly_small_text_region_training_mask,
east.text_region_boundary_training_mask,
score_map_loss_weight, small_text_weight),
rbox_loss(east.overly_small_text_region_training_mask,
east.text_region_boundary_training_mask,
small_text_weight, east.target_score_map)
],
loss_weights=[1., 1.],
optimizer='adam')
model_json = east.model.to_json()
with open(FLAGS.checkpoint_path + '/model.json', 'w') as json_file:
json_file.write(model_json)
history = parallel_model.fit_generator(
train_data_generator,
epochs=FLAGS.max_epochs,
steps_per_epoch=train_samples_count / FLAGS.batch_size,
workers=FLAGS.nb_workers,
use_multiprocessing=True,
max_queue_size=10,
callbacks=callbacks,
verbose=1)
saver = tf.train.Saver()
saver.save(train_sess, 'checkpoints')
# one_hot_encode_labels = torch.squeeze(one_hot_encode_labels,0)
# one_hot_encode_labels = one_hot_encode_labels.permute(0,4,1,2,3).contiguous()
# loss = dice_loss(output_2,one_hot_encode_labels)
# optimizer = optim.Adam(model.parameters())
# optimizer.zero_grad()
# loss.backward()
# optimizer.step()
# print('Done')
image = torch.from_numpy(batch['data']).float().cuda()
label = torch.from_numpy(batch['seg']).cuda().long()
labels_for_conf = batch['seg']
output_1, output_2 = model(image)
one_hot_encode_labels = F.one_hot(label, n_classes)
one_hot_encode_labels = one_hot_encode_labels.permute(
0, 4, 1, 2, 3).contiguous()
loss = dice_loss(output_2, one_hot_encode_labels)
conf_matrix = confusion_matrix(
torch.argmax(output_2, 1).view(-1).cpu().detach().numpy(),
labels_for_conf.view(-1).cpu().detach().numpy())
TPR, TNR, PPV, FPR, FNR, ACC = get_metrics(conf_matrix)
accuracy = accuracy_score(
labels_for_conf.view(-1).cpu().detach().numpy(),
torch.argmax(output_2, 1).view(-1).cpu().detach().numpy())
mean_accuracy.append(accuracy)
logger.info(
'TPR == {} | \nTNR == {} | \nPRCSN == {} | \nFPR == {}\n | \nFNR == {} | \nACC == {}.'
.format(TPR, TNR, PPV, FPR, FNR, ACC))
logger.info('Accuracy = {}'.format(accuracy))
losses.append(loss.item())
optimizer = optim.Adam(model.parameters())
optimizer.zero_grad()
def main():
if use_multiinput_architecture is False:
if modeltype == 'unet':
model = UNet(n_classes=n_classes,
padding=True,
depth=model_depth,
up_mode='upconv',
batch_norm=True,
residual=False).double().to(device)
elif modeltype == 'resunet':
model = UNet(n_classes=n_classes,
padding=True,
depth=model_depth,
up_mode='upconv',
batch_norm=True,
residual=True).double().to(device)
elif use_multiinput_architecture is True:
if modeltype == 'unet':
model = Attention_UNet(
n_classes=n_classes,
padding=True,
up_mode='upconv',
batch_norm=True,
residual=False,
wf=wf,
use_attention=use_attention).double().to(device)
elif modeltype == 'resunet':
model = Attention_UNet(
n_classes=n_classes,
padding=True,
up_mode='upconv',
batch_norm=True,
residual=True,
wf=wf,
use_attention=use_attention).double().to(device)
checkpoint = torch.load(args.model_path,
map_location=lambda storage, loc: storage)
model.load_state_dict(checkpoint['model_state_dict'])
valid_loader = dataloader_cxr.DataLoader(data_path,
dataloader_type=dataloader_type,
batchsize=batch_size,
device=device,
image_resolution=image_resolution)
if len(valid_loader.data_list) % batch_size == 0:
total_idx_valid = len(valid_loader.data_list) // batch_size
else:
total_idx_valid = len(valid_loader.data_list) // batch_size + 1
model.eval()
prediction_array = np.zeros((len(valid_loader.data_list),
image_resolution[0], image_resolution[1]))
if valid_loader.dataloader_type != "test":
input_mask_array = np.zeros((len(valid_loader.data_list),
image_resolution[0], image_resolution[1]))
valid_count = 0
valid_dice_score = 0.0
valid_dice_score_0 = 0.0
valid_dice_score_1 = 0.0
for idx in range(total_idx_valid):
with torch.no_grad():
if valid_loader.dataloader_type != "test":
batch_images_input, batch_label_input = valid_loader[idx]
else:
batch_images_input = valid_loader[idx]
output = model(batch_images_input)
if use_multiinput_architecture is False:
if len(valid_loader.data_list) % batch_size == 0:
temp_image = torch.max(
output, 1)[1].detach().cpu().numpy().astype(np.bool)
prediction_array[idx * batch_size:(idx + 1) *
batch_size] = remove_small_regions(
temp_image,
0.02 * np.prod(image_resolution))
if valid_loader.dataloader_type != "test":
input_mask_array[idx * batch_size:(
idx + 1) * batch_size] = batch_label_input.detach(
).cpu().numpy().astype(np.uint8)
else:
if idx == len(valid_loader.data_list) // batch_size:
temp_image = torch.max(
output,
1)[1].detach().cpu().numpy().astype(np.bool)
prediction_array[idx *
batch_size:] = remove_small_regions(
temp_image,
0.02 * np.prod(image_resolution))
if valid_loader.dataloader_type != "test":
input_mask_array[
idx * batch_size:] = batch_label_input.detach(
).cpu().numpy().astype(np.uint8)
else:
temp_image = torch.max(
output,
1)[1].detach().cpu().numpy().astype(np.bool)
prediction_array[idx * batch_size:(idx + 1) *
batch_size] = remove_small_regions(
temp_image,
0.02 * np.prod(image_resolution))
if valid_loader.dataloader_type != "test":
input_mask_array[idx * batch_size:(
idx +
1) * batch_size] = batch_label_input.detach(
).cpu().numpy().astype(np.uint8)
else:
if len(valid_loader.data_list) % batch_size == 0:
temp_image = torch.max(output[-1],
1)[1].detach().cpu().numpy().astype(
np.bool)
prediction_array[idx * batch_size:(idx + 1) *
batch_size] = remove_small_regions(
temp_image,
0.02 * np.prod(image_resolution))
if valid_loader.dataloader_type != "test":
input_mask_array[idx * batch_size:(
idx + 1) * batch_size] = batch_label_input.detach(
).cpu().numpy().astype(np.uint8)
else:
if idx == len(valid_loader.data_list) // batch_size:
temp_image = torch.max(
output[-1],
1)[1].detach().cpu().numpy().astype(np.bool)
prediction_array[idx *
batch_size:] = remove_small_regions(
temp_image,
0.02 * np.prod(image_resolution))
if valid_loader.dataloader_type != "test":
input_mask_array[
idx * batch_size:] = batch_label_input.detach(
).cpu().numpy().astype(np.uint8)
else:
temp_image = torch.max(
output[-1],
1)[1].detach().cpu().numpy().astype(np.bool)
prediction_array[idx * batch_size:(idx + 1) *
batch_size] = remove_small_regions(
temp_image,
0.02 * np.prod(image_resolution))
if valid_loader.dataloader_type != "test":
input_mask_array[idx * batch_size:(
idx +
1) * batch_size] = batch_label_input.detach(
).cpu().numpy().astype(np.uint8)
if valid_loader.dataloader_type != "test":
if use_multiinput_architecture is False:
loss = losses.dice_loss(output, batch_label_input)
else:
loss = losses.dice_loss(output[-1], batch_label_input)
valid_count += batch_images_input.shape[0]
if use_multiinput_architecture is False:
score = losses.dice_score(output, batch_label_input)
else:
score = losses.dice_score(output[-1], batch_label_input)
valid_dice_score += (score.sum().item() / score.size(0)
) * batch_images_input.shape[0]
valid_dice_score_0 += score[0].item(
) * batch_images_input.shape[0]
valid_dice_score_1 += score[1].item(
) * batch_images_input.shape[0]
if valid_loader.dataloader_type != "test":
valid_dice_score = valid_dice_score / valid_count
valid_dice_score_0 = valid_dice_score_0 / valid_count
valid_dice_score_1 = valid_dice_score_1 / valid_count
if generate_mask is True:
for i, files in enumerate(valid_loader.data_list):
temp_mask = prediction_array[i].astype(int)
temp_mask = ndimage.zoom(
temp_mask,
np.asarray(valid_loader.original_size_array[files]) /
np.asarray(temp_mask.shape),
order=0)
io.imsave(save_path + '/Pred_mask_' + files.split('.')[0] + '.png',
temp_mask)
if valid_loader.dataloader_type != "test":
print('Valid Dice Score: ', valid_dice_score, ' Valid Dice Score 0: ',
valid_dice_score_0, ' Valid Dice Score 1: ', valid_dice_score_1)
var_input = Variable(var_input, requires_grad=True)
with torch.set_grad_enabled(
phase in ['train', 'train_ancillary']):
if model_type == 'Unet':
preds = model(var_input)
elif model_type == 'BB_Unet':
var_bbox = batch["bboxes"].reshape(
-1, 1, size, size)
var_bbox = torch.tensor(var_bbox,
dtype=torch.float)
var_bbox = Variable(var_bbox, requires_grad=True)
preds = model(var_input, var_bbox)
loss = losses.dice_loss(preds.to('cpu'), gt_samples)
print('loss for batch {} on epoch {} is {}'.format(
i, epoch, loss))
predicted_output = threshold(preds.to('cpu'))
predicted_output = predicted_output.type(torch.float32)
d_metric, dict = dice_score(predicted_output.to('cpu'),
gt_samples.to('cpu'))
if phase in ['train_ancillary', 'train']:
train_loss_total += loss
tr_dice_total += d_metric
optimizer.zero_grad()
loss.backward()
optimizer.step()
def main():
if use_multiinput_architecture is False:
if modeltype == 'unet':
model = UNet(n_classes=n_classes,
padding=True,
depth=model_depth,
up_mode='upsample',
batch_norm=True,
residual=False).double().to(device)
elif modeltype == 'resunet':
model = UNet(n_classes=n_classes,
padding=True,
depth=model_depth,
up_mode='upsample',
batch_norm=True,
residual=True).double().to(device)
elif use_multiinput_architecture is True:
if modeltype == 'unet':
model = Attention_UNet(
n_classes=n_classes,
padding=True,
up_mode='upconv',
batch_norm=True,
residual=False,
wf=wf,
use_attention=use_attention).double().to(device)
elif modeltype == 'resunet':
model = Attention_UNet(
n_classes=n_classes,
padding=True,
up_mode='upconv',
batch_norm=True,
residual=True,
wf=wf,
use_attention=use_attention).double().to(device)
checkpoint = torch.load(args.model_path,
map_location=lambda storage, loc: storage)
model.load_state_dict(checkpoint['model_state_dict'])
valid_loader = dataloader_cxr.DataLoader(data_path,
dataloader_type=dataloader_type,
batchsize=batch_size,
device=device,
image_resolution=image_resolution,
invert=invert,
remove_wires=remove_wires)
if len(valid_loader.data_list) % batch_size == 0:
total_idx_valid = len(valid_loader.data_list) // batch_size
else:
total_idx_valid = len(valid_loader.data_list) // batch_size + 1
model.eval()
prediction_array = np.zeros((len(valid_loader.data_list),
image_resolution[0], image_resolution[1]))
if valid_loader.dataloader_type != "test":
input_mask_array = np.zeros((len(valid_loader.data_list),
image_resolution[0], image_resolution[1]))
valid_count = 0
valid_dice_score = 0.0
if 0 in classes:
valid_dice_score_0 = 0.0
if 1 in classes:
valid_dice_score_1 = 0.0
for idx in range(total_idx_valid):
with torch.no_grad():
if valid_loader.dataloader_type != "test":
batch_images_input, batch_label_input = valid_loader[idx]
else:
batch_images_input = valid_loader[idx]
output = model(batch_images_input)
if use_multiinput_architecture is False:
if len(valid_loader.data_list) % batch_size == 0:
temp_image = torch.max(
output, 1)[1].detach().cpu().numpy().astype(np.bool)
prediction_array[idx * batch_size:(idx + 1) *
batch_size] = remove_small_regions(
temp_image,
0.02 * np.prod(image_resolution))
if valid_loader.dataloader_type != "test":
input_mask_array[idx * batch_size:(
idx + 1) * batch_size] = batch_label_input.detach(
).cpu().numpy().astype(np.uint8)
else:
if idx == len(valid_loader.data_list) // batch_size:
temp_image = torch.max(
output,
1)[1].detach().cpu().numpy().astype(np.bool)
prediction_array[idx *
batch_size:] = remove_small_regions(
temp_image,
0.02 * np.prod(image_resolution))
if valid_loader.dataloader_type != "test":
input_mask_array[
idx * batch_size:] = batch_label_input.detach(
).cpu().numpy().astype(np.uint8)
else:
temp_image = torch.max(
output,
1)[1].detach().cpu().numpy().astype(np.bool)
prediction_array[idx * batch_size:(idx + 1) *
batch_size] = remove_small_regions(
temp_image,
0.02 * np.prod(image_resolution))
if valid_loader.dataloader_type != "test":
input_mask_array[idx * batch_size:(
idx +
1) * batch_size] = batch_label_input.detach(
).cpu().numpy().astype(np.uint8)
else:
if len(valid_loader.data_list) % batch_size == 0:
temp_image = torch.max(output[-1],
1)[1].detach().cpu().numpy().astype(
np.bool)
prediction_array[idx * batch_size:(idx + 1) *
batch_size] = remove_small_regions(
temp_image,
0.02 * np.prod(image_resolution))
if valid_loader.dataloader_type != "test":
input_mask_array[idx * batch_size:(
idx + 1) * batch_size] = batch_label_input.detach(
).cpu().numpy().astype(np.uint8)
else:
if idx == len(valid_loader.data_list) // batch_size:
temp_image = torch.max(
output[-1],
1)[1].detach().cpu().numpy().astype(np.bool)
prediction_array[idx *
batch_size:] = remove_small_regions(
temp_image,
0.02 * np.prod(image_resolution))
if valid_loader.dataloader_type != "test":
input_mask_array[
idx * batch_size:] = batch_label_input.detach(
).cpu().numpy().astype(np.uint8)
else:
temp_image = torch.max(
output[-1],
1)[1].detach().cpu().numpy().astype(np.bool)
prediction_array[idx * batch_size:(idx + 1) *
batch_size] = remove_small_regions(
temp_image,
0.02 * np.prod(image_resolution))
if valid_loader.dataloader_type != "test":
input_mask_array[idx * batch_size:(
idx +
1) * batch_size] = batch_label_input.detach(
).cpu().numpy().astype(np.uint8)
if valid_loader.dataloader_type != "test":
if use_multiinput_architecture is False:
loss = losses.dice_loss(output, batch_label_input,
exclude_0)
else:
loss = losses.dice_loss(output[-1], batch_label_input,
exclude_0)
valid_count += batch_images_input.shape[0]
if use_multiinput_architecture is False:
score = losses.dice_score(output, batch_label_input,
exclude_0)
else:
score = losses.dice_score(output[-1], batch_label_input,
exclude_0)
valid_dice_score += (score.sum().item() / score.size(0)
) * batch_images_input.shape[0]
if 0 in classes and 1 in classes and len(classes) == 2 and len(
clubbed) == 0:
valid_dice_score_0 += score[0].item(
) * batch_images_input.shape[0]
valid_dice_score_1 += score[1].item(
) * batch_images_input.shape[0]
if valid_loader.dataloader_type != "test":
valid_dice_score = valid_dice_score / valid_count
if 0 in classes:
valid_dice_score_0 = valid_dice_score_0 / valid_count
if 1 in classes:
valid_dice_score_1 = valid_dice_score_1 / valid_count
if generate_mask is True:
for i, files in enumerate(valid_loader.data_list):
temp_mask = prediction_array[i].astype(int)
temp_mask = ndimage.zoom(
temp_mask,
np.asarray(valid_loader.original_size_array[files]) /
np.asarray(temp_mask.shape),
order=0)
temp_mask = sitk.GetImageFromArray(temp_mask)
temp_mask = sitk.Cast(temp_mask, sitk.sitkUInt8)
resampler = sitk.ResampleImageFilter()
resampler.SetReferenceImage(temp_mask)
resampler.SetOutputSpacing(valid_loader.spacing[files])
resampler.SetSize(valid_loader.size_[files])
resampler.SetInterpolator(sitk.sitkNearestNeighbor)
temp_mask = resampler.Execute(temp_mask)
temp_mask.SetOrigin(valid_loader.origin[files])
#temp_name = ''
#for j in range(len(files.split('.'))-1):
#if files.split('.')[j] != 'nii':
#temp_name = temp_name + files.split('.')[j] + '.'
temp_name = files[:-4]
sitk.WriteImage(temp_mask,
save_path + '/Pred_mask_' + temp_name + '.nii.gz')
# sitk.WriteImage(temp_mask, save_path + '/Pred_mask_' + files.split('.')[0] + '.nii.gz')
# io.imsave(save_path + '/Pred_mask_' + files.split('.')[0] + '.png', temp_mask)
# if valid_loader.dataloader_type != "test":
# temp_img_plus_mask = prediction_array[i].astype(int) + input_mask_array[i]*2
# temp_img_plus_mask = ndimage.zoom(temp_img_plus_mask, np.asarray(valid_loader.original_size_array[files]) / np.asarray(temp_img_plus_mask.shape), order=0)
# temp_img_plus_mask = sitk.GetImageFromArray(temp_img_plus_mask)
# resampler = sitk.ResampleImageFilter()
# resampler.SetReferenceImage(temp_img_plus_mask)
# resampler.SetOutputSpacing(valid_loader.spacing[files])
# resampler.SetSize(valid_loader.size_[files])
# resampler.SetInterpolator(sitk.sitkNearestNeighbor)
# temp_img_plus_mask = resampler.Execute(temp_img_plus_mask)
# temp_img_plus_mask.SetOrigin(valid_loader.origin[files])
# sitk.WriteImage(temp_img_plus_mask, save_path + '/Pred_merged_mask_' + files.split('.')[0] + '.nii.gz')
if valid_loader.dataloader_type != "test":
if 0 in classes and 1 in classes and len(clubbed) == 0:
print('Valid Dice Score: ', valid_dice_score,
' Valid Dice Score 0: ', valid_dice_score_0,
' Valid Dice Score 1: ', valid_dice_score_1)
def validate(val_loader, model, criterion, epoch, args):
batch_time = AverageMeter('Time', ':6.3f')
losses = AverageMeter('Loss', ':.4e')
top1 = AverageMeter('Acc@1', ':6.2f')
dice = AverageMeter('Dice', ':6.2f')
progress = ProgressMeter(len(val_loader), [batch_time, losses, top1, dice],
prefix='Test: ')
# switch to evaluate mode
model.eval()
with torch.no_grad():
end = time.time()
for i, (images, target) in enumerate(val_loader):
if args.debug and i > 10: break
if args.gpu is not None:
images = images.cuda(args.gpu, non_blocking=True)
target = target.cuda(args.gpu, non_blocking=True)
# compute output
coarse_output, output = model(images)
loss = criterion(output, target)
if args.need_vis:
for t, (in_dbg, out_dbg, out_gt,
in_dbg1) in get_vis(images, output, target):
cv2.imwrite("./tmp/{}_{}_targ.png".format(i, t), out_gt)
cv2.imwrite("./tmp/{}_{}_in.png".format(i, t),
cv2.cvtColor(in_dbg, cv2.COLOR_BGR2RGB))
cv2.imwrite("./tmp/{}_{}_out.png".format(i, t), out_dbg)
cv2.imwrite("./tmp/{}_{}_crop.png".format(i, t),
cv2.cvtColor(in_dbg1, cv2.COLOR_BGRA2RGBA))
# measure accuracy and record loss
acc1 = accuracy(output.permute(0, 2, 3,
1).reshape(-1, args.num_classes),
target.view(-1),
topk=(1, ))
losses.update(loss.item(),
output.size(0) * output.size(2) * output.size(3))
top1.update(acc1[0][0].item(),
output.size(0) * output.size(2) * output.size(3))
dice.update(1 - dice_loss(output, target).item(), output.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % args.print_freq == 0:
progress.display(i)
# TODO: this should also be done with the ProgressMeter
print(' * Acc@1 {top1.avg:.3f} Dice {dice.avg:.3f}'.format(top1=top1,
dice=dice))
if epoch != 'e':
writer.add_scalar('acc1/acc1_val', top1.avg, epoch)
writer.add_scalar('dice/dice_val', dice.avg, epoch)
_, (in_dbg, out_dbg, out_gt,
in_dbg1) = get_vis(images, output, target)[0]
writer.add_image("a_input", in_dbg.transpose(2, 0, 1), epoch)
writer.add_image("b_predicted", out_dbg[..., None].transpose(2, 0, 1),
epoch)
writer.add_image("c_ground truth",
out_gt[..., None].transpose(2, 0, 1), epoch)
writer.add_image("d_cropped", in_dbg1.transpose(2, 0, 1), epoch)
return dice.avg
iter = tf.data.Iterator.from_structure(train_ds.output_types,
train_ds.output_shapes)
full_init = iter.make_initializer(ds)
train_init = iter.make_initializer(train_ds)
test_init = iter.make_initializer(test_ds)
input = iter.get_next()
logits, latent = model.model(input['image'], training)
# very imbalanced classification, requires loss robust to high class imbalance
# I use combination of balanced cross entropy and dice loss
loss = sum([
tf.reduce_mean(
losses.balanced_sigmoid_cross_entropy_with_logits(
labels=input['image'], logits=logits)),
tf.reduce_mean(losses.dice_loss(labels=input['image'], logits=logits)),
tf.losses.get_regularization_loss()
])
optimizer = tf.train.AdamOptimizer(learning_rate)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
train_step = optimizer.minimize(loss, global_step)
metrics, update_metrics = build_metrics(loss,
labels=input['image'],
logits=logits)
summary = build_summary(metrics, input['image'], logits, learning_rate)
locals_init = tf.local_variables_initializer()
saver = tf.train.Saver()
