diff --git a/eval.py b/eval.py index 50974ff..ec6b9f1 100644 --- a/eval.py +++ b/eval.py @@ -1,118 +1,110 @@ -import os -import sys -from collections import OrderedDict -from os import walk +#!/usr/bin/env python +# coding: utf-8 +# +# Author: Kazuto Nakashima +# URL: http://kazuto1011.github.io +# Created: 2017-11-03 + + +import argparse +import os.path as osp import cv2 -import matplotlib.pyplot as plt import numpy as np -import scipy import torch -import torch.nn as nn import torch.nn.functional as F +import yaml +from tensorboardX import SummaryWriter from torch.autograd import Variable - -from models import DeepLab -import torchvision.models as models -from docopt import docopt - -docstr = """Evaluate ResNet-DeepLab trained on scenes (VOC 2012),a total of 21 labels including background - -Usage: - evalpyt.py [options] - -Options: - -h, --help Print this message - --visualize view outputs of each sketch - --snapPrefix= Snapshot [default: VOC12_scenes_] - --testGTpath= Ground truth path prefix [default: data/gt/] - --testIMpath= Sketch images path prefix [default: data/img/] - --NoLabels= The number of different labels in training data, VOC has 21 labels, including background [default: 21] - --gpu0= GPU number [default: 0] -""" - -args = docopt(docstr, version='v0.1') -print args - - -def get_iou(pred, gt): - if pred.shape != gt.shape: - print 'pred shape', pred.shape, 'gt shape', gt.shape - assert(pred.shape == gt.shape) - gt = gt.astype(np.float32) - pred = pred.astype(np.float32) - - max_label = int(args['--NoLabels']) - 1 # labels from 0,1, ... 20(for VOC) - count = np.zeros((max_label + 1,)) - for j in range(max_label + 1): - x = np.where(pred == j) - p_idx_j = set(zip(x[0].tolist(), x[1].tolist())) - x = np.where(gt == j) - GT_idx_j = set(zip(x[0].tolist(), x[1].tolist())) - # pdb.set_trace() - n_jj = set.intersection(p_idx_j, GT_idx_j) - u_jj = set.union(p_idx_j, GT_idx_j) - - if len(GT_idx_j) != 0: - count[j] = float(len(n_jj)) / float(len(u_jj)) - - result_class = count - Aiou = np.sum(result_class[:]) / float(len(np.unique(gt))) - - return Aiou - - -gpu0 = int(args['--gpu0']) -im_path = args['--testIMpath'] -model = DeepLab(int(args['--NoLabels'])) -model.eval() -counter = 0 -model.cuda(gpu0) -snapPrefix = args['--snapPrefix'] -gt_path = args['--testGTpath'] -img_list = open('data/list/val.txt').readlines() - -# TODO set the (different iteration)models that you want to evaluate on. Models are saved during training after every 1000 iters by default. -for iter in range(1, 21): - saved_state_dict = torch.load(os.path.join( - 'data/snapshots/', snapPrefix + str(iter) + '000.pth')) - if counter == 0: - print snapPrefix - counter += 1 - model.load_state_dict(saved_state_dict) - - pytorch_list = [] - for i in img_list: - img = np.zeros((513, 513, 3)) - - img_temp = cv2.imread(os.path.join( - im_path, i[:-1] + '.jpg')).astype(float) - img_original = img_temp - img_temp[:, :, 0] = img_temp[:, :, 0] - 104.008 - img_temp[:, :, 1] = img_temp[:, :, 1] - 116.669 - img_temp[:, :, 2] = img_temp[:, :, 2] - 122.675 - img[:img_temp.shape[0], :img_temp.shape[1], :] = img_temp - gt = cv2.imread(os.path.join(gt_path, i[:-1] + '.png'), 0) - gt[gt == 255] = 0 - - output = model(Variable(torch.from_numpy(img[np.newaxis, :].transpose( - 0, 3, 1, 2)).float(), volatile=True).cuda(gpu0)) - interp = nn.UpsamplingBilinear2d(size=(513, 513)) - output = interp(output[3]).cpu().data[0].numpy() - output = output[:, :img_temp.shape[0], :img_temp.shape[1]] - - output = output.transpose(1, 2, 0) +from torchnet.meter import MovingAverageValueMeter +from tqdm import tqdm + +from libs.datasets import get_dataset +from libs.models import DeepLab +from libs.utils import CrossEntropyLoss2d, scores + + +def main(args): + # Configuration + with open(args.config) as f: + config = yaml.load(f) + + image_size = (config['image']['size']['test'], + config['image']['size']['test']) + n_classes = config['dataset'][args.dataset]['n_classes'] + + # Dataset + dataset = get_dataset(args.dataset)( + root=config['dataset'][args.dataset]['root'], + split='test', + image_size=image_size, + scale=False, + flip=False, + preload=False + ) + + # DataLoader + loader = torch.utils.data.DataLoader( + dataset=dataset, + batch_size=args.batch_size, + num_workers=config['num_workers'], + shuffle=False + ) + loader_iter = iter(loader) + + checkpoint = torch.load(args.checkpoint, + map_location=lambda storage, + loc: storage) + state_dict = checkpoint['weight'] + print('Result after {} iterations'.format(checkpoint['iteration'])) + + # Model + model = DeepLab(n_classes=n_classes) + model.load_state_dict(state_dict) + model.eval() + if args.cuda: + model.cuda() + + targets, outputs = [], [] + for data, target in tqdm(loader, total=len(loader), + leave=False, dynamic_ncols=True): + # Image + data = data.cuda() if args.cuda else data + data = Variable(data, volatile=True) + + # Forward propagation + output = model(data) + output = F.upsample(output[3], size=image_size, mode='bilinear') + output = output[0].cpu().data.numpy().transpose(1, 2, 0) output = np.argmax(output, axis=2) - if args['--visualize']: - plt.subplot(3, 1, 1) - plt.imshow(img_original) - plt.subplot(3, 1, 2) - plt.imshow(gt) - plt.subplot(3, 1, 3) - plt.imshow(output) - plt.show() - - iou_pytorch = get_iou(output, gt) - pytorch_list.append(iou_pytorch) - - print 'pytorch', iter, np.sum(np.asarray(pytorch_list)) / len(pytorch_list) + target = target.numpy() + + for o, t in zip(output, target): + outputs.append(o) + targets.append(t) + + score, class_iou = scores(targets, outputs, n_class=n_classes) + + for k, v in score.items(): + print k, v + + for i in range(n_classes): + print i, class_iou[i] + + +if __name__ == '__main__': + # Parsing arguments + parser = argparse.ArgumentParser(description='') + parser.add_argument('--no_cuda', action='store_true', default=False) + parser.add_argument('--dataset', nargs='?', type=str, default='cocostuff') + parser.add_argument('--config', type=str, default='config/default.yaml') + parser.add_argument('--checkpoint', type=str, default=None) + parser.add_argument('--batch_size', type=int, default=1) + + args = parser.parse_args() + args.cuda = not args.no_cuda and torch.cuda.is_available() + + for arg in vars(args): + print('{0:20s}: {1}'.format(arg.rjust(20), getattr(args, arg))) + + main(args) diff --git a/libs/utils/__init__.py b/libs/utils/__init__.py index d58a0ad..fd26545 100644 --- a/libs/utils/__init__.py +++ b/libs/utils/__init__.py @@ -1,2 +1,2 @@ from libs.utils.loss import CrossEntropyLoss2d -from libs.utils.lr_scheduler import poly_lr_scheduler \ No newline at end of file +from libs.utils.metric import scores \ No newline at end of file diff --git a/libs/utils/lr_scheduler.py b/libs/utils/lr_scheduler.py deleted file mode 100644 index e6aab94..0000000 --- a/libs/utils/lr_scheduler.py +++ /dev/null @@ -1,7 +0,0 @@ -def poly_lr_scheduler(optimizer, init_lr, iter, lr_decay_iter=10, max_iter=30000, power=0.9): - if iter % lr_decay_iter or iter > max_iter: - return None - - new_lr = init_lr * (1 - iter / max_iter)**power - optimizer.params_groups[0]['lr'] = new_lr - optimizer.params_groups[1]['lr'] = 10 * new_lr \ No newline at end of file diff --git a/libs/utils/metric.py b/libs/utils/metric.py new file mode 100644 index 0000000..6c45dee --- /dev/null +++ b/libs/utils/metric.py @@ -0,0 +1,37 @@ +# Originally written by wkentaro +# https://github.com/wkentaro/pytorch-fcn/blob/master/torchfcn/utils.py + +import numpy as np + + +def _fast_hist(label_true, label_pred, n_class): + mask = (label_true >= 0) & (label_true < n_class) + hist = np.bincount( + n_class * label_true[mask].astype(int) + + label_pred[mask], minlength=n_class**2).reshape(n_class, n_class) + return hist + + +def scores(label_trues, label_preds, n_class): + """Returns accuracy score evaluation result. + - overall accuracy + - mean accuracy + - mean IU + - fwavacc + """ + hist = np.zeros((n_class, n_class)) + for lt, lp in zip(label_trues, label_preds): + hist += _fast_hist(lt.flatten(), lp.flatten(), n_class) + acc = np.diag(hist).sum() / hist.sum() + acc_cls = np.diag(hist) / hist.sum(axis=1) + acc_cls = np.nanmean(acc_cls) + iu = np.diag(hist) / (hist.sum(axis=1) + hist.sum(axis=0) - np.diag(hist)) + mean_iu = np.nanmean(iu) + freq = hist.sum(axis=1) / hist.sum() + fwavacc = (freq[freq > 0] * iu[freq > 0]).sum() + cls_iu = dict(zip(range(n_class), iu)) + + return {'Overall Acc: \t': acc, + 'Mean Acc : \t': acc_cls, + 'FreqW Acc : \t': fwavacc, + 'Mean IoU : \t': mean_iu, }, cls_iu diff --git a/train.py b/train.py index ddc3863..61d1906 100644 --- a/train.py +++ b/train.py @@ -75,10 +75,11 @@ def main(args): dataset = get_dataset(args.dataset)( root=config['dataset'][args.dataset]['root'], split='train', - image_size=(321, 321), + image_size=(config['image']['size']['train'], + config['image']['size']['train']), scale=True, flip=True, - # preload=True + preload=True ) # DataLoader @@ -186,7 +187,6 @@ def main(args): if iteration % len(loader) == 0: loader_iter = iter(loader) - torch.save( {'iteration': iteration - 1, 'weight': model.state_dict()},