eval_zeroshot.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import argparse
import numpy as np
import os.path as osp
import time
from collections import OrderedDict

import pandas as pd
import torch
import torchvision.transforms as transforms
import torchvision.transforms._transforms_video as transforms_video
from sklearn.metrics import confusion_matrix

from lavila.data import datasets
from lavila.data.video_transforms import Permute, SpatialCrop, TemporalCrop
from lavila.models import models
from lavila.models.utils import inflate_positional_embeds
from lavila.utils import distributed as dist_utils
from lavila.utils.evaluation import accuracy, get_mean_accuracy
from lavila.utils.evaluation_egomcq import egomcq_accuracy_metrics
from lavila.utils.evaluation_ek100mir import (calculate_k_counts, calculate_IDCG, calculate_mAP, calculate_nDCG)
from lavila.utils.evaluation_charades import charades_map
from lavila.utils.preprocess import generate_label_map, generate_tokenizer


def get_args_parser():
    parser = argparse.ArgumentParser(description='LAVILA 0-shot evaluations', add_help=False)
    parser.add_argument('--dataset', default='ek100_mir', type=str,
                        choices=['ek100_cls', 'ek100_mir', 'charades_ego', 'egtea', 'ego4d_mcq'])
    parser.add_argument('--root',
                        default='datasets/EK100/video_ht256px/',
                        type=str, help='path to dataset root')
    parser.add_argument('--metadata-val',
                        default='datasets/EK100/epic-kitchens-100-annotations/retrieval_annotations/EPIC_100_retrieval_test.csv',
                        type=str, help='path to metadata file (val set)')
    parser.add_argument('--relevancy-path',
                        default='datasets/EK100/epic-kitchens-100-annotations/retrieval_annotations/relevancy/caption_relevancy_EPIC_100_retrieval_test.pkl',
                        type=str, help='path to relevancy matrix (val set)')
    parser.add_argument('--output-dir', default='./', type=str, help='output dir')
    parser.add_argument('--num-crops', default=1, type=int, help='number of crops in transforms')
    parser.add_argument('--num-clips', default=1, type=int, help='number of clips (for untrimmed videos, eg. Charades)')
    parser.add_argument('--clip-length', default=4, type=int, help='clip length')
    parser.add_argument('--clip-stride', default=16, type=int, help='clip stride')
    parser.add_argument('--sparse-sample', action='store_true', help='switch to sparse sampling')
    parser.add_argument('--batch-size', default=16, type=int, help='batch_size')
    parser.add_argument('--cls-use-template', action='store_true', help='use prompt in 0-shot classification')
    parser.add_argument('--print-freq', default=100, type=int)
    parser.add_argument('-j', '--workers', default=10, type=int, metavar='N',
                        help='number of data loading workers per process')
    parser.add_argument('--resume', default='', type=str, help='path to latest checkpoint')
    parser.add_argument('--use-half', action='store_true')
    return parser


def main(args):
    if args.resume:
        ckpt_path = args.resume
    elif osp.isfile(osp.join(args.output_dir, 'checkpoint_best.pt')):
        ckpt_path = osp.join(args.output_dir, 'checkpoint_best.pt')
    else:
        raise Exception('no checkpoint found')

    ckpt = torch.load(ckpt_path, map_location='cpu')

    # create model
    state_dict = OrderedDict()
    for k, v in ckpt['state_dict'].items():
        state_dict[k.replace('module.', '')] = v

    old_args = ckpt['args']
    print('=> creating model: {}'.format(old_args.model))
    model = getattr(models, old_args.model)(
        text_use_cls_token=old_args.use_cls_token,
        project_embed_dim=old_args.project_embed_dim,
        gated_xattn=False if 'gated_xattn' not in old_args else old_args.gated_xattn,
        timesformer_gated_xattn=False if 'timesformer_gated_xattn' not in old_args else old_args.timesformer_gated_xattn,
        timesformer_freeze_space=False if 'timesformer_freeze_space' not in old_args else old_args.timesformer_freeze_space,
        freeze_lm_vclm=False if 'freeze_lm_vclm' not in old_args else old_args.freeze_lm_vclm,
        freeze_visual_vclm=False if 'freeze_visual_vclm' not in old_args else old_args.freeze_visual_vclm,
        num_frames=args.clip_length,
        drop_path_rate=0,
    )
    model.cuda()
    if 'TIMESFORMER' in old_args.model or 'EGOVLP' in old_args.model:
        # inflate weight
        print('=> inflating PE in models due to different frame numbers')
        state_dict = inflate_positional_embeds(
            model.state_dict(), state_dict,
            num_frames=args.clip_length,
            load_temporal_fix='bilinear',
        )
    model.load_state_dict(state_dict, strict=True)
    print("=> loaded resume checkpoint '{}' (epoch {}, best_metric = {})".format(args.resume, ckpt['epoch'], ckpt['best_acc1']))

    torch.backends.cudnn.benchmark = True

    if args.dataset in ['ek100_cls', 'charades_ego', 'egtea']:
        labels, mapping_vn2act = generate_label_map(args.dataset)
    else:
        mapping_vn2act = None
    tokenizer = generate_tokenizer(old_args.model)
    crop_size = 224 if '336PX' not in old_args.model else 336
    if args.num_crops == 1 and args.num_clips == 1:
        val_transform = transforms.Compose([
            Permute([3, 0, 1, 2]),  # T H W C -> C T H W
            transforms.Resize(crop_size),
            transforms.CenterCrop(crop_size),
            (transforms_video.NormalizeVideo(mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375]) if ('OPENAI' not in old_args.model) else
             transforms_video.NormalizeVideo(mean=[108.3272985, 116.7460125, 104.09373615000001], std=[68.5005327, 66.6321579, 70.32316305])),
        ])
    else:
        val_transform = transforms.Compose([
            Permute([3, 0, 1, 2]),  # T H W C -> C T H W
            transforms.Resize(crop_size),
            (transforms_video.NormalizeVideo(mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375]) if ('OPENAI' not in old_args.model) else
             transforms_video.NormalizeVideo(mean=[108.3272985, 116.7460125, 104.09373615000001], std=[68.5005327, 66.6321579, 70.32316305])),
            TemporalCrop(frames_per_clip=args.clip_length, stride=args.clip_length),
            SpatialCrop(crop_size=crop_size, num_crops=args.num_crops),
            ])

    val_dataset = datasets.get_downstream_dataset(
        val_transform, tokenizer, args, subset='val', label_mapping=mapping_vn2act,
    )

    val_loader = torch.utils.data.DataLoader(
        val_dataset, batch_size=args.batch_size, shuffle=False,
        num_workers=args.workers, pin_memory=True, drop_last=False)

    if args.cls_use_template:
        templates = ['#C C {}', '#C {}']
    else:
        templates = ['{}']

    if args.dataset in ['ek100_cls', 'charades_ego', 'egtea']:
        preds, targets = validate_zeroshot(val_loader, templates, labels, model, tokenizer)
    if args.dataset == 'ek100_cls':
        if args.use_half:
            preds = preds.float()
        top1, top5 = accuracy(preds, targets, topk=(1, 5))
        print('top1 = {:.3f}'.format(top1.item()))
        print('top5 = {:.3f}'.format(top5.item()))
    elif args.dataset == 'charades_ego':
        preds, targets = preds.numpy(), targets.numpy()
        m_ap, _, _ = charades_map(preds, targets)
        print('mAP = {:.3f}'.format(m_ap))
    elif args.dataset == 'egtea':
        preds, targets = preds.numpy(), targets.numpy()
        print(preds.shape, targets.shape)
        cm = confusion_matrix(targets, preds.argmax(axis=1))
        mean_class_acc, acc = get_mean_accuracy(cm)
        print('Mean Acc. = {:.3f}, Top-1 Acc. = {:.3f}'.format(mean_class_acc, acc))

    if args.dataset == 'ek100_mir':
        val_dataset = datasets.VideoCaptionDatasetCLIP(
            'ek100_mir',
            args.root,
            args.metadata_val,
            transform=val_transform, is_training=False,
            tokenizer=tokenizer,
            clip_length=args.clip_length,
            clip_stride=args.clip_stride,
            sparse_sample=False
        )
        val_loader = torch.utils.data.DataLoader(
            val_dataset, batch_size=args.batch_size, shuffle=False,
            num_workers=args.workers, pin_memory=True, drop_last=False
        )
        similarity_matrix = get_similarity_matrix(val_loader, model, print_freq=args.print_freq, use_half=args.use_half)
        similarity_matrix = (similarity_matrix + 1) / 2
        video_id = pd.read_csv(args.metadata_val).values[:, 0]
        text_id = pd.read_csv(args.metadata_val.replace("test.csv", "test_sentence.csv")).values[:, 0]
        indexes = [video_id.tolist().index(elem) for elem in text_id]
        similarity_matrix = similarity_matrix[:, indexes]
        print(similarity_matrix.shape)
        rel_matrix = pd.read_pickle(args.relevancy_path)
        vis_map = calculate_mAP(similarity_matrix, rel_matrix)
        txt_map = calculate_mAP(similarity_matrix.T, rel_matrix.T)
        print('mAP: V->T: {:.3f} T->V: {:.3f} AVG: {:.3f}'.format(vis_map, txt_map, (vis_map + txt_map) / 2))
        vis_k_counts = calculate_k_counts(rel_matrix)
        txt_k_counts = calculate_k_counts(rel_matrix.T)
        vis_IDCG = calculate_IDCG(rel_matrix, vis_k_counts)
        txt_IDCG = calculate_IDCG(rel_matrix.T, txt_k_counts)
        vis_nDCG = calculate_nDCG(similarity_matrix, rel_matrix, k_counts=vis_k_counts, IDCG=vis_IDCG)
        txt_nDCG = calculate_nDCG(similarity_matrix.T, rel_matrix.T, k_counts=txt_k_counts, IDCG=txt_IDCG)
        print('nDCG: V->T: {:.3f} T->V: {:.3f} AVG: {:.3f}'.format(vis_nDCG, txt_nDCG, (vis_nDCG + txt_nDCG) / 2))

    if args.dataset == 'ego4d_mcq':
        val_dataset = datasets.VideoCaptionDatasetMCQ(
            args.dataset,
            args.root,
            args.metadata_val,
            transform=val_transform, is_training=False,
            tokenizer=tokenizer,
            clip_length=args.clip_length,
            clip_stride=args.clip_stride,
            sparse_sample=False,
        )
        val_loader = torch.utils.data.DataLoader(
            val_dataset, batch_size=args.batch_size, shuffle=False,
            num_workers=args.workers, pin_memory=True, drop_last=False
        )
        validate_mcq(val_loader, model, use_half=args.use_half)


def validate_zeroshot(val_loader, templates, labels, model, tokenizer):
    model.eval()
    if args.use_half:
        model = model.half()
    all_outputs = []
    all_targets = []
    all_vis_features = []
    print('=> encoding captions')
    with torch.no_grad():
        text_features = []
        for label in labels:
            if isinstance(label, list):
                texts = [tmpl.format(lbl) for tmpl in templates for lbl in label]
            else:
                texts = [tmpl.format(label) for tmpl in templates]
            texts = tokenizer(texts)
            if isinstance(texts, tuple):
                # Bert-style tokenizer will output both ids and mask
                texts, masks = texts
                texts = texts.cuda(non_blocking=True)
                masks = masks.cuda(non_blocking=True)
            else:
                texts = texts.cuda(non_blocking=True)
                masks = None
            texts = texts.view(-1, 77).contiguous()
            masks = masks.view(-1, 77).contiguous() if masks is not None else None
            if masks is not None:
                class_embeddings = dist_utils.get_model(model).encode_text(texts, attention_mask=masks)
            else:
                class_embeddings = dist_utils.get_model(model).encode_text(texts)
            class_embeddings = class_embeddings / class_embeddings.norm(dim=-1, keepdim=True)

            class_embeddings = class_embeddings.mean(dim=0)
            class_embeddings = class_embeddings / class_embeddings.norm(dim=-1, keepdim=True)

            text_features.append(class_embeddings)
        text_features = torch.stack(text_features, dim=0)

        print('=> start forwarding')
        end_time = time.time()
        for i, (images, target) in enumerate(val_loader):
            if i % args.print_freq == 0:
                print('finish batch {}/{} in {} sec'.format(i, len(val_loader), time.time() - end_time))
                end_time = time.time()
            if isinstance(images, torch.Tensor):
                images = images.cuda(non_blocking=True)
                if args.use_half:
                    images = images.half()
                target = target.cuda(non_blocking=True)

                # encode images
                image_features = dist_utils.get_model(model).encode_image(images)
                image_features = image_features / image_features.norm(dim=-1, keepdim=True)
                all_vis_features.append(image_features)
                # cosine similarity as logits
                logits_per_image = image_features @ text_features.t()
                # logits_per_image = torch.softmax(logits_per_image, dim=1)
            else:
                target = target.cuda(non_blocking=True)
                images_list = images
                logits_all_clips = []
                for images in images_list:
                    images = images.cuda(non_blocking=True)
                    if args.use_half:
                        images = images.half()
                    image_features = dist_utils.get_model(model).encode_image(images)
                    image_features = image_features / image_features.norm(dim=-1, keepdim=True)
                    logits_per_image = image_features @ text_features.t()
                    logits_all_clips.append(logits_per_image)

                logits_all_clips = torch.stack(logits_all_clips, dim=0)
                logits_per_image = logits_all_clips.max(0).values
                # logits_per_image = logits_all_clips.mean(0)
                logits_per_image = torch.softmax(logits_per_image, dim=1)

            all_outputs.append(logits_per_image.cpu())
            all_targets.append(target.cpu())

    return torch.cat(all_outputs), torch.cat(all_targets)


def get_similarity_matrix(val_loader, model, print_freq=100, use_half=False):
    model.eval()
    if use_half:
        model = model.half()
    all_text_embed = []
    all_video_embed = []
    with torch.no_grad():
        print('=> encoding visual and textual')
        for i, inputs in enumerate(val_loader):
            if i % print_freq == 0:
                print('finish batch {}/{}'.format(i, len(val_loader)))
            frames = inputs[0].cuda(non_blocking=True)
            if use_half:
                frames = frames.half()
            texts = inputs[1].cuda(non_blocking=True)
            if len(inputs) == 4:
                masks = inputs[2].cuda(non_blocking=True)
            else:
                masks = None

            # encode images
            image_features = dist_utils.get_model(model).encode_image(frames)
            image_features = image_features / image_features.norm(dim=-1, keepdim=True)
            all_video_embed.append(image_features.cpu().numpy())

            if texts.ndim == 3:
                is_multiple_narrations = True
                texts = texts.view(-1, texts.shape[-1])
            else:
                is_multiple_narrations = False
            if masks is not None:
                text_features = dist_utils.get_model(model).encode_text(texts, attention_mask=masks)
            else:
                text_features = dist_utils.get_model(model).encode_text(texts)
            text_features = text_features / text_features.norm(dim=-1, keepdim=True)
            all_text_embed.append(text_features.cpu().numpy())

        all_text_embed = np.vstack(all_text_embed)
        all_video_embed = np.vstack(all_video_embed)
        similarity_matrix = np.matmul(all_video_embed, all_text_embed.T)
        if is_multiple_narrations:
            similarity_matrix = similarity_matrix.reshape(all_video_embed.shape[0], all_video_embed.shape[0], -1)

    return similarity_matrix


def validate_mcq(val_loader, model, use_half=False):
    model.eval()
    if use_half:
        model.half()
    with torch.no_grad():
        print('=> start forwarding')
        all_preds = []
        all_gts = []
        all_types = []
        end_time = time.time()
        for i, inputs in enumerate(val_loader):
            if i % args.print_freq == 0:
                print('finish batch {}/{} in {} sec'.format(i, len(val_loader), time.time() - end_time))
                end_time = time.time()
            texts_query = inputs[0].cuda(non_blocking=True)
            frames_options = inputs[1].cuda(non_blocking=True)
            if use_half:
                frames_options = frames_options.half()
            answer = inputs[3]
            q_type = inputs[4]
            if len(inputs) == 7:
                masks_query = inputs[5].cuda(non_blocking=True)
            else:
                masks_query = None

            batch_size = frames_options.shape[0]

            frames_options = frames_options.view(-1, *frames_options.shape[2:])
            image_features = dist_utils.get_model(model).encode_image(frames_options)
            image_features = image_features.view(batch_size, -1, *image_features.shape[1:])

            if masks_query is not None:
                query_features = dist_utils.get_model(model).encode_text(texts_query, attention_mask=masks_query)
            else:
                query_features = dist_utils.get_model(model).encode_text(texts_query)

            all_gts.append(answer)
            all_types.append(q_type)
            for j in range(batch_size):
                similarity_matrix = torch.matmul(query_features[j], image_features[j].T)
                similarity_matrix = similarity_matrix.cpu().detach()
                all_preds.append(similarity_matrix)
        all_preds = torch.stack(all_preds)
        all_gts = torch.cat(all_gts)
        all_types = torch.cat(all_types)
        metrics = egomcq_accuracy_metrics(all_preds, all_gts, all_types)
        print(metrics)


if __name__ == '__main__':
    parser = argparse.ArgumentParser('lavila 0-shot evaluations', parents=[get_args_parser()])
    args = parser.parse_args()
    main(args)