demo.py

### demo.py
# Define model classes for inference.
###
import json
import numpy as np
import os
import pandas as pd

import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
import torchvision.transforms as transforms
import torchvision.transforms._transforms_video as transforms_video
from sklearn.metrics import confusion_matrix
from einops import rearrange
from transformers import BertTokenizer

from svitt.model import SViTT
from svitt.datasets import VideoClassyDataset
from svitt.video_transforms import Permute
from svitt.config import load_cfg, setup_config
from svitt.evaluation_charades import charades_map
from svitt.evaluation import get_mean_accuracy


class VideoModel(nn.Module):
    """ Base model for video understanding based on SViTT architecture. """
    def __init__(self, config):
        """ Initializes the model.
        Parameters:
            config: config file
        """
        super(VideoModel, self).__init__()
        self.cfg = load_cfg(config)
        self.model = self.build_model()
        use_gpu = torch.cuda.is_available()
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        if use_gpu:
            self.model = self.model.to(self.device)
        self.templates = ['{}']
        self.dataset = self.cfg['data']['dataset']
        self.eval()

    def build_model(self):
        cfg = self.cfg
        if cfg['model'].get('pretrain', False):
            ckpt_path = cfg['model']['pretrain']
        else:
            raise Exception('no checkpoint found')
        
        if cfg['model'].get('config', False):
            config_path = cfg['model']['config']
        else:
            raise Exception('no model config found')
        
        self.model_cfg = setup_config(config_path)
        self.tokenizer = BertTokenizer.from_pretrained(self.model_cfg.text_encoder)
        model = SViTT(config=self.model_cfg, tokenizer=self.tokenizer)

        print(f"Loading checkpoint from {ckpt_path}")
        checkpoint = torch.load(ckpt_path, map_location="cpu")
        state_dict = checkpoint["model"]

        # fix for zero-shot evaluation
        for key in list(state_dict.keys()):
            if "bert" in key:
                encoder_key = key.replace("bert.", "")
                state_dict[encoder_key] = state_dict[key]
                    
        if torch.cuda.is_available():
            model.cuda()
                        
        model.load_state_dict(state_dict, strict=False)

        return model
    
    

    def eval(self):
        cudnn.benchmark = True
        for p in self.model.parameters():
            p.requires_grad = False
        self.model.eval()


class VideoCLSModel(VideoModel):
    """ Video model for video classification tasks (Charades-Ego, EGTEA). """
    def __init__(self, config):
        super(VideoCLSModel, self).__init__(config)
        self.labels, self.mapping_vn2act = self.gen_label_map()
        self.text_features = self.get_text_features()

    def gen_label_map(self):
        labelmap = self.cfg.get('label_map', 'meta/charades_ego/label_map.json')
        if os.path.isfile(labelmap):
            print(f"=> Loading label maps from {labelmap}")
            meta = json.load(open(labelmap, 'r'))
            labels, mapping_vn2act = meta['labels'], meta['mapping_vn2act']
        else:
            from svitt.preprocess import generate_label_map
            labels, mapping_vn2act = generate_label_map(self.dataset)
            meta = {'labels': labels, 'mapping_vn2act': mapping_vn2act}
            meta_dir = f'meta/{self.dataset}'
            if not os.path.exists(meta_dir):
                os.makedirs(meta_dir)
            json.dump(meta, open(f'{meta_dir}/label_map.json', 'w'))
            print(f"=> Label map is generated and saved to {meta_dir}/label_map.json")

        return labels, mapping_vn2act
        
    def load_data(self, idx=None): 
        print(f"=> Creating dataset")
        cfg, dataset = self.cfg, self.dataset
        data_cfg = cfg['data']
        crop_size = 224 
        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=[108.3272985, 116.7460125, 104.09373615000001], 
                std=[68.5005327, 66.6321579, 70.32316305],
            ),
        ])

        if idx is None:
            metadata_val = data_cfg['metadata_val']
        else:
            metadata_val = data_cfg['metadata_val'].format(idx)
        if dataset in ['charades_ego', 'egtea']:
            val_dataset = VideoClassyDataset(
                dataset, 
                data_cfg['root'], 
                metadata_val,
                transform=val_transform, 
                is_training=False,
                label_mapping=self.mapping_vn2act, 
                is_trimmed=False,
                num_clips=1, 
                clip_length=data_cfg['clip_length'], 
                clip_stride=data_cfg['clip_stride'],
                sparse_sample=data_cfg['sparse_sample'],
            )
        else:
            raise NotImplementedError

        val_loader = torch.utils.data.DataLoader(
            val_dataset, batch_size=8, shuffle=False,
            num_workers=4, pin_memory=True, sampler=None, drop_last=False
        )

        return val_loader
        
    @torch.no_grad()
    def get_text_features(self):
        print('=> Extracting text features')
        embeddings = self.tokenizer(
            self.labels, 
            padding="max_length", 
            truncation=True,
            max_length=self.model_cfg.max_txt_l.video, 
            return_tensors="pt",
        ).to(self.device)
        _, class_embeddings = self.model.encode_text(embeddings)
        return class_embeddings

    @torch.no_grad()
    def forward(self, idx=None):
        print('=> Start forwarding')
        val_loader = self.load_data(idx)
        all_outputs = []
        all_targets = []
        for i, values in enumerate(val_loader):
            images = values[0]
            target = values[1]

            images = images.to(self.device)

            # encode images
            images = rearrange(images, 'b c k h w ->  b k c h w')
            dims = images.shape
            images = images.reshape(-1, 4, dims[-3], dims[-2], dims[-1])
            
            image_features, _ = self.model.encode_image(images)
            
            if image_features.ndim == 3:
                image_features = rearrange(image_features, '(b k) n d -> b (k n) d', b=1)
            else:
                image_features = rearrange(image_features, '(b k) d -> b k d', b=1)
            
            # cosine similarity as logits
            similarity = self.model.get_sim(image_features, self.text_features)[0]

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

        all_outputs = torch.cat(all_outputs)
        all_targets = torch.cat(all_targets)

        return all_outputs, all_targets

    @torch.no_grad()
    def predict(self, idx=0):
        all_outputs, all_targets = self.forward(idx)
        preds, targets = all_outputs.numpy(), all_targets.numpy()
        #sel = np.where(np.cumsum(sorted(preds[0].tolist(), reverse=True)) > 0.06)[0][0]
        sel = 5
        df = pd.DataFrame(self.labels)
        pred_action = df.iloc[preds[0].argsort()[-sel:]].values.tolist()
        gt_action = df.iloc[np.where(targets[0])[0]].values.tolist()
        pred_action = sorted([x[0] for x in pred_action])
        gt_action = sorted([x[0] for x in gt_action])
        return pred_action, gt_action      

    @torch.no_grad()
    def evaluate(self):
        all_outputs, all_targets = self.forward()
        preds, targets = all_outputs.numpy(), all_targets.numpy()
        if self.dataset == 'charades_ego':
            m_ap, _, m_aps = charades_map(preds, targets)
            print('mAP = {:.3f}'.format(m_ap))
        elif self.dataset == 'egtea':
            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))
        else:
            raise NotImplementedError