from __future__ import annotations import argparse import csv import hashlib import importlib.metadata import json import logging import math import platform import re import shutil import subprocess import sys import time import urllib.request from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path from typing import Any import cv2 import mediapipe as mp import numpy as np import opensmile import torch from PIL import Image from transformers import AutoImageProcessor, AutoModel, AutoModelForCTC, AutoTokenizer from .audit import _read_labels TEXT_MODEL_ID = "google-bert/bert-base-uncased" CTC_MODEL_ID = "facebook/wav2vec2-base-960h" VISION_MODEL_ID = "facebook/deit-tiny-patch16-224" FACE_MODEL_URL = ( "https://storage.googleapis.com/mediapipe-models/face_landmarker/" "face_landmarker/float16/latest/face_landmarker.task" ) AUDIO_SAMPLE_RATE = 16_000 VISION_FPS = 5.0 VISION_DIM = 52 VISION_GLOBAL_DIM = 192 LOGGER = logging.getLogger("q1.extract") @dataclass class ExtractedSample: arrays: dict[str, np.ndarray] info: dict[str, Any] def _package_version(name: str) -> str: try: return importlib.metadata.version(name) except importlib.metadata.PackageNotFoundError: return "not-installed" def _tool_version(executable: str) -> str: candidate = shutil.which(executable) if candidate is None and executable == "uv": local_uv = Path.home() / ".local" / "bin" / "uv" candidate = str(local_uv) if local_uv.is_file() else None if candidate is None: return "not-found" try: flag = "--version" if executable == "uv" else "-version" result = subprocess.run([candidate, flag], capture_output=True, text=True, check=True) output = result.stdout or result.stderr return output.splitlines()[0].strip() if output else "unknown" except (OSError, subprocess.CalledProcessError) as error: return f"unavailable: {type(error).__name__}" def _safe_id(value: str) -> str: return re.sub(r"[^A-Za-z0-9_.-]+", "_", value) def _ffprobe_duration(video_path: Path) -> float: result = subprocess.run( [ "ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "default=noprint_wrappers=1:nokey=1", str(video_path), ], check=True, capture_output=True, text=True, ) try: return max(float(result.stdout.strip()), 0.0) except ValueError: return 0.0 def _read_audio(video_path: Path) -> np.ndarray: command = [ "ffmpeg", "-nostdin", "-v", "error", "-i", str(video_path), "-vn", "-ac", "1", "-ar", str(AUDIO_SAMPLE_RATE), "-f", "f32le", "pipe:1", ] result = subprocess.run(command, check=True, capture_output=True) waveform = np.frombuffer(result.stdout, dtype=" dict[str, np.ndarray]: return { "text_words": np.empty((0,), dtype="U1"), "text_features": np.empty((0, hidden_size), dtype=np.float16), } def _ctc_viterbi( log_probs: np.ndarray, targets: list[int], blank_id: int, ) -> tuple[np.ndarray, float] | None: """Viterbi-decode a CTC target into a frame-to-extended-target path.""" if not targets or log_probs.ndim != 2: return None if log_probs.shape[0] < len(targets): return None extended = np.full(2 * len(targets) + 1, blank_id, dtype=np.int64) extended[1::2] = np.asarray(targets, dtype=np.int64) states = extended.size frames = log_probs.shape[0] previous = np.full(states, -np.inf, dtype=np.float32) previous[0] = log_probs[0, blank_id] previous[1] = log_probs[0, extended[1]] back = np.zeros((frames, states), dtype=np.uint8) back[0, 1] = 1 state_indices = np.arange(states) for frame in range(1, frames): candidates = np.full((3, states), -np.inf, dtype=np.float32) candidates[0] = previous candidates[1, 1:] = previous[:-1] skip_allowed = (state_indices >= 2) & (extended != blank_id) skip_allowed[2:] &= extended[2:] != extended[:-2] candidates[2, skip_allowed] = previous[state_indices[skip_allowed] - 2] choice = candidates.argmax(axis=0).astype(np.uint8) best = candidates[choice, state_indices] previous = best + log_probs[frame, extended] back[frame] = choice final_state = states - 1 if states > 1 and previous[-2] > previous[-1]: final_state -= 1 final_score = float(previous[final_state]) if not math.isfinite(final_score): return None path = np.empty(frames, dtype=np.int32) state = final_state path[-1] = state for frame in range(frames - 1, 0, -1): state -= int(back[frame, state]) path[frame - 1] = state return path, final_score / frames class FeatureModels: def __init__(self, device: torch.device, vision_model_path: Path, cache_dir: Path | None = None): self.device = device LOGGER.info("Loading text model %s", TEXT_MODEL_ID) self.text_tokenizer = AutoTokenizer.from_pretrained( TEXT_MODEL_ID, use_fast=True, cache_dir=str(cache_dir) if cache_dir else None ) self.text_model = AutoModel.from_pretrained( TEXT_MODEL_ID, cache_dir=str(cache_dir) if cache_dir else None ).to(device).eval() self.text_dim = int(self.text_model.config.hidden_size) LOGGER.info("Loading CTC model %s", CTC_MODEL_ID) self.ctc_tokenizer = AutoTokenizer.from_pretrained( CTC_MODEL_ID, cache_dir=str(cache_dir) if cache_dir else None ) self.ctc_model = AutoModelForCTC.from_pretrained( CTC_MODEL_ID, cache_dir=str(cache_dir) if cache_dir else None ).to(device).eval() self.blank_id = int(self.ctc_tokenizer.pad_token_id) self.delimiter_id = int( self.ctc_tokenizer.convert_tokens_to_ids( self.ctc_tokenizer.word_delimiter_token or "|" ) ) self.ctc_vocab = self.ctc_tokenizer.get_vocab() self.stride_samples = int(np.prod(self.ctc_model.config.conv_stride)) self.receptive_samples = self._receptive_field( self.ctc_model.config.conv_kernel, self.ctc_model.config.conv_stride, ) LOGGER.info("Loading general vision model %s", VISION_MODEL_ID) self.vision_processor = AutoImageProcessor.from_pretrained( VISION_MODEL_ID, cache_dir=str(cache_dir) if cache_dir else None ) self.vision_model = AutoModel.from_pretrained( VISION_MODEL_ID, cache_dir=str(cache_dir) if cache_dir else None ).to(device).eval() self.vision_dim = int(self.vision_model.config.hidden_size) if self.vision_dim != VISION_GLOBAL_DIM: raise ValueError(f"expected {VISION_GLOBAL_DIM}-dimensional DeiT features, received {self.vision_dim}") self.vision_model_path = vision_model_path self._ensure_vision_model() LOGGER.info("Feature models are ready on %s", device) @staticmethod def _receptive_field(kernels: list[int], strides: list[int]) -> int: receptive, jump = 1, 1 for kernel, stride in zip(kernels, strides): receptive += (kernel - 1) * jump jump *= stride return receptive def _ensure_vision_model(self) -> None: self.vision_model_path.parent.mkdir(parents=True, exist_ok=True) if self.vision_model_path.exists() and self.vision_model_path.stat().st_size > 1_000_000: return LOGGER.info("Downloading MediaPipe Face Landmarker model") temporary = self.vision_model_path.with_suffix(self.vision_model_path.suffix + ".tmp") urllib.request.urlretrieve(FACE_MODEL_URL, temporary) if temporary.stat().st_size < 1_000_000: temporary.unlink(missing_ok=True) raise RuntimeError("downloaded Face Landmarker model is unexpectedly small") temporary.replace(self.vision_model_path) def text_features(self, transcript: str) -> tuple[list[str], np.ndarray]: words = transcript.split() if not words: return [], np.empty((0, self.text_dim), dtype=np.float32) embeddings = np.zeros((len(words), self.text_dim), dtype=np.float32) # Chunk by original words so every supplied transcript token is retained. chunk_size = 180 with torch.inference_mode(): for offset in range(0, len(words), chunk_size): chunk = words[offset:offset + chunk_size] encoded = self.text_tokenizer( chunk, is_split_into_words=True, add_special_tokens=True, truncation=False, return_tensors="pt", ) word_ids = encoded.word_ids(batch_index=0) model_inputs = {key: value.to(self.device) for key, value in encoded.items()} hidden = self.text_model(**model_inputs).last_hidden_state[0].float().cpu().numpy() bins: list[list[int]] = [[] for _ in chunk] for token_index, word_index in enumerate(word_ids): if word_index is not None and 0 <= word_index < len(chunk): bins[word_index].append(token_index) for local_word_index, token_indices in enumerate(bins): if token_indices: embeddings[offset + local_word_index] = hidden[token_indices].mean(axis=0) return words, embeddings def force_align(self, words: list[str], waveform: np.ndarray, duration_s: float) -> tuple[np.ndarray, np.ndarray, np.ndarray, bool]: count = len(words) if count == 0: return ( np.empty((0, 2), dtype=np.float32), np.empty((0,), dtype=np.float32), np.empty((0,), dtype=np.bool_), False, ) normalized = [re.sub(r"[^a-z']", "", word.lower().replace("’", "'")) for word in words] targets: list[int] = [] target_word: list[int] = [] word_target_indices: list[list[int]] = [[] for _ in words] for word_index, word in enumerate(normalized): if word_index: targets.append(self.delimiter_id) target_word.append(-1) for character in word: token_id = self.ctc_vocab.get(character) if token_id is None: token_id = int(self.ctc_tokenizer.unk_token_id) word_target_indices[word_index].append(len(targets)) targets.append(int(token_id)) target_word.append(word_index) valid = np.zeros(count, dtype=np.bool_) confidence = np.zeros(count, dtype=np.float32) intervals = np.zeros((count, 2), dtype=np.float32) ctc_success = False if targets and waveform.size: with torch.inference_mode(): input_values = torch.from_numpy(waveform).to(self.device).unsqueeze(0) logits = self.ctc_model(input_values=input_values).logits[0] log_probs = torch.log_softmax(logits.float(), dim=-1).cpu().numpy() decoded = _ctc_viterbi(log_probs, targets, self.blank_id) if decoded is not None: path, _ = decoded ctc_success = True frame_s = self.stride_samples / AUDIO_SAMPLE_RATE center_s = self.receptive_samples / (2 * AUDIO_SAMPLE_RATE) for word_index, target_indices in enumerate(word_target_indices): states = np.asarray([2 * index + 1 for index in target_indices], dtype=np.int32) frame_indices = np.flatnonzero(np.isin(path, states)) if frame_indices.size: first, last = int(frame_indices[0]), int(frame_indices[-1]) start = first * frame_s + center_s - frame_s / 2 end = (last + 1) * frame_s + center_s - frame_s / 2 intervals[word_index] = (max(0.0, start), min(duration_s, max(start, end))) valid[word_index] = True char_scores = [] for target_index in target_indices: char_frames = np.flatnonzero(path == 2 * target_index + 1) if char_frames.size: token_id = targets[target_index] char_scores.extend(log_probs[char_frames, token_id].tolist()) if char_scores: confidence[word_index] = float(np.exp(np.mean(char_scores))) if not bool(valid.all()): good = np.flatnonzero(valid) if not good.size: centers = (np.arange(count, dtype=np.float32) + 0.5) * duration_s / count for index, center in enumerate(centers): intervals[index] = (center, center) else: cursor = 0 while cursor < count: if valid[cursor]: cursor += 1 continue run_start = cursor while cursor < count and not valid[cursor]: cursor += 1 run_end = cursor - 1 left = float(intervals[run_start - 1, 1]) if run_start > 0 else 0.0 right = float(intervals[cursor, 0]) if cursor < count else duration_s right = max(left, right) run_length = run_end - run_start + 1 centers = np.linspace(left, right, run_length + 2, dtype=np.float32)[1:-1] for index, center in zip(range(run_start, run_end + 1), centers): # A point fallback preserves order without inventing a word duration. # The false validity mask makes the fallback explicit downstream. intervals[index] = (center, center) # Keep returned token order monotone even if low-confidence CTC spans touch. intervals[:, 0] = np.clip(intervals[:, 0], 0.0, duration_s) intervals[:, 1] = np.clip(np.maximum(intervals[:, 1], intervals[:, 0]), 0.0, duration_s) intervals[:, 0] = np.maximum.accumulate(intervals[:, 0]) intervals[:, 1] = np.maximum.accumulate(np.maximum(intervals[:, 1], intervals[:, 0])) return intervals, confidence, valid, ctc_success @staticmethod def audio_features(waveform: np.ndarray) -> tuple[np.ndarray, np.ndarray, list[str]]: smile = opensmile.Smile( feature_set=opensmile.FeatureSet.eGeMAPSv02, feature_level=opensmile.FeatureLevel.LowLevelDescriptors, sampling_rate=AUDIO_SAMPLE_RATE, ) frame = smile.process_signal(waveform, sampling_rate=AUDIO_SAMPLE_RATE) values = frame.to_numpy(dtype=np.float32, copy=True) valid = np.isfinite(values).all(axis=1) values = np.nan_to_num(values, nan=0.0, posinf=0.0, neginf=0.0) index = frame.index if getattr(index, "nlevels", 1) >= 2: starts = index.get_level_values(0) ends = index.get_level_values(1) starts_s = np.asarray([item.total_seconds() for item in starts], dtype=np.float32) ends_s = np.asarray([item.total_seconds() for item in ends], dtype=np.float32) times = (starts_s + ends_s) / 2 else: times = np.asarray([item.total_seconds() for item in index], dtype=np.float32) return times, values, list(map(str, frame.columns)), valid def vision_features( self, video_path: Path ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, list[str], int]: base_options = mp.tasks.BaseOptions(model_asset_path=str(self.vision_model_path)) options = mp.tasks.vision.FaceLandmarkerOptions( base_options=base_options, running_mode=mp.tasks.vision.RunningMode.VIDEO, num_faces=1, output_face_blendshapes=True, output_facial_transformation_matrixes=False, ) cap = cv2.VideoCapture(str(video_path)) if not cap.isOpened(): raise RuntimeError(f"OpenCV could not open {video_path}") fps = float(cap.get(cv2.CAP_PROP_FPS) or 0.0) if not math.isfinite(fps) or fps <= 0: fps = 30.0 stride_frames = max(1, int(round(fps / VISION_FPS))) features: list[np.ndarray] = [] times: list[float] = [] valid: list[bool] = [] names: list[str] = [] sampled_rgb: list[np.ndarray] = [] frames_seen = 0 last_timestamp_ms = -1 with mp.tasks.vision.FaceLandmarker.create_from_options(options) as landmarker: while True: ok, bgr = cap.read() if not ok: break frame_index = frames_seen frames_seen += 1 if frame_index % stride_frames != 0: continue timestamp_s = float(cap.get(cv2.CAP_PROP_POS_MSEC)) / 1000.0 if not math.isfinite(timestamp_s) or timestamp_s < 0: timestamp_s = frame_index / fps timestamp_ms = max(last_timestamp_ms + 1, int(round(timestamp_s * 1000))) last_timestamp_ms = timestamp_ms rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) sampled_rgb.append(rgb) image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np.ascontiguousarray(rgb)) result = landmarker.detect_for_video(image, timestamp_ms) categories = result.face_blendshapes[0] if result.face_blendshapes else [] if categories: current_names = [category.category_name for category in categories] if not names: names = current_names lookup = {category.category_name: float(category.score) for category in categories} vector = np.asarray([lookup.get(name, 0.0) for name in names], dtype=np.float32) features.append(vector) valid.append(True) else: if not names: names = [f"blendshape_{index:02d}" for index in range(VISION_DIM)] features.append(np.zeros((len(names),), dtype=np.float32)) valid.append(False) times.append(timestamp_s) cap.release() if not features: return ( np.empty((0,), dtype=np.float32), np.empty((0, VISION_DIM), dtype=np.float32), np.empty((0,), dtype=np.bool_), np.empty((0, self.vision_dim), dtype=np.float32), np.empty((0,), dtype=np.bool_), [f"blendshape_{index:02d}" for index in range(VISION_DIM)], frames_seen, ) matrix = np.stack(features).astype(np.float32, copy=False) if matrix.shape[1] != VISION_DIM: raise ValueError(f"expected {VISION_DIM} MediaPipe blendshapes, received {matrix.shape[1]}") image_embeddings: list[np.ndarray] = [] with torch.inference_mode(): for start in range(0, len(sampled_rgb), 16): batch_images = [Image.fromarray(frame) for frame in sampled_rgb[start:start + 16]] encoded = self.vision_processor(images=batch_images, return_tensors="pt") encoded = {key: value.to(self.device) for key, value in encoded.items()} hidden = self.vision_model(**encoded).last_hidden_state[:, 0, :] image_embeddings.append(hidden.float().cpu().numpy()) global_features = np.concatenate(image_embeddings, axis=0).astype(np.float32, copy=False) return ( np.asarray(times, dtype=np.float32), matrix, np.asarray(valid, dtype=np.bool_), global_features, np.ones((len(times),), dtype=np.bool_), names, frames_seen, ) def _mean_by_intervals( times: np.ndarray, values: np.ndarray, valid: np.ndarray, intervals: np.ndarray, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: dimensions = values.shape[1] if values.ndim == 2 else 0 aligned = np.zeros((len(intervals), dimensions), dtype=np.float32) aligned_valid = np.zeros((len(intervals),), dtype=np.bool_) fallback = np.zeros((len(intervals),), dtype=np.bool_) valid_positions = np.flatnonzero(valid) if not valid_positions.size: return aligned, aligned_valid, fallback for index, (start, end) in enumerate(intervals): inside = valid_positions[(times[valid_positions] >= start) & (times[valid_positions] < end)] if inside.size: aligned[index] = values[inside].mean(axis=0) aligned_valid[index] = True else: center = (start + end) / 2 nearest = valid_positions[int(np.argmin(np.abs(times[valid_positions] - center)))] aligned[index] = values[nearest] aligned_valid[index] = True fallback[index] = True return aligned, aligned_valid, fallback def extract_sample(models: FeatureModels, video_path: Path, record: dict[str, Any]) -> ExtractedSample: duration_s = _ffprobe_duration(video_path) errors: dict[str, str] = {} warnings: list[str] = [] waveform = np.empty((0,), dtype=np.float32) try: waveform = _read_audio(video_path) except Exception as error: # Keep the sample row and report modality failure. errors["audio_decode"] = f"{type(error).__name__}: {error}" try: words, text_values = models.text_features(str(record.get("text") or "")) except Exception as error: errors["text_features"] = f"{type(error).__name__}: {error}" words, text_values = [], np.empty((0, models.text_dim), dtype=np.float32) if not words: warnings.append("transcript_empty_or_unreadable") try: word_intervals, word_confidence, word_align_valid, ctc_success = models.force_align( words, waveform, duration_s ) except Exception as error: errors["forced_alignment"] = f"{type(error).__name__}: {error}" word_intervals = np.empty((len(words), 2), dtype=np.float32) if words: edges = np.linspace(0.0, duration_s, len(words) + 1, dtype=np.float32) word_intervals = np.column_stack((edges[:-1], edges[1:])) word_confidence = np.zeros((len(words),), dtype=np.float32) word_align_valid = np.zeros((len(words),), dtype=np.bool_) ctc_success = False if words and not bool(word_align_valid.all()): warnings.append(f"ctc_words_fallback:{int((~word_align_valid).sum())}") audio_columns: list[str] = [] try: if waveform.size: audio_times, audio_values, audio_columns, audio_valid = models.audio_features(waveform) else: raise ValueError("no decoded audio waveform") except Exception as error: errors["audio_features"] = f"{type(error).__name__}: {error}" audio_times = np.empty((0,), dtype=np.float32) audio_values = np.empty((0, 0), dtype=np.float32) audio_valid = np.empty((0,), dtype=np.bool_) try: ( vision_times, face_values, face_valid, vision_values, vision_valid, vision_names, decoded_frames, ) = models.vision_features(video_path) except Exception as error: errors["vision_features"] = f"{type(error).__name__}: {error}" vision_times = np.empty((0,), dtype=np.float32) face_values = np.empty((0, VISION_DIM), dtype=np.float32) face_valid = np.empty((0,), dtype=np.bool_) vision_values = np.empty((0, models.vision_dim), dtype=np.float32) vision_valid = np.empty((0,), dtype=np.bool_) vision_names = [f"blendshape_{index:02d}" for index in range(VISION_DIM)] decoded_frames = 0 if words and word_intervals.shape[0] == len(words): if audio_values.shape[1] if audio_values.ndim == 2 else 0: audio_word, audio_word_valid, audio_word_fallback = _mean_by_intervals( audio_times, audio_values, audio_valid, word_intervals ) else: audio_word = np.empty((len(words), 0), dtype=np.float32) audio_word_valid = np.zeros((len(words),), dtype=np.bool_) audio_word_fallback = np.zeros((len(words),), dtype=np.bool_) vision_word, vision_word_valid, vision_word_fallback = _mean_by_intervals( vision_times, vision_values, vision_valid, word_intervals ) face_word, face_word_valid, face_word_fallback = _mean_by_intervals( vision_times, face_values, face_valid, word_intervals ) else: audio_word = np.empty((0, audio_values.shape[1] if audio_values.ndim == 2 else 0), dtype=np.float32) audio_word_valid = np.empty((0,), dtype=np.bool_) audio_word_fallback = np.empty((0,), dtype=np.bool_) vision_word = np.empty((0, models.vision_dim), dtype=np.float32) vision_word_valid = np.empty((0,), dtype=np.bool_) vision_word_fallback = np.empty((0,), dtype=np.bool_) face_word = np.empty((0, VISION_DIM), dtype=np.float32) face_word_valid = np.empty((0,), dtype=np.bool_) face_word_fallback = np.empty((0,), dtype=np.bool_) if not np.any(face_valid): warnings.append("no_face_blendshapes_detected") sample_key = f"{record['video_id']}/{record['clip_id']}" arrays = { "text_words": np.asarray(words, dtype=f"U{max((len(word) for word in words), default=1)}"), "text_features": text_values.astype(np.float16), "word_intervals_s": word_intervals.astype(np.float32), "word_alignment_confidence": word_confidence.astype(np.float32), "word_alignment_valid": word_align_valid.astype(np.bool_), "audio_times_s": audio_times.astype(np.float32), "audio_features": audio_values.astype(np.float16), "audio_valid": audio_valid.astype(np.bool_), "audio_feature_names": np.asarray(audio_columns, dtype="U96"), "vision_times_s": vision_times.astype(np.float32), "vision_features": vision_values.astype(np.float16), "vision_valid": vision_valid.astype(np.bool_), "vision_feature_names": np.asarray([f"vit_cls_{index:03d}" for index in range(models.vision_dim)], dtype="U32"), "face_blendshape_features": face_values.astype(np.float16), "face_blendshape_valid": face_valid.astype(np.bool_), "face_blendshape_names": np.asarray(vision_names, dtype="U64"), "audio_word_features": audio_word.astype(np.float16), "audio_word_valid": audio_word_valid.astype(np.bool_), "audio_word_nearest_fallback": audio_word_fallback.astype(np.bool_), "vision_word_features": vision_word.astype(np.float16), "vision_word_valid": vision_word_valid.astype(np.bool_), "vision_word_nearest_fallback": vision_word_fallback.astype(np.bool_), "face_word_features": face_word.astype(np.float16), "face_word_valid": face_word_valid.astype(np.bool_), "face_word_nearest_fallback": face_word_fallback.astype(np.bool_), } info = { "sample_id": sample_key, "video_id": str(record["video_id"]), "clip_id": str(record["clip_id"]), "text": str(record.get("text") or ""), "label": record.get("label"), "annotation": str(record.get("annotation") or ""), "video_path": str(video_path), "duration_s": duration_s, "ctc_inference_succeeded": ctc_success, "word_count": len(words), "word_alignment_valid_count": int(word_align_valid.sum()), "audio_raw_length": int(len(audio_times)), "audio_valid_length": int(audio_valid.sum()), "audio_feature_dim": int(audio_values.shape[1]) if audio_values.ndim == 2 else 0, "vision_raw_length": int(len(vision_times)), "vision_valid_length": int(vision_valid.sum()), "vision_feature_dim": models.vision_dim, "face_blendshape_valid_length": int(face_valid.sum()), "face_blendshape_feature_dim": VISION_DIM, "vision_decoded_source_frames": int(decoded_frames), "text_feature_dim": models.text_dim, "alignment_granularity": "transcript word; Audio/Vision pooled over forced CTC word interval", "errors": errors, "warnings": warnings, } return ExtractedSample(arrays=arrays, info=info) def _save_npz(path: Path, arrays: dict[str, np.ndarray]) -> None: path.parent.mkdir(parents=True, exist_ok=True) temporary = path.with_suffix(path.suffix + ".tmp") with temporary.open("wb") as file: np.savez_compressed(file, **arrays) temporary.replace(path) def _sha256(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as file: for chunk in iter(lambda: file.read(1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() def _write_csv(path: Path, rows: list[dict[str, Any]], fieldnames: list[str]) -> None: path.parent.mkdir(parents=True, exist_ok=True) with path.open("w", encoding="utf-8-sig", newline="") as file: writer = csv.DictWriter(file, fieldnames=fieldnames, extrasaction="ignore") writer.writeheader() writer.writerows(rows) def _create_typical_figure( info: dict[str, Any], npz_path: Path, video_path: Path, output_path: Path ) -> None: import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt with np.load(npz_path, allow_pickle=False) as data: words = data["text_words"].astype(str) intervals = data["word_intervals_s"].astype(np.float32) audio_times = data["audio_times_s"].astype(np.float32) audio = data["audio_features"].astype(np.float32) audio_names = data["audio_feature_names"].astype(str) vision_times = data["vision_times_s"].astype(np.float32) vision_valid = data["vision_valid"].astype(bool) if not len(words): return mid = len(words) // 2 start_word = max(0, mid - 6) end_word = min(len(words), start_word + 12) start_word = max(0, end_word - 12) chosen = np.arange(start_word, end_word) if not chosen.size: return left = float(intervals[chosen[0], 0]) right = float(intervals[chosen[-1], 1]) if right <= left: left, right = 0.0, max(info["duration_s"], 0.1) figure = plt.figure(figsize=(15, 10), constrained_layout=True) layout = figure.add_gridspec(3, 1, height_ratios=(1.0, max(2.0, len(chosen) * 0.25), 1.5)) audio_axis = figure.add_subplot(layout[0, 0]) word_axis = figure.add_subplot(layout[1, 0], sharex=audio_axis) frame_axis = figure.add_subplot(layout[2, 0]) if audio.shape[0] and audio.shape[1]: keep = (audio_times >= left) & (audio_times <= right) selected_audio = audio[keep] selected_times = audio_times[keep] if selected_audio.size: normalized = (selected_audio - np.nanmedian(selected_audio, axis=0, keepdims=True)) scales = np.nanstd(normalized, axis=0, keepdims=True) normalized = normalized / np.where(scales < 1e-6, 1.0, scales) count = min(12, normalized.shape[1]) image = audio_axis.imshow( normalized[:, :count].T, aspect="auto", origin="lower", interpolation="nearest", extent=(float(selected_times[0]), float(selected_times[-1]), -0.5, count - 0.5), cmap="coolwarm", vmin=-2.5, vmax=2.5, ) labels = audio_names[:count] audio_axis.set_yticks(np.arange(count), labels=labels, fontsize=7) figure.colorbar(image, ax=audio_axis, label="per-clip z-score", fraction=0.02, pad=0.01) audio_axis.set_xlim(left, right) audio_axis.set_title("eGeMAPSv02 low-level speech features on the selected word interval") audio_axis.set_ylabel("audio descriptor") audio_axis.tick_params(axis="x", labelbottom=True) word_axis.set_ylim(-0.7, len(chosen) - 0.3) word_axis.set_yticks(chosen - chosen[0], labels=[words[index] for index in chosen], fontsize=8) for row_index, word_index in enumerate(chosen): word_start, word_end = intervals[word_index] word_axis.broken_barh( [(float(word_start), max(0.015, float(word_end - word_start)))], (row_index - 0.32, 0.64), facecolors="#2878B5", alpha=0.8, ) word_axis.text(float(word_start), row_index, f" {word_start:.2f}–{word_end:.2f}s", va="center", fontsize=7) word_axis.set_xlim(left, right) word_axis.set_xlabel("time from clip start (seconds)") word_axis.set_title("Transcript words mapped to their forced CTC audio intervals") word_axis.grid(axis="x", alpha=0.25) cap = cv2.VideoCapture(str(video_path)) candidate_times = vision_times[(vision_times >= left) & (vision_times <= right)] if not candidate_times.size: candidate_times = np.linspace(left, right, min(6, max(1, len(chosen)))) if candidate_times.size > 6: candidate_times = candidate_times[np.linspace(0, len(candidate_times) - 1, 6).round().astype(int)] frame_images: list[tuple[float, np.ndarray]] = [] for timestamp in candidate_times: cap.set(cv2.CAP_PROP_POS_MSEC, float(timestamp) * 1000) ok, bgr = cap.read() if ok: rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) frame_images.append((float(timestamp), rgb)) cap.release() frame_axis.axis("off") if frame_images: subgrid = layout[2, 0].subgridspec(1, len(frame_images), wspace=0.03) for index, (timestamp, image) in enumerate(frame_images): axis = figure.add_subplot(subgrid[0, index]) axis.imshow(image) axis.set_title(f"video t={timestamp:.2f}s", fontsize=9) axis.axis("off") else: frame_axis.text(0.5, 0.5, "No decodable video frames in this interval", ha="center", va="center") figure.suptitle( f"Q1 typical sample: {info['sample_id']}\n" f"Transcript snippet: {' '.join(words[chosen])}\n" "Text words → forced speech intervals → sampled video frames", fontsize=13, ) output_path.parent.mkdir(parents=True, exist_ok=True) figure.savefig(output_path, dpi=160) plt.close(figure) def run_extraction( video_root: Path, label_file: Path, output_dir: Path, model_dir: Path, cache_dir: Path | None, start: int = 0, limit: int | None = None, resume: bool = False, ) -> dict[str, Any]: records = _read_labels(label_file) selected = records[start:] if limit is None else records[start:start + limit] if not selected: raise ValueError("the requested sample range is empty") features_dir = output_dir / "features" logs_dir = output_dir / "logs" reports_dir = output_dir / "reports" features_dir.mkdir(parents=True, exist_ok=True) logs_dir.mkdir(parents=True, exist_ok=True) reports_dir.mkdir(parents=True, exist_ok=True) log_path = logs_dir / "extraction.log" LOGGER.setLevel(logging.INFO) LOGGER.handlers.clear() formatter = logging.Formatter("%(asctime)s %(levelname)s %(message)s") stream = logging.StreamHandler(sys.stdout) stream.setFormatter(formatter) file_handler = logging.FileHandler(log_path, encoding="utf-8") file_handler.setFormatter(formatter) LOGGER.addHandler(stream) LOGGER.addHandler(file_handler) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") LOGGER.info("Using %s", device) models = FeatureModels(device, model_dir / "face_landmarker.task", cache_dir=cache_dir) log_jsonl_path = logs_dir / "samples.jsonl" existing: dict[str, dict[str, Any]] = {} if resume and log_jsonl_path.exists(): for line in log_jsonl_path.read_text(encoding="utf-8").splitlines(): try: entry = json.loads(line) existing[entry["sample_id"]] = entry except (json.JSONDecodeError, KeyError): continue completed: list[dict[str, Any]] = [] sample_log = log_jsonl_path.open("a" if resume else "w", encoding="utf-8") try: for index, record in enumerate(selected, start=start): video_id, clip_id = str(record["video_id"]), str(record["clip_id"]) sample_id = f"{video_id}/{clip_id}" output_path = features_dir / f"{_safe_id(video_id)}__{_safe_id(clip_id)}.npz" if resume and sample_id in existing and output_path.exists(): info = existing[sample_id] completed.append(info) LOGGER.info("[%d/%d] resume existing %s", index + 1, len(records), sample_id) continue video_path = video_root / video_id / f"{clip_id}.mp4" begin = time.perf_counter() if not video_path.is_file(): info = { "sample_id": sample_id, "video_id": video_id, "clip_id": clip_id, "text": str(record.get("text") or ""), "label": record.get("label"), "annotation": str(record.get("annotation") or ""), "video_path": str(video_path), "duration_s": 0.0, "errors": {"input": "source video file missing"}, "warnings": [], "text_feature_dim": models.text_dim, "audio_feature_dim": 0, "vision_feature_dim": models.vision_dim, "face_blendshape_valid_length": 0, "face_blendshape_feature_dim": VISION_DIM, "word_count": 0, "word_alignment_valid_count": 0, "audio_raw_length": 0, "audio_valid_length": 0, "vision_raw_length": 0, "vision_valid_length": 0, "alignment_granularity": "transcript word; Audio/Vision pooled over forced CTC word interval", } empty = { **_empty_text_arrays(models.text_dim), "word_intervals_s": np.empty((0, 2), dtype=np.float32), "word_alignment_confidence": np.empty((0,), dtype=np.float32), "word_alignment_valid": np.empty((0,), dtype=np.bool_), "audio_times_s": np.empty((0,), dtype=np.float32), "audio_features": np.empty((0, 0), dtype=np.float16), "audio_valid": np.empty((0,), dtype=np.bool_), "audio_feature_names": np.empty((0,), dtype="U1"), "vision_times_s": np.empty((0,), dtype=np.float32), "vision_features": np.empty((0, models.vision_dim), dtype=np.float16), "vision_valid": np.empty((0,), dtype=np.bool_), "vision_feature_names": np.asarray([f"vit_cls_{i:03d}" for i in range(models.vision_dim)], dtype="U32"), "face_blendshape_features": np.empty((0, VISION_DIM), dtype=np.float16), "face_blendshape_valid": np.empty((0,), dtype=np.bool_), "face_blendshape_names": np.asarray([f"blendshape_{i:02d}" for i in range(VISION_DIM)], dtype="U32"), "audio_word_features": np.empty((0, 0), dtype=np.float16), "audio_word_valid": np.empty((0,), dtype=np.bool_), "audio_word_nearest_fallback": np.empty((0,), dtype=np.bool_), "vision_word_features": np.empty((0, models.vision_dim), dtype=np.float16), "vision_word_valid": np.empty((0,), dtype=np.bool_), "vision_word_nearest_fallback": np.empty((0,), dtype=np.bool_), "face_word_features": np.empty((0, VISION_DIM), dtype=np.float16), "face_word_valid": np.empty((0,), dtype=np.bool_), "face_word_nearest_fallback": np.empty((0,), dtype=np.bool_), } _save_npz(output_path, empty) else: try: extracted = extract_sample(models, video_path, record) info = extracted.info _save_npz(output_path, extracted.arrays) except Exception as error: LOGGER.exception("Sample-level extraction failed for %s", sample_id) info = { "sample_id": sample_id, "video_id": video_id, "clip_id": clip_id, "text": str(record.get("text") or ""), "label": record.get("label"), "annotation": str(record.get("annotation") or ""), "video_path": str(video_path), "duration_s": 0.0, "errors": {"sample": f"{type(error).__name__}: {error}"}, "warnings": [], "text_feature_dim": models.text_dim, "audio_feature_dim": 0, "vision_feature_dim": models.vision_dim, "face_blendshape_valid_length": 0, "face_blendshape_feature_dim": VISION_DIM, "word_count": 0, "word_alignment_valid_count": 0, "audio_raw_length": 0, "audio_valid_length": 0, "vision_raw_length": 0, "vision_valid_length": 0, "alignment_granularity": "transcript word; Audio/Vision pooled over forced CTC word interval", } # A failed sample still has a corresponding record instead of being omitted. _save_npz(output_path, { **_empty_text_arrays(models.text_dim), "word_intervals_s": np.empty((0, 2), dtype=np.float32), "word_alignment_confidence": np.empty((0,), dtype=np.float32), "word_alignment_valid": np.empty((0,), dtype=np.bool_), "audio_times_s": np.empty((0,), dtype=np.float32), "audio_features": np.empty((0, 0), dtype=np.float16), "audio_valid": np.empty((0,), dtype=np.bool_), "audio_feature_names": np.empty((0,), dtype="U1"), "vision_times_s": np.empty((0,), dtype=np.float32), "vision_features": np.empty((0, models.vision_dim), dtype=np.float16), "vision_valid": np.empty((0,), dtype=np.bool_), "vision_feature_names": np.asarray([f"vit_cls_{i:03d}" for i in range(models.vision_dim)], dtype="U32"), "face_blendshape_features": np.empty((0, VISION_DIM), dtype=np.float16), "face_blendshape_valid": np.empty((0,), dtype=np.bool_), "face_blendshape_names": np.asarray([f"blendshape_{i:02d}" for i in range(VISION_DIM)], dtype="U32"), "audio_word_features": np.empty((0, 0), dtype=np.float16), "audio_word_valid": np.empty((0,), dtype=np.bool_), "audio_word_nearest_fallback": np.empty((0,), dtype=np.bool_), "vision_word_features": np.empty((0, models.vision_dim), dtype=np.float16), "vision_word_valid": np.empty((0,), dtype=np.bool_), "vision_word_nearest_fallback": np.empty((0,), dtype=np.bool_), "face_word_features": np.empty((0, VISION_DIM), dtype=np.float16), "face_word_valid": np.empty((0,), dtype=np.bool_), "face_word_nearest_fallback": np.empty((0,), dtype=np.bool_), }) info["processing_seconds"] = round(time.perf_counter() - begin, 3) info["feature_file"] = str(output_path) info["feature_file_bytes"] = output_path.stat().st_size info["feature_file_sha256"] = _sha256(output_path) sample_log.write(json.dumps(info, ensure_ascii=False, default=str) + "\n") sample_log.flush() completed.append(info) errors = info.get("errors", {}) LOGGER.info( "[%d/%d] %s duration=%.3fs words=%d audio=%d vision=%d errors=%s time=%.1fs", index + 1, len(records), sample_id, info.get("duration_s", 0.0), info.get("word_count", 0), info.get("audio_raw_length", 0), info.get("vision_raw_length", 0), ",".join(errors) or "none", info["processing_seconds"], ) finally: sample_log.close() detailed_rows: list[dict[str, Any]] = [] for info in completed: for modality in ("text", "audio", "vision"): if modality == "text": length, valid_length, dimension = ( info.get("word_count", 0), info.get("word_alignment_valid_count", 0), info.get("text_feature_dim", 0) ) granularity = "transcript word" elif modality == "audio": length, valid_length, dimension = ( info.get("audio_raw_length", 0), info.get("audio_valid_length", 0), info.get("audio_feature_dim", 0) ) granularity = "eGeMAPSv02 LLD frame (~10 ms)" else: length, valid_length, dimension = ( info.get("vision_raw_length", 0), info.get("vision_valid_length", 0), info.get("vision_feature_dim", VISION_DIM) ) granularity = f"sampled frame ({VISION_FPS:g} fps), ViT CLS; optional 52-d face blendshape" detailed_rows.append({ "sample_id": info.get("sample_id"), "video_id": info.get("video_id"), "clip_id": info.get("clip_id"), "modality": modality, "effective_duration_s": info.get("duration_s", 0.0), "feature_dimension": dimension, "sequence_length": length, "valid_length": valid_length, "padding_rule": "no on-disk padding; per-sample variable length", "alignment_granularity": granularity, "source_time_mapping": "seconds from clip start; text tokens map to word_intervals_s", "feature_file": info.get("feature_file", ""), "status": "error" if info.get("errors") else "ok_with_warning" if info.get("warnings") else "ok", "errors": json.dumps(info.get("errors", {}), ensure_ascii=False), "warnings": json.dumps(info.get("warnings", []), ensure_ascii=False), }) _write_csv( reports_dir / "feature_summary_100x3.csv", detailed_rows, ["sample_id", "video_id", "clip_id", "modality", "effective_duration_s", "feature_dimension", "sequence_length", "valid_length", "padding_rule", "alignment_granularity", "source_time_mapping", "feature_file", "status", "errors", "warnings"], ) _write_csv( reports_dir / "sample_summary.csv", completed, ["sample_id", "video_id", "clip_id", "duration_s", "word_count", "word_alignment_valid_count", "text_feature_dim", "audio_raw_length", "audio_valid_length", "audio_feature_dim", "vision_raw_length", "vision_valid_length", "vision_feature_dim", "face_blendshape_valid_length", "face_blendshape_feature_dim", "ctc_inference_succeeded", "alignment_granularity", "feature_file", "feature_file_bytes", "errors", "warnings"], ) manifest = { "created_at_utc": datetime.now(timezone.utc).isoformat(), "command": sys.argv, "python": sys.version, "platform": platform.platform(), "uv_version": _tool_version("uv"), "device": str(device), "cuda_available": bool(torch.cuda.is_available()), "torch_cuda_build": torch.version.cuda, "cudnn_version": torch.backends.cudnn.version(), "gpu_name": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None, "external_tools": { "ffmpeg": _tool_version("ffmpeg"), "ffprobe": _tool_version("ffprobe"), }, "package_versions": { name: _package_version(name) for name in ( "torch", "torchvision", "transformers", "tokenizers", "mediapipe", "opencv-python-headless", "opensmile", "numpy", "soundfile", "openpyxl", "matplotlib", "Pillow", ) }, "models": { "text": { "id": TEXT_MODEL_ID, "revision": getattr(models.text_model.config, "_commit_hash", None), "representation": "mean of final-layer BERT subtokens per supplied whitespace token", }, "forced_alignment": { "id": CTC_MODEL_ID, "revision": getattr(models.ctc_model.config, "_commit_hash", None), "algorithm": "CTC Viterbi over supplied transcript; clip-relative seconds", }, "audio": {"set": "eGeMAPSv02", "level": "LowLevelDescriptors", "sample_rate_hz": AUDIO_SAMPLE_RATE}, "vision": { "model": VISION_MODEL_ID, "revision": getattr(models.vision_model.config, "_commit_hash", None), "representation": "DeiT final-layer CLS embedding for every sampled frame", "sample_rate_fps": VISION_FPS, "supplemental_face_feature": "MediaPipe Face Landmarker 52 blendshape scores with face_blendshape_valid mask", }, }, "alignment": { "granularity": "word-level", "primary_rule": "forced CTC transcript word intervals; mean Audio LLD and DeiT frame embeddings inside interval", "empty_window_rule": "nearest valid source frame; fallback flags are retained in NPZ", "missing_face_rule": "DeiT frame embeddings remain valid; face blendshapes are zero with face_blendshape_valid=false", "storage": "compressed NPZ; values float16, times float32, masks bool; no padding in per-sample files", }, "face_landmarker_asset": { "path": str(models.vision_model_path), "sha256": _sha256(models.vision_model_path), "size_bytes": models.vision_model_path.stat().st_size, }, "inputs": {"video_root": str(video_root), "label_file": str(label_file), "label_sha256": _sha256(label_file)}, "outputs": {"feature_dir": str(features_dir), "summary": str(reports_dir / "feature_summary_100x3.csv"), "sample_table": str(reports_dir / "sample_summary.csv"), "sample_log": str(log_jsonl_path)}, "requested_sample_range": {"start": start, "limit": limit, "selected_count": len(selected)}, "completed_count": len(completed), "sample_errors": sum(bool(item.get("errors")) for item in completed), "sample_warnings": sum(bool(item.get("warnings")) for item in completed), "aggregate_feature_bytes": sum(int(item.get("feature_file_bytes", 0)) for item in completed), } manifest_path = reports_dir / "run_manifest.json" manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2, default=str), encoding="utf-8") candidates = [ item for item in completed if item.get("vision_valid_length", 0) > 0 and item.get("word_alignment_valid_count", 0) > 0 ] if candidates: typical = max( candidates, key=lambda item: ( min(item.get("vision_valid_length", 0), 12) + min(item.get("word_alignment_valid_count", 0), 12), -abs(item.get("duration_s", 0.0) - 12.0), ), ) _create_typical_figure( typical, Path(typical["feature_file"]), Path(typical["video_path"]), reports_dir / f"typical_sample_{_safe_id(typical['video_id'])}__{_safe_id(typical['clip_id'])}.png", ) manifest["typical_sample_id"] = typical["sample_id"] manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2, default=str), encoding="utf-8") return manifest def main() -> int: project_dir = Path(__file__).resolve().parents[1] repo_dir = project_dir.parent default_video_root = repo_dir / "E题数据" / "附件1-数据集原始多模态样本" / "MOSEI数据集部分原始视频-100条" parser = argparse.ArgumentParser(description="Extract timestamped Q1 Text, Audio, and Vision features from Attachment 1.") parser.add_argument("--video-root", type=Path, default=default_video_root) parser.add_argument("--labels", type=Path, default=default_video_root / "label-100.xlsx") parser.add_argument("--output-dir", type=Path, default=project_dir / "outputs") parser.add_argument("--model-dir", type=Path, default=project_dir / "models") parser.add_argument("--cache-dir", type=Path, default=None) parser.add_argument("--start", type=int, default=0, help="zero-based row offset in label workbook") parser.add_argument("--limit", type=int, default=None, help="number of rows; omit for all remaining rows") parser.add_argument("--resume", action="store_true", help="reuse already completed sample files/log entries") args = parser.parse_args() manifest = run_extraction( args.video_root, args.labels, args.output_dir, args.model_dir, args.cache_dir, start=args.start, limit=args.limit, resume=args.resume, ) print(json.dumps(manifest, ensure_ascii=False, indent=2, default=str)) return 0 if manifest["sample_errors"] == 0 else 1 if __name__ == "__main__": raise SystemExit(main())