csf_emb.py

import os
import cv2
import numpy as np
import pandas as pd
import tensorflow as tf
from mtcnn import MTCNN
from tensorflow.keras.applications import EfficientNetB0
from tensorflow.keras.applications.efficientnet import preprocess_input
from tensorflow.keras import layers, Model
from moviepy import VideoFileClip
import librosa
from transformers import Wav2Vec2Processor, Wav2Vec2Model

# -----------------------
# user data loading
# -----------------------
df_test = pd.read_csv("data/annotations/test_gt.csv")
df_train = pd.read_csv("data/annotations/training_gt.csv")
df_vali = pd.read_csv("data/annotations/validation_gt.csv")

df_temp = df_test  # swap to df_train or df_vali when needed
video_filenames = df_temp["VideoName"].tolist()
y_labels = df_temp[["ValueOpenness", "ValueConscientiousness",
                    "ValueExtraversion", "ValueAgreeableness",
                    "ValueNeurotisicm"]].values.astype(np.float32)
train_folder = "data/test/"

print("videos:", len(video_filenames))
print("y shape:", y_labels.shape)



# -----------------------
# Text pipeline (Whisper -> BERT -> CSF)
# -----------------------
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import BertTokenizer, BertModel

# reuse VideoTranscriber class from your code above
# if not in scope, paste the VideoTranscriber class first.

# CSF PyTorch module (from your snippet, wrapped as class)
class CSF(nn.Module):
    def __init__(self, bert_dim=768, hid=256):
        super().__init__()
        self.lstm = nn.LSTM(bert_dim, hid, num_layers=2,
                            batch_first=True, bidirectional=True)
        self.att_fc = nn.Linear(hid * 2, hid)
        self.att_out = nn.Linear(hid, 1)
        self.mlp = nn.Sequential(
            nn.Linear(hid * 2, 512),
            nn.ReLU(),
            nn.Linear(512, 512)
        )

    def forward(self, bert_x):
        # bert_x: (batch, seq_len, 768)
        h, _ = self.lstm(bert_x)                # (batch, seq_len, hid*2)
        a = torch.tanh(self.att_fc(h))          # (batch, seq_len, hid)
        a = self.att_out(a)                     # (batch, seq_len, 1)
        a = F.softmax(a, dim=1)                 # attention weights
        f_sent = (a * h).sum(dim=1)             # (batch, hid*2)
        f_text = self.mlp(f_sent)               # (batch, 512)
        return f_text

# init devices and models
dev = 'cuda' if torch.cuda.is_available() else 'cpu'
vt = VideoTranscriber(model_size='medium', device=dev)  # uses whisper
tk = BertTokenizer.from_pretrained('bert-base-uncased')
bm = BertModel.from_pretrained('bert-base-uncased').to(dev)
csf = CSF(bert_dim=768, hid=256).to(dev)
bm.eval()
csf.eval()

# helper: get bert token embeddings for a piece of text
def bert_embs(text, max_len=128):
    with torch.no_grad():
        t = tk(text, return_tensors='pt', truncation=True,
               padding='max_length', max_length=max_len)
        inp = {k: v.to(dev) for k, v in t.items()}
        out = bm(**inp).last_hidden_state        # (1, seq_len, 768)
    return out.cpu()                            # return on cpu tensor

# helper: run CSF and get 512-d numpy
def text_to_vec(bert_hs):
    # bert_hs: cpu tensor (1, seq_len, 768)
    x = bert_hs.to(dev)
    with torch.no_grad():
        v = csf(x)                             # (1,512)
    return v.squeeze(0).cpu().numpy()         # (512,)

# CSV save loop for text embeddings
txt_csv = "text_embeddings_test.csv"
if os.path.exists(txt_csv):
    d0 = pd.read_csv(txt_csv, usecols=['video_name'])
    done_txt = set(d0['video_name'].tolist())
else:
    done_txt = set()

# use video_filenames and train_folder from main script
for vf in video_filenames:
    if vf in done_txt:
        print("skip text", vf)
        continue
    path = os.path.join(train_folder, vf)
    if not os.path.exists(path):
        print("no file", vf)
        continue

    # transcribe
    try:
        res = vt.transcribe_video(path, language='en', verbose=False)
    except Exception as e:
        print("whisper err", vf, e)
        continue

    if not res or 'text' not in res:
        print("no text", vf)
        continue

    txt = res['text'].strip()
    if len(txt) == 0:
        print("empty text", vf)
        continue

    # convert to bert embeddings and CSF vector
    bh = bert_embs(txt)        # (1, seq, 768) cpu tensor
    vec = text_to_vec(bh)      # (512,)

    # save
    row = pd.DataFrame([[vf] + vec.tolist()],
                       columns=['video_name'] + [f'dim_{i}' for i in range(len(vec))])
    row.to_csv(txt_csv, mode='a', header=not os.path.exists(txt_csv), index=False)
    print("saved text", vf)

print("text done")