fusion.py

import os
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import pandas as pd
import numpy as np

# -----------------------
# IMD Fusion (audio + video) -> 1024d
# -----------------------
class IMDFusion(nn.Module):
    def __init__(self, feature_dim=512, num_heads=4, max_iterations=3):
        super().__init__()
        self.F = feature_dim
        self.H = num_heads
        self.hd = self.F // self.H
        self.iters = max_iterations

        self.Wq = nn.Linear(self.F, self.F)
        self.Wk = nn.Linear(self.F, self.F)
        self.Wv = nn.Linear(self.F, self.F)

        self.stop = nn.Sequential(
            nn.Linear(self.F * 2, 64),
            nn.ReLU(),
            nn.Linear(64, 1),
            nn.Sigmoid()
        )

        self.ln = nn.LayerNorm(self.F)
        self.scale = math.sqrt(self.hd)

    def _cross(self, qx, kx):
        B = qx.size(0)
        Q = self.Wq(qx).view(B, self.H, 1, self.hd)    # [B,H,1,hd]
        K = self.Wk(kx).view(B, self.H, 1, self.hd)    # [B,H,1,hd]
        V = self.Wv(kx).view(B, self.H, self.hd)       # [B,H,hd]

        # compute scores via dot product along hd
        scores = (Q * K).sum(-1, keepdim=True) / self.scale  # [B,H,1,1]
        att = F.softmax(scores, dim=-2)                      # [B,H,1,1]
        att = att.squeeze(-2)                                # [B,H,1]
        out = (att * V).view(B, -1)                          # [B, F]
        return out

    def forward(self, v, a):
        vcur = v
        acur = a
        for i in range(self.iters):
            vnew = self.ln(vcur + self._cross(vcur, acur))
            anew = self.ln(acur + self._cross(acur, vcur))
            vcur, acur = vnew, anew

            s_in = torch.cat([vcur, acur], dim=1)
            p = self.stop(s_in)
            if torch.mean(p) > 0.9:
                break

        return torch.cat([vcur, acur], dim=1)  # [B, 1024]


# -----------------------
# Cross attention + iterative for 3 modalities -> 1536d
# -----------------------
class CrossAttention(nn.Module):
    def __init__(self, dim=512, heads=4):
        super().__init__()
        self.h = heads
        self.dim = dim
        self.hd = dim // heads
        self.q = nn.Linear(dim, dim)
        self.k = nn.Linear(dim, dim)
        self.v = nn.Linear(dim, dim)
        self.out = nn.Linear(dim, dim)

    def forward(self, fa, fb):
        B = fa.size(0)
        q = self.q(fa).view(B, 1, self.h, self.hd).transpose(1,2)  # [B,H,1,hd]
        k = self.k(fb).view(B, 1, self.h, self.hd).transpose(1,2)
        v = self.v(fb).view(B, 1, self.h, self.hd).transpose(1,2)

        scores = torch.matmul(q, k.transpose(-2,-1)) / math.sqrt(self.hd)  # [B,H,1,1]
        att = F.softmax(scores, dim=-1)
        out = torch.matmul(att, v)  # [B,H,1,hd]
        out = out.transpose(1,2).contiguous().view(B, 1, self.dim)
        out = self.out(out).squeeze(1)
        return out

class IterativeModalityDialogue(nn.Module):
    def __init__(self, dim=512, iters=3):
        super().__init__()
        self.iters = iters
        self.v_a = CrossAttention(dim)
        self.a_t = CrossAttention(dim)
        self.t_v = CrossAttention(dim)
        self.stop = nn.Linear(dim * 3, 1)

    def forward(self, fv, fa, ft):
        v, a, t = fv, fa, ft
        for i in range(self.iters):
            v = v + self.v_a(v, a)
            a = a + self.a_t(a, t)
            t = t + self.t_v(t, v)

            cat = torch.cat([v, a, t], dim=1)
            g = torch.sigmoid(self.stop(cat))
            if (g > 0.9).all() and i > 0:
                break
        return torch.cat([v, a, t], dim=1)  # [B, 1536]


# -----------------------
# utils: load embeddings and align by video_name
# -----------------------
def load_emb(path):
    df = pd.read_csv(path)
    names = df.iloc[:, 0].astype(str).tolist()
    feats = df.iloc[:, 1:].astype(np.float32).values
    return names, feats

def align_by_name(vn, an, tn=None):
    # vn, an, tn are lists of names
    if tn is None:
        common = set(vn).intersection(set(an))
    else:
        common = set(vn).intersection(set(an)).intersection(set(tn))
    common = sorted(list(common))
    return common

def pick_feats(names, src_names, src_feats):
    idx = [src_names.index(n) for n in names]
    return src_feats[idx]


# -----------------------
# main fuse runner
# -----------------------
def fuse_main(mode="av", v_path="video_embeddings_vali.csv",
              a_path="audio_embeddings_vali.csv",
              t_path="text_embeddings_vali.csv",
              out_path=None, bs=64):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    v_names, v_feats = load_emb(v_path)
    a_names, a_feats = load_emb(a_path)

    if mode == "av":
        names = align_by_name(v_names, a_names)
        v_sel = pick_feats(names, v_names, v_feats)
        a_sel = pick_feats(names, a_names, a_feats)

        v_t = torch.tensor(v_sel, dtype=torch.float32).to(device)
        a_t = torch.tensor(a_sel, dtype=torch.float32).to(device)

        model = IMDFusion().to(device)
        model.eval()

        out_list = []
        with torch.no_grad():
            for i in range(0, v_t.size(0), bs):
                vb = v_t[i:i+bs]
                ab = a_t[i:i+bs]
                o = model(vb, ab)  # [B,1024]
                out_list.append(o.cpu().numpy())
        out_all = np.vstack(out_list)
        cols = [f"dim_{i}" for i in range(out_all.shape[1])]
        df_out = pd.DataFrame(out_all, columns=cols)
        df_out.insert(0, "video_name", names)
        if out_path is None:
            out_path = "fused_av.csv"
        df_out.to_csv(out_path, index=False)
        print("Saved", out_path)
        return df_out

    elif mode == "avt":
        # need text file too
        t_names, t_feats = load_emb(t_path)
        names = align_by_name(v_names, a_names, t_names)
        v_sel = pick_feats(names, v_names, v_feats)
        a_sel = pick_feats(names, a_names, a_feats)
        t_sel = pick_feats(names, t_names, t_feats)

        v_t = torch.tensor(v_sel, dtype=torch.float32).to(device)
        a_t = torch.tensor(a_sel, dtype=torch.float32).to(device)
        t_t = torch.tensor(t_sel, dtype=torch.float32).to(device)

        model = IterativeModalityDialogue().to(device)
        model.eval()

        out_list = []
        with torch.no_grad():
            for i in range(0, v_t.size(0), bs):
                vb = v_t[i:i+bs]
                ab = a_t[i:i+bs]
                tb = t_t[i:i+bs]
                o = model(vb, ab, tb)  # [B,1536]
                out_list.append(o.cpu().numpy())
        out_all = np.vstack(out_list)
        cols = [f"dim_{i}" for i in range(out_all.shape[1])]
        df_out = pd.DataFrame(out_all, columns=cols)
        df_out.insert(0, "video_name", names)
        if out_path is None:
            out_path = "fused_avt.csv"
        df_out.to_csv(out_path, index=False)
        print("Saved", out_path)
        return df_out

    else:
        raise ValueError("mode must be 'av' or 'avt'.")


# -----------------------
# example usage
# -----------------------
if __name__ == "__main__":
    # mode = "av" for video+audio fusion -> fused_av.csv (1024 dims)
    # mode = "avt" for video+audio+text fusion -> fused_avt.csv (1536 dims)
    df1 = fuse_main(mode="av", v_path="video_embeddings_vali.csv",
                    a_path="audio_embeddings_vali.csv", bs=64)

    # to run 3-modality fusion, uncomment
    # df2 = fuse_main(mode="avt",
    #                 v_path="video_embeddings_vali.csv",
    #                 a_path="audio_embeddings_vali.csv",
    #                 t_path="text_embeddings_vali.csv",
    #                 bs=64)