DenseAV/denseav/aggregators.py

from abc import abstractmethod

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from tqdm import tqdm

from DenseAV.denseav.constants import *


@torch.jit.script
def masked_mean(x: torch.Tensor, mask: torch.Tensor, dim: int):
    mask = mask.to(x)
    return (x * mask).sum(dim, keepdim=True) / mask.sum(dim, keepdim=True).clamp_min(.001)


@torch.jit.script
def masked_max(x: torch.Tensor, mask: torch.Tensor, dim: int):
    mask = mask.to(torch.bool)
    eps = 1e7
    # eps = torch.finfo(x.dtype).max
    return (x - (~mask) * eps).max(dim, keepdim=True).values


def masked_lse(x: torch.Tensor, mask: torch.Tensor, dim: int, temp):
    x = x.to(torch.float32)
    mask = mask.to(torch.float32)
    x_masked = (x - (1 - mask) * torch.finfo(x.dtype).max)
    return (torch.logsumexp(x_masked * temp, dim, keepdim=True) - torch.log(mask.sum(dim, keepdim=True))) / temp


class BaseAggregator(torch.nn.Module):

    def __init__(self, nonneg_sim, mask_silence, num_heads, head_agg, use_cls):
        super().__init__()

        self.nonneg_sim = nonneg_sim
        self.mask_silence = mask_silence
        self.num_heads = num_heads
        self.head_agg = head_agg
        self.use_cls = use_cls

    @abstractmethod
    def _agg_sim(self, sim, mask):
        pass

    def prepare_sims(self, sim, mask, agg_sim, agg_heads):
        sim_size = sim.shape
        assert len(mask.shape) == 2
        assert len(sim_size) in {6, 7}, f"sim has wrong number of dimensions: {sim.shape}"
        pairwise = len(sim_size) == 6

        if self.mask_silence:
            mask = mask
        else:
            mask = torch.ones_like(mask)

        if self.nonneg_sim:
            sim = sim.clamp_min(0)

        if pairwise:
            head_dim = 1
        else:
            head_dim = 2

        if self.head_agg == "max_elementwise" and agg_heads:
            sim = sim.max(head_dim, keepdim=True).values

        if agg_sim:
            sim = self._agg_sim(sim, mask)

        if agg_heads:
            if self.head_agg == "sum" or self.head_agg == "max_elementwise":
                sim = sim.sum(head_dim)
            elif self.head_agg == "max":
                sim = sim.max(head_dim).values
            else:
                raise ValueError(f"Unknown head_agg: {self.head_agg}")

        return sim

    def _get_full_sims(self, preds, raw, agg_sim, agg_heads):
        if agg_sim or agg_heads or raw:
            assert (agg_sim or agg_heads) != raw, "Cannot have raw on at the same time as agg_sim or agg_heads"

        audio_feats = preds[AUDIO_FEATS]
        audio_mask = preds[AUDIO_MASK]
        image_feats = preds[IMAGE_FEATS]

        b1, c2, f, t1 = audio_feats.shape
        b2, t2 = audio_mask.shape
        d, c1, h, w = image_feats.shape
        assert b1 == b2 and c1 == c2 and t1 == t2
        assert c1 % self.num_heads == 0
        new_c = c1 // self.num_heads
        audio_feats = audio_feats.reshape(b1, self.num_heads, new_c, f, t1)
        image_feats = image_feats.reshape(d, self.num_heads, new_c, h, w)
        raw_sims = torch.einsum(
            "akcft,vkchw->avkhwft",
            audio_feats.to(torch.float32),
            image_feats.to(torch.float32))

        if self.use_cls:
            audio_cls = preds[AUDIO_CLS].reshape(b1, self.num_heads, new_c)
            image_cls = preds[IMAGE_CLS].reshape(d, self.num_heads, new_c)
            cls_sims = torch.einsum(
                "akc,vkc->avk",
                audio_cls.to(torch.float32),
                image_cls.to(torch.float32))
            raw_sims += cls_sims.reshape(b1, d, self.num_heads, 1, 1, 1, 1)

        if raw:
            return raw_sims
        else:
            return self.prepare_sims(raw_sims, audio_mask, agg_sim, agg_heads)

    def get_pairwise_sims(self, preds, raw, agg_sim, agg_heads):
        if agg_sim or agg_heads or raw:
            assert (agg_sim or agg_heads) != raw, "Cannot have raw on at the same time as agg_sim or agg_heads"

        audio_feats = preds[AUDIO_FEATS]
        audio_mask = preds[AUDIO_MASK]
        image_feats = preds[IMAGE_FEATS]

        a1, c1, f, t1 = audio_feats.shape
        a2, t2 = audio_mask.shape

        assert c1 % self.num_heads == 0
        new_c = c1 // self.num_heads
        audio_feats = audio_feats.reshape(a1, self.num_heads, new_c, f, t1)

        if len(image_feats.shape) == 5:
            print("Using similarity for video, should only be called during plotting")
            v, vt, c2, h, w = image_feats.shape
            image_feats = image_feats.reshape(v, vt, self.num_heads, new_c, h, w)
            raw_sims = torch.einsum(
                "bkcft,bskchw,bt->bskhwft",
                audio_feats.to(torch.float32),
                image_feats.to(torch.float32),
                audio_mask.to(torch.float32))

            if self.use_cls:
                print(preds[AUDIO_CLS].shape)
                audio_cls = preds[AUDIO_CLS].reshape(v, self.num_heads, new_c)
                image_cls = preds[IMAGE_CLS].reshape(v, vt, self.num_heads, new_c)
                cls_sims = torch.einsum(
                    "bkc,bskc->bsk",
                    audio_cls.to(torch.float32),
                    image_cls.to(torch.float32))
                raw_sims += cls_sims.reshape(v, vt, self.num_heads, 1, 1, 1, 1)


        elif len(image_feats.shape) == 4:
            v, c2, h, w = image_feats.shape
            image_feats = image_feats.reshape(v, self.num_heads, new_c, h, w)
            raw_sims = torch.einsum(
                "bkcft,bkchw,bt->bkhwft",
                audio_feats.to(torch.float32),
                image_feats.to(torch.float32),
                audio_mask.to(torch.float32))

            if self.use_cls:
                print(preds[AUDIO_CLS].shape)
                audio_cls = preds[AUDIO_CLS].reshape(v, self.num_heads, new_c)
                image_cls = preds[IMAGE_CLS].reshape(v, self.num_heads, new_c)
                cls_sims = torch.einsum(
                    "bkc,bkc->bk",
                    audio_cls.to(torch.float32),
                    image_cls.to(torch.float32))
                raw_sims += cls_sims.reshape(v, self.num_heads, 1, 1, 1, 1)
        else:
            raise ValueError(f"Improper image shape: {image_feats.shape}")

        assert a1 == a2 and c2 == c2 and t1 == t2

        if raw:
            return raw_sims
        else:
            return self.prepare_sims(raw_sims, audio_mask, agg_sim, agg_heads)

    def forward(self, preds, agg_heads):
        return self._get_full_sims(
            preds, raw=False, agg_sim=True, agg_heads=agg_heads)

    def forward_batched(self, preds, agg_heads, batch_size):
        new_preds = {k: v for k, v in preds.items()}
        big_image_feats = new_preds.pop(IMAGE_FEATS)
        if self.use_cls:
            big_image_cls = new_preds.pop(IMAGE_CLS)

        n = big_image_feats.shape[0]
        n_steps = math.ceil(n / batch_size)
        outputs = []
        for step in tqdm(range(n_steps), "Calculating Sim", leave=False):
            new_preds[IMAGE_FEATS] = big_image_feats[step * batch_size:(step + 1) * batch_size].cuda()
            if self.use_cls:
                new_preds[IMAGE_CLS] = big_image_cls[step * batch_size:(step + 1) * batch_size].cuda()

            sim = self.forward(new_preds, agg_heads=agg_heads)
            outputs.append(sim.cpu())
        return torch.cat(outputs, dim=1)


class ImageThenAudioAggregator(BaseAggregator):

    def __init__(self, image_agg_type, audio_agg_type, nonneg_sim, mask_silence, num_heads, head_agg, use_cls):
        super().__init__(nonneg_sim, mask_silence, num_heads, head_agg, use_cls)
        if image_agg_type == "max":
            self.image_agg = lambda x, dim: x.max(dim=dim, keepdim=True).values
        elif image_agg_type == "avg":
            self.image_agg = lambda x, dim: x.mean(dim=dim, keepdim=True)
        else:
            raise ValueError(f"Unknown image_agg_type {image_agg_type}")

        if audio_agg_type == "max":
            self.time_agg = masked_max
        elif audio_agg_type == "avg":
            self.time_agg = masked_mean
        else:
            raise ValueError(f"Unknown audio_agg_type {audio_agg_type}")

        self.freq_agg = lambda x, dim: x.mean(dim=dim, keepdim=True)

    def _agg_sim(self, sim, mask):
        sim_shape = sim.shape
        new_mask_shape = [1] * len(sim_shape)
        new_mask_shape[0] = sim_shape[0]
        new_mask_shape[-1] = sim_shape[-1]
        mask = mask.reshape(new_mask_shape)
        sim = self.image_agg(sim, -3)
        sim = self.image_agg(sim, -4)
        sim = self.freq_agg(sim, -2)
        sim = self.time_agg(sim, mask, -1)
        return sim.squeeze(-1).squeeze(-1).squeeze(-1).squeeze(-1)


class PairedAggregator(BaseAggregator):

    def __init__(self, nonneg_sim, mask_silence, num_heads, head_agg, use_cls):
        super().__init__(nonneg_sim, mask_silence, num_heads, head_agg, use_cls)
        self.image_agg_max = lambda x, dim: x.max(dim=dim, keepdim=True).values
        self.image_agg_mean = lambda x, dim: x.mean(dim=dim, keepdim=True)

        self.time_agg_max = masked_max
        self.time_agg_mean = masked_mean

        self.freq_agg = lambda x, dim: x.mean(dim=dim, keepdim=True)

    def _agg_sim(self, sim, mask):
        sim_shape = sim.shape
        new_mask_shape = [1] * len(sim_shape)
        new_mask_shape[0] = sim_shape[0]
        new_mask_shape[-1] = sim_shape[-1]
        mask = mask.reshape(new_mask_shape)

        sim_1 = self.image_agg_max(sim, -3)
        sim_1 = self.image_agg_max(sim_1, -4)
        sim_1 = self.freq_agg(sim_1, -2)
        sim_1 = self.time_agg_mean(sim_1, mask, -1)

        sim_2 = self.freq_agg(sim, -2)
        sim_2 = self.time_agg_max(sim_2, mask, -1)
        sim_2 = self.image_agg_mean(sim_2, -3)
        sim_2 = self.image_agg_mean(sim_2, -4)

        sim = 1 / 2 * (sim_1 + sim_2)

        return sim.squeeze(-1).squeeze(-1).squeeze(-1).squeeze(-1)



class CAVMAEAggregator(BaseAggregator):

    def __init__(self, *args, **kwargs):
        super().__init__(False, False, 1, "sum", False)

    def _get_full_sims(self, preds, raw, agg_sim, agg_heads):
        if agg_sim:
            audio_feats = preds[AUDIO_FEATS]
            image_feats = preds[IMAGE_FEATS]
            pool_audio_feats = F.normalize(audio_feats.mean(dim=[-1, -2]), dim=1)
            pool_image_feats = F.normalize(image_feats.mean(dim=[-1, -2]), dim=1)
            sims = torch.einsum(
                "bc,dc->bd",
                pool_audio_feats.to(torch.float32),
                pool_image_feats.to(torch.float32))
            if agg_heads:
                return sims
            else:
                return sims.unsqueeze(-1)

        else:
            return BaseAggregator._get_full_sims(self, preds, raw, agg_sim, agg_heads)

    def get_pairwise_sims(self, preds, raw, agg_sim, agg_heads):
        if agg_sim:
            audio_feats = preds[AUDIO_FEATS]
            image_feats = preds[IMAGE_FEATS]
            pool_audio_feats = F.normalize(audio_feats.mean(dim=[-1, -2]), dim=1)
            pool_image_feats = F.normalize(image_feats.mean(dim=[-1, -2]), dim=1)
            sims = torch.einsum(
                "bc,bc->b",
                pool_audio_feats.to(torch.float32),
                pool_image_feats.to(torch.float32))
            if agg_heads:
                return sims
            else:
                return sims.unsqueeze(-1)

        else:
            return BaseAggregator.get_pairwise_sims(self, preds, raw, agg_sim, agg_heads)


class ImageBindAggregator(BaseAggregator):

    def __init__(self, num_heads, *args, **kwargs):
        super().__init__(False, False, num_heads, "sum", False)

    def _get_full_sims(self, preds, raw, agg_sim, agg_heads):
        if agg_sim:
            sims = torch.einsum(
                "bc,dc->bd",
                preds[AUDIO_CLS].to(torch.float32),
                preds[IMAGE_CLS].to(torch.float32))
            if agg_heads:
                return sims
            else:
                sims = sims.unsqueeze(-1)
                return sims.repeat(*([1] * (sims.dim() - 1)), self.num_heads)


        else:
            return BaseAggregator._get_full_sims(self, preds, raw, agg_sim, agg_heads)

    def get_pairwise_sims(self, preds, raw, agg_sim, agg_heads):
        if agg_sim:
            sims = torch.einsum(
                "bc,dc->b",
                preds[AUDIO_CLS].to(torch.float32),
                preds[IMAGE_CLS].to(torch.float32))
            if agg_heads:
                return sims
            else:
                sims = sims.unsqueeze(-1)
                return sims.repeat(*([1] * (sims.dim() - 1)), self.num_heads)

        else:
            return BaseAggregator.get_pairwise_sims(self, preds, raw, agg_sim, agg_heads)

    def forward_batched(self, preds, agg_heads, batch_size):
        return self.forward(preds, agg_heads)


class SimPool(nn.Module):
    def __init__(self, dim, num_heads=1, qkv_bias=False, qk_scale=None, gamma=None, use_beta=False):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = qk_scale or head_dim ** -0.5

        self.norm_patches = nn.LayerNorm(dim, eps=1e-6)

        self.wq = nn.Linear(dim, dim, bias=qkv_bias)
        self.wk = nn.Linear(dim, dim, bias=qkv_bias)

        if gamma is not None:
            self.gamma = torch.tensor([gamma])
            if use_beta:
                self.beta = nn.Parameter(torch.tensor([0.0]))
        self.eps = torch.tensor([1e-6])

        self.gamma = gamma
        self.use_beta = use_beta

    def prepare_input(self, x):
        if len(x.shape) == 3:  # Transformer
            # Input tensor dimensions:
            # x: (B, N, d), where B is batch size, N are patch tokens, d is depth (channels)
            B, N, d = x.shape
            gap_cls = x.mean(-2)  # (B, N, d) -> (B, d)
            gap_cls = gap_cls.unsqueeze(1)  # (B, d) -> (B, 1, d)
            return gap_cls, x
        if len(x.shape) == 4:  # CNN
            # Input tensor dimensions:
            # x: (B, d, H, W), where B is batch size, d is depth (channels), H is height, and W is width
            B, d, H, W = x.shape
            gap_cls = x.mean([-2, -1])  # (B, d, H, W) -> (B, d)
            x = x.reshape(B, d, H * W).permute(0, 2, 1)  # (B, d, H, W) -> (B, d, H*W) -> (B, H*W, d)
            gap_cls = gap_cls.unsqueeze(1)  # (B, d) -> (B, 1, d)
            return gap_cls, x
        else:
            raise ValueError(f"Unsupported number of dimensions in input tensor: {len(x.shape)}")

    def forward(self, x):
        self.eps = self.eps.to(x.device)
        # Prepare input tensor and perform GAP as initialization
        gap_cls, x = self.prepare_input(x)

        # Prepare queries (q), keys (k), and values (v)
        q, k, v = gap_cls, self.norm_patches(x), self.norm_patches(x)

        # Extract dimensions after normalization
        Bq, Nq, dq = q.shape
        Bk, Nk, dk = k.shape
        Bv, Nv, dv = v.shape

        # Check dimension consistency across batches and channels
        assert Bq == Bk == Bv
        assert dq == dk == dv

        # Apply linear transformation for queries and keys then reshape
        qq = self.wq(q).reshape(Bq, Nq, self.num_heads, dq // self.num_heads).permute(0, 2, 1,
                                                                                      3)  # (Bq, Nq, dq) -> (B, num_heads, Nq, dq/num_heads)
        kk = self.wk(k).reshape(Bk, Nk, self.num_heads, dk // self.num_heads).permute(0, 2, 1,
                                                                                      3)  # (Bk, Nk, dk) -> (B, num_heads, Nk, dk/num_heads)

        vv = v.reshape(Bv, Nv, self.num_heads, dv // self.num_heads).permute(0, 2, 1,
                                                                             3)  # (Bv, Nv, dv) -> (B, num_heads, Nv, dv/num_heads)

        # Compute attention scores
        attn = (qq @ kk.transpose(-2, -1)) * self.scale
        # Apply softmax for normalization
        attn = attn.softmax(dim=-1)

        # If gamma scaling is used
        if self.gamma is not None:
            # Apply gamma scaling on values and compute the weighted sum using attention scores
            x = torch.pow(attn @ torch.pow((vv - vv.min() + self.eps), self.gamma),
                          1 / self.gamma)  # (B, num_heads, Nv, dv/num_heads) -> (B, 1, 1, d)
            # If use_beta, add a learnable translation
            if self.use_beta:
                x = x + self.beta
        else:
            # Compute the weighted sum using attention scores
            x = (attn @ vv).transpose(1, 2).reshape(Bq, Nq, dq)

        return x.squeeze()



class SimPoolAggregator(BaseAggregator):

    def __init__(self, num_heads, dim, *args, **kwargs):
        super().__init__(False, False, num_heads, "sum", False)
        self.pool = SimPool(dim, gamma=1.25)

    def _get_full_sims(self, preds, raw, agg_sim, agg_heads):
        if agg_sim:
            device = self.pool.wq.weight.data.device
            pooled_audio = self.pool(preds[AUDIO_FEATS].to(torch.float32).to(device))
            pooled_image = self.pool(preds[IMAGE_FEATS].to(torch.float32).to(device))

            sims = torch.einsum(
                "bc,dc->bd",
                pooled_audio,
                pooled_image)
            if agg_heads:
                return sims
            else:
                sims = sims.unsqueeze(-1)
                return sims.repeat(*([1] * (sims.dim() - 1)), self.num_heads)


        else:
            return BaseAggregator._get_full_sims(self, preds, raw, agg_sim, agg_heads)

    def get_pairwise_sims(self, preds, raw, agg_sim, agg_heads):
        if agg_sim:
            device = self.pool.wq.weight.data.device
            pooled_audio = self.pool(preds[AUDIO_FEATS].to(torch.float32).to(device))
            pooled_image = self.pool(preds[IMAGE_FEATS].to(torch.float32).to(device))

            sims = torch.einsum(
                "bc,dc->b",
                pooled_audio,
                pooled_image)
            if agg_heads:
                return sims
            else:
                sims = sims.unsqueeze(-1)
                return sims.repeat(*([1] * (sims.dim() - 1)), self.num_heads)

        else:
            return BaseAggregator.get_pairwise_sims(self, preds, raw, agg_sim, agg_heads)

    def forward_batched(self, preds, agg_heads, batch_size):
        return self.forward(preds, agg_heads)



def get_aggregator(sim_agg_type, nonneg_sim, mask_silence, num_heads, head_agg, use_cls, dim):
    shared_args = dict(
        nonneg_sim=nonneg_sim,
        mask_silence=mask_silence,
        num_heads=num_heads,
        head_agg=head_agg,
        use_cls=use_cls,
    )

    if sim_agg_type == "paired":
        agg1 = PairedAggregator(**shared_args)
    elif sim_agg_type == "misa":
        agg1 = ImageThenAudioAggregator("max", "avg", **shared_args)
    elif sim_agg_type == "mima":
        agg1 = ImageThenAudioAggregator("max", "max", **shared_args)
    elif sim_agg_type == "sisa":
        agg1 = ImageThenAudioAggregator("avg", "avg", **shared_args)
    elif sim_agg_type == "cavmae":
        agg1 = CAVMAEAggregator()
    elif sim_agg_type == "imagebind":
        agg1 = ImageBindAggregator(num_heads=shared_args["num_heads"])
    elif sim_agg_type == "simpool":
        agg1 = SimPoolAggregator(num_heads=shared_args["num_heads"], dim=dim)
    else:
        raise ValueError(f"Unknown loss_type {sim_agg_type}")

    return agg1