test_oti.py

'''
在ucf101、hmdb51数据集上进行测试
'''
import os
import sys
import time
import argparse
from timm import create_model
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torch.nn.parallel import DistributedDataParallel
import torch.distributed as dist
import torch.backends.cudnn as cudnn
from torch.cuda.amp import GradScaler
import torchvision
import torch.optim as optim
from utils.utils import init_distributed_mode, AverageMeter, reduce_tensor, accuracy
from utils.logger import setup_logger
import clip

from pathlib import Path
import yaml
import pprint
from dotmap import DotMap
import numpy as np
import pickle

import datetime
import shutil
from contextlib import suppress
from modules import cswin
from datasets.transforms import GroupScale, GroupCenterCrop, Stack, ToTorchFormatTensor, GroupNormalize, GroupOverSample, GroupFullResSample
from datasets.dataset_test import Video_dataset
from modules.oti_zsvr import video_header, OTI

from utils.Augmentation import get_augmentation, randAugment
from utils.solver import _lr_scheduler
from modules.text_prompt import text_prompt

'''
参数的传入
'''


def get_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument('--config', type=str, help='global config file')
    parser.add_argument('--weights', type=str, default=None)
    parser.add_argument('--log_time', default='front_lay')
    parser.add_argument('--dist_url', default='env://',
                        help='url used to set up distributed training')
    parser.add_argument('--world_size', default=1, type=int,
                        help='number of distributed processes')
    parser.add_argument("--local_rank", type=int,
                        help='local rank for DistributedDataParallel')
    parser.add_argument("--gpu_count",type=int,default=1,help="")
    parser.add_argument(
        "--precision",
        choices=["amp", "fp16", "fp32"],
        default="amp",
        help="Floating point precition."
    )
    parser.add_argument('--test_crops', type=int, default=3)
    parser.add_argument('--test_clips', type=int, default=3)
    parser.add_argument('--dense', default=False, action="store_true",
                        help='use multiple clips for test')
    args = parser.parse_args()
    return args


'''
去掉参数key中的module.
'''


def update_dict(dict):
    new_dict = {}
    for k, v in dict.items():
        new_dict[k.replace('module.', '')] = v
    return new_dict

'''
获取cswin模型最后norm层的维度
'''
def get_cswinmodel_pa(model):
    ks=[]
    vs=[]
    for k,v in model.named_parameters():
        ks.append(k)
        vs.append(v)
    print("The number of parameters of the cs_model is {}".format(len(ks)))
    return vs[-1].shape[0]

'''
主程序
'''


def main(args):
    init_distributed_mode(args)
    ## config
    with open(args.config, 'r') as f:
        config = yaml.load(f, Loader=yaml.FullLoader)
    config = DotMap(config)
    ## device
    device = "cpu"
    if torch.cuda.is_available():
        device = "cuda"
        cudnn.benchmark = True
    '''
    the save dir about test file
    '''
    save_root = os.path.join('/mnt/dolphinfs/hdd_pool/docker/user/hadoop-mtcv/zhuyan29/zsvideo_us/test_records', config['data']['dataset'], config['network']['arch'],args.log_time)
    if dist.get_rank() == 0:
        Path(save_root).mkdir(parents=True, exist_ok=True)
    '''
    load clip model
    '''
    clip_model, clip_state_dict = clip.load(config.network.arch,
                                            device='cpu', jit=False,
                                            internal_modeling=config.network.tm,
                                            T=config.data.num_segments,
                                            dropout=config.network.drop_out,
                                            emb_dropout=config.network.emb_dropout,
                                            pretrain=config.network.init,
                                            joint_st=config.network.joint_st)
    if args.precision == "amp" or args.precision == "fp32":
        clip_model = clip_model.float()

    ## crop
    input_mean = [0.48145466, 0.4578275, 0.40821073]
    input_std = [0.26862954, 0.26130258, 0.27577711]

    # rescale size
    if 'something' in config.data.dataset:
        scale_size = (240, 320)
    else:
        scale_size = 256 if config.data.input_size == 224 else config.data.input_size

    # crop size
    input_size = config.data.input_size

    # control the spatial crop
    if args.test_crops == 1:  # one crop
        cropping = torchvision.transforms.Compose([
            GroupScale(scale_size),
            GroupCenterCrop(input_size),
        ])
    elif args.test_crops == 3:  # do not flip, so only 3 crops (left right center)
        cropping = torchvision.transforms.Compose([
            GroupFullResSample(
                crop_size=input_size,
                scale_size=scale_size,
                flip=False)
        ])
    elif args.test_crops == 5:  # do not flip, so only 5 crops
        cropping = torchvision.transforms.Compose([
            GroupOverSample(
                crop_size=input_size,
                scale_size=scale_size,
                flip=False)
        ])
    elif args.test_crops == 10:
        cropping = torchvision.transforms.Compose([
            GroupOverSample(
                crop_size=input_size,
                scale_size=scale_size,
            )
        ])
    else:
        raise ValueError("Only 1, 3, 5, 10 crops are supported while we got {}".format(args.test_crops))
    ## val data
    val_data = Video_dataset(
        config.data.val_root, config.data.val_list, config.data.label_list,
        random_shift=False, num_segments=config.data.num_segments,
        modality=config.data.modality,
        image_tmpl=config.data.image_tmpl,
        test_mode=True,
        transform=torchvision.transforms.Compose([
            cropping,
            Stack(roll=False),
            ToTorchFormatTensor(div=True),
            GroupNormalize(input_mean, input_std),
        ]),
        dense_sample=args.dense,
        test_clips=args.test_clips,
        new_length=config.data.seg_length)

    val_sampler = torch.utils.data.distributed.DistributedSampler(val_data)
    val_loader = DataLoader(val_data,
                            batch_size=config.data.batch_size, num_workers=config.data.workers,
                            sampler=val_sampler, pin_memory=True, drop_last=False)

    #
    classes, _, text_dict = text_prompt(val_data)  # classes: 400*77
    n_class = text_dict[0].size(0)
    clip_model.eval()
    with torch.no_grad():
        classes_features = clip_model.encode_text(classes)
    '''
    model 
    '''
    video_head = video_header(config.network.sim_header, clip_state_dict)
    model_full = OTI(clip_model,video_head,config.data.num_segments)
    '''
    load parameter
    '''
    if os.path.isfile(args.weights):
        checkpoint = torch.load(args.weights, map_location='cpu')
        if dist.get_rank() == 0:
            print('load model: epoch {}'.format(checkpoint['epoch']))

        model_full.load_state_dict(update_dict(checkpoint['model_state_dict']),strict=True)
        del checkpoint

    if args.distributed:
        model_full = DistributedDataParallel(model_full.cuda(), device_ids=[args.gpu], find_unused_parameters=True)
#     print('model loading completed')
    '''
    预测
    '''
    prec1 = validate(
        val_loader, device,
        model_full, config, classes_features, args.test_crops, args.test_clips,save_root)
    print(prec1)
    return


def get_orthogonal_num_ration(now_tensor,best_tensor,true_label,cls):
    l2sums=torch.norm(best_tensor,dim=1)**2
    l2sums=torch.unsqueeze(l2sums,dim=1)
    maps=(now_tensor*best_tensor)/l2sums
    maps*=best_tensor 
    gain=now_tensor-maps
    accuracies=[]
    # res_ucf101=[]
    for i in range(0,11):
        ratio=i/10
        af_otf=best_tensor+gain*ratio
        pres=(af_otf@cls.T).softmax(dim=-1).topk(1).indices
        accuracy=sum(pres==true_label).numpy()[0]/gain.shape[0]
        accuracy=accuracy*100
        accuracies.append(accuracy)
    maxweight=accuracies.index(max(accuracies))/10
    best_af_otf=best_tensor+gain*maxweight
    accuracies=[round(i,1) for i in accuracies]
    return max(accuracies),accuracies,maxweight,best_af_otf


def validate(val_loader, device, model, config, text_features, test_crops, test_clips,save_root):
    top1 = AverageMeter()
    top5 = AverageMeter()
    front_top1=AverageMeter()
    front_top5=AverageMeter()
    model.eval()
    proc_start_time = time.time()

    sim_logits = []  #
    labels = []  #
    i_features = []
    predics_top5 = []
    front_features=[]
    
    with torch.no_grad():
        n_class = text_features.size(0)

        for i, (images, class_id) in enumerate(val_loader):
#             print(image.shape)
            batch_size = class_id.numel()
            num_crop = test_crops

            num_crop *= test_clips  # 4 clips for testing when using dense sample

            class_id = class_id.to(device)
            text_features = text_features.to(device)
#             n_seg = config.data.num_segments
#             image = image.view((-1, n_seg, 3) + image.size()[-2:])
            images = images.view((-1, config.data.num_segments * config.data.seg_length, 3) + images.size()[-2:])
            b, t, c, h, w = images.size()
#             print(b, t, c, h, w )
            # image_input = image.to(device).view(-1, c, h, w)
            images = images.to(device)
            video_emb,video_emb_front = model.module.encode_image(images)
#             print("image_features is {}".format(image_features.shape))
            cnt_time = time.time() - proc_start_time
    
            video_emb = video_emb.reshape(batch_size, num_crop, -1).mean(1)  # bs dim
            video_emb_front=video_emb_front.reshape(batch_size, num_crop, -1).mean(1)
            
            video_emb /= video_emb.norm(dim=-1, keepdim=True)
            video_emb_front/=video_emb_front.norm(dim=-1, keepdim=True)
            text_features /= text_features.norm(dim=-1, keepdim=True)
            
            similarity = (100.0 * video_emb @ text_features.T)
            similarity = similarity.view(batch_size, -1, n_class).softmax(dim=-1)
            similarity = similarity.mean(dim=1, keepdim=False)  # bs 200
            
            front_similarity = (100.0 * video_emb_front @ text_features.T)
            front_similarity = front_similarity.view(batch_size, -1, n_class).softmax(dim=-1)
            front_similarity = front_similarity.mean(dim=1, keepdim=False)  # bs 200
            
            
            prec5 = similarity.topk(5)
            prec5 = prec5.indices

            ########## gathering
            i_features.append(concat_all_gather(video_emb))
            front_features.append(concat_all_gather(video_emb_front))
            sim_logits.append(concat_all_gather(similarity))
            labels.append(concat_all_gather(class_id))
            predics_top5.append(concat_all_gather(prec5))
            ##########

            prec = accuracy(similarity, class_id, topk=(1, 5))
            prec1 = reduce_tensor(prec[0])
            prec5 = reduce_tensor(prec[1])
            top1.update(prec1.item(), class_id.size(0))
            top5.update(prec5.item(), class_id.size(0))

            
            front_prec = accuracy(front_similarity, class_id, topk=(1, 5))
            front_prec1 = reduce_tensor(front_prec[0])
            front_prec5 = reduce_tensor(front_prec[1])
            front_top1.update(front_prec1.item(), class_id.size(0))
            front_top5.update(front_prec5.item(), class_id.size(0))
            
            if i % config.logging.print_freq == 0 and dist.get_rank() == 0:
                runtime = float(cnt_time) / (i + 1) / (batch_size * dist.get_world_size())
                print(
                    ('Test: [{0}/{1}], average {runtime:.4f} sec/video \t'
                     'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
                     'Prec@5 {top5.val:.3f} ({top5.avg:.3f})\t'
                     'front_Prec@1 {front_top1.val:.3f} ({front_top1.avg:.3f})\t'
                     'front_Prec@5 {front_top5.val:.3f} ({front_top5.avg:.3f})\t'.format(
                        i, len(val_loader), runtime=runtime, top1=top1, top5=top5,front_top1=front_top1,front_top5=front_top5)))

    if dist.get_rank() == 0:
        print('-----Full-classes Evaluation------')
        print('Overall Top1 {:.03f}% Top5 {:.03f}%'.format(top1.avg, top5.avg))

        ## half-classes evaluation
        sim, la = sim_logits[0], labels[0]
        vid_feat = i_features[0]
        pre_lab5 = predics_top5[0]
        front_vid_feat=front_features[0]
        for i in range(1, len(sim_logits)):
            sim = torch.cat((sim, sim_logits[i]), 0)
            la = torch.cat((la, labels[i]), 0)
            vid_feat = torch.cat((vid_feat, i_features[i]), 0)
            pre_lab5 = torch.cat((pre_lab5, predics_top5[i]), 0)
            front_vid_feat=torch.cat((front_vid_feat, front_features[i]), 0)
        las=torch.unsqueeze(la,1)
#         print(pre_lab5.shape,la.shape)
        labels=torch.cat((pre_lab5.cpu(),las.cpu()),1)
        save_pickle(labels, vid_feat.cpu(),front_vid_feat.cpu(), text_features.cpu(), save_root)
        bestaccu,accuracies,maxweight,best_af_otf=get_orthogonal_num_ration(vid_feat,front_vid_feat,las,text_features)
        print('Best_prec@1 using orthogonal temporal interpolation feature is {}, the weight is {}.\n Prec@1 under different interpolation weights are {}.'.format(bestaccu,maxweight,accuracies))
        
        best_af_otf_acc_split,best_af_otf_acc_split_top5=multi_split_test(best_af_otf.cpu(), text_features.cpu(), la.cpu())
        accuracy_split,accuracy_split_std=np.mean(best_af_otf_acc_split), np.std(best_af_otf_acc_split)
        print('-----Half-classes Evaluation af_otf-----')
        print('Top1: mean {:.03f}%, std {:.03f}%'.format(accuracy_split, accuracy_split_std))
        

        # acc_split, acc_split_top5 = multi_split_test(vid_feat.cpu(), text_features.cpu(), la.cpu())
        # accuracy_split, accuracy_split_std = np.mean(acc_split), np.std(acc_split)
        # accuracy_split_top5, accuracy_split_top5_std = np.mean(acc_split_top5), np.std(acc_split_top5)
        
        # front_acc_split, front_acc_split_top5 = multi_split_test(front_vid_feat.cpu(), text_features.cpu(), la.cpu())
        # front_accuracy_split, front_accuracy_split_std = np.mean(front_acc_split), np.std(front_acc_split)
        # front_accuracy_split_top5, front_accuracy_split_top5_std = np.mean(front_acc_split_top5), np.std(front_acc_split_top5)
        # print('-----Half-classes Evaluation after layer-----')
        # print('Top1: mean {:.03f}%, std {:.03f}%'.format(accuracy_split, accuracy_split_std))
        # # print('Top5: mean {:.03f}%, std {:.03f}%'.format(accuracy_split_top5, accuracy_split_top5_std))
        
        # print('-----Half-classes Evaluation before layer-----')
        # print('Top1: mean {:.03f}%, std {:.03f}%'.format(front_accuracy_split, front_accuracy_split_std))
        # print('Top5: mean {:.03f}%, std {:.03f}%'.format(front_accuracy_split_top5, front_accuracy_split_top5_std))

    return top1.avg


# utils
@torch.no_grad()
def concat_all_gather(tensor):
    """
    Performs all_gather operation on the provided tensors.
    *** Warning ***: torch.distributed.all_gather has no gradient.
    """
    tensors_gather = [torch.ones_like(tensor)
                      for _ in range(torch.distributed.get_world_size())]
    torch.distributed.all_gather(tensors_gather, tensor, async_op=False)

    output = torch.cat(tensors_gather, dim=0)
    return output.cpu()
def save_pickle(labels, visul_video,front_vid_feat, text_features, save_root):
    flabel = open(os.path.join(save_root, 'pre_labels'), 'wb')
    pickle.dump(labels, flabel)
    flabel.close()
    ffeas = open(os.path.join(save_root, 'video_features'), 'wb')
    pickle.dump(visul_video, ffeas)
    ffeas.close()
    ffeas = open(os.path.join(save_root, 'before_video_features'), 'wb')
    pickle.dump(front_vid_feat, ffeas)
    ffeas.close()
    ffeas = open(os.path.join(save_root, 'cls_feature'), 'wb')
    pickle.dump(text_features, ffeas)
    ffeas.close()



def compute_accuracy(vis_emb, text_emb, label):
    n_class = len(text_emb)
    n_samples = len(vis_emb)
    similarity = (100.0 * vis_emb @ text_emb.T)
    similarity = similarity.view(n_samples, -1, n_class).softmax(dim=-1)
    similarity = similarity.mean(dim=1, keepdim=False)  # b 101
    prec = accuracy(similarity, label, topk=(1, 5))
    return prec[0], prec[1]



def multi_split_test(vis_embs, text_embs, true_label):
    # vis_embs: [10000, 768]
    # text_embs: [101, 768]
    # true_label: [10000,]
    full_acc1, full_acc5 = compute_accuracy(vis_embs, text_embs, true_label)

    # Calculate accuracy per split
    # Only when the model has been trained on a different dataset
    true_label = true_label.numpy()
    accuracy_split, accuracy_split_top5 = np.zeros(10), np.zeros(10)
    for split in range(len(accuracy_split)):
        np.random.seed(split)
        sel_classes = np.random.permutation(len(text_embs))[:len(text_embs) // 2]  # the half of categories
        sel = [l in sel_classes for l in true_label]  # 
        subclasses = np.unique(true_label[sel])  ## 
        #### the mapping of labels
        tl = np.array([int(np.where(l == subclasses)[0]) for l in true_label[sel]])  ## true_label[sel] 
        tl = torch.from_numpy(tl)  
        acc, acc5 = compute_accuracy(vis_embs[sel], text_embs[subclasses], tl)
        accuracy_split[split] = acc
        accuracy_split_top5[split] = acc5

    return accuracy_split, accuracy_split_top5


if __name__ == '__main__':
    args = get_parser()
    main(args)