1083 lines
53 KiB
Python
1083 lines
53 KiB
Python
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="<f4").copy()
|
||
if waveform.size == 0:
|
||
raise ValueError("ffmpeg returned an empty audio stream")
|
||
return np.nan_to_num(waveform, nan=0.0, posinf=0.0, neginf=0.0)
|
||
|
||
|
||
def _empty_text_arrays(hidden_size: int) -> 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())
|