| import datasets |
| import os |
| import torch |
| import torch.nn as nn |
| from torch.utils.data import DataLoader |
| import numpy as np |
| from os.path import join as pjoin |
| from tqdm import tqdm |
| import torch.nn.functional as F |
| import argparse |
| import matplotlib.pyplot as plt |
| from thop import profile |
|
|
| from models.AYANet import AYANet |
| from misc.metric_tool import ConfuseMatrixMeter |
| from misc.logger_tool import Logger |
| import misc.utils as utils |
| from misc.utils import de_norm |
|
|
|
|
| class CDEval(): |
| def __init__(self, arguments): |
| self.args = arguments |
|
|
| self.device = torch.device("cuda:%s" % self.args.gpu_ids[0] if torch.cuda.is_available() and len(self.args.gpu_ids)>0 |
| else "cpu") |
|
|
| self.net_G = AYANet(self.args.encoder_arch, self.args.decoder_arch) |
|
|
| self.dataset_eval_loader = DataLoader(datasets.bcd_eval(pjoin(self.args.datadir), self.args.test_split), |
| batch_size=self.args.batch_size, |
| shuffle=False) |
|
|
| self.running_metric = ConfuseMatrixMeter(n_class=2) |
|
|
| self.ckpt_list = [] |
| if self.args.checkpoint != 'All': |
| ckpt_path = pjoin(self.args.checkpointdir, self.args.checkpoint) |
| self.ckpt_list.append(ckpt_path) |
| else: |
| for el in os.listdir(self.args.checkpointdir): |
| if el.endswith('.pt'): |
| ckpt_path = pjoin(self.args.checkpointdir, el) |
| self.ckpt_list.append(ckpt_path) |
|
|
| logger_path = os.path.join(self.args.checkpointdir, 'log_test.txt') |
| self.logger = Logger(logger_path) |
| self.logger.write_dict_str(args.__dict__) |
|
|
| self.epoch_acc = 0 |
| self.best_val_acc = 0.0 |
| self.best_epoch_id = 0 |
|
|
| self.batch = None |
| self.batch_id = 0 |
|
|
|
|
| def _load_checkpoint(self): |
|
|
| try: |
| self.logger.write('loading checkpoint...\n') |
| |
| checkpoint = torch.load(self.checkpoint, map_location=self.device) |
|
|
| self.net_G.load_state_dict(checkpoint['model_G_state_dict']) |
|
|
| self.net_G.to(self.device) |
|
|
| |
| self.best_val_acc = checkpoint['best_val_acc'] |
| self.best_epoch_id = checkpoint['best_epoch_id'] |
|
|
| self.logger.write('Eval Historical_best_acc = %.4f (at epoch %d)\n' % |
| (self.best_val_acc, self.best_epoch_id)) |
| self.logger.write('\n') |
|
|
| except: |
| raise FileNotFoundError('no such checkpoint %s' % self.checkpoint) |
|
|
|
|
| def _clear_cache(self): |
| self.running_metric.clear() |
|
|
| def _update_metric(self): |
| target = self.batch['L'].to(self.device).detach() |
| G_pred = self.G_pred.detach() |
|
|
| G_pred = torch.argmax(G_pred, dim=1) |
|
|
| current_score = self.running_metric.update_cm(pr=G_pred.cpu().numpy(), gt=target.cpu().numpy()) |
| return current_score |
|
|
| def _visualize_pred(self): |
| pred = torch.argmax(self.G_pred, dim=1, keepdim=True) |
| pred_vis = pred * 255 |
| return pred_vis |
|
|
| def _collect_running_batch_states(self): |
|
|
| running_acc = self._update_metric() |
|
|
| m = len(self.dataset_eval_loader) |
|
|
| if np.mod(self.batch_id, 100) == 1: |
| message = '[%d,%d], running_mf1: %.5f\n' %\ |
| (self.batch_id, m, running_acc) |
| self.logger.write(message) |
|
|
| img_t0, img_t1 = torch.split(self.batch['I'], 3, 1) |
| vis_input = utils.make_numpy_grid(de_norm(img_t0)) |
| vis_input2 = utils.make_numpy_grid(de_norm(img_t1)) |
|
|
| vis_pred = utils.make_numpy_grid(self._visualize_pred()) |
|
|
| vis_gt = utils.make_numpy_grid(self.batch['L']) |
| val, ct = np.unique(vis_gt, return_counts=True) |
|
|
| if (1 in val): |
| if ct[np.where(val==1)] > 16384: |
| |
| diff = vis_pred - vis_gt |
|
|
| mask_tp = np.all(diff == [254, 254, 254], axis=-1) |
| mask_fp = np.all(diff == [255, 255, 255], axis=-1) |
| mask_fn = np.all(diff == [-1, -1, -1], axis=-1) |
| mask_tn = np.all(diff == [0, 0, 0], axis=-1) |
|
|
| diff[mask_tp] = [1, 1, 1] |
| diff[mask_fp] = [0.4, 1, 1] |
| diff[mask_fn] = [1, 0.4, 0.4] |
|
|
| vis = np.concatenate([vis_input, vis_input2, vis_pred, vis_gt, diff], axis=0) |
| vis = np.clip(vis, a_min=0.0, a_max=1.0) |
| file_name = os.path.join( |
| self.args.visdir, 'eval_' + str(self.batch_id)+'.jpg') |
| plt.imsave(file_name, vis) |
|
|
| def _collect_epoch_states(self): |
|
|
| scores_dict = self.running_metric.get_scores() |
|
|
| np.save(os.path.join(self.args.checkpointdir, 'scores_dict.npy'), scores_dict) |
|
|
| self.epoch_acc = scores_dict['mf1'] |
|
|
| message = '' |
| for k, v in scores_dict.items(): |
| message += '%s: %.5f ' % (k, v) |
| self.logger.write('%s\n' % message) |
|
|
| self.logger.write('\n') |
|
|
| ckpt = self.checkpoint.split('/')[-1] |
| message += ('\n' + str(ckpt) + '\n') |
| with open(os.path.join(self.args.checkpointdir, 'report.txt' % (self.epoch_acc)), mode='a') as file: |
| file.write(message) |
|
|
|
|
| def _forward(self, inputs, label): |
| data = {} |
| data['I'] = inputs |
| data['L'] = label |
| self.batch = data |
| |
| inputs_img = self.batch['I'].to(self.device) |
| |
| |
| |
| |
| |
|
|
| self.G_pred = self.net_G(inputs_img) |
| |
|
|
|
|
| def eval(self): |
| for ckpt in self.ckpt_list: |
| self.checkpoint = ckpt |
| self._load_checkpoint() |
|
|
| self.logger.write('Begin evaluation...\n') |
| self._clear_cache() |
| self.net_G.eval() |
|
|
| for self.batch_id, (inputs_test, mask_test) in enumerate(tqdm(self.dataset_eval_loader)): |
| with torch.no_grad(): |
| self._forward(inputs_test, mask_test) |
| self._collect_running_batch_states() |
| self._collect_epoch_states() |
|
|
|
|
| if __name__ =='__main__': |
| parser = argparse.ArgumentParser(description='START EVALUATING...') |
| parser.add_argument('--gpu_ids', default=0) |
| parser.add_argument('--dataset', type=str, default='bcd', required=True) |
| parser.add_argument('--test_split', type=str, default='test', required=True) |
| parser.add_argument('--datadir',required=True) |
| parser.add_argument('--resultdir',required=True) |
| parser.add_argument('--checkpointroot', type=str, required=True) |
| parser.add_argument('--visroot', type=str, required=True) |
| parser.add_argument('--checkpoint', required=True) |
| parser.add_argument('--encoder-arch', type=str, required=True) |
| parser.add_argument('--decoder-arch', type=str, required=True) |
| parser.add_argument('--batch-size', type=int, default=1) |
| parser.add_argument('--multi-gpu', action='store_true', help='processing with multi-gpus') |
| parser.add_argument('--project_name', type=str) |
|
|
| args = parser.parse_args() |
| args.checkpointdir = pjoin(args.checkpointroot, args.project_name) |
| os.makedirs(args.checkpointdir, exist_ok=True) |
|
|
| args.visdir = pjoin(args.visroot, args.project_name) |
| os.makedirs(args.visdir, exist_ok=True) |
|
|
| if args.dataset == 'bcd': |
| eval = CDEval(args) |
| eval.eval() |
| else: |
| print('Error: Cannot identify the dataset...(dataset: bcd)') |
| exit(-1) |
|
|
