insightface/parsing/dml_csr/test.py

#!/usr/bin/env python
# -*- encoding: utf-8 -*-
"""
@Author     :   Qingping Zheng
@Contact    :   qingpingzheng2014@gmail.com
@File       :   datasets.py
@Time       :   10/01/21 00:00 PM
@Desc       :   
@License    :   Licensed under the Apache License, Version 2.0 (the "License"); 
@Copyright  :   Copyright 2015 The Authors. All Rights Reserved.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import argparse
import cv2
import numpy as np
import matplotlib.pyplot as plt
import os
import torch
import torch.backends.cudnn as cudnn
import torchvision.transforms as transforms

from copy import deepcopy
from inplace_abn import InPlaceABN
from dataset import datasets
from networks import dml_csr
from utils import miou

torch.multiprocessing.set_start_method("spawn", force=True)

DATA_DIRECTORY = './datasets/Helen'
IGNORE_LABEL = 255
NUM_CLASSES = 20
SNAPSHOT_DIR = './snapshots/'
INPUT_SIZE = (473,473)


def get_arguments():
    """Parse all the arguments provided from the CLI.
    
    Returns:
      A list of parsed arguments.
    """
    parser = argparse.ArgumentParser(description="DML_CSR Network")
    parser.add_argument("--batch-size", type=int, default=1,
                        help="Number of images sent to the network in one step.")
    parser.add_argument("--data-dir", type=str, default=DATA_DIRECTORY,
                        help="Path to the directory containing the PASCAL VOC dataset.")
    parser.add_argument("--out-dir", type=str, default=DATA_DIRECTORY,
                        help="Path to the directory containing the PASCAL VOC dataset.")
    parser.add_argument("--dataset", type=str, default='val',
                        help="Path to the file listing the images in the dataset.")
    parser.add_argument("--ignore-label", type=int, default=IGNORE_LABEL,
                        help="The index of the label to ignore during the training.")
    parser.add_argument("--num-classes", type=int, default=NUM_CLASSES,
                        help="Number of classes to predict (including background).")
    parser.add_argument("--restore-from", type=str,
                        help="Where restore model parameters from.")
    parser.add_argument("--gpu", type=str, default='7',
                        help="choose gpu device.")
    parser.add_argument("--input-size", type=str, default=INPUT_SIZE,
                        help="Comma-separated string with height and width of images.")
    parser.add_argument("--local_rank", type=int, default=0,
                        help="choose gpu numbers") 
    parser.add_argument('--dist-backend', default='nccl', type=str,
                        help='distributed backend')
    parser.add_argument("--model_type", type=int, default=0,
                        help="choose model type") 
    return parser.parse_args()


def valid(model, valloader, input_size, num_samples, dir=None, dir_edge=None, dir_img=None):

    height = input_size[0]
    width  = input_size[1]
    with torch.autograd.profiler.profile(enabled=True, use_cuda=True, \
        record_shapes=False, profile_memory=False) as prof:
        model.eval()
        parsing_preds = np.zeros((num_samples, height, width), dtype=np.uint8)
        scales = np.zeros((num_samples, 2), dtype=np.float32)
        centers = np.zeros((num_samples, 2), dtype=np.int32)

        idx = 0
        interp = torch.nn.Upsample(size=(height, width), mode='bilinear', align_corners=True)

        with torch.no_grad():
            for index, batch in enumerate(valloader):
                image, meta = batch
                num_images = image.size(0)
                if index % 10 == 0:
                    print('%d  processd' % (index * num_images))

                c = meta['center'].numpy()
                s = meta['scale'].numpy()
                scales[idx:idx + num_images, :] = s[:, :]
                centers[idx:idx + num_images, :] = c[:, :]

                results = model(image.cuda())
                outputs = results

                if isinstance(results, list):
                    outputs = results[0]

                if isinstance(outputs, list):
                    for k, output in enumerate(outputs):
                        parsing = output
                        nums = len(parsing)
                        parsing = interp(parsing).data.cpu().numpy()
                        parsing = parsing.transpose(0, 2, 3, 1)  # NCHW NHWC
                        parsing_preds[idx:idx + nums, :, :] = np.asarray(np.argmax(parsing, axis=3), dtype=np.uint8)
                        idx += nums
                else:
                    parsing = outputs
                    parsing = interp(parsing).data.cpu().numpy()
                    parsing = parsing.transpose(0, 2, 3, 1)  # NCHW NHWC
                    parsing_preds[idx:idx + num_images, :, :] = np.asarray(np.argmax(parsing, axis=3), dtype=np.uint8)

                    if dir is not None:
                        for i in range(len(meta['name'])):
                            cv2.imwrite(os.path.join(dir, meta['name'][i] + '.png'), np.asarray(np.argmax(parsing, axis=3))[i])
                    idx += num_images
        parsing_preds = parsing_preds[:num_samples, :, :]

    return parsing_preds, scales, centers


def main():
    """Create the model and start the evaluation process."""

    args = get_arguments()

    os.environ["CUDA_VISIBLE_DEVICES"]=args.gpu
    gpus = [int(i) for i in args.gpu.split(',')]

    print(args.gpu)

    h, w = map(int, args.input_size.split(','))
    
    input_size = (h, w)

    cudnn.benchmark = True
    cudnn.enabled = True

    model = dml_csr.DML_CSR(args.num_classes, InPlaceABN, False)
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])

    transform = transforms.Compose([
        transforms.ToTensor(),
        normalize,
    ])

    dataset = datasets.FaceDataSet(args.data_dir, args.dataset, \
        crop_size=input_size, transform=transform)
    num_samples = len(dataset)

    valloader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size, \
        shuffle=False, pin_memory=True)

    restore_from = args.restore_from
    print(restore_from)
    state_dict = torch.load(restore_from,map_location='cuda:0')
    model.load_state_dict(state_dict)
        
    model.cuda()
    model.eval()

    save_path =  os.path.join(args.out_dir, args.dataset, 'parsing')
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    parsing_preds, scales, centers = valid(model, valloader, input_size, num_samples, save_path)
    mIoU, f1 = miou.compute_mean_ioU(parsing_preds, scales, centers, args.num_classes, args.data_dir, \
        input_size, args.dataset, reverse=True)

    print(mIoU)
    print(f1)

if __name__ == '__main__':
    main()