231 lines
11 KiB
Python
231 lines
11 KiB
Python
"""Read and derive Q1 V2 feature views from the native-source artifacts."""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import re
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
from scipy import sparse
|
|
|
|
Q1_DIR = Path(__file__).resolve().parent
|
|
FEATURE_DIR = Q1_DIR / "features_v2"
|
|
MODALITIES = ("text", "audio", "vision")
|
|
FEATURE_DIMS = {"text": 768, "audio": 74, "vision": 35, "speech": 768}
|
|
|
|
|
|
def load_sample(sample_id: str, feature_dir: Path = FEATURE_DIR) -> dict[str, Any]:
|
|
"""Load one sample by `video_id/clip_id` and return arrays plus metadata."""
|
|
manifest = feature_dir / "manifest_q1.jsonl"
|
|
found = None
|
|
for line in manifest.read_text(encoding="utf-8").splitlines():
|
|
row = json.loads(line)
|
|
if row["sample_id"] == sample_id:
|
|
found = row
|
|
break
|
|
if found is None:
|
|
raise KeyError(f"sample_id not found: {sample_id}")
|
|
path = Q1_DIR / found["feature_path"]
|
|
with np.load(path, allow_pickle=False) as archive:
|
|
result = {name: archive[name] for name in archive.files}
|
|
result["_path"] = path
|
|
result["_manifest"] = found
|
|
result["_meta"] = json.loads(str(result["meta_json"]))
|
|
return result
|
|
|
|
|
|
def csr_from_sample(sample: dict[str, Any], prefix: str) -> sparse.csr_matrix:
|
|
"""Reconstruct a stored sparse alignment/query map."""
|
|
shape = tuple(int(v) for v in sample[f"{prefix}_shape"])
|
|
return sparse.csr_matrix(
|
|
(
|
|
sample[f"{prefix}_data"].astype(np.float32),
|
|
sample[f"{prefix}_indices"].astype(np.int32),
|
|
sample[f"{prefix}_indptr"].astype(np.int32),
|
|
),
|
|
shape=shape,
|
|
)
|
|
|
|
|
|
def _native(sample: dict[str, Any], modality: str) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
|
|
if modality == "text":
|
|
return (
|
|
sample["native_text_features"].astype(np.float32),
|
|
np.broadcast_to(sample["native_text_observed"][:, None], sample["native_text_features"].shape),
|
|
sample["native_text_intervals"].astype(np.float32),
|
|
sample["native_text_quality_effective"].astype(np.float32),
|
|
sample["native_text_quality_available"].astype(bool),
|
|
)
|
|
if modality in ("audio", "vision"):
|
|
return (
|
|
sample[f"native_{modality}_features"].astype(np.float32),
|
|
sample[f"native_{modality}_mask"].astype(bool),
|
|
sample[f"native_{modality}_intervals"].astype(np.float32),
|
|
sample[f"native_{modality}_quality"].astype(np.float32),
|
|
np.asarray(sample[f"native_{modality}_quality_available"], dtype=bool),
|
|
)
|
|
if modality == "speech":
|
|
meta = sample["_meta"]
|
|
safe = lambda value: re.sub(r"[^A-Za-z0-9_.-]+", "_", str(value))
|
|
path = Q1_DIR / "cache" / "native" / f"{safe(meta['video_id'])}__{safe(meta['clip_id'])}.npz"
|
|
if not path.is_file():
|
|
raise FileNotFoundError(f"B4 speech source cache is unavailable: {path}")
|
|
with np.load(path, allow_pickle=False) as archive:
|
|
values = archive["speech_features"].astype(np.float32)
|
|
times = sample["native_ctc_times"].astype(np.float32)
|
|
if len(values) != len(times):
|
|
raise ValueError("B4 speech rows do not match the stored CTC timestamps")
|
|
step = 320.0 / 16000.0
|
|
intervals = np.column_stack((np.maximum(0.0, times - step / 2), np.minimum(float(meta["duration_s"]), times + step / 2))).astype(np.float32)
|
|
return values, np.isfinite(values), intervals, np.ones(len(values), np.float32), np.zeros(len(values), bool)
|
|
raise ValueError(f"unknown modality: {modality}")
|
|
|
|
|
|
def _aggregate(
|
|
values: np.ndarray,
|
|
observed: np.ndarray,
|
|
source_intervals: np.ndarray,
|
|
quality: np.ndarray,
|
|
target_intervals: np.ndarray,
|
|
quality_available: np.ndarray | None = None,
|
|
) -> dict[str, np.ndarray]:
|
|
"""Quality-weight B0 means while keeping physical coverage and availability separate."""
|
|
n, d = len(target_intervals), values.shape[1]
|
|
mean = np.zeros((n, d), np.float32)
|
|
var = np.zeros((n, d), np.float32)
|
|
count = np.zeros((n, d), np.uint8)
|
|
coverage = np.zeros((n, d), np.float32)
|
|
qbar = np.ones((n, d), np.float32)
|
|
qavail = np.zeros((n, d), np.float32)
|
|
available = np.zeros(len(values), bool) if quality_available is None else np.asarray(quality_available, bool)
|
|
mask = np.zeros((n, d), bool)
|
|
for i, (left, right) in enumerate(target_intervals):
|
|
width = float(right - left)
|
|
if width <= 0 or not len(source_intervals):
|
|
continue
|
|
overlap = np.maximum(
|
|
0.0,
|
|
np.minimum(source_intervals[:, 1], right) - np.maximum(source_intervals[:, 0], left),
|
|
)
|
|
physical = overlap[:, None] * observed
|
|
weighted = physical * quality[:, None]
|
|
total = weighted.sum(axis=0)
|
|
valid = total > 0
|
|
if valid.any():
|
|
weighted_sum = (values * weighted).sum(axis=0)
|
|
weighted_sq = (np.square(values) * weighted).sum(axis=0)
|
|
mean[i, valid] = (weighted_sum[valid] / total[valid]).astype(np.float32)
|
|
var[i, valid] = np.maximum(0.0, weighted_sq[valid] / total[valid] - mean[i, valid] ** 2)
|
|
mask[i, valid] = True
|
|
physical_total = physical.sum(axis=0)
|
|
physical_valid = physical_total > 0
|
|
qbar[i, physical_valid] = weighted.sum(axis=0)[physical_valid] / physical_total[physical_valid]
|
|
qavail[i, physical_valid] = (physical * available[:, None]).sum(axis=0)[physical_valid] / physical_total[physical_valid]
|
|
coverage[i] = np.minimum(1.0, physical_total / width)
|
|
count[i] = np.minimum(255, ((overlap[:, None] > 0) & observed).sum(axis=0)).astype(np.uint8)
|
|
return {
|
|
"x": mean.astype(np.float16),
|
|
"var": var.astype(np.float32),
|
|
"count": count,
|
|
"coverage": coverage,
|
|
"mask": mask,
|
|
"quality_mean": qbar.astype(np.float32),
|
|
"quality_available_fraction": qavail.astype(np.float32),
|
|
}
|
|
|
|
|
|
def get_view(
|
|
sample: dict[str, Any],
|
|
modality: str,
|
|
view: str = "sec",
|
|
context_s: float | None = None,
|
|
geometry_aware: bool = False,
|
|
) -> dict[str, np.ndarray]:
|
|
"""Return `sec`, `phase50`, `word`, `multi`, or `posterior` features.
|
|
|
|
`multi` is always aggregated from native rows around each 0.1 s bin midpoint.
|
|
For it, set `context_s` to one of 0.1, 0.3, or 0.7.
|
|
"""
|
|
if modality not in (*MODALITIES, "speech"):
|
|
raise ValueError(f"unknown modality: {modality}")
|
|
if view in ("sec", "phase50", "word"):
|
|
bounds = sample[f"views_{view}_time_bounds_s"].astype(np.float32)
|
|
if modality == "speech":
|
|
values, observed, intervals, quality, available = _native(sample, modality)
|
|
return {**_aggregate(values, observed, intervals, quality, bounds, available), "time_bounds_s": bounds}
|
|
prefix = f"views_{view}_{modality}_"
|
|
result = {
|
|
"x": sample[prefix + "x"],
|
|
"var": sample[prefix + "var"],
|
|
"count": sample[prefix + "count"],
|
|
"coverage": sample[prefix + "coverage_u8"].astype(np.float32) / 255.0,
|
|
"mask": sample[prefix + "mask"].astype(bool),
|
|
"quality_mean": sample[prefix + "qbar"].astype(np.float32),
|
|
"quality_available_fraction": sample[prefix + "quality_available_fraction"].astype(np.float32),
|
|
"time_bounds_s": bounds,
|
|
}
|
|
if view == "sec" and modality == "vision" and geometry_aware:
|
|
result["x"] = sample["views_sec_vision_b1_x"]
|
|
result["mask"] = sample["views_sec_vision_b1_mask"].astype(bool)
|
|
return result
|
|
if view == "multi":
|
|
if context_s not in (0.1, 0.3, 0.7):
|
|
raise ValueError("context_s must be one of 0.1, 0.3, or 0.7")
|
|
centers = sample["views_sec_time_bounds_s"].astype(np.float32).mean(axis=1)
|
|
duration = float(sample["_meta"]["duration_s"])
|
|
bounds = np.column_stack((np.maximum(0.0, centers - context_s / 2), np.minimum(duration, centers + context_s / 2))).astype(np.float32)
|
|
values, observed, intervals, quality, available = _native(sample, modality)
|
|
return {**_aggregate(values, observed, intervals, quality, bounds, available), "time_bounds_s": bounds}
|
|
if view == "posterior":
|
|
if modality != "text":
|
|
raise ValueError("posterior projection is defined for text only")
|
|
return posterior_text_view(sample)
|
|
raise ValueError(f"unknown view: {view}")
|
|
|
|
|
|
def posterior_text_view(sample: dict[str, Any]) -> dict[str, np.ndarray]:
|
|
"""Project CTC forward-backward word occupancy directly to the 0.1 s grid."""
|
|
edges = sample["views_sec_time_bounds_s"].astype(np.float32)
|
|
words = sample["native_text_features"].astype(np.float32)
|
|
word_ok = sample["native_text_observed"].astype(bool)
|
|
occupancy = sample["native_ctc_occupancy"].astype(np.float32)
|
|
times = sample["native_ctc_times"].astype(np.float32)
|
|
if occupancy.ndim != 2 or occupancy.shape[1] != len(words):
|
|
raise ValueError("CTC occupancy word dimension does not match the native text rows")
|
|
step = 320.0 / 16000.0
|
|
half = step / 2
|
|
duration = float(sample["_meta"]["duration_s"])
|
|
source = np.column_stack((np.maximum(0.0, times - half), np.minimum(duration, times + half))).astype(np.float32)
|
|
n, d = len(edges), words.shape[1]
|
|
result = np.zeros((n, d), np.float32)
|
|
var = np.zeros((n, d), np.float32)
|
|
count = np.zeros(n, np.uint8)
|
|
coverage = np.zeros(n, np.float32)
|
|
mask = np.zeros(n, bool)
|
|
mass_by_word = np.zeros((n, len(words)), np.float32)
|
|
for i, (left, right) in enumerate(edges):
|
|
overlap = np.maximum(0.0, np.minimum(source[:, 1], right) - np.maximum(source[:, 0], left))
|
|
mass_by_word[i] = (occupancy * overlap[:, None]).sum(axis=0)
|
|
w = mass_by_word[i] * word_ok
|
|
total = float(w.sum())
|
|
if total > 0:
|
|
normalized = w / total
|
|
result[i] = normalized @ words
|
|
centered = words - result[i]
|
|
var[i] = (normalized[:, None] * np.square(centered)).sum(axis=0)
|
|
mask[i] = True
|
|
count[i] = min(255, int(np.count_nonzero(w > 0)))
|
|
coverage[i] = min(1.0, total / max(float(right - left), 1e-8))
|
|
return {
|
|
"x": result.astype(np.float16),
|
|
"var": var,
|
|
"count": np.broadcast_to(count[:, None], result.shape).copy(),
|
|
"coverage": np.broadcast_to(coverage[:, None], result.shape).copy(),
|
|
"mask": np.broadcast_to(mask[:, None], result.shape).copy(),
|
|
"quality_mean": np.ones_like(result, dtype=np.float32),
|
|
"quality_available_fraction": np.zeros_like(result, dtype=np.float32),
|
|
"word_posterior_mass": mass_by_word,
|
|
"time_bounds_s": edges,
|
|
}
|