210 lines
7.6 KiB
Python
210 lines
7.6 KiB
Python
#!/usr/bin/env python
|
|
# -*- encoding: utf-8 -*-
|
|
|
|
"""
|
|
@Author : Peike Li
|
|
@Contact : peike.li@yahoo.com
|
|
@File : evaluate.py
|
|
@Time : 8/4/19 3:36 PM
|
|
@Desc :
|
|
@License : This source code is licensed under the license found in the
|
|
LICENSE file in the root directory of this source tree.
|
|
"""
|
|
|
|
import os
|
|
import argparse
|
|
import numpy as np
|
|
import torch
|
|
|
|
from torch.utils import data
|
|
from tqdm import tqdm
|
|
from PIL import Image as PILImage
|
|
import torchvision.transforms as transforms
|
|
import torch.backends.cudnn as cudnn
|
|
|
|
import networks
|
|
from datasets.datasets import LIPDataValSet
|
|
from utils.miou import compute_mean_ioU
|
|
from utils.transforms import BGR2RGB_transform
|
|
from utils.transforms import transform_parsing
|
|
|
|
|
|
def get_arguments():
|
|
"""Parse all the arguments provided from the CLI.
|
|
|
|
Returns:
|
|
A list of parsed arguments.
|
|
"""
|
|
parser = argparse.ArgumentParser(description="Self Correction for Human Parsing")
|
|
|
|
# Network Structure
|
|
parser.add_argument("--arch", type=str, default='resnet101')
|
|
# Data Preference
|
|
parser.add_argument("--data-dir", type=str, default='./data/LIP')
|
|
parser.add_argument("--batch-size", type=int, default=1)
|
|
parser.add_argument("--input-size", type=str, default='473,473')
|
|
parser.add_argument("--num-classes", type=int, default=20)
|
|
parser.add_argument("--ignore-label", type=int, default=255)
|
|
parser.add_argument("--random-mirror", action="store_true")
|
|
parser.add_argument("--random-scale", action="store_true")
|
|
# Evaluation Preference
|
|
parser.add_argument("--log-dir", type=str, default='./log')
|
|
parser.add_argument("--model-restore", type=str, default='./log/checkpoint.pth.tar')
|
|
parser.add_argument("--gpu", type=str, default='0', help="choose gpu device.")
|
|
parser.add_argument("--save-results", action="store_true", help="whether to save the results.")
|
|
parser.add_argument("--flip", action="store_true", help="random flip during the test.")
|
|
parser.add_argument("--multi-scales", type=str, default='1', help="multiple scales during the test")
|
|
return parser.parse_args()
|
|
|
|
|
|
def get_palette(num_cls):
|
|
""" Returns the color map for visualizing the segmentation mask.
|
|
Args:
|
|
num_cls: Number of classes
|
|
Returns:
|
|
The color map
|
|
"""
|
|
n = num_cls
|
|
palette = [0] * (n * 3)
|
|
for j in range(0, n):
|
|
lab = j
|
|
palette[j * 3 + 0] = 0
|
|
palette[j * 3 + 1] = 0
|
|
palette[j * 3 + 2] = 0
|
|
i = 0
|
|
while lab:
|
|
palette[j * 3 + 0] |= (((lab >> 0) & 1) << (7 - i))
|
|
palette[j * 3 + 1] |= (((lab >> 1) & 1) << (7 - i))
|
|
palette[j * 3 + 2] |= (((lab >> 2) & 1) << (7 - i))
|
|
i += 1
|
|
lab >>= 3
|
|
return palette
|
|
|
|
|
|
def multi_scale_testing(model, batch_input_im, crop_size=[473, 473], flip=True, multi_scales=[1]):
|
|
flipped_idx = (15, 14, 17, 16, 19, 18)
|
|
if len(batch_input_im.shape) > 4:
|
|
batch_input_im = batch_input_im.squeeze()
|
|
if len(batch_input_im.shape) == 3:
|
|
batch_input_im = batch_input_im.unsqueeze(0)
|
|
|
|
interp = torch.nn.Upsample(size=crop_size, mode='bilinear', align_corners=True)
|
|
ms_outputs = []
|
|
for s in multi_scales:
|
|
interp_im = torch.nn.Upsample(scale_factor=s, mode='bilinear', align_corners=True)
|
|
scaled_im = interp_im(batch_input_im)
|
|
parsing_output = model(scaled_im)
|
|
parsing_output = parsing_output[0][-1]
|
|
output = parsing_output[0]
|
|
if flip:
|
|
flipped_output = parsing_output[1]
|
|
flipped_output[14:20, :, :] = flipped_output[flipped_idx, :, :]
|
|
output += flipped_output.flip(dims=[-1])
|
|
output *= 0.5
|
|
output = interp(output.unsqueeze(0))
|
|
ms_outputs.append(output[0])
|
|
ms_fused_parsing_output = torch.stack(ms_outputs)
|
|
ms_fused_parsing_output = ms_fused_parsing_output.mean(0)
|
|
ms_fused_parsing_output = ms_fused_parsing_output.permute(1, 2, 0) # HWC
|
|
parsing = torch.argmax(ms_fused_parsing_output, dim=2)
|
|
parsing = parsing.data.cpu().numpy()
|
|
ms_fused_parsing_output = ms_fused_parsing_output.data.cpu().numpy()
|
|
return parsing, ms_fused_parsing_output
|
|
|
|
|
|
def main():
|
|
"""Create the model and start the evaluation process."""
|
|
args = get_arguments()
|
|
multi_scales = [float(i) for i in args.multi_scales.split(',')]
|
|
gpus = [int(i) for i in args.gpu.split(',')]
|
|
assert len(gpus) == 1
|
|
if not args.gpu == 'None':
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
|
|
|
|
cudnn.benchmark = True
|
|
cudnn.enabled = True
|
|
|
|
h, w = map(int, args.input_size.split(','))
|
|
input_size = [h, w]
|
|
|
|
model = networks.init_model(args.arch, num_classes=args.num_classes, pretrained=None)
|
|
|
|
IMAGE_MEAN = model.mean
|
|
IMAGE_STD = model.std
|
|
INPUT_SPACE = model.input_space
|
|
print('image mean: {}'.format(IMAGE_MEAN))
|
|
print('image std: {}'.format(IMAGE_STD))
|
|
print('input space:{}'.format(INPUT_SPACE))
|
|
if INPUT_SPACE == 'BGR':
|
|
print('BGR Transformation')
|
|
transform = transforms.Compose([
|
|
transforms.ToTensor(),
|
|
transforms.Normalize(mean=IMAGE_MEAN,
|
|
std=IMAGE_STD),
|
|
|
|
])
|
|
if INPUT_SPACE == 'RGB':
|
|
print('RGB Transformation')
|
|
transform = transforms.Compose([
|
|
transforms.ToTensor(),
|
|
BGR2RGB_transform(),
|
|
transforms.Normalize(mean=IMAGE_MEAN,
|
|
std=IMAGE_STD),
|
|
])
|
|
|
|
# Data loader
|
|
lip_test_dataset = LIPDataValSet(args.data_dir, 'val', crop_size=input_size, transform=transform, flip=args.flip)
|
|
num_samples = len(lip_test_dataset)
|
|
print('Totoal testing sample numbers: {}'.format(num_samples))
|
|
testloader = data.DataLoader(lip_test_dataset, batch_size=args.batch_size, shuffle=False, pin_memory=True)
|
|
|
|
# Load model weight
|
|
state_dict = torch.load(args.model_restore)['state_dict']
|
|
from collections import OrderedDict
|
|
new_state_dict = OrderedDict()
|
|
for k, v in state_dict.items():
|
|
name = k[7:] # remove `module.`
|
|
new_state_dict[name] = v
|
|
model.load_state_dict(new_state_dict)
|
|
model.cuda()
|
|
model.eval()
|
|
|
|
sp_results_dir = os.path.join(args.log_dir, 'sp_results')
|
|
if not os.path.exists(sp_results_dir):
|
|
os.makedirs(sp_results_dir)
|
|
|
|
palette = get_palette(20)
|
|
parsing_preds = []
|
|
scales = np.zeros((num_samples, 2), dtype=np.float32)
|
|
centers = np.zeros((num_samples, 2), dtype=np.int32)
|
|
with torch.no_grad():
|
|
for idx, batch in enumerate(tqdm(testloader)):
|
|
image, meta = batch
|
|
if (len(image.shape) > 4):
|
|
image = image.squeeze()
|
|
im_name = meta['name'][0]
|
|
c = meta['center'].numpy()[0]
|
|
s = meta['scale'].numpy()[0]
|
|
w = meta['width'].numpy()[0]
|
|
h = meta['height'].numpy()[0]
|
|
scales[idx, :] = s
|
|
centers[idx, :] = c
|
|
parsing, logits = multi_scale_testing(model, image.cuda(), crop_size=input_size, flip=args.flip,
|
|
multi_scales=multi_scales)
|
|
if args.save_results:
|
|
parsing_result = transform_parsing(parsing, c, s, w, h, input_size)
|
|
parsing_result_path = os.path.join(sp_results_dir, im_name + '.png')
|
|
output_im = PILImage.fromarray(np.asarray(parsing_result, dtype=np.uint8))
|
|
output_im.putpalette(palette)
|
|
output_im.save(parsing_result_path)
|
|
|
|
parsing_preds.append(parsing)
|
|
assert len(parsing_preds) == num_samples
|
|
mIoU = compute_mean_ioU(parsing_preds, scales, centers, args.num_classes, args.data_dir, input_size)
|
|
print(mIoU)
|
|
return
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|