test.py

import torch
import torch.nn.functional as F
import torch.optim as optim
import model
from tensorboard_logger import log_value
import utils.wsad_utils as utils
import numpy as np
from torch.autograd import Variable
from eval.classificationMAP import getClassificationMAP as cmAP
from eval.eval_detection import ANETdetection
import wsad_dataset
from eval.detectionMAP import getDetectionMAP as dmAP
import scipy.io as sio
import pdb
from tensorboard_logger import Logger
import multiprocessing as mp
import options
import proposal_methods as PM
import pandas as pd
from collections import defaultdict

torch.set_default_tensor_type('torch.cuda.FloatTensor')


@torch.no_grad()
def test(itr, dataset, args, model, logger, device, pool):
    model.eval()
    done = False
    instance_logits_stack = []
    element_logits_stack = {}
    logits_stack={}
    attn_stack={}
    video_list_stack=[]
    labels_stack = []

    back_norms = []
    front_norms = []
    ind = 0
    proposals = []
    attn_dict = defaultdict(list)
    while not done:
        if dataset.currenttestidx % (len(dataset.testidx) // 5) == 0:
            print('Testing test data point %d of %d' % (dataset.currenttestidx, len(dataset.testidx)))

        data_dict=dataset.load_data(is_training=False)
        features,labels,vn,done=data_dict['feat'],data_dict['lab'],data_dict['vn'],data_dict['done']
        # features, labels, vn, done = dataset.load_data(is_training=False)

        seq_len = [features.shape[0]]
        if seq_len == 0:
            continue
        features = torch.from_numpy(features).float().to(device).unsqueeze(0)
        with torch.no_grad():
            outputs = model(Variable(features), is_training=False, seq_len=seq_len)

            _, element_logits, atn_supp, atn_drop, element_atn, x_atn_logit = outputs
            if 'Thumos' in args.dataset_name:
                proposals.append(PM.multiple_threshold_hamnet(vn,args, element_logits, element_atn))
            else:
                proposals.append(PM.multiple_threshold_hamnet_ant(vn,args, element_logits, element_atn))

            logits = element_logits.squeeze(0)
        tmp = F.softmax(torch.mean(torch.topk(logits, k=int(np.ceil(len(features) / 8)), dim=0)[0], dim=0),
                        dim=0).cpu().data.numpy()
        # logits = logits.cpu().data.numpy()

        instance_logits_stack.append(tmp)
        element_logits_stack[vn]=(element_logits).detach().cpu().numpy()
        logits_stack[vn]=logits.detach().cpu().numpy()
        attn_stack[vn]=element_atn.detach().cpu().numpy()
        video_list_stack.append(vn)
        labels_stack.append(labels)

    # np.save('temp/attn.npy',attn_dict)
    # np.save('temp/video_list_ANT13.npy',video_list_stack)
    # np.save('temp/element_logits_stack_ANT13.npy',element_logits_stack)
    # np.save('temp/attn_stack_ANT13.npy',attn_stack)
    # np.save('temp/labels_stack_ANT13.npy',labels_stack)

    # np.save('temp/pred_wo_SCL.npy',logits_stack)
    # np.save('temp/p_MIL_0.npy',element_logits_stack)
    # np.save('temp/atten_wo_SCL.npy',np.array(attn_stack))
    # np.save('temp/atten_Cls_VCL.npy',np.array(attn_stack))

    # update in 2022/3/2
    # to generate the scores for visualization
    # 1. to save the predictions of Full method
    # np.save('temp/attn_Full_1.npy',attn_stack)
    # np.save('temp/pred_Full_1.npy',logits_stack)


    instance_logits_stack = np.array(instance_logits_stack)
    labels_stack = np.array(labels_stack)
    proposals = pd.concat(proposals).reset_index(drop=True)

    # CVPR2020
    if 'Thumos14' in args.dataset_name:
        iou = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
        dmap_detect = ANETdetection(dataset.path_to_annotations, iou, args=args)
        dmap_detect.ground_truth.to_csv('temp/groundtruth.csv')
        dmap_detect.prediction = proposals
        dmap = dmap_detect.evaluate()
    else:
        iou = [0.5, 0.55, 0.60, 0.65, 0.70, 0.75, 0.80, 0.85, 0.90, 0.95]

        dmap_detect = ANETdetection(dataset.path_to_annotations, iou, args=args, subset='validation')
        dmap_detect.ground_truth.to_csv('temp/groundtruth_ANT13.csv')
        dmap_detect.prediction = proposals
        dmap = dmap_detect.evaluate()
    # ECCV2018
    # dmap, iou = dmAP(element_logits_stack, dataset.path_to_annotations, args,pool)

    if args.dataset_name == 'Thumos14':
        test_set = sio.loadmat('test_set_meta.mat')['test_videos'][0]
        for i in range(np.shape(labels_stack)[0]):
            if test_set[i]['background_video'] == 'YES':
                labels_stack[i, :] = np.zeros_like(labels_stack[i, :])


    ############ for Snippet-level classification ######################
    groundTruthGroupByVideoId = dmap_detect.ground_truth.groupby('video-id')
    # to evaluate the snippet-level classification
    import copy
    from sklearn.metrics import accuracy_score
    def calSnippetClassificationAcc(pred, vns, labels):
        labels = copy.deepcopy(labels)

        preds = []
        for i, vn in enumerate(vns):
            preds.append(pred[vn].squeeze())

        preds = np.concatenate(preds, axis=0)
        labels = np.concatenate(labels, axis=0)

        args_preds = np.argmax(preds, axis=-1)
        args_labels = np.argmax(labels, axis=-1)

        return accuracy_score(args_labels, args_preds)

    def softmax(pred):
        return np.exp(pred) / np.exp(pred).sum(axis=-1, keepdims=True)

    def calBinaryClassificationAcc(pred,vns,labels,background=False):
        labels = copy.deepcopy(labels)

        preds = []
        for i, vn in enumerate(vns):
            preds.append(pred[vn].squeeze())

        preds = np.concatenate(preds, axis=0)
        labels = np.concatenate(labels, axis=0)

        if background:
            preds=1-softmax(preds)
            preds=preds[...,-1]
        foreground_labels=1-labels[...,-1]

        # generate preds to 0-1
        fore_preds=np.array([p>0.5 for p in preds]).astype(float)
        # translate the foreground_labels and foreground_preds
        return ((fore_preds * foreground_labels).sum() +
                ((1 - fore_preds) * (1 - foreground_labels)).sum()) \
               / foreground_labels.shape[0]

    from sklearn.metrics import roc_auc_score
    def calBinaryAUC(pred,vns,labels,background=False):
        labels = copy.deepcopy(labels)

        preds = []
        for i, vn in enumerate(vns):
            preds.append(pred[vn].squeeze())

        preds = np.concatenate(preds, axis=0)
        labels = np.concatenate(labels, axis=0)

        if background:
            preds = 1 - softmax(preds)
            preds = preds[..., -1]
        foreground_labels = 1 - labels[..., -1]

        return roc_auc_score(foreground_labels.astype(int), preds)

    snippet_labels=[]
    for vn in video_list_stack:
        p = logits_stack[vn].squeeze()
        # [t,c+1]
        snippet_label = np.zeros_like(p)
        snippet_label[:, -1] = 1

        if 'ActivityNet1.3' not in args.dataset_name:
            gt = groundTruthGroupByVideoId.get_group(vn.decode())
        else:
            gt = groundTruthGroupByVideoId.get_group(vn)
        for idx, this_pred in gt.iterrows():
            snippet_label[this_pred['t-start']:this_pred["t-end"], this_pred['label']] = 1
            snippet_label[this_pred['t-start']:this_pred["t-end"], -1] = 0
        snippet_labels.append(snippet_label)

    pred_cmap = calSnippetClassificationAcc(logits_stack, video_list_stack, snippet_labels)
    pred_bin_cmap=calBinaryClassificationAcc(logits_stack,video_list_stack,snippet_labels,background=True)
    att_bin_cmap=calBinaryClassificationAcc(attn_stack,video_list_stack,snippet_labels,background=False)
    pred_bin_auc=calBinaryAUC(logits_stack,video_list_stack,snippet_labels,background=True)
    att_bin_auc=calBinaryAUC(attn_stack,video_list_stack,snippet_labels,background=False)
    print('snippet-level classification mAP:{}'.format(pred_cmap),
          'snippet_binary_classification mAP:{}'.format(pred_bin_cmap),
          'snippet_binary_attention mAP:{}'.format(att_bin_cmap),
          'snippet_binary_classification ACU:{}'.format(pred_bin_auc),
          'snippet_binary_attention AUC:{}'.format(att_bin_auc),
          )
    #########################################################


    cmap = cmAP(instance_logits_stack, labels_stack)
    print('Classification map %f' % cmap)
    print('||'.join(['map @ {} = {:.3f} '.format(iou[i], dmap[i] * 100) for i in range(len(iou))]))
    print('mAP Avg ALL: {:.3f}'.format(sum(dmap) / len(iou) * 100))
    # print('Detection map @ %2f = %2f, map @ %2f = %2f, map @ %2f = %2f, map @ %2f = %2f, map @ %2f = %2f' %(iou[0], dmap[0],iou[1], dmap[1],iou[2], dmap[2],iou[3], dmap[3],iou[4], dmap[4]))

    logger.log_value('Test Classification mAP', cmap, itr)
    for item in list(zip(dmap, iou)):
        logger.log_value('Test Detection mAP @ IoU = ' + str(item[1]), item[0], itr)
    utils.write_to_file(args.dataset_name, dmap, cmap, itr)

    return iou, dmap


if __name__ == '__main__':
    args = options.parser.parse_args()
    device = torch.device("cuda")
    dataset = getattr(wsad_dataset, args.dataset)(args)

    model = getattr(model, args.use_model)(dataset.feature_size, dataset.num_class, opt=args).to(device)
    model.load_state_dict(torch.load('./ckpt/best_' + args.model_name + '.pkl'))
    logger = Logger('./logs/test_' + args.model_name)
    pool = mp.Pool(5)

    iou, dmap = test(-1, dataset, args, model, logger, device, pool)
    print('mAP Avg 0.1-0.5: {}, mAP Avg 0.1-0.7: {}, mAP Avg ALL: {}'.format(np.mean(dmap[:5]) * 100,
                                                                             np.mean(dmap[:7]) * 100,
                                                                             np.mean(dmap) * 100))