""" 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)