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))