Create test.py

Author: QingpingZheng

parsing/dml_csr/test.py

@@ -0,0 +1,182 @@
+#!/usr/bin/env python
+# -*- encoding: utf-8 -*-
+"""
+@Author     :   Qingping Zheng
+@Contact    :   [email protected]
+@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()