scripts/feature_extractors.py

"""
Pretrained, frozen encoders for each modality.

Each `extract_*` function returns a fixed-size numpy embedding for one clip.
These models are NOT trained in this project -- only used as feature extractors.
The fusion head (fusion_model.py) is the part you actually train.
"""
import argparse
import os
import numpy as np
import pandas as pd
import torch
from tqdm import tqdm

EMBED_DIM = 256  # common projected size for every modality, set in fusion_model.py too

# ----------------------------------------------------------------------------
# Lazy-loaded global models (loaded once, reused across clips)
# ----------------------------------------------------------------------------
_audio_model = None
_audio_processor = None
_text_tokenizer = None
_text_model = None
_whisper_model = None
_face_extractor = None
_face_model = None


def _device():
    return "cuda" if torch.cuda.is_available() else "cpu"


def _load_audio_model():
    global _audio_model, _audio_processor
    if _audio_model is None:
        from transformers import Wav2Vec2Processor, Wav2Vec2Model
        name = "audeering/wav2vec2-large-robust-12-ft-emotion-msp-dim"
        _audio_processor = Wav2Vec2Processor.from_pretrained(name)
        _audio_model = Wav2Vec2Model.from_pretrained(name).to(_device()).eval()
    return _audio_model, _audio_processor


def _load_text_model():
    global _text_tokenizer, _text_model
    if _text_model is None:
        from transformers import AutoTokenizer, AutoModel
        try:
            name = "mental/mental-bert-base-uncased"
            _text_tokenizer = AutoTokenizer.from_pretrained(name)
            _text_model = AutoModel.from_pretrained(name).to(_device()).eval()
        except Exception:
            # fallback if mental-bert isn't reachable in your environment
            name = "roberta-base"
            _text_tokenizer = AutoTokenizer.from_pretrained(name)
            _text_model = AutoModel.from_pretrained(name).to(_device()).eval()
    return _text_model, _text_tokenizer


def _load_whisper():
    global _whisper_model
    if _whisper_model is None:
        import whisper
        _whisper_model = whisper.load_model("base")
    return _whisper_model


def _load_face_model():
    global _face_extractor, _face_model
    if _face_model is None:
        from transformers import AutoImageProcessor, AutoModel
        name = "trpakov/vit-face-expression"
        _face_extractor = AutoImageProcessor.from_pretrained(name)
        _face_model = AutoModel.from_pretrained(name).to(_device()).eval()
    return _face_model, _face_extractor


# ----------------------------------------------------------------------------
# Projection: raw encoder hidden size -> common EMBED_DIM
# (a single untrained linear layer per modality is fine since the fusion head
#  learns on top of it; alternatively just mean-pool/truncate to EMBED_DIM)
# ----------------------------------------------------------------------------
def _project(vec: np.ndarray, dim: int = EMBED_DIM) -> np.ndarray:
    if vec.shape[-1] == dim:
        return vec
    if vec.shape[-1] > dim:
        # simple average-pool down to target dim
        factor = vec.shape[-1] // dim
        usable = factor * dim
        return vec[:usable].reshape(dim, factor).mean(axis=1)
    # pad if smaller
    pad = np.zeros(dim - vec.shape[-1], dtype=vec.dtype)
    return np.concatenate([vec, pad])


# ----------------------------------------------------------------------------
# Public extraction functions
# ----------------------------------------------------------------------------
def extract_audio_embedding(audio_path: str) -> np.ndarray:
    import librosa
    model, processor = _load_audio_model()
    wav, sr = librosa.load(audio_path, sr=16000)
    inputs = processor(wav, sampling_rate=16000, return_tensors="pt").input_values
    with torch.no_grad():
        out = model(inputs.to(_device())).last_hidden_state  # [1, T, H]
    pooled = out.mean(dim=1).squeeze(0).cpu().numpy()
    return _project(pooled)


def transcribe_audio(audio_path: str) -> str:
    whisper_model = _load_whisper()
    result = whisper_model.transcribe(
        audio_path,
        no_speech_threshold=0.6,
        logprob_threshold=-1.0,
        condition_on_previous_text=True
    )
    return result["text"].strip()


def extract_text_embedding(text: str) -> np.ndarray:
    model, tokenizer = _load_text_model()
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=256).to(_device())
    with torch.no_grad():
        out = model(**inputs).last_hidden_state  # [1, T, H]
    pooled = out.mean(dim=1).squeeze(0).cpu().numpy()
    return _project(pooled)


def extract_face_embedding(video_path: str, max_frames: int = 16) -> np.ndarray:
    import cv2
    model, extractor = _load_face_model()
    cap = cv2.VideoCapture(video_path)
    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) or max_frames
    step = max(total // max_frames, 1)

    frame_embeds = []
    idx = 0
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        if idx % step == 0:
            rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            inputs = extractor(images=rgb, return_tensors="pt").to(_device())
            with torch.no_grad():
                out = model(**inputs).last_hidden_state.mean(dim=1).squeeze(0).cpu().numpy()
            frame_embeds.append(out)
        idx += 1
        if len(frame_embeds) >= max_frames:
            break
    cap.release()

    if not frame_embeds:
        return np.zeros(EMBED_DIM, dtype=np.float32)
    pooled = np.mean(frame_embeds, axis=0)
    return _project(pooled)


def extract_all(audio_path: str, video_path: str):
    """Returns (audio_emb, text_emb, face_emb, transcript) for one clip."""
    audio_emb = extract_audio_embedding(audio_path)
    transcript = transcribe_audio(audio_path)
    text_emb = extract_text_embedding(transcript)
    face_emb = extract_face_embedding(video_path)
    return audio_emb, text_emb, face_emb, transcript


# ----------------------------------------------------------------------------
# Batch extraction over a manifest.csv -> embeddings.npz
# ----------------------------------------------------------------------------
def build_embeddings_file(manifest_path: str, out_path: str):
    df = pd.read_csv(manifest_path)
    audio_embs, text_embs, face_embs, labels, clip_ids = [], [], [], [], []

    for _, row in tqdm(df.iterrows(), total=len(df), desc="Extracting embeddings"):
        a, t, f, _ = extract_all(row["audio_path"], row["video_path"])
        audio_embs.append(a)
        text_embs.append(t)
        face_embs.append(f)
        labels.append(row["label"])
        clip_ids.append(row["clip_id"])

    np.savez(
        out_path,
        audio=np.stack(audio_embs),
        text=np.stack(text_embs),
        face=np.stack(face_embs),
        labels=np.array(labels),
        clip_ids=np.array(clip_ids),
    )
    print(f"Saved {len(df)} clips -> {out_path}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--manifest", required=True)
    parser.add_argument("--out", required=True)
    args = parser.parse_args()
    os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
    build_embeddings_file(args.manifest, args.out)