def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
# initialization
print("Input arguments:")
for key, val in vars(args).items():
print("{:16} {}".format(key, val))
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
model = get_model(num_classes=args.num_classes)
# if not os.path.exists(args.save_dir):
# os.makedirs(args.save_dir)
palette = get_lip_palette()
restore_from = args.restore_from
saved_state_dict = torch.load(restore_from)
model.load_state_dict(saved_state_dict)
model.eval()
model.cuda()
testloader = data.DataLoader(TestGenerator(args.root,
args.data_list,
crop_size=args.crop_size),
batch_size=1,
shuffle=False,
pin_memory=True)
confusion_matrix = np.zeros((args.num_classes, args.num_classes))
for index, batch in enumerate(testloader):
if index % 100 == 0:
print('%d images have been proceeded' % index)
image, label, ori_size, name = batch
ori_size = ori_size[0].numpy()
output = predict(model,
image.numpy(),
(np.asscalar(ori_size[0]), np.asscalar(ori_size[1])),
is_mirror=args.is_mirror,
scales=args.eval_scale)
seg_pred = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
# output_im = PILImage.fromarray(seg_pred)
# output_im.putpalette(palette)
# output_im.save(args.save_dir + name[0] + '.png')
seg_gt = np.asarray(label[0].numpy(), dtype=np.int)
ignore_index = seg_gt != 255
seg_gt = seg_gt[ignore_index]
seg_pred = seg_pred[ignore_index]
confusion_matrix += get_confusion_matrix(seg_gt, seg_pred,
args.num_classes)
pos = confusion_matrix.sum(1)
res = confusion_matrix.sum(0)
tp = np.diag(confusion_matrix)
pixel_accuracy = tp.sum() / pos.sum()
mean_accuracy = (tp / np.maximum(1.0, pos)).mean()
IU_array = (tp / np.maximum(1.0, pos + res - tp))
mean_IU = IU_array.mean()
# get_confusion_matrix_plot()
print('Pixel accuracy: %f \n' % pixel_accuracy)
print('Mean accuracy: %f \n' % mean_accuracy)
print('Mean IU: %f \n' % mean_IU)
for index, IU in enumerate(IU_array):
print('%f ', IU)
def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
# initialization
print("Input arguments:")
for key, val in vars(args).items():
print("{:16} {}".format(key, val))
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
# if not os.path.exists(args.save_dir):
# os.makedirs(args.save_dir)
# obtain the color map
palette = get_atr_palette()
# conduct model & load pre-trained weights
model = get_model(num_classes=args.num_classes)
restore_from = args.restore_from
saved_state_dict = torch.load(restore_from)
model.load_state_dict(saved_state_dict)
model.eval()
model.cuda()
# data loader
testloader = data.DataLoader(TestGenerator(args.root, args.data_list, crop_size=args.crop_size),
batch_size=1, shuffle=False, pin_memory=True)
confusion_matrix = np.zeros((args.num_classes, args.num_classes))
for index, batch in enumerate(testloader):
if index % 10 == 0:
print('%d images have been proceeded' % index)
image, label, ori_size, name = batch
ori_size = ori_size[0].numpy()
output = predict(model, image.numpy(), (np.asscalar(ori_size[0]), np.asscalar(ori_size[1])),
is_mirror=args.is_mirror, scales=args.eval_scale)
seg_pred = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
# output_im = PILImage.fromarray(seg_pred)
# output_im.putpalette(palette)
# output_im.save(args.save_dir + name[0] + '.png')
# seg_gt = np.asarray(label[0].numpy(), dtype=np.int)
# ignore_index = seg_gt != 255
# seg_gt = seg_gt[ignore_index]
# seg_pred = seg_pred[ignore_index]
#
# confusion_matrix = get_confusion_matrix(seg_gt, seg_pred, args.num_classes)
#
# pos = confusion_matrix.sum(1)
# res = confusion_matrix.sum(0)
# tp = np.diag(confusion_matrix)
# non_zero_idx = res != 0
# pos_idx = pos[non_zero_idx]
# res_idx = res[non_zero_idx]
# tp_idx = tp[non_zero_idx]
#
# iou_array = tp_idx / np.maximum(1.0, pos_idx + res_idx - tp_idx)
# mean_iou = np.nanmean(iou_array)
#
# output_im = Image.fromarray(seg_pred)
# output_im.putpalette(palette)
# output_im.save(os.path.join(args.save_dir, str(mean_iou)[2:5] + '_' + name[0] + '.png'))
seg_gt = np.asarray(label[0].numpy(), dtype=np.int)
ignore_index = seg_gt != 255
seg_gt = seg_gt[ignore_index]
seg_pred = seg_pred[ignore_index]
confusion_matrix += get_confusion_matrix(seg_gt, seg_pred, args.num_classes)
pos = confusion_matrix.sum(1) # TP + FP
res = confusion_matrix.sum(0) # p
tp = np.diag(confusion_matrix)
pixel_accuracy = tp.sum() / pos.sum() # mean Acc
fg_accuracy = tp[1:].sum() / pos[1:].sum() # foreground Acc
mean_accuracy = (tp / np.maximum(1.0, pos)).mean() # mean Precision
# cal_list = 0
# for i in range(len(res)):
# if res[i] != 0:
# cal_list += (tp[i] / res[i])
# mean_recall = cal_list / len(res) # mean Recall
mean_recall = (tp / res).mean() # mean Recall
f1_score = 2 * mean_accuracy * mean_recall / (mean_accuracy + mean_recall)
accuracy = (tp / np.maximum(1.0, pos))
recall = (tp / res)
cls_f1_score = 2 * accuracy * recall / (accuracy + recall)
# mean_recall = (tp / (res+1e-10)).mean() # mean Recall
# f1_score = 2 * mean_accuracy * mean_recall / (mean_accuracy + mean_recall + 1e-10)
# accuracy = (tp / (np.maximum(1.0, pos)+1e-10))
# recall = (tp / (res+1e-10))
# cls_f1_score = 2 * accuracy * recall / (accuracy + recall + 1e-10)
print('pixelAcc. --> {}'.format(pixel_accuracy))
print('foregroundAcc. --> {}'.format(fg_accuracy))
print('meanPrecision --> {}'.format(mean_accuracy))
print('meanRecall --> {}'.format(mean_recall))
print('meanF1-score --> {}'.format(f1_score))
for i in range(args.num_classes):
print('cls {}, F1-score --> {}'.format(i, cls_f1_score[i]))
def main(args):
# initialization
print("Input arguments:")
for key, val in vars(args).items():
print("{:16} {}".format(key, val))
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
writer = SummaryWriter(log_dir=os.path.join(args.log_dir, args.method))
random.seed(args.seed)
torch.manual_seed(args.seed)
cudnn.enabled = True
cudnn.benchmark = True
# conduct seg network
seg_model = get_model(num_classes=args.num_classes)
saved_state_dict = torch.load(args.restore_from)
new_params = seg_model.state_dict().copy()
if args.init:
for i in saved_state_dict:
i_parts = i.split('.')
if not i_parts[0] == 'fc':
new_params['encoder.' + '.'.join(i_parts[:])] = saved_state_dict[i]
seg_model.load_state_dict(new_params)
print('loading params w/o fc')
else:
seg_model.load_state_dict(saved_state_dict)
print('loading params all')
model = DataParallelModel(seg_model)
model.float()
model.cuda()
# define dataloader
train_loader = data.DataLoader(DataGenerator(root=args.root, list_path=args.lst,
crop_size=args.crop_size, training=True),
batch_size=args.batch_size, shuffle=True, num_workers=4, pin_memory=False)
val_loader = data.DataLoader(DataGenerator(root=args.val_root, list_path=args.val_lst,
crop_size=args.crop_size, training=False),
batch_size=args.batch_size, shuffle=False, num_workers=4, pin_memory=False)
# define criterion & optimizer
criterion = ABRLovaszLoss(ignore_index=args.ignore_label, only_present=True)
criterion = DataParallelCriterion(criterion).cuda()
optimizer = optim.SGD(
[{'params': filter(lambda p: p.requires_grad, seg_model.parameters()), 'lr': args.learning_rate}],
lr=args.learning_rate, momentum=0.9, weight_decay=5e-4)
# key points
best_val_mIoU = 0
best_val_pixAcc = 0
start = time.time()
for epoch in range(0, args.epochs):
print('\n{} | {}'.format(epoch, args.epochs - 1))
# training
_ = train(model, train_loader, epoch, criterion, optimizer, writer)
# validation
if epoch %10 ==0 or epoch > args.epochs-10:
val_pixacc, val_miou = validation(model, val_loader, epoch, writer)
# save model
if val_pixacc > best_val_pixAcc:
best_val_pixAcc = val_pixacc
if val_miou > best_val_mIoU:
best_val_mIoU = val_miou
model_dir = os.path.join(args.snapshot_dir, args.method + '_miou.pth')
torch.save(seg_model.state_dict(), model_dir)
print('Model saved to %s' % model_dir)
os.rename(model_dir, os.path.join(args.snapshot_dir, args.method + '_miou'+str(best_val_mIoU)+'.pth'))
print('Complete using', time.time() - start, 'seconds')
print('Best pixAcc: {} | Best mIoU: {}'.format(best_val_pixAcc, best_val_mIoU))