test.py

# System libs
import os, time
import argparse
from utils import Evaluator
from distutils.version import LooseVersion
from collections import OrderedDict
# Numerical libs
import numpy as np
import torch
import torch.nn as nn
#from scipy.io import loadmat
import csv
# Our libs
from dataset import TestDataset
from models import ModelBuilder, SegmentationModule
from utils import colorEncode, find_recursive, setup_logger
from lib.nn import user_scattered_collate, async_copy_to
from lib.utils import as_numpy
from PIL import Image
from tqdm import tqdm
from config import cfg
import pickle as pkl
_palette=[0, 0, 0, 128, 0, 0, 0, 128, 0, 128, 128, 0, 0, 0, 128, 128, 0, 128, 0, 128, 128, 128, 128, 128, 64, 0, 0, 191, 0, 0, 64, 128, 0, 191, 128, 0, 64, 0, 128, 191, 0, 128, 64, 128, 128, 191, 128, 128, 0, 64, 0, 128, 64, 0, 0, 191, 0, 128, 191, 0, 0, 64, 128, 128, 64, 128, 22, 22, 22, 23, 23, 23, 24, 24, 24, 25, 25, 25, 26, 26, 26, 27, 27, 27, 28, 28, 28, 29, 29, 29, 30, 30, 30, 31, 31, 31, 32, 32, 32, 33, 33, 33, 34, 34, 34, 35, 35, 35, 36, 36, 36, 37, 37, 37, 38, 38, 38, 39, 39, 39, 40, 40, 40, 41, 41, 41, 42, 42, 42, 43, 43, 43, 44, 44, 44, 45, 45, 45, 46, 46, 46, 47, 47, 47, 48, 48, 48, 49, 49, 49, 50, 50, 50, 51, 51, 51, 52, 52, 52, 53, 53, 53, 54, 54, 54, 55, 55, 55, 56, 56, 56, 57, 57, 57, 58, 58, 58, 59, 59, 59, 60, 60, 60, 61, 61, 61, 62, 62, 62, 63, 63, 63, 64, 64, 64, 65, 65, 65, 66, 66, 66, 67, 67, 67, 68, 68, 68, 69, 69, 69, 70, 70, 70, 71, 71, 71, 72, 72, 72, 73, 73, 73, 74, 74, 74, 75, 75, 75, 76, 76, 76, 77, 77, 77, 78, 78, 78, 79, 79, 79, 80, 80, 80, 81, 81, 81, 82, 82, 82, 83, 83, 83, 84, 84, 84, 85, 85, 85, 86, 86, 86, 87, 87, 87, 88, 88, 88, 89, 89, 89, 90, 90, 90, 91, 91, 91, 92, 92, 92, 93, 93, 93, 94, 94, 94, 95, 95, 95, 96, 96, 96, 97, 97, 97, 98, 98, 98, 99, 99, 99, 100, 100, 100, 101, 101, 101, 102, 102, 102, 103, 103, 103, 104, 104, 104, 105, 105, 105, 106, 106, 106, 107, 107, 107, 108, 108, 108, 109, 109, 109, 110, 110, 110, 111, 111, 111, 112, 112, 112, 113, 113, 113, 114, 114, 114, 115, 115, 115, 116, 116, 116, 117, 117, 117, 118, 118, 118, 119, 119, 119, 120, 120, 120, 121, 121, 121, 122, 122, 122, 123, 123, 123, 124, 124, 124, 125, 125, 125, 126, 126, 126, 127, 127, 127, 128, 128, 128, 129, 129, 129, 130, 130, 130, 131, 131, 131, 132, 132, 132, 133, 133, 133, 134, 134, 134, 135, 135, 135, 136, 136, 136, 137, 137, 137, 138, 138, 138, 139, 139, 139, 140, 140, 140, 141, 141, 141, 142, 142, 142, 143, 143, 143, 144, 144, 144, 145, 145, 145, 146, 146, 146, 147, 147, 147, 148, 148, 148, 149, 149, 149, 150, 150, 150, 151, 151, 151, 152, 152, 152, 153, 153, 153, 154, 154, 154, 155, 155, 155, 156, 156, 156, 157, 157, 157, 158, 158, 158, 159, 159, 159, 160, 160, 160, 161, 161, 161, 162, 162, 162, 163, 163, 163, 164, 164, 164, 165, 165, 165, 166, 166, 166, 167, 167, 167, 168, 168, 168, 169, 169, 169, 170, 170, 170, 171, 171, 171, 172, 172, 172, 173, 173, 173, 174, 174, 174, 175, 175, 175, 176, 176, 176, 177, 177, 177, 178, 178, 178, 179, 179, 179, 180, 180, 180, 181, 181, 181, 182, 182, 182, 183, 183, 183, 184, 184, 184, 185, 185, 185, 186, 186, 186, 187, 187, 187, 188, 188, 188, 189, 189, 189, 190, 190, 190, 191, 191, 191, 192, 192, 192, 193, 193, 193, 194, 194, 194, 195, 195, 195, 196, 196, 196, 197, 197, 197, 198, 198, 198, 199, 199, 199, 200, 200, 200, 201, 201, 201, 202, 202, 202, 203, 203, 203, 204, 204, 204, 205, 205, 205, 206, 206, 206, 207, 207, 207, 208, 208, 208, 209, 209, 209, 210, 210, 210, 211, 211, 211, 212, 212, 212, 213, 213, 213, 214, 214, 214, 215, 215, 215, 216, 216, 216, 217, 217, 217, 218, 218, 218, 219, 219, 219, 220, 220, 220, 221, 221, 221, 222, 222, 222, 223, 223, 223, 224, 224, 224, 225, 225, 225, 226, 226, 226, 227, 227, 227, 228, 228, 228, 229, 229, 229, 230, 230, 230, 231, 231, 231, 232, 232, 232, 233, 233, 233, 234, 234, 234, 235, 235, 235, 236, 236, 236, 237, 237, 237, 238, 238, 238, 239, 239, 239, 240, 240, 240, 241, 241, 241, 242, 242, 242, 243, 243, 243, 244, 244, 244, 245, 245, 245, 246, 246, 246, 247, 247, 247, 248, 248, 248, 249, 249, 249, 250, 250, 250, 251, 251, 251, 252, 252, 252, 253, 253, 253, 254, 254, 254, 255, 255, 255]
#colors = loadmat('data/color150.mat')['colors']
names = {}


def visualize_result(data, pred, cfg):
    (img, info) = data

    # print predictions in descending order
    pred = np.int32(pred)
    pixs = pred.size
    uniques, counts = np.unique(pred, return_counts=True)
    print("Predictions in [{}]:".format(info))
    for idx in np.argsort(counts)[::-1]:
        name = names[uniques[idx] + 1]
        ratio = counts[idx] / pixs * 100
        if ratio > 0.1:
            print("  {}: {:.2f}%".format(name, ratio))

    # colorize prediction
    pred_color = colorEncode(pred, colors).astype(np.uint8)

    # aggregate images and save
    im_vis = np.concatenate((img, pred_color), axis=1)

    img_name = info.split('/')[-1]
    Image.fromarray(im_vis).save(
        os.path.join(cfg.TEST.result, img_name.replace('.jpg', '.png')))


def test(segmentation_module, loader, gpu,args,evaluator,eval_video,video):
    segmentation_module.eval()

    for i,data in enumerate(loader):
        # process data
        batch_data ={}
        if args.split !='test':
            imgs, gts,gtnames = data
            gts = gts.cuda(args.start_gpu)
            batch_data['seg_label'] = gts
        else:
            imgs, gtnames = data
      
        imgs = imgs.cuda(args.start_gpu)
        batch_data['img_data']= imgs
     
        segSize = (imgs.size(2),
                   imgs.size(3))

        with torch.no_grad():
            scores = segmentation_module(batch_data, segSize=segSize)
            pred = torch.argmax(scores, dim=1)
            pred = pred.data.cpu().numpy()


            # Add batch sample into evaluator
            if args.split !='test':
                target = gts.squeeze(1).cpu().numpy()
                evaluator.add_batch(target, pred)
                eval_video.add_batch(target,pred)
            if args.is_save:
                for j in range(pred.shape[0]):
                    imgpred = pred[j]
                    imgpred = Image.fromarray(imgpred.astype('uint8')).convert('P')
                    imgpred.putpalette(_palette)
                    if not os.path.exists(os.path.join(args.saveroot,video)):
                        os.makedirs(os.path.join(args.saveroot,video))
                    imgpred.save(os.path.join(args.saveroot,video,gtnames[j]))
                    print('video:{} image:{} saved.'.format(video,gtnames[j]))
             #############




        # visualization


def main(cfg, gpu,args):



    num_class =args.num_class 
    torch.cuda.set_device(gpu)

    # Network Builders
    net_encoder = ModelBuilder.build_encoder(
        arch=cfg.MODEL.arch_encoder,
        fc_dim=cfg.MODEL.fc_dim,
        weights=cfg.MODEL.weights_encoder)
    net_decoder = ModelBuilder.build_decoder(
        arch=cfg.MODEL.arch_decoder,
        fc_dim=cfg.MODEL.fc_dim,
        num_class=num_class,
        weights=cfg.MODEL.weights_decoder,
        use_softmax=True)

    crit = nn.NLLLoss(ignore_index=-1)

    segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)

    to_load = torch.load(args.load,map_location=torch.device("cuda:"+str(args.start_gpu)))
    new_state_dict = OrderedDict()
    for k, v in to_load.items():
        name = k[7:] # remove `module.`，表面从第7个key值字符取到最后一个字符，正好去掉了module.
        new_state_dict[name] = v #新字典的key值对应的value为一一对应的值。

    segmentation_module.load_state_dict(new_state_dict)
    print('load model parameters')

    segmentation_module.cuda(args.start_gpu)
    with open(os.path.join(args.dataroot,args.split+'.txt')) as f:
            lines=f.readlines()
            videolists = [line[:-1] for line in lines]


    # Dataset and Loader
    evaluator = Evaluator(num_class)
    eval_video = Evaluator(num_class)
    evaluator.reset()
    eval_video.reset()
    total_vmIOU=0.0
    total_vfwIOU=0.0
    total_video = len(videolists)
    v = []
    n = []
    for video in videolists:
        eval_video.reset()
        dataset_test = TestDataset(
            args.dataroot,
            video,args)
        loader_test = torch.utils.data.DataLoader(
            dataset_test,
            batch_size=args.batchsize,
            shuffle=False,
            num_workers=5,
            drop_last=False)
    # Main loop
        test(segmentation_module, loader_test, gpu,args,evaluator,eval_video,video)
        if args.split !='test':
            v_mIOU =eval_video.Mean_Intersection_over_Union()
            v.append(v)
            n.append(video)
            print(video, v_mIOU)
            total_vmIOU += v_mIOU
            v_fwIOU = eval_video.Frequency_Weighted_Intersection_over_Union()
    
            total_vfwIOU += v_fwIOU
    if args.split !='test':
        total_vmIOU  = total_vmIOU/total_video
        total_vfwIOU = total_vfwIOU/total_video
    
        Acc = evaluator.Pixel_Accuracy()
        Acc_class = evaluator.Pixel_Accuracy_Class()
        mIoU = evaluator.Mean_Intersection_over_Union()
        FWIoU = evaluator.Frequency_Weighted_Intersection_over_Union()
        print("Acc:{}, Acc_class:{}, mIoU:{}, fwIoU: {}, video mIOU: {}, video fwIOU: {}".format(Acc, Acc_class, mIoU, FWIoU,total_vmIOU,total_vfwIOU))

    print('Inference done!')
    


if __name__ == '__main__':
    assert LooseVersion(torch.__version__) >= LooseVersion('0.4.0'), \
        'PyTorch>=0.4.0 is required'

    def str2bool(v):
        if isinstance(v, bool):
           return v
        if v.lower() in ('yes', 'true', 't', 'y', '1'):
            return True
        elif v.lower() in ('no', 'false', 'f', 'n', '0'):
            return False
        else:
            raise argparse.ArgumentTypeError('Boolean value expected.')

    parser = argparse.ArgumentParser(
        description="PyTorch Semantic Segmentation Testing"
    )
    parser.add_argument(
        "--cfg",
        default="config/ade20k-hrnetv2.yaml",
        metavar="FILE",
        help="path to config file",
        type=str,
    )
    parser.add_argument(
        "opts",
        help="Modify config options using the command-line",
        default=None,
        nargs=argparse.REMAINDER,
    )
 
    parser.add_argument("--start_gpu",type=int,default=0)
    parser.add_argument("--num_class",type=int,default=124)
    parser.add_argument("--dataroot",type=str,default='')
    parser.add_argument("--saveroot",type=str,default='')
    parser.add_argument("--load_en",type=str,default='')
    parser.add_argument("--load_de",type=str,default='')
    parser.add_argument("--load",type=str,default='')
    parser.add_argument("--batchsize",type=int,default=4)
    parser.add_argument("--split",type=str,default='val')
    parser.add_argument("--is_save",type=str2bool,default=False)
    parser.add_argument("--lesslabel",type=str2bool,default=False)
    parser.add_argument("--use_720p",type=str2bool,default=False)
    args = parser.parse_args()

    cfg.merge_from_file(args.cfg)
    cfg.merge_from_list(args.opts)
    # cfg.freeze()

    logger = setup_logger(distributed_rank=0)   # TODO
    logger.info("Loaded configuration file {}".format(args.cfg))
    logger.info("Running with config:\n{}".format(cfg))

    cfg.MODEL.arch_encoder = cfg.MODEL.arch_encoder.lower()
    cfg.MODEL.arch_decoder = cfg.MODEL.arch_decoder.lower()

    # absolute paths of model weights
    cfg.MODEL.weights_encoder = ''
    cfg.MODEL.weights_decoder = ''

    print(cfg.MODEL.weights_encoder)

    # generate testing image list
    if not os.path.isdir(args.saveroot):
        os.makedirs(args.saveroot)

    main(cfg, args.start_gpu,args)
    print(args)