diff --git a/.gitignore b/.gitignore index ecbb082..7cf1785 100644 --- a/.gitignore +++ b/.gitignore @@ -11,8 +11,6 @@ __pycache__/ /manifests/ /outputs/ /target_frame/ -/third_party/EgoSpeak/ -/third_party/MuVAP/ # Local editors and environments. .idea/ diff --git a/PROJECT_CONTEXT.md b/PROJECT_CONTEXT.md index df8b343..211d3ea 100644 --- a/PROJECT_CONTEXT.md +++ b/PROJECT_CONTEXT.md @@ -73,6 +73,40 @@ result. See the immutable [`result record`](results/easycom_track_gru_residual_4ep_2026-09-07/README.md) and [`implementation note`](docs/easycom_track_gru_residual.md). +- `causal-person-adapter-v2` is now implemented but unmeasured. It preserves an + exact matched pooled `G+A+V` base, reuses the same causal anonymous track + histories, replaces the blocking zero scalar with a sample-dependent vector + gate, and zero-initializes only the final residual-logit projection. Its + isolated runbook is + [`docs/easycom_causal_person_adapter_v2.md`](docs/easycom_causal_person_adapter_v2.md). + This is an optimization-repair/incremental screen, not evidence that track + identity matters. +- `causal-person-adapter-egospeak-gru` is an unmeasured controlled ablation of + V2. It changes only the per-track visual temporal stem to EgoSpeak + MiniROADA's `Linear-LN-ReLU-Dropout-GRU-ReLU` ordering while retaining V2's + pooled base, G+A person query, set attention, dynamic gate and zero-output + residual. It is an EgoSpeak-style per-track adaptation, not a reproduction of + EgoSpeak's global 5 Hz RGB+audio GRU. See its + [`runbook`](docs/easycom_egospeak_style_person_gru.md). +- `highrate-person-av-attention` is now implemented but unmeasured. A separate + extractor runs the published, frozen MuVAP frontend at 25 Hz: mixed audio is + encoded once as a shared CPC stream, while each face track contributes a + causal visual state plus one AV active-speaker score. At each 5 Hz right edge + `G+U` queries that variable person set. The first screen loads and freezes + the matched trained global-Omni + traditional-audio base and learns a + zero-initialized person-logit residual. This is global Omni context plus + traditional high-rate audio/person perception; + it is not the full MuVAP forecasting model and does not yet include online + detection. See the + [`runbook`](docs/easycom_highrate_person_av.md). +- A MuVAP-only Stage-3 baseline is now implemented as a clean-room paper + reconstruction. It freezes the released 25 Hz mono role-VAP and causal ASD + backbones, lets `Z_VAP` query the variable `Z_ASD` person set at each frame, + applies a one-layer causal temporal Transformer, and trains only that fusion + plus the EgoSpeak 10-horizon head. It deliberately excludes Omni and is not + described as upstream Stage-3 code because that entry point remains + unpublished. See the + [`runbook`](docs/easycom_muvap_stage3_baseline.md). ### Measured oracle-information screen @@ -305,3 +339,6 @@ Do not reorganize them physically without compatibility wrappers. - Remote paths differ (`/data/...` and `/data1/...`). Commands should expose dataset/model roots as variables rather than hard-code one machine. - Large models, datasets, checkpoints and extracted features live outside Git. + The explicitly audited 72 MiB MuVAP ASD-CPC release and 7 MiB CPC frontend + are a deliberate exception: they are vendored with their exact source and + provenance under `third_party/MuVAP` for reproducible extraction. diff --git a/docs/README.md b/docs/README.md index 8f0e3ea..c220be9 100644 --- a/docs/README.md +++ b/docs/README.md @@ -25,6 +25,10 @@ tests. | [`ego4d_encoder_only.md`](ego4d_encoder_only.md) | Preparing and running the Ego4D encoder-only comparison | | [`ego4d_local_baselines.md`](ego4d_local_baselines.md) | Previewing Ego4D boxes, extracting local features and running 8-GPU jobs | | [`easycom_track_gru_residual.md`](easycom_track_gru_residual.md) | Matched pooled-to-track GRU experiment, causal identity contract and L20 commands | +| [`easycom_causal_person_adapter_v2.md`](easycom_causal_person_adapter_v2.md) | Dynamic-gated, zero-output causal person-track residual and matched L20 screen | +| [`easycom_egospeak_style_person_gru.md`](easycom_egospeak_style_person_gru.md) | Controlled V2 temporal-stem ablation using EgoSpeak MiniROADA's layer order | +| [`easycom_highrate_person_av.md`](easycom_highrate_person_av.md) | Frozen 25 Hz causal ASD person evidence fused with 5 Hz global Omni context | +| [`easycom_muvap_stage3_baseline.md`](easycom_muvap_stage3_baseline.md) | MuVAP-only frozen VAP+ASD Stage-3 paper reconstruction on EgoSpeak/EasyCom targets | | [`easycom_oracle_information.md`](easycom_oracle_information.md) | Causal oracle per-person speaking-history upper bound and Session-12 screening commands | ## Design history and implementation options diff --git a/docs/easycom_causal_person_adapter_v2.md b/docs/easycom_causal_person_adapter_v2.md new file mode 100644 index 0000000..c02deed --- /dev/null +++ b/docs/easycom_causal_person_adapter_v2.md @@ -0,0 +1,113 @@ +# EasyCom causal person adapter V2 + +## Status and question + +Implemented and CPU-testable as `--local-fusion causal-person-adapter-v2`; +GPU results have not yet been measured. The mode is an optimization-corrected, +incremental experiment after the matched pooled `G+A+V` baseline. It does not +yet establish that track identity or visual grounding improves wearer floor +behavior. + +The experiment asks whether strictly causal, per-track visual histories contain +information beyond equal person/frame pooling. It deliberately preserves the +legacy `track-gru-residual` mode so its earlier null screen remains reproducible. +The controlled `causal-person-adapter-egospeak-gru` ablation keeps this complete +V2 fusion path and changes only the per-track temporal stem; see +[`easycom_egospeak_style_person_gru.md`](easycom_egospeak_style_person_gru.md). + +## Architecture + +For 5 Hz row `t`, the media cutoff is `(t+1)/5` seconds. The local loader selects +only person frames strictly earlier than that cutoff. Each currently visible +anonymous track contributes at most `K` observations: + +```text +p_i,t = shared_GRU(project(v_i,t-K+1:t)) +q_t = project(G_t) + project(A_t) +c_t = MHA(q_t, [NULL, p_1,t, ..., p_N,t]) +l_t = MLP([c_t, projected_A_t, person_valid, audio_valid]) +gate = sigmoid(MLP([pooled_base_hidden_t, l_t])) +logits_t = pooled_logits_t + person_valid * W_delta(gate * l_t) +``` + +`W_delta.weight` and `W_delta.bias` are zero initialized. After copying a +matched pooled checkpoint, V2 therefore starts with logits exactly equal to +the pooled model. The first backward pass updates `W_delta`; subsequent passes +send gradients into the shared GRU, attention and dynamic gate. This fixes the +legacy mode's zero scalar gate, which blocked all upstream branch gradients at +initialization. + +Track IDs are routing metadata only and never become embeddings. Parameters are +shared across people, the attention uses a real-person mask plus a learned NULL +person, source-clip/max-gap boundaries reset links, and changing padded values +or permuting the current person order cannot change predictions. At the 2 FPS +local grid, this runbook uses `K=4`, covering at most two seconds. Repeated 5 Hz +queries +recompute/read the same causal history; they do not append duplicate visual +observations. + +This immediate V2 keeps pooled visual evidence in the exact base. If it passes, +a later replacement ablation must compare `G+A+person-track` without pooled V; +otherwise the final method cannot claim unconditional pooling is unnecessary. + +## L20 smoke test + +The smoke test uses only training sessions and writes isolated outputs: + +```bash +EPOCHS=1 \ +WARMUP_EPOCHS=0 \ +SESSIONS="4 5" \ +DEV_SESSIONS=5 \ +SCREEN_SESSIONS=5 \ +CHECKPOINT_ROOT=checkpoints/easycom_person_adapter_v2_smoke \ +RESULT_ROOT=results/easycom_person_adapter_v2_smoke \ +CUDA_VISIBLE_DEVICES=0 \ +bash run_easycom_causal_person_adapter_v2.sh all 0 +``` + +Confirm that: + +- the pooled checkpoint is accepted as an exact base; +- `person_adapter_diagnostics.residual_output_weight_norm` becomes nonzero; +- `delta_logit_mean_abs` becomes nonzero; +- shapes and valid counts match the pooled run; +- no Session 1--3 output is produced. + +## Three-seed Session-12 screen + +Run one seed per GPU in separate foreground terminals. Filenames are disjoint, +so the runs may share the same roots: + +```bash +CUDA_VISIBLE_DEVICES=0 bash run_easycom_causal_person_adapter_v2.sh all 0 +CUDA_VISIBLE_DEVICES=1 bash run_easycom_causal_person_adapter_v2.sh all 1 +CUDA_VISIBLE_DEVICES=2 bash run_easycom_causal_person_adapter_v2.sh all 2 +``` + +Summarize after all processes finish: + +```bash +python summarize_local_baselines.py \ + --root results/easycom_causal_person_adapter_v2 \ + --dataset easycom \ + --protocol causal_valid_only \ + --paired-deltas +``` + +Do not run the official test split while choosing this architecture. Once the +method, seed-independent settings and event protocol are frozen, test inference +is explicit: + +```bash +CUDA_VISIBLE_DEVICES=0 \ +bash run_easycom_causal_person_adapter_v2.sh test 0 1 2 +``` + +## Interpretation boundary + +Dense AP/mAP is an architecture screen. The event evaluator still requires +competing-risk censoring before response/yield metrics become the paper +protocol. A publishable visual claim additionally requires correct-person, +past-shifted, wrong-clip, NULL-person and parameter-matched controls. V2 alone +does not predict active-speaker identity or addressee state. diff --git a/docs/easycom_egospeak_style_person_gru.md b/docs/easycom_egospeak_style_person_gru.md new file mode 100644 index 0000000..53e988d --- /dev/null +++ b/docs/easycom_egospeak_style_person_gru.md @@ -0,0 +1,92 @@ +# EasyCom EgoSpeak-style person-track GRU ablation + +## Status and purpose + +Implemented as +`--local-fusion causal-person-adapter-egospeak-gru`; GPU results are not yet +measured. This is a controlled temporal-stem ablation of +`causal-person-adapter-v2`, not a reproduction of EgoSpeak MiniROADA. + +The comparison changes only how each anonymous person's causal visual history +is embedded before the existing GRU/set-attention path: + +| Component | Native V2 | EgoSpeak-style ablation | +|---|---|---| +| track input | `LayerNorm -> Linear -> GELU` | `Linear -> LayerNorm -> ReLU -> Dropout` | +| temporal layer | one-layer shared causal GRU | the same one-layer shared causal GRU | +| GRU output | final hidden state | `ReLU(final hidden state)` | +| all downstream fusion | unchanged | unchanged | + +Both modes use the same raw local visual features, causal track histories, +history length, G+A query, masked multi-person attention, direct local-audio +state, dynamic vector gate, zero-initialized residual output, pooled base +checkpoint, loss, split, optimizer and seed. The actual trainable parameter +count is printed and stored because moving LayerNorm across the projection +causes a small parameter-count difference. With the default 2048D visual input +and 256D person state, the EgoSpeak-style stem has 3,584 fewer affine +LayerNorm parameters than native V2; all Linear and GRU matrix shapes match. + +## Relation to EgoSpeak + +EgoSpeak MiniROADA applies +`Linear -> LayerNorm -> ReLU -> Dropout -> one-layer GRU -> ReLU` to one +**global 5 Hz RGB+audio sequence**. Its EasyCom configuration uses a 128-step +window, 2048D embedding and 1024D hidden state, followed by its own anticipation +expansion and shared classifier. + +This repository uses a clean-room PyTorch implementation of only that temporal +stem and applies it independently to each anonymous **visual person track**. +It intentionally retains V2's 256D screen, two-second person history and +residual forecasting head so the experiment isolates temporal parameterization. +The correct name in results and the paper is **EgoSpeak-style per-track GRU**, +not “EgoSpeak reproduction.” The audited reference is +[`third_party/EgoSpeak/MiniROAD/model/rnn/rnn.py`](../third_party/EgoSpeak/MiniROAD/model/rnn/rnn.py). + +## L20 smoke test + +Use only training sessions and isolated output roots: + +```bash +EPOCHS=1 \ +WARMUP_EPOCHS=0 \ +SESSIONS="4 5" \ +DEV_SESSIONS=5 \ +SCREEN_SESSIONS=5 \ +CHECKPOINT_ROOT=checkpoints/easycom_egospeak_person_gru_smoke \ +RESULT_ROOT=results/easycom_egospeak_person_gru_smoke \ +CUDA_VISIBLE_DEVICES=0 \ +bash run_easycom_causal_person_adapter_v2.sh all 0 +``` + +This runs the matched pooled base, native V2 and EgoSpeak-style ablation. Check +that all three use the same split and base-checkpoint hash, both adapter modes +start from the pooled logits, and their residual-output/delta diagnostics become +nonzero after training. + +## Three-seed development screen + +If matched pooled/V2 checkpoints already exist, run only the new ablation: + +```bash +SUMMARIZE=0 CUDA_VISIBLE_DEVICES=0 bash run_easycom_causal_person_adapter_v2.sh egospeak-gru 0 +SUMMARIZE=0 CUDA_VISIBLE_DEVICES=1 bash run_easycom_causal_person_adapter_v2.sh egospeak-gru 1 +SUMMARIZE=0 CUDA_VISIBLE_DEVICES=2 bash run_easycom_causal_person_adapter_v2.sh egospeak-gru 2 +``` + +Otherwise, run `all` instead of `egospeak-gru` to create each matched pooled +checkpoint and both adapters. Do not run official Sessions 1--3 while selecting +the temporal stem. + +After all jobs finish: + +```bash +python summarize_local_baselines.py \ + --root results/easycom_causal_person_adapter_v2 \ + --dataset easycom \ + --protocol causal_valid_only \ + --paired-deltas +``` + +The relevant paired rows are `egospeak_gru - pooled` and +`egospeak_gru - person_adapter_v2`. Dense AP/mAP is only the architecture +screen; wearer event claims still wait for the frozen competing-risk evaluator. diff --git a/docs/easycom_highrate_person_av.md b/docs/easycom_highrate_person_av.md new file mode 100644 index 0000000..fcdbc81 --- /dev/null +++ b/docs/easycom_highrate_person_av.md @@ -0,0 +1,185 @@ +# Global Omni + shared 25 Hz audio + per-person evidence + +## What this implementation does + +The first version deliberately keeps the two jobs separate: + +1. Existing Qwen2.5-Omni states provide the global scene/conversation context + `G_t` at the verified 5 Hz task grid. +2. MuVAP's frozen CPC frontend encodes the dense session's mixed 16 kHz audio + once into a shared causal 25 Hz stream `U_t`; it does not reset at the + artificial one-minute file boundaries. +3. MuVAP's causal visual track encoder produces one visual state per visible + person every 40 ms. Its audio-visual active-speaker readout contributes one + speaking logit, so a person token is `[visual_state, ASD_score]`; the shared + audio vector is not duplicated inside every person token. +4. At training time `G_t + U_t` queries the variable-sized person set using + masked attention. +5. A zero-initialized residual readout maps the selected person evidence into + a correction of the ordinary `G_t + U_t` logits. With matching shared + weights, the initial model is exactly the audio-only baseline; the person + branch must earn every change during training. +6. The ten-horizon floor head predicts the same 0.2--2.0 second EgoSpeak target. + +In compact form: + +```text +full scene + mixed audio -> frozen Omni -> G_t --------+--> G+U base logits -----+ +mixed audio -> frozen CPC once -> U_t -----------------+ sum + ^ +25 Hz face track -> causal visual state --+ | +U_t + visual state -> ASD speaking score -+-> {P_t,i} -> zero-init residual ----+ + G_t+U_t queries person set +``` + +The person branch uses MuVAP's released **ASD module**, not the unreleased/full +MuVAP multiparty forecasting module. This is a traditional high-rate +active-speaker frontend combined with an Omni global representation. + +## Causality and rate contract + +- MuVAP ASD runs at 25 Hz. Model state `k` consumes the complete interval + `[k/25, (k+1)/25)` and is timestamped by its right edge `(k+1)/25`. +- Every fifth ASD state ends on a 5 Hz task cutoff. Equality is causal here: + the saved state at `0.2s` has consumed samples strictly before `0.2s`. +- MuVAP-style chunks score 10 seconds at a time and prepend at most 2 seconds + of past context. Chunks overlap only into the past. +- A source-clip boundary or missing face observation resets a local face track. + The shared audio stream remains continuous on the same concatenated dense + timeline used by the targets and global Omni features. +- The released MuVAP data loader's whole-recording RMS normalization is not + reused: at an early cutoff that statistic would depend on future audio. The + extractor uses fixed float-PCM scaling with pointwise clipping only. Record + this causal deployment/domain-shift choice when reporting the experiment. +- Speech labels, wearer targets and future frames are never extractor inputs. +- Provided boxes isolate the representation/fusion question. They are not an + online detector result and must be described that way in the paper. + +## External dependency + +The audited MuVAP checkout and published ASD-CPC release are vendored under +`third_party/MuVAP`, including an immutable `.upstream-revision` record. For a +fresh checkout of this project, no separate MuVAP clone or weight download is +needed. The following commands document how that vendored snapshot was +obtained and are only needed when intentionally refreshing it: + +```bash +git clone https://github.com/Haotian-Qi/MuVAP.git third_party/MuVAP +git -C third_party/MuVAP checkout a1ef9c8 +python third_party/MuVAP/tools/fetch_frontends.py cpc + +huggingface-cli download Haotian-Qi/MuVAP \ + --include 'asd-cpc/*' \ + --local-dir third_party/MuVAP/weights +``` + +The extractor imports the pinned vendored checkout directly; it does not need +an editable MuVAP install. Check that `torch`, `torchaudio`, `einops` and +`yaml` import in the existing Qwen container instead of upgrading its full +dependency stack. + +The extractor checks the source revision and hashes the selected release +weights into every output contract. If a newer MuVAP revision is intentionally +tested, give its exact commit with `--muvap-revision` and use a new output root. + +## EasyCom smoke extraction + +Run a 10-second partial smoke first. A partial sidecar is intentionally +rejected by training. + +```bash +export DATASET_ROOT=/data1/yfzx/wangx/Datasets/easycom/extracted_easycom +export FEATURE_ROOT=/data1/yfzx/wangx/Datasets/easycom/omni_features/open_stream_2fps_8s_last4_stack +export HIGHRATE_LOCAL_ROOT=/data1/yfzx/wangx/Datasets/easycom/omni_features/local_muvap_asd_25hz_v2 +export ASD_WEIGHTS=third_party/MuVAP/weights/asd-cpc + +CUDA_VISIBLE_DEVICES=0 python extract_highrate_person_av_features.py \ + --dataset easycom \ + --data-root "$DATASET_ROOT" \ + --feature-root "$FEATURE_ROOT" \ + --output-root features/highrate_person_av_smoke \ + --samples 12 \ + --muvap-root third_party/MuVAP \ + --asd-weights "$ASD_WEIGHTS" \ + --max-steps 50 \ + --device cuda:0 \ + --dtype bfloat16 +``` + +Audit the result: + +```bash +python - <<'PY' +import torch + +p = torch.load("features/highrate_person_av_smoke/Session_12.pt", map_location="cpu", weights_only=False) +print("format:", p["format"]) +print("partial:", p["partial"]) +print("times:", p["times"][:3].tolist(), "...", p["times"][-1].item()) +print("person features:", tuple(p["persons"]["features"].shape)) +print("shared audio:", tuple(p["audio"].shape)) +print("frame offsets:", tuple(p["persons"]["frame_offsets"].shape)) +print("person-row coverage:", float(p["visual_valid"].float().mean())) +print("uses speech labels:", p["config"]["uses_speech_labels"]) +print("causal policy:", p["config"]["causal_policy"]) +PY +``` + +Expected invariants are `causal-highrate-person-av-v2`, `partial=True`, times +`[0.2, 0.4, ...]`, shared audio width 256, person width 257 (`256D causal +visual state + 1 AV activity logit`), and `uses_speech_labels=False`. + +## Full EasyCom extraction + +Run one command per GPU/terminal. The shard assignments are disjoint and the +destination contains real files, not symbolic links: + +```bash +CUDA_VISIBLE_DEVICES=0 python extract_highrate_person_av_features.py \ + --dataset easycom --data-root "$DATASET_ROOT" --feature-root "$FEATURE_ROOT" \ + --output-root "$HIGHRATE_LOCAL_ROOT" --samples 1 2 3 4 5 6 7 8 9 10 11 12 \ + --num-shards 8 --shard-index 0 --muvap-root third_party/MuVAP \ + --asd-weights "$ASD_WEIGHTS" --device cuda:0 --dtype bfloat16 +``` + +Repeat with physical GPU and shard index `1` through `7`. Do not add +`--overwrite` unless the exact file has been inspected and a replacement is +intentional. + +## First matched training screen + +The screen trains on Sessions 4--11 and uses Session 12 for checkpoint +selection. It does not inspect Sessions 1--3. Run different seeds on +different GPUs: + +```bash +CUDA_VISIBLE_DEVICES=0 bash run_easycom_highrate_person_av.sh all 0 +CUDA_VISIBLE_DEVICES=1 bash run_easycom_highrate_person_av.sh all 1 +CUDA_VISIBLE_DEVICES=2 bash run_easycom_highrate_person_av.sh all 2 +``` + +This compares four matched conditions: + +- `global`: frozen Omni global state `G` only; +- `qwen_audio`: existing `G+A_qwen` strong baseline; +- `traditional_audio`: global Omni `G` plus one shared CPC stream `U`; +- `highrate`: load that seed's trained `G+U` checkpoint, freeze the exact base, + and learn only the attention/residual over frozen 25 Hz causal ASD person + states. This first screen therefore asks whether the new person evidence + adds information, without allowing the base to drift. + +The `all` command runs `traditional_audio` before `highrate`, so the initializer +exists. If running separately, run `traditional_audio` first for the same seed. + +Only after selecting the architecture on Session 12 should the fixed +checkpoints be run on the official test set: + +```bash +CUDA_VISIBLE_DEVICES=0 bash run_easycom_highrate_person_av.sh test 0 +CUDA_VISIBLE_DEVICES=1 bash run_easycom_highrate_person_av.sh test 1 +CUDA_VISIBLE_DEVICES=2 bash run_easycom_highrate_person_av.sh test 2 +``` + +Dense AP/mAP is a screening/comparability result. The event evaluator still +needs competing-risk censoring before claims about false entries, responses or +yields are final. diff --git a/docs/easycom_muvap_stage3_baseline.md b/docs/easycom_muvap_stage3_baseline.md new file mode 100644 index 0000000..57ad5f2 --- /dev/null +++ b/docs/easycom_muvap_stage3_baseline.md @@ -0,0 +1,211 @@ +# MuVAP-only Stage-3 reconstruction on EasyCom/EgoSpeak targets + +## Question and claim boundary + +This baseline asks whether MuVAP's specialized acoustic turn-taking and +person-grounded active-speaker representations can solve the EgoSpeak future +wearer-state task without an Omni representation. + +The official MuVAP repository currently releases VAP and ASD training code and +weights, but not its multiparty Stage-3 implementation. Therefore this code is +a clean-room reconstruction from MuVAP Sections 4.3 and 5.3, not an execution +of upstream Stage-3 code. + +The frozen inputs are: + +- `Z_VAP`: the published 25 Hz mono, role-relative CPC VAP embedding; +- `Z_ASD_i`: the published causal ASD fused embedding for visible track `i`. + +The matched control uses the identical projections, frame-attention block, +temporal Transformer and head, but masks every observed `Z_ASD_i` token. It is +named `vap_only`; the full condition is named `stage3`. Comparing these two +conditions isolates aligned person evidence from a simple capacity increase. + +The trainable path is: + +```text +frozen mixed audio -> frozen role-VAP -> Z_VAP(t) ---- query ----+ + frame attention +frozen mixed audio + provided face track i | + -> frozen ASD -> {Z_ASD_i(t)} -- key/value+ + -> causal temporal Transformer + -> EgoSpeak 10-horizon activity head +``` + +Independent Linear-LayerNorm projections map both inputs to 256 dimensions. +At every 25 Hz frame, `Z_VAP` queries the variable person set. A one-layer, +four-head causal temporal Transformer then processes the last 8 seconds and +the state at each 5 Hz right edge predicts the ten EgoSpeak horizons from 0.2 +to 2.0 seconds. + +Differences from the unpublished original Stage 3 must remain explicit: + +- supervision is EgoSpeak's `background/wearer/others` future activity, not + MuVAP's GVAP plus per-visible-person SVAP objectives; +- missing tracks are masked rather than padded as permanent known identities; +- provided EasyCom boxes isolate representation quality and are not an online + face detector result; +- the first strict-causal experiment does not apply whole-recording RMS + normalization because that would expose a future-dependent statistic; +- the temporal memory is fixed to 8 seconds for a reproducible streaming + contract. + +## Required releases + +The repository vendors the pinned MuVAP source, CPC frontend and the two small +published releases consumed here. After pulling the experiment commit, verify: + +```bash +test -s third_party/MuVAP/weights/asd-cpc/weights.pt +test -s third_party/MuVAP/weights/vap-role-cpc/weights.pt +test -s third_party/MuVAP/models/assets/checkpoints/cpc/60k_epoch4-d0f474de.pt +``` + +Expected paths: + +```text +third_party/MuVAP/weights/asd-cpc/{config.yaml,weights.pt} +third_party/MuVAP/weights/vap-role-cpc/{config.yaml,weights.pt} +third_party/MuVAP/models/assets/checkpoints/cpc/60k_epoch4-d0f474de.pt +``` + +On the L20 environment, run the CPU model/data-contract tests before touching +the full dataset: + +```bash +python -m unittest test_muvap_egospeak_stage3 +``` + +## Extraction smoke test + +The reference root supplies only the already-audited 5 Hz timestamps and raw +source inventory. Its Omni `states` are never copied or loaded by the training +code. + +```bash +export DATASET_ROOT=/data1/yfzx/wangx/Datasets/easycom/extracted_easycom +export REFERENCE_ROOT=/data1/yfzx/wangx/Datasets/easycom/omni_features/open_stream_2fps_8s_last4_stack +export MUVAP_STAGE3_ROOT=/data1/yfzx/wangx/Datasets/easycom/omni_features/muvap_stage3_cpc_25hz_v1 + +CUDA_VISIBLE_DEVICES=0 python extract_muvap_egospeak_features.py \ + --dataset easycom \ + --data-root "$DATASET_ROOT" \ + --reference-feature-root "$REFERENCE_ROOT" \ + --output-root features/muvap_stage3_smoke \ + --samples 12 \ + --muvap-root third_party/MuVAP \ + --vap-weights third_party/MuVAP/weights/vap-role-cpc \ + --asd-weights third_party/MuVAP/weights/asd-cpc \ + --max-steps 50 \ + --device cuda:0 \ + --dtype bfloat16 +``` + +The smoke should report 250 MuVAP frames for 50 target states, finite +`vap=[250,D_vap]`, a ragged person tensor with width `D_asd`, exact task frame +indices `[4,9,...,249]`, and `partial=True`. Training intentionally rejects +this partial file. + +## Full extraction + +Use a fresh output directory. Eight terminals can run disjoint real-file +shards; no symbolic links are used. Replace both `GPU` and `SHARD` with the +same value from 0 through 7: + +```bash +GPU=0 +SHARD=0 +CUDA_VISIBLE_DEVICES="$GPU" python extract_muvap_egospeak_features.py \ + --dataset easycom \ + --data-root "$DATASET_ROOT" \ + --reference-feature-root "$REFERENCE_ROOT" \ + --output-root "$MUVAP_STAGE3_ROOT" \ + --samples 1 2 3 4 5 6 7 8 9 10 11 12 \ + --num-shards 8 \ + --shard-index "$SHARD" \ + --muvap-root third_party/MuVAP \ + --vap-weights third_party/MuVAP/weights/vap-role-cpc \ + --asd-weights third_party/MuVAP/weights/asd-cpc \ + --device cuda:0 \ + --dtype bfloat16 +``` + +Do not add `--overwrite` without inspecting the exact destination file. + +## CPU contract audit + +```bash +python - <<'PY' +from pathlib import Path +from muvap_egospeak_stage3 import load_muvap_session + +feature_root = Path("/data1/yfzx/wangx/Datasets/easycom/omni_features/muvap_stage3_cpc_25hz_v1") +target_root = Path("/data1/yfzx/wangx/Datasets/easycom/easycom_target_frame") +contract = None +for sid in range(1, 13): + data = load_muvap_session(feature_root, target_root, sid, expected_contract=contract) + contract = data.contract + print( + data.name, + "task=", len(data.targets), + "25Hz=", len(data.vap), + "VAP=", data.vap_dim, + "ASD=", data.asd_dim, + "max_people=", data.people.shape[1], + ) +PY +``` + +## One-epoch L20 smoke + +Use only training sessions. This verifies forward/backward/checkpoint and does +not touch Sessions 1--3: + +```bash +CUDA_VISIBLE_DEVICES=0 python train_muvap_egospeak_stage3.py \ + --feature-root "$MUVAP_STAGE3_ROOT" \ + --target-root /data1/yfzx/wangx/Datasets/easycom/easycom_target_frame \ + --output checkpoints/muvap_stage3_smoke.pth \ + --sessions 4 5 \ + --dev-sessions 5 \ + --epochs 1 \ + --batch-size 4 \ + --history-seconds 8 \ + --max-train-examples 64 \ + --device cuda:0 +``` + +Use a fresh checkpoint name if repeating the smoke. + +## Three-seed experiment + +First screen the full Stage-3 reconstruction and its capacity-matched VAP-only +control on Session 12. Run one seed per GPU/terminal; the examples below use +six GPUs concurrently: + +```bash +CUDA_VISIBLE_DEVICES=0 PERSON_EVIDENCE=aligned bash run_easycom_muvap_stage3.sh all 0 +CUDA_VISIBLE_DEVICES=1 PERSON_EVIDENCE=aligned bash run_easycom_muvap_stage3.sh all 1 +CUDA_VISIBLE_DEVICES=2 PERSON_EVIDENCE=aligned bash run_easycom_muvap_stage3.sh all 2 +CUDA_VISIBLE_DEVICES=3 PERSON_EVIDENCE=none bash run_easycom_muvap_stage3.sh all 0 +CUDA_VISIBLE_DEVICES=4 PERSON_EVIDENCE=none bash run_easycom_muvap_stage3.sh all 1 +CUDA_VISIBLE_DEVICES=5 PERSON_EVIDENCE=none bash run_easycom_muvap_stage3.sh all 2 +``` + +`all` trains on Sessions 4--11, selects the best of five epochs on Session 12, +and writes the Session-12 screen. It does not touch the official test set. +After fixing the architecture and hyperparameters from training/Session 12, +run the official test once: + +```bash +CUDA_VISIBLE_DEVICES=0 PERSON_EVIDENCE=aligned bash run_easycom_muvap_stage3.sh test 0 +CUDA_VISIBLE_DEVICES=1 PERSON_EVIDENCE=aligned bash run_easycom_muvap_stage3.sh test 1 +CUDA_VISIBLE_DEVICES=2 PERSON_EVIDENCE=aligned bash run_easycom_muvap_stage3.sh test 2 +CUDA_VISIBLE_DEVICES=3 PERSON_EVIDENCE=none bash run_easycom_muvap_stage3.sh test 0 +CUDA_VISIBLE_DEVICES=4 PERSON_EVIDENCE=none bash run_easycom_muvap_stage3.sh test 1 +CUDA_VISIBLE_DEVICES=5 PERSON_EVIDENCE=none bash run_easycom_muvap_stage3.sh test 2 +``` + +Dense AP/mAP is currently the primary screening output. Floor-event results +remain timing-derived proxies until competing-risk censoring is frozen in +`evaluate_floor_state_changes.py`. diff --git a/docs/progressive_feature_fusion_route_2026-09-02.md b/docs/progressive_feature_fusion_route_2026-09-02.md index b1d6826..31acffa 100644 --- a/docs/progressive_feature_fusion_route_2026-09-02.md +++ b/docs/progressive_feature_fusion_route_2026-09-02.md @@ -231,6 +231,29 @@ ten-horizon mean mAP 配对增量仅为 `+0.01306 +/- 0.02368` 个百分点， residual 最后输出层零初始化，并在 dev 上记录 delta-logit、概率变化和非零梯度；详见 [`结果记录`](../results/easycom_track_gru_residual_4ep_2026-09-07/README.md)。 +修正版已实现为 `causal-person-adapter-v2`，但尚无 GPU 结果。它仍从 matched pooled +`G+A+V` checkpoint 精确初始化并默认冻结 base；每条当前 track 经共享因果 GRU 后由 G+A +query 读取，再与直接 local audio 构成局部 state。模型使用逐样本向量 gate，并仅把最终 +residual-logit projection 的 weight/bias 置零。这样初始 logits 与 pooled baseline +逐元素相同，第一次更新 residual output，后续更新才可进入 GRU、attention 和动态 gate。 +运行契约与命令见 +[`easycom_causal_person_adapter_v2.md`](easycom_causal_person_adapter_v2.md)。它保留 pooled V， +因此若结果通过，还必须补一个移除 pooled V 的 replacement ablation，才能决定最终方法。 + +为减少“GRU 具体写法”这一实验变量，另实现 +`causal-person-adapter-egospeak-gru`：保留 V2 的所有下游融合与优化设计，只把逐人输入改成 +EgoSpeak MiniROADA 风格的 `Linear-LN-ReLU-Dropout-GRU-ReLU`。这不是 EgoSpeak 的全局 +5 Hz RGB+audio GRU 复现，而是人物 track 上的受控 temporal-stem 消融；运行说明见 +[`easycom_egospeak_style_person_gru.md`](easycom_egospeak_style_person_gru.md)。 + +为直接检验“低频 Qwen crop 是否遗漏唇音动态”，另实现独立的 +`highrate-person-av-attention` 分支。它不把同一张 2 FPS 人脸重复送进 GRU：MuVAP CPC 在 +25 Hz 上只编码一次 mixed audio 得到共享 `U`，每条人脸轨迹另产生 causal visual state 和 +一个 AV speaking score，只在每个 5 Hz 右边界保存。随后由 `G+U` query 做集合注意力，并 +用零初始化 residual 修正已经训练且冻结的 matched `G+U` logits。它首先回答传统高频 +audio/person 证据能否弥补全局 Omni，而不是检验 identity continuity。实现和命令见 +[`easycom_highrate_person_av.md`](easycom_highrate_person_av.md)。 + ## 8. F4：Person-conditioned audio/Omni history F4 让每个匿名人物分别读取已有的 5 Hz 因果历史： diff --git a/easycom_head.py b/easycom_head.py index 3b04441..0683e1e 100644 --- a/easycom_head.py +++ b/easycom_head.py @@ -66,8 +66,18 @@ CURRENT_FRAME_PERSON_SET_FUSION_MODES = ( "oracle-frame-set-attention", "audio-person-frame-set", "gated-person-state-adapter", + "highrate-person-av-attention", +) +HIGHRATE_PERSON_AV_FUSION_MODES = ("highrate-person-av-attention",) +TRACK_PERSON_FUSION_MODES = ( + "track-gru-residual", + "causal-person-adapter-v2", + "causal-person-adapter-egospeak-gru", +) +PERSON_ADAPTER_V2_MODES = ( + "causal-person-adapter-v2", + "causal-person-adapter-egospeak-gru", ) -TRACK_PERSON_FUSION_MODES = ("track-gru-residual",) PERSON_SET_FUSION_MODES = ( *CURRENT_FRAME_PERSON_SET_FUSION_MODES, *TRACK_PERSON_FUSION_MODES, @@ -79,7 +89,16 @@ LOCAL_FUSION_SCOPES = { "oracle-frame-set-attention": "causal-oracle-visible-person-speaking-history", "audio-person-frame-set": "f5-screening-current-frame-no-person-history", "gated-person-state-adapter": "f6-screening-post-thinker-not-internal", + "highrate-person-av-attention": ( + "global-omni-traditional-audio-base-plus-zero-init-25hz-person-residual" + ), "track-gru-residual": "causal-anonymous-track-gru-pooled-baseline-residual", + "causal-person-adapter-v2": ( + "causal-anonymous-track-gru-dynamic-gated-logit-residual" + ), + "causal-person-adapter-egospeak-gru": ( + "causal-anonymous-track-egospeak-style-gru-dynamic-gated-logit-residual" + ), } DatasetName = Literal["easycom", "ego4d"] @@ -211,15 +230,36 @@ class TurnPredictionHead(nn.Module): self.audio_person_gate: nn.Parameter | None = None self.person_state_adapter: nn.Module | None = None self.person_state_adapter_gate: nn.Parameter | None = None + self.person_track_input_embedding: nn.Module | None = None self.person_track_gru: nn.GRU | None = None + self.person_track_output_activation: nn.Module | None = None self.person_track_delta: nn.Module | None = None self.person_track_residual_gate: nn.Parameter | None = None + self.person_track_local_state: nn.Module | None = None + self.person_track_dynamic_gate: nn.Module | None = None + self.person_track_residual_output: nn.Linear | None = None + self.highrate_person_residual: nn.Linear | None = None + self._last_person_adapter_diagnostics: dict[str, float] = {} if local_fusion in PERSON_SET_FUSION_MODES: - self.person_input_projection = nn.Sequential( - nn.LayerNorm(local_visual_dim), - nn.Linear(local_visual_dim, person_attention_dim), - nn.GELU(), - ) + if local_fusion == "causal-person-adapter-egospeak-gru": + # Clean-room adaptation of EgoSpeak MiniROADA's temporal stem + # (MiniROAD/model/rnn/rnn.py): Linear -> LayerNorm -> ReLU -> + # Dropout -> one-layer GRU -> ReLU. EgoSpeak applies it to a + # global RGB+audio sequence; this ablation applies the same + # layer order independently to each anonymous visual track. + self.person_track_input_embedding = nn.Sequential( + nn.Linear(local_visual_dim, person_attention_dim), + nn.LayerNorm(person_attention_dim), + nn.ReLU(), + nn.Dropout(dropout), + ) + self.person_track_output_activation = nn.ReLU() + else: + self.person_input_projection = nn.Sequential( + nn.LayerNorm(local_visual_dim), + nn.Linear(local_visual_dim, person_attention_dim), + nn.GELU(), + ) self.person_state_query = nn.Sequential( nn.LayerNorm(hidden_size), nn.Linear(hidden_size, person_attention_dim), @@ -263,25 +303,75 @@ class TurnPredictionHead(nn.Module): nn.Linear(person_attention_dim, hidden_size), ) self.person_state_adapter_gate = nn.Parameter(torch.zeros(())) - elif local_fusion == "track-gru-residual": + elif local_fusion in TRACK_PERSON_FUSION_MODES: # Apply one shared causal GRU independently to every anonymous # face track. The regular pooled G+A+V head remains the exact - # base path; this branch can only add a zero-gated correction. + # base path for every track experiment. self.person_track_gru = nn.GRU( person_attention_dim, person_attention_dim, num_layers=1, batch_first=True, ) - delta_input = hidden_size // 2 + person_attention_dim + 1 - self.person_track_delta = nn.Sequential( - nn.LayerNorm(delta_input), - nn.Linear(delta_input, hidden_size // 2), - nn.GELU(), - nn.Dropout(dropout), - nn.Linear(hidden_size // 2, num_horizons * num_classes), + if local_fusion == "track-gru-residual": + # Legacy screen: the scalar zero gate is intentionally + # preserved so old checkpoints and the measured null result + # remain reproducible. + delta_input = hidden_size // 2 + person_attention_dim + 1 + self.person_track_delta = nn.Sequential( + nn.LayerNorm(delta_input), + nn.Linear(delta_input, hidden_size // 2), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(hidden_size // 2, num_horizons * num_classes), + ) + self.person_track_residual_gate = nn.Parameter(torch.zeros(())) + elif local_fusion in PERSON_ADAPTER_V2_MODES: + # V2: person-track context and direct local audio form a + # floor-evidence state. A sample-dependent vector gate + # controls its correction to the exact pooled baseline. + # Zero-initialize only the final logit projection: unlike a + # zero scalar multiplier, it starts exactly at the baseline + # without permanently starving the branch of gradients. + base_feature_dim = hidden_size // 2 + projected_audio_dim = local_projection_dim or local_audio_dim + local_state_input = ( + person_attention_dim + projected_audio_dim + 2 + ) + self.person_track_local_state = nn.Sequential( + nn.LayerNorm(local_state_input), + nn.Linear(local_state_input, base_feature_dim), + nn.GELU(), + nn.Dropout(dropout), + ) + self.person_track_dynamic_gate = nn.Sequential( + nn.LayerNorm(base_feature_dim * 2), + nn.Linear(base_feature_dim * 2, base_feature_dim), + ) + self.person_track_residual_output = nn.Linear( + base_feature_dim, + num_horizons * num_classes, + ) + nn.init.zeros_(self.person_track_residual_output.weight) + nn.init.zeros_(self.person_track_residual_output.bias) + else: + raise AssertionError( + f"Unhandled track fusion mode: {local_fusion}" + ) + elif local_fusion in HIGHRATE_PERSON_AV_FUSION_MODES: + # Keep global Omni + shared traditional audio as the complete + # base path. The high-rate ASD + # person set contributes only a zero-initialized logit + # residual, so step zero is exactly the same function as the + # audio-only head with matching shared weights. This gives a + # controlled answer to whether person evidence is useful. + residual_input = hidden_size // 2 + person_attention_dim + 1 + self.highrate_person_residual = nn.Linear( + residual_input, + num_horizons * num_classes, ) - self.person_track_residual_gate = nn.Parameter(torch.zeros(())) + nn.init.zeros_(self.highrate_person_residual.weight) + nn.init.zeros_(self.highrate_person_residual.bias) else: fused_size += person_attention_dim + 1 self.net = nn.Sequential( @@ -413,17 +503,20 @@ class TurnPredictionHead(nn.Module): local: torch.Tensor, person_histories: torch.Tensor | None, history_mask: torch.Tensor | None, - ) -> tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Encode each causal track with a shared GRU, then query the set.""" if states.ndim != 2: raise ValueError("Track GRU fusion expects batched [B,D] states") if local.shape != (len(states), self.local_input_size): raise ValueError("Track GRU fusion requires pooled [visual,audio,valid] input") + if person_histories is None or person_histories.ndim != 4: + raise ValueError( + "Track GRU fusion requires person histories " + "[B,P,person_history_frames,visual_dim]" + ) if ( - person_histories is None - or person_histories.ndim != 4 - or person_histories.shape[0] != len(states) + person_histories.shape[0] != len(states) or person_histories.shape[2] != self.person_history_frames or person_histories.shape[3] != self.local_visual_dim ): @@ -431,49 +524,22 @@ class TurnPredictionHead(nn.Module): "Track GRU fusion requires person histories " "[B,P,person_history_frames,visual_dim]" ) - if ( - history_mask is None - or history_mask.dtype != torch.bool - or history_mask.shape != person_histories.shape[:3] - ): - raise ValueError("Track GRU fusion requires a boolean history mask [B,P,L]") - if not torch.isfinite(person_histories).all(): - raise ValueError("Person track histories contain NaN/Inf") - if bool((history_mask[:, :, 1:] & ~history_mask[:, :, :-1]).any()): - raise ValueError("Person track histories must be left aligned") - assert self.person_input_projection is not None assert self.person_state_query is not None assert self.person_audio_query is not None assert self.person_attention is not None assert self.null_person is not None - assert self.person_track_gru is not None - - batch, people, length, _ = person_histories.shape - person_valid = history_mask.any(dim=-1) - if people: - projected = self.person_input_projection(person_histories) - flat = projected.reshape(batch * people, length, -1) - lengths = history_mask.sum(dim=-1).reshape(-1) - safe_lengths = lengths.clamp_min(1).to(device="cpu") - packed = nn.utils.rnn.pack_padded_sequence( - flat, - safe_lengths, - batch_first=True, - enforce_sorted=False, - ) - _, hidden = self.person_track_gru(packed) - track_states = hidden[-1].reshape(batch, people, -1) - track_states = track_states * person_valid[..., None] - else: - track_states = states.new_zeros( - (batch, 0, self.person_attention_dim) - ) + track_states, person_valid = self._encode_person_track_histories( + person_histories, + history_mask, + ) + batch = len(states) visual_width = self.local_visual_dim + 1 audio_values = local[:, visual_width : visual_width + self.local_audio_dim] audio_valid = local[:, visual_width + self.local_audio_dim :] if audio_valid.shape != (batch, 1): raise AssertionError("Unexpected pooled AV local feature layout") + audio_projected = self.local_projections["audio"](audio_values) * audio_valid query = self.person_state_query(states) query = query + self.person_audio_query(audio_values) * audio_valid memory = torch.cat( @@ -494,7 +560,81 @@ class TurnPredictionHead(nn.Module): need_weights=False, ) any_person = person_valid.any(dim=1, keepdim=True).to(states.dtype) - return context[:, 0, :] * any_person, any_person + return ( + context[:, 0, :] * any_person, + any_person, + audio_projected, + audio_valid, + ) + + def _encode_person_track_histories( + self, + person_histories: torch.Tensor, + history_mask: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return one shared-GRU state per causal anonymous person track.""" + + if ( + person_histories.ndim != 4 + or person_histories.shape[2] != self.person_history_frames + or person_histories.shape[3] != self.local_visual_dim + ): + raise ValueError( + "Track GRU fusion requires person histories " + "[B,P,person_history_frames,visual_dim]" + ) + if ( + history_mask is None + or history_mask.dtype != torch.bool + or history_mask.shape != person_histories.shape[:3] + ): + raise ValueError( + "Track GRU fusion requires a boolean history mask [B,P,L]" + ) + if not torch.isfinite(person_histories).all(): + raise ValueError("Person track histories contain NaN/Inf") + if bool((history_mask[:, :, 1:] & ~history_mask[:, :, :-1]).any()): + raise ValueError("Person track histories must be left aligned") + assert self.person_track_gru is not None + assert ( + self.person_input_projection is not None + or self.person_track_input_embedding is not None + ) + + batch, people, length, _ = person_histories.shape + person_valid = history_mask.any(dim=-1) + if not people: + return ( + person_histories.new_zeros( + (batch, 0, self.person_attention_dim) + ), + person_valid, + ) + + if self.person_track_input_embedding is not None: + projected = self.person_track_input_embedding(person_histories) + else: + assert self.person_input_projection is not None + projected = self.person_input_projection(person_histories) + flat = projected.reshape(batch * people, length, -1) + lengths = history_mask.sum(dim=-1).reshape(-1) + safe_lengths = lengths.clamp_min(1).to(device="cpu") + packed = nn.utils.rnn.pack_padded_sequence( + flat, + safe_lengths, + batch_first=True, + enforce_sorted=False, + ) + initial_hidden = flat.new_zeros( + self.person_track_gru.num_layers, + batch * people, + self.person_track_gru.hidden_size, + ) + _, hidden = self.person_track_gru(packed, initial_hidden) + track_states = hidden[-1].reshape(batch, people, -1) + if self.person_track_output_activation is not None: + track_states = self.person_track_output_activation(track_states) + return track_states * person_valid[..., None], person_valid def _shared_features( self, @@ -506,7 +646,23 @@ class TurnPredictionHead(nn.Module): """Return the activity MLP representation shared by both readouts.""" states = self._fuse_layers(states) - if self.local_fusion in CURRENT_FRAME_PERSON_SET_FUSION_MODES: + if self.local_fusion in HIGHRATE_PERSON_AV_FUSION_MODES: + # Do not inject person values into the base representation. The + # person set is selected separately in _activity_logits and can + # only alter the prediction through its explicit residual. + if local is None or local.shape != ( + len(states), + self.local_audio_dim + 1, + ): + raise ValueError( + "High-rate person AV fusion requires aligned " + "[audio,valid] local features" + ) + audio_values = local[:, : self.local_audio_dim] + audio_valid = local[:, self.local_audio_dim :] + audio = self.local_projections["audio"](audio_values) * audio_valid + states = torch.cat((states, audio, audio_valid), dim=-1) + elif self.local_fusion in CURRENT_FRAME_PERSON_SET_FUSION_MODES: if local is None: raise ValueError("Person-set fusion requires local audio features") audio, audio_valid, person_context, person_valid = self._frame_set_context( @@ -569,26 +725,118 @@ class TurnPredictionHead(nn.Module): person_mask, ) flat_logits = self.net[-1](features) - if self.local_fusion == "track-gru-residual": + if self.local_fusion in HIGHRATE_PERSON_AV_FUSION_MODES: if local is None: - raise ValueError("Track GRU residual requires pooled local AV") - assert self.person_track_delta is not None - assert self.person_track_residual_gate is not None + raise ValueError("High-rate person AV fusion requires local audio") fused_states = self._fuse_layers(states) - context, person_valid = self._track_gru_context( + _, _, context, person_valid = self._frame_set_context( fused_states, local, person_features, person_mask, ) - delta = self.person_track_delta( + assert self.highrate_person_residual is not None + delta = self.highrate_person_residual( torch.cat((features, context, person_valid), dim=-1) ) - flat_logits = flat_logits + ( - torch.tanh(self.person_track_residual_gate) - * delta - * person_valid + flat_logits = flat_logits + delta * person_valid + with torch.no_grad(): + valid_rows = person_valid[:, 0].bool() + diagnostics = { + "valid_person_rate": float(person_valid.mean().cpu()), + "residual_output_weight_norm": float( + self.highrate_person_residual.weight.norm().cpu() + ), + "residual_output_bias_norm": float( + self.highrate_person_residual.bias.norm().cpu() + ), + } + if bool(valid_rows.any()): + valid_delta = delta[valid_rows] + diagnostics.update( + { + "delta_logit_mean_abs": float( + valid_delta.abs().mean().cpu() + ), + "delta_logit_max_abs": float( + valid_delta.abs().max().cpu() + ), + } + ) + self._last_person_adapter_diagnostics = diagnostics + elif self.local_fusion in TRACK_PERSON_FUSION_MODES: + if local is None: + raise ValueError("Person-track fusion requires pooled local AV") + fused_states = self._fuse_layers(states) + context, person_valid, audio, audio_valid = self._track_gru_context( + fused_states, + local, + person_features, + person_mask, ) + if self.local_fusion == "track-gru-residual": + assert self.person_track_delta is not None + assert self.person_track_residual_gate is not None + delta = self.person_track_delta( + torch.cat((features, context, person_valid), dim=-1) + ) + flat_logits = flat_logits + ( + torch.tanh(self.person_track_residual_gate) + * delta + * person_valid + ) + elif self.local_fusion in PERSON_ADAPTER_V2_MODES: + assert self.person_track_local_state is not None + assert self.person_track_dynamic_gate is not None + assert self.person_track_residual_output is not None + local_state = self.person_track_local_state( + torch.cat( + (context, audio, person_valid, audio_valid), + dim=-1, + ) + ) + gate = torch.sigmoid( + self.person_track_dynamic_gate( + torch.cat((features, local_state), dim=-1) + ) + ) + delta = self.person_track_residual_output(gate * local_state) + flat_logits = flat_logits + delta * person_valid + with torch.no_grad(): + valid_rows = person_valid[:, 0].bool() + diagnostics = { + "valid_person_rate": float(person_valid.mean().cpu()), + "residual_output_weight_norm": float( + self.person_track_residual_output.weight.norm().cpu() + ), + "residual_output_bias_norm": float( + self.person_track_residual_output.bias.norm().cpu() + ), + } + if bool(valid_rows.any()): + valid_gate = gate[valid_rows] + valid_delta = delta[valid_rows] + diagnostics.update( + { + "dynamic_gate_mean": float(valid_gate.mean().cpu()), + "dynamic_gate_std": float( + valid_gate.std(unbiased=False).cpu() + ), + "dynamic_gate_min": float(valid_gate.min().cpu()), + "dynamic_gate_max": float(valid_gate.max().cpu()), + "delta_logit_mean_abs": float( + valid_delta.abs().mean().cpu() + ), + "delta_logit_max_abs": float( + valid_delta.abs().max().cpu() + ), + } + ) + self._last_person_adapter_diagnostics = diagnostics + else: + raise AssertionError( + f"Unhandled track fusion mode: {self.local_fusion}" + ) return features, flat_logits def forward( @@ -660,6 +908,35 @@ class TurnPredictionHead(nn.Module): ) return result + def person_adapter_diagnostics(self) -> dict[str, float]: + """Return V2's latest dynamic/residual diagnostics for audit logs.""" + + if self.local_fusion not in ( + *PERSON_ADAPTER_V2_MODES, + *HIGHRATE_PERSON_AV_FUSION_MODES, + ): + return {} + result = dict(self._last_person_adapter_diagnostics) + if self.person_track_residual_output is not None: + result.setdefault( + "residual_output_weight_norm", + float(self.person_track_residual_output.weight.detach().norm().cpu()), + ) + result.setdefault( + "residual_output_bias_norm", + float(self.person_track_residual_output.bias.detach().norm().cpu()), + ) + if self.highrate_person_residual is not None: + result.setdefault( + "residual_output_weight_norm", + float(self.highrate_person_residual.weight.detach().norm().cpu()), + ) + result.setdefault( + "residual_output_bias_norm", + float(self.highrate_person_residual.bias.detach().norm().cpu()), + ) + return result + def _validate_loss_inputs( logits: torch.Tensor, @@ -1297,16 +1574,27 @@ def load_session_data( raise ValueError( "Person-set fusion requires --local-modalities av" ) - from local_turn_features import load_frame_set_features + if local_fusion in HIGHRATE_PERSON_AV_FUSION_MODES: + from local_turn_features import load_highrate_person_av_features - local, local_people, contract = load_frame_set_features( - local_path, - payload, - sample_id=name, - dataset=dataset, - expected_contract=local_contract, - oracle_analysis=local_fusion == "oracle-frame-set-attention", - ) + local, local_people, contract = load_highrate_person_av_features( + local_path, + payload, + sample_id=name, + dataset=dataset, + expected_contract=local_contract, + ) + else: + from local_turn_features import load_frame_set_features + + local, local_people, contract = load_frame_set_features( + local_path, + payload, + sample_id=name, + dataset=dataset, + expected_contract=local_contract, + oracle_analysis=local_fusion == "oracle-frame-set-attention", + ) else: from local_turn_features import load_local_features diff --git a/extract_highrate_person_av_features.py b/extract_highrate_person_av_features.py new file mode 100644 index 0000000..a4e1948 --- /dev/null +++ b/extract_highrate_person_av_features.py @@ -0,0 +1,846 @@ +"""Extract frozen 25 Hz causal per-person AV states for Omni floor prediction. + +The global Qwen2.5-Omni states are not recomputed. A released MuVAP ASD model +provides one shared traditional CPC audio stream and one causal visual/activity +state per provided face track on a 25 Hz timeline. Only states whose right edge +lands on the verified 5 Hz task grid are saved. ``highrate-person-av-attention`` +therefore compares ``global Omni + traditional audio`` against the same frozen +base plus person-specific visual/speaking evidence. + +No speech/activity label is read by this extractor. Provided boxes are an +experimental isolation condition, not an online detector claim. +""" + +from __future__ import annotations + +import argparse +import hashlib +import math +import os +import re +import subprocess +import sys +import tempfile +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from tqdm.auto import tqdm + +from easycom_head import resolve_feature_path +from extract_turn_local_features import ( + EGO4D_MAX_TERMINAL_FRAME_OVERRUN, + box_coordinate_offset, + build_media_and_boxes, + load_ego4d_boxes, + select_samples, + validate_source_timeline, +) +from local_turn_features import ( + HIGHRATE_PERSON_AV_FORMAT, + crop_box, + reference_alignment, +) + + +MODEL_FPS = 25.0 +OUTPUT_HZ = 5.0 +FACE_SIZE = 112 +SAMPLE_RATE = 16_000 +SAMPLES_PER_MODEL_FRAME = SAMPLE_RATE // int(MODEL_FPS) +PERSON_FEATURE_KIND = "muvap-asd-causal-visual-plus-activity-logit" +AUDIO_FEATURE_KIND = "muvap-cpc-causal-audio-state" +PERSON_FEATURE_EXTRA_DIMS = 1 + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + while block := stream.read(1024 * 1024): + digest.update(block) + return digest.hexdigest() + + +def _release_paths(path: Path) -> tuple[Path, Path]: + path = path.expanduser().resolve() + weights = path / "weights.pt" if path.is_dir() else path + config = weights.parent / "config.yaml" + if not weights.is_file(): + raise FileNotFoundError(weights) + if not config.is_file(): + raise FileNotFoundError( + f"MuVAP release config must sit beside the weights: {config}" + ) + return weights, config + + +class FrozenMuVAPASDEncoder: + """Load the published causal ASD model without copying its source code.""" + + def __init__( + self, + *, + muvap_root: Path, + revision: str, + release: Path, + device: str, + dtype: str, + ) -> None: + root = muvap_root.expanduser().resolve() + if not re.fullmatch(r"[0-9a-fA-F]{7,40}", revision): + raise ValueError( + "--muvap-revision must be a 7--40 character hexadecimal Git " + "commit prefix" + ) + required = root / "models" / "asd.py" + if not required.is_file(): + raise FileNotFoundError( + f"MuVAP source is missing: {required}. See third_party/README.md" + ) + revision_file = root / ".upstream-revision" + if revision_file.is_file(): + actual_revision = revision_file.read_text().strip().lower() + else: + completed = subprocess.run( + ["git", "-C", str(root), "rev-parse", "HEAD"], + check=False, + capture_output=True, + text=True, + ) + actual_revision = ( + completed.stdout.strip().lower() + if completed.returncode == 0 + else "" + ) + if ( + not re.fullmatch(r"[0-9a-f]{40}", actual_revision) + or not actual_revision.startswith(revision.lower()) + ): + raise RuntimeError( + f"MuVAP checkout is {actual_revision or 'unreadable'}, expected " + f"revision prefix {revision!r}. See third_party/README.md and " + f"{revision_file}" + ) + self.source_revision = actual_revision + self.weights_path, config_path = _release_paths(release) + try: + import yaml + except ImportError as error: + raise RuntimeError("PyYAML is required to read the MuVAP release") from error + + # MuVAP uses top-level imports such as ``models.asd``. Keep the + # dependency external/ignored and verify that Python resolves it from + # the explicitly requested checkout. + sys.path.insert(0, str(root)) + try: + from models.asd import AudioVisualASD + from models.release import load_weights + except Exception as error: + raise RuntimeError( + f"Cannot import MuVAP from {root}; ensure the dependencies in " + "requirements.txt are installed and the pinned checkout is complete" + ) from error + imported = Path(sys.modules[AudioVisualASD.__module__].__file__).resolve() + if root not in imported.parents: + raise RuntimeError( + f"Python resolved {AudioVisualASD.__module__} from {imported}, not {root}" + ) + + payload = yaml.safe_load(config_path.read_text()) or {} + config = payload.get("asd") + if not isinstance(config, dict): + raise ValueError(f"MuVAP release has no `asd` config: {config_path}") + audio_spec = config.get("audio_encoder", "cpc") + audio_name = audio_spec if isinstance(audio_spec, str) else audio_spec.get("name") + if audio_name != "cpc": + raise ValueError( + "This extractor currently freezes the published 25 Hz CPC ASD " + f"release; got audio_encoder={audio_name!r}" + ) + for block in ("cross", "temporal"): + if config.get(block, {}).get("causal", True) is not True: + raise ValueError( + f"MuVAP ASD {block} stack must be causal for floor prediction" + ) + self.device = torch.device(device) + if self.device.type == "cuda" and not torch.cuda.is_available(): + raise RuntimeError(f"CUDA requested but unavailable: {device}") + self.autocast_dtype = { + "float32": None, + "float16": torch.float16, + "bfloat16": torch.bfloat16, + }[dtype] + try: + model = AudioVisualASD(config) + except FileNotFoundError as error: + raise RuntimeError( + "MuVAP's frozen CPC frontend is missing. Run " + f"`python {root / 'tools' / 'fetch_frontends.py'} cpc` first." + ) from error + self.model = load_weights(model, self.weights_path).to(self.device).eval() + self.model.requires_grad_(False) + self.audio_dim = int(config["cross"]["d_model"]) + self.visual_dim = int(config["cross"]["d_model"]) + self.fused_dim = int(config["temporal"]["d_model"]) + self.feature_dim = self.visual_dim + PERSON_FEATURE_EXTRA_DIMS + self.release_config = config + self.release_sha256 = _sha256(self.weights_path) + + def _autocast(self): + return torch.autocast( + device_type=self.device.type, + dtype=self.autocast_dtype, + enabled=self.autocast_dtype is not None, + ) + + @torch.inference_mode() + def encode_audio(self, waveform: np.ndarray, num_frames: int) -> torch.Tensor: + """Encode mixed audio once as causal 25 Hz CPC states.""" + + expected_samples = num_frames * SAMPLES_PER_MODEL_FRAME + if waveform.dtype != np.float32 or waveform.shape != (expected_samples,): + raise ValueError( + f"Audio must be float32 [{expected_samples}], got {waveform.shape}" + ) + audio = torch.from_numpy(np.ascontiguousarray(waveform)).view(1, 1, -1) + audio = audio.to(self.device, dtype=torch.float32) + with self._autocast(): + embedding = self.model.norm_audio( + self.model.encoder_proj(self.model.audio_encoder(audio)) + ) + if embedding.shape != (1, num_frames, self.audio_dim): + raise RuntimeError( + f"Unexpected MuVAP audio shape: {tuple(embedding.shape)}" + ) + result = embedding[0].float().cpu() + if not torch.isfinite(result).all(): + raise RuntimeError("MuVAP produced invalid shared audio states") + return result + + @torch.inference_mode() + def encode_person( + self, + faces: np.ndarray, + audio_states: torch.Tensor, + ) -> torch.Tensor: + """Encode one face track; keep visual state plus one AV activity logit.""" + + if ( + faces.dtype != np.uint8 + or faces.ndim != 3 + or faces.shape[1:] != (FACE_SIZE, FACE_SIZE) + ): + raise ValueError(f"Faces must be uint8 [T,112,112], got {faces.shape}") + if ( + not isinstance(audio_states, torch.Tensor) + or audio_states.shape != (len(faces), self.audio_dim) + or not audio_states.is_floating_point() + or not torch.isfinite(audio_states).all() + ): + raise ValueError( + "Person ASD requires finite aligned shared audio states " + f"[T,{self.audio_dim}]" + ) + visual = torch.from_numpy(np.ascontiguousarray(faces)).unsqueeze(0) + visual = visual.to(self.device, dtype=torch.float32) + audio_embedding = audio_states[None].to(self.device, dtype=torch.float32) + with self._autocast(): + visual_embedding = self.model.norm_visual( + self.model.visual_encoder(visual) + ) + audio_cross = self.model.cross_v2a( + audio_embedding, + src=visual_embedding, + ) + visual_cross = self.model.cross_a2v( + visual_embedding, + src=audio_embedding, + ) + audio_logits = self.model.audio_head(audio_cross) + visual_logits = self.model.visual_head(visual_cross) + activity_logits = (audio_logits + visual_logits) / 2 + result = torch.cat((visual_embedding, activity_logits), dim=-1)[0] + result = result.float().cpu() + if result.shape != (len(faces), self.feature_dim) or not torch.isfinite(result).all(): + raise RuntimeError("MuVAP produced invalid person AV states") + return result + + @torch.inference_mode() + def encode_person_fused( + self, + faces: np.ndarray, + audio_states: torch.Tensor, + ) -> torch.Tensor: + """Return the released ASD backbone's fused ``Z_ASD`` state. + + The ordinary local-feature experiment intentionally stores a visual + state and a scalar speaking score. MuVAP Stage 3 instead consumes the + fused output of the causal ASD temporal stack. Keeping this as a + separate method prevents the two feature contracts from being + silently interchanged. + """ + + if ( + faces.dtype != np.uint8 + or faces.ndim != 3 + or faces.shape[1:] != (FACE_SIZE, FACE_SIZE) + ): + raise ValueError(f"Faces must be uint8 [T,112,112], got {faces.shape}") + if ( + not isinstance(audio_states, torch.Tensor) + or audio_states.shape != (len(faces), self.audio_dim) + or not audio_states.is_floating_point() + or not torch.isfinite(audio_states).all() + ): + raise ValueError( + "Person ASD requires finite aligned shared audio states " + f"[T,{self.audio_dim}]" + ) + visual = torch.from_numpy(np.ascontiguousarray(faces)).unsqueeze(0) + visual = visual.to(self.device, dtype=torch.float32) + audio_embedding = audio_states[None].to(self.device, dtype=torch.float32) + with self._autocast(): + visual_embedding = self.model.norm_visual( + self.model.visual_encoder(visual) + ) + audio_cross = self.model.cross_v2a( + audio_embedding, + src=visual_embedding, + ) + visual_cross = self.model.cross_a2v( + visual_embedding, + src=audio_embedding, + ) + fused = self.model.transformer( + torch.cat((audio_cross, visual_cross), dim=-1) + ) + result = fused[0].float().cpu() + if result.shape != (len(faces), self.fused_dim): + raise RuntimeError( + f"Unexpected MuVAP fused ASD shape: {tuple(result.shape)}" + ) + if not torch.isfinite(result).all(): + raise RuntimeError("MuVAP produced invalid fused ASD states") + return result + + +def _square_gray_face(crop: Any) -> np.ndarray: + """Match MuVAP's 112x112 grayscale face geometry without distortion.""" + + from PIL import Image + + gray = crop.convert("L") + size = max(gray.size) + canvas = Image.new("L", (size, size), color=128) + canvas.paste(gray, ((size - gray.width) // 2, (size - gray.height) // 2)) + resampling = getattr(Image, "Resampling", Image).BILINEAR + return np.asarray(canvas.resize((FACE_SIZE, FACE_SIZE), resampling), dtype=np.uint8) + + +def _chunked_audio_encode( + encoder: FrozenMuVAPASDEncoder, + waveform: np.ndarray, + *, + num_frames: int, + output_frames: int, + context_frames: int, +) -> torch.Tensor: + """Encode shared mixed audio causally and cover every 25 Hz frame once.""" + + if output_frames <= 0 or context_frames < 0: + raise ValueError("ASD chunk/context frame counts are invalid") + if len(waveform) != num_frames * SAMPLES_PER_MODEL_FRAME: + raise ValueError("Audio/high-rate frame lengths differ") + pieces = [] + for output_start in range(0, num_frames, output_frames): + input_start = max(0, output_start - context_frames) + output_end = min(num_frames, output_start + output_frames) + sample_start = input_start * SAMPLES_PER_MODEL_FRAME + sample_end = output_end * SAMPLES_PER_MODEL_FRAME + encoded = encoder.encode_audio( + waveform[sample_start:sample_end], + output_end - input_start, + ) + pieces.append(encoded[output_start - input_start :]) + result = torch.cat(pieces) if pieces else torch.empty(0, encoder.audio_dim) + if result.shape != (num_frames, encoder.audio_dim): + raise RuntimeError("Chunked audio output did not cover every input frame once") + return result + + +def _chunked_person_encode( + encoder: FrozenMuVAPASDEncoder, + faces: np.ndarray, + audio_states: torch.Tensor, + *, + output_frames: int, + context_frames: int, +) -> torch.Tensor: + """Encode one causal face track against the already shared audio stream.""" + + if output_frames <= 0 or context_frames < 0: + raise ValueError("ASD chunk/context frame counts are invalid") + if audio_states.shape != (len(faces), encoder.audio_dim): + raise ValueError("Face/shared-audio high-rate lengths differ") + pieces = [] + for output_start in range(0, len(faces), output_frames): + input_start = max(0, output_start - context_frames) + output_end = min(len(faces), output_start + output_frames) + encoded = encoder.encode_person( + faces[input_start:output_end], + audio_states[input_start:output_end], + ) + pieces.append(encoded[output_start - input_start :]) + result = torch.cat(pieces) if pieces else torch.empty(0, encoder.feature_dim) + if result.shape != (len(faces), encoder.feature_dim): + raise RuntimeError("Chunked person output did not cover every input frame once") + return result + + +def _model_tick_bounds(clip: Any, end: float) -> tuple[int, int]: + start_value = float(clip.dense_start) * MODEL_FPS + stop_value = min(float(clip.dense_end), end) * MODEL_FPS + start = round(start_value) + stop = round(stop_value) + if abs(start_value - start) > 1e-6 or abs(stop_value - stop) > 1e-6: + raise ValueError( + "Source clip boundaries must align to the 25 Hz ASD lattice; " + f"got [{clip.dense_start}, {min(float(clip.dense_end), end)}]" + ) + return start, stop + + +def _collect_clip_tracks( + media: Any, + annotations: list[dict[int, list[dict[str, Any]]]], + clip_index: int, + start_tick: int, + stop_tick: int, + args: argparse.Namespace, +) -> list[dict[str, Any]]: + """Decode one source clip once and return consecutive anonymous tracks.""" + + active: dict[str, dict[str, Any]] = {} + segments: list[dict[str, Any]] = [] + clip = media.clips[clip_index] + for tick in tqdm( + range(start_tick, stop_tick), + desc=f"ASD faces {clip.path.name}", + unit="frame", + disable=args.no_progress, + leave=False, + ): + centre = (tick + 0.5) / MODEL_FPS + mapped_clip, frame_index, source_pts = media._source_frame_for_time(centre) + if mapped_clip != clip_index: + raise RuntimeError("25 Hz interval crossed a source-clip boundary") + image = media.frame_at_or_before(centre) + offset_x, offset_y = box_coordinate_offset(clip, args.box_coordinate_space) + seen: set[str] = set() + for person in annotations[clip_index].get(frame_index, []): + person_id = str(person.get("id", "")).strip() + if not person_id or person_id.lower() in {"-1", "none", "null", "unknown"}: + # Unknown detections cannot be joined across time. Give the + # one observation its own stable-within-extraction key. + person_id = f"unknown:{tick}:{len(seen)}" + if person_id in seen: + raise ValueError( + f"Duplicate person ID {person_id!r} at {centre:.3f}s" + ) + seen.add(person_id) + box = crop_box( + person["xyxy"], + *image.size, + offset_x=offset_x, + offset_y=offset_y, + expansion=args.crop_expansion, + ) + if box is None: + continue + record = { + "id": person_id, + "box": list(box), + "source_clip_index": clip_index, + "source_frame_index": frame_index, + "source_time": float(clip.dense_start + source_pts), + "asd_model_tick": tick, + "asd_interval_end": (tick + 1) / MODEL_FPS, + } + if record["source_time"] >= record["asd_interval_end"] - 1e-8: + raise RuntimeError( + "High-rate face sample is not strictly earlier than its " + f"ASD right edge: source={record['source_time']}, " + f"right_edge={record['asd_interval_end']}" + ) + current = active.get(person_id) + if current is None or current["ticks"][-1] + 1 != tick: + if current is not None: + segments.append(current) + current = {"id": person_id, "ticks": [], "faces": [], "records": []} + active[person_id] = current + current["ticks"].append(tick) + current["faces"].append(_square_gray_face(image.crop(box))) + current["records"].append(record) + segments.extend(active.values()) + return segments + + +def _audio_for_tick_range(media: Any, start_tick: int, stop_tick: int) -> np.ndarray: + if start_tick < 0 or stop_tick <= start_tick: + raise ValueError("ASD audio tick range must be nonempty and increasing") + waveform = media.audio_slice( + start_tick / MODEL_FPS, + stop_tick / MODEL_FPS, + ).astype(np.float32, copy=False) + wanted = (stop_tick - start_tick) * SAMPLES_PER_MODEL_FRAME + result = np.zeros(wanted, dtype=np.float32) + used = min(wanted, len(waveform)) + if used: + # AAC may decode a little outside nominal float PCM bounds. Clipping + # each observed sample is causal; content-dependent clip/prefix RMS + # normalization is intentionally not used because it would make the + # same past waveform depend on a later prediction cutoff. + result[:used] = np.clip(waveform[:used], -1.0, 1.0) + return result + + +def extract_highrate_person_states( + media: Any, + annotations: list[dict[int, list[dict[str, Any]]]], + encoder: FrozenMuVAPASDEncoder, + times: torch.Tensor, + args: argparse.Namespace, +) -> dict[str, Any]: + """Produce one variable-cardinality person set at every 5 Hz right edge.""" + + rows: list[list[tuple[torch.Tensor, dict[str, Any]]]] = [ + [] for _ in range(len(times)) + ] + shared_audio = torch.zeros( + len(times), + encoder.audio_dim, + dtype=torch.float32, + ) + shared_audio_valid = torch.zeros(len(times), dtype=torch.bool) + required_end = len(times) / OUTPUT_HZ + output_frames = round(args.asd_output_seconds * MODEL_FPS) + context_frames = round(args.asd_context_seconds * MODEL_FPS) + total_ticks = len(times) * int(MODEL_FPS / OUTPUT_HZ) + session_waveform = _audio_for_tick_range(media, 0, total_ticks) + session_audio_states = _chunked_audio_encode( + encoder, + session_waveform, + num_frames=total_ticks, + output_frames=output_frames, + context_frames=context_frames, + ) + shared_audio.copy_( + session_audio_states[ + int(MODEL_FPS / OUTPUT_HZ) - 1 :: int(MODEL_FPS / OUTPUT_HZ) + ] + ) + shared_audio_valid.fill_(True) + for clip_index, clip in enumerate(media.clips): + if clip.dense_start >= required_end - 1e-8: + break + start_tick, stop_tick = _model_tick_bounds(clip, required_end) + segments = _collect_clip_tracks( + media, + annotations, + clip_index, + start_tick, + stop_tick, + args, + ) + progress = tqdm( + segments, + desc=f"ASD tracks {clip.path.name}", + unit="track", + disable=args.no_progress, + leave=False, + ) + for segment in progress: + ticks = segment["ticks"] + faces = np.stack(segment["faces"]) + positions = torch.tensor(ticks, dtype=torch.long) + if bool(((positions < 0) | (positions >= len(session_audio_states))).any()): + raise RuntimeError("Person track extends beyond its shared audio stream") + states = _chunked_person_encode( + encoder, + faces, + session_audio_states[positions], + output_frames=output_frames, + context_frames=context_frames, + ) + for offset, tick in enumerate(ticks): + if (tick + 1) % int(MODEL_FPS / OUTPUT_HZ): + continue + row = (tick + 1) // int(MODEL_FPS / OUTPUT_HZ) - 1 + if not 0 <= row < len(rows): + continue + record = segment["records"][offset] + cutoff = float(times[row]) + if ( + abs(float(record["asd_interval_end"]) - cutoff) > 1e-7 + or float(record["source_time"]) >= cutoff + ): + raise RuntimeError( + "High-rate person observation is not strictly causal at " + f"row {row}: source={record['source_time']}, " + f"interval_end={record['asd_interval_end']}, cutoff={cutoff}" + ) + rows[row].append((states[offset].to(torch.bfloat16), record)) + + if not bool(shared_audio_valid.all()): + missing = torch.nonzero(~shared_audio_valid).flatten().tolist() + raise RuntimeError( + f"Shared traditional audio misses 5 Hz rows: {missing[:20]}" + ) + + features = [] + records = [] + offsets = [0] + pooled = torch.zeros(len(rows), encoder.feature_dim, dtype=torch.float32) + valid = torch.zeros(len(rows), dtype=torch.bool) + for row_index, observations in enumerate(rows): + if observations: + values = torch.stack([item[0] for item in observations]) + features.append(values) + records.extend(item[1] for item in observations) + pooled[row_index] = values.float().mean(0) + valid[row_index] = True + offsets.append(offsets[-1] + len(observations)) + flat = ( + torch.cat(features) + if features + else torch.empty(0, encoder.feature_dim, dtype=torch.bfloat16) + ) + return { + "audio": shared_audio.to(torch.bfloat16), + "audio_valid": shared_audio_valid, + "visual": pooled.to(torch.bfloat16), + "visual_valid": valid, + "persons": { + "features": flat, + # These are interval right edges. Equality with the task cutoff is + # explicitly allowed by the high-rate loader. + "frame_times": times.to(torch.float64).clone(), + "frame_offsets": torch.tensor(offsets, dtype=torch.long), + "records": records, + }, + } + + +def requested_config( + args: argparse.Namespace, + encoder: FrozenMuVAPASDEncoder, +) -> dict[str, Any]: + return { + "encoder": "MuVAP AudioVisualASD", + "encoder_source_revision": encoder.source_revision, + "encoder_release_sha256": encoder.release_sha256, + "encoder_release_config": encoder.release_config, + "dtype": args.dtype, + "model_fps": MODEL_FPS, + "output_hz": OUTPUT_HZ, + "asd_output_seconds": args.asd_output_seconds, + "asd_context_seconds": args.asd_context_seconds, + "face_size": [FACE_SIZE, FACE_SIZE], + "person_feature_kind": PERSON_FEATURE_KIND, + "person_visual_dim": encoder.visual_dim, + "person_activity_logit_dim": PERSON_FEATURE_EXTRA_DIMS, + "audio_feature_kind": AUDIO_FEATURE_KIND, + "visual_dim": encoder.feature_dim, + "audio_dim": encoder.audio_dim, + "audio_source": "single-shared-traditional-cpc-stream", + "person_tokens_duplicate_audio": False, + "asd_audio_preprocessing": "causal-float-pcm-clip-only-no-rms-normalization", + "asd_temporal_attention": "causal", + "box_kind": args.box_kind if args.dataset == "easycom" else "ego4d-provided-tracks", + "box_coordinate_space": args.box_coordinate_space, + "crop_expansion": args.crop_expansion, + "box_protocol": "provided-boxes-not-online-detector", + "causal_policy": "right-edge-causal-highrate-asd", + "uses_speech_labels": False, + } + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--dataset", choices=("easycom", "ego4d"), default="easycom") + parser.add_argument("--data-root", type=Path, required=True) + parser.add_argument("--feature-root", type=Path, required=True, help="Global Omni .pt root") + parser.add_argument("--output-root", type=Path, required=True) + selection = parser.add_mutually_exclusive_group(required=True) + selection.add_argument("--samples", nargs="+") + selection.add_argument("--manifest", type=Path) + parser.add_argument("--split", choices=("all", "train", "test"), default="all") + parser.add_argument("--start-index", type=int, default=1) + parser.add_argument("--end-index", type=int) + parser.add_argument("--limit", type=int) + parser.add_argument("--num-shards", type=int, default=1) + parser.add_argument("--shard-index", type=int, default=0) + parser.add_argument("--ego4d-annotations", type=Path, nargs="+") + parser.add_argument( + "--ego4d-max-terminal-overrun-frames", + type=int, + default=EGO4D_MAX_TERMINAL_FRAME_OVERRUN, + ) + parser.add_argument("--muvap-root", type=Path, default=Path("third_party/MuVAP")) + parser.add_argument("--muvap-revision", default="a1ef9c8") + parser.add_argument("--asd-weights", type=Path, required=True) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--dtype", choices=("float32", "float16", "bfloat16"), default="bfloat16") + parser.add_argument("--asd-output-seconds", type=float, default=10.0) + parser.add_argument("--asd-context-seconds", type=float, default=2.0) + parser.add_argument("--box-kind", choices=("face", "head"), default="face") + parser.add_argument("--box-coordinate-space", choices=("camera", "encoded"), default=None) + parser.add_argument("--crop-expansion", type=float, default=1.0) + parser.add_argument("--decoder-backend", choices=("decord", "torchcodec", "pynvcodec"), default="decord") + parser.add_argument("--decoder-threads", type=int, default=1) + parser.add_argument("--decoder-gpu-id", type=int, default=0) + parser.add_argument("--max-steps", type=int, help="Smoke only; partial output is rejected by training") + parser.add_argument("--overwrite", action="store_true") + parser.add_argument("--dry-run", action="store_true") + parser.add_argument("--no-progress", action="store_true") + args = parser.parse_args(argv) + if args.manifest is not None and args.dataset != "ego4d": + parser.error("--manifest is only supported for Ego4D") + if args.manifest is None and ( + args.split != "all" or args.start_index != 1 or args.end_index is not None or args.limit is not None + ): + parser.error("--split/--start-index/--end-index/--limit require --manifest") + if args.start_index < 1 or (args.end_index is not None and args.end_index < args.start_index): + parser.error("Invalid manifest range") + if args.limit is not None and args.limit <= 0: + parser.error("--limit must be positive") + if args.num_shards <= 0 or not 0 <= args.shard_index < args.num_shards: + parser.error("Require --num-shards > 0 and 0 <= --shard-index < --num-shards") + if args.max_steps is not None and args.max_steps <= 0: + parser.error("--max-steps must be positive") + if not math.isfinite(args.asd_output_seconds) or args.asd_output_seconds <= 0: + parser.error("--asd-output-seconds must be finite and positive") + if not math.isfinite(args.asd_context_seconds) or args.asd_context_seconds < 0: + parser.error("--asd-context-seconds must be finite and nonnegative") + if round(args.asd_output_seconds * MODEL_FPS) <= 0: + parser.error("--asd-output-seconds is shorter than one model frame") + if args.box_coordinate_space is None: + args.box_coordinate_space = "encoded" if args.dataset == "easycom" else "camera" + # ``build_media_and_boxes`` consumes this compatibility field. + args.video_fps = MODEL_FPS + return args + + +def main(argv: list[str] | None = None) -> None: + args = parse_args(argv) + names = select_samples(args) + if not names: + print(f"Shard {args.shard_index}/{args.num_shards}: no assigned clips") + return + print(f"Shard {args.shard_index}/{args.num_shards}: {len(names)} assigned samples") + references = { + name: resolve_feature_path(args.feature_root, name, dataset=args.dataset) + for name in names + } + for name in names: + required = [references[name]] + if args.dataset == "ego4d": + required.append(args.data_root / "v1" / "clips" / f"{name}.mp4") + for path in required: + if not path.is_file(): + raise FileNotFoundError(path) + destination = args.output_root / f"{name}.pt" + if destination.exists() and not args.overwrite: + raise FileExistsError( + f"{destination} exists; use a new root or explicit --overwrite" + ) + if args.dry_run: + for name in names: + reference = torch.load(references[name], map_location="cpu", weights_only=False) + times, _ = reference_alignment(reference) + steps = min(len(times), args.max_steps or len(times)) + print(f"{name}: would extract {steps} states ({steps / OUTPUT_HZ:g}s)") + return + + paths: list[Path] = [] + if args.dataset == "ego4d": + paths = args.ego4d_annotations or [ + args.data_root / "v1" / "annotations" / f"av_{split}.json" + for split in ("train", "val") + ] + ego_boxes = load_ego4d_boxes(paths, names) if args.dataset == "ego4d" else {} + encoder = FrozenMuVAPASDEncoder( + muvap_root=args.muvap_root, + revision=args.muvap_revision, + release=args.asd_weights, + device=args.device, + dtype=args.dtype, + ) + args.output_root.mkdir(parents=True, exist_ok=True) + expected_contract = None + for position, name in enumerate(names, 1): + print(f"[{position}/{len(names)}] {name}", flush=True) + reference = torch.load(references[name], map_location="cpu", weights_only=False) + times, indices = reference_alignment(reference) + full_steps = len(times) + if args.max_steps is not None: + times, indices = times[: args.max_steps], indices[: args.max_steps] + media, annotations = build_media_and_boxes( + args, + name, + len(times) / OUTPUT_HZ, + ego_boxes, + ) + validate_source_timeline(reference, media) + output = extract_highrate_person_states( + media, + annotations, + encoder, + times, + args, + ) + config = requested_config(args, encoder) + contract = {"format": HIGHRATE_PERSON_AV_FORMAT, "modalities": "av", "config": config} + if expected_contract is not None and contract != expected_contract: + raise ValueError("High-rate encoder/audio contract differs across samples") + expected_contract = contract + output.update( + { + "format": HIGHRATE_PERSON_AV_FORMAT, + "dataset": args.dataset, + "sample_id": name, + "times": times, + "label_indices": indices, + "partial": len(times) != full_steps, + "reference_features": str(references[name].resolve()), + "direct_media": media.metadata(), + "config": config, + } + ) + if args.dataset == "ego4d": + output["box_frame_audit"] = ego_boxes[name]["frame_audit"] + destination = args.output_root / f"{name}.pt" + with tempfile.NamedTemporaryFile( + dir=args.output_root, + prefix=f".{name}.", + suffix=".tmp", + delete=False, + ) as handle: + temporary = Path(handle.name) + try: + torch.save(output, temporary) + os.replace(temporary, destination) + finally: + temporary.unlink(missing_ok=True) + count = len(output["persons"]["features"]) + coverage = float(output["visual_valid"].float().mean()) + print( + f"Saved {destination}; states={len(times)}, person observations={count}, " + f"person-row coverage={coverage:.1%}, partial={output['partial']}", + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/extract_muvap_egospeak_features.py b/extract_muvap_egospeak_features.py new file mode 100644 index 0000000..02cadac --- /dev/null +++ b/extract_muvap_egospeak_features.py @@ -0,0 +1,579 @@ +"""Extract frozen MuVAP VAP/ASD states for an EgoSpeak Stage-3 baseline. + +This is a clean-room adaptation of the architecture described in MuVAP Sec. +4.3/5.3. The upstream repository has not released its Stage-3 trainer, so the +feature contract is deliberately named ``paper-reconstruction`` rather than +claiming to reproduce unpublished code. + +The extractor never reads speech targets. A verified global-feature sidecar +is used only for its dense 5 Hz timestamps and source-media inventory; none of +its Omni tensors are copied into the output. MuVAP's released role-VAP and ASD +backbones are frozen and evaluated on a 25 Hz, strictly causal timeline. +""" + +from __future__ import annotations + +import argparse +import hashlib +import math +import os +import re +import subprocess +import sys +import tempfile +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from tqdm.auto import tqdm + +from easycom_head import resolve_feature_path +from extract_highrate_person_av_features import ( + FACE_SIZE, + MODEL_FPS, + OUTPUT_HZ, + SAMPLE_RATE, + SAMPLES_PER_MODEL_FRAME, + FrozenMuVAPASDEncoder, + _audio_for_tick_range, + _chunked_audio_encode, + _collect_clip_tracks, + _model_tick_bounds, + _release_paths, +) +from extract_turn_local_features import ( + EGO4D_MAX_TERMINAL_FRAME_OVERRUN, + build_media_and_boxes, + load_ego4d_boxes, + select_samples, + validate_source_timeline, +) +from local_turn_features import reference_alignment +from muvap_egospeak_stage3 import MUVAP_EGOSPEAK_INPUT_FORMAT + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + while block := stream.read(1024 * 1024): + digest.update(block) + return digest.hexdigest() + + +def _verified_root(root: Path, revision: str) -> tuple[Path, str]: + root = root.expanduser().resolve() + if not re.fullmatch(r"[0-9a-fA-F]{7,40}", revision): + raise ValueError("--muvap-revision must be a hexadecimal commit prefix") + if not (root / "models" / "vap.py").is_file(): + raise FileNotFoundError(f"MuVAP source is incomplete: {root}") + revision_file = root / ".upstream-revision" + if revision_file.is_file(): + actual = revision_file.read_text().strip().lower() + else: + completed = subprocess.run( + ["git", "-C", str(root), "rev-parse", "HEAD"], + check=False, + capture_output=True, + text=True, + ) + actual = completed.stdout.strip().lower() if completed.returncode == 0 else "" + if not re.fullmatch(r"[0-9a-f]{40}", actual) or not actual.startswith( + revision.lower() + ): + raise RuntimeError( + f"MuVAP checkout is {actual or 'unreadable'}, expected {revision!r}" + ) + return root, actual + + +class FrozenMuVAPVAPEncoder: + """Load the published 25 Hz mono role-VAP and expose ``Z_VAP``.""" + + def __init__( + self, + *, + muvap_root: Path, + revision: str, + release: Path, + device: str, + dtype: str, + ) -> None: + root, self.source_revision = _verified_root(muvap_root, revision) + self.weights_path, config_path = _release_paths(release) + try: + import yaml + except ImportError as error: + raise RuntimeError("PyYAML is required to read MuVAP releases") from error + + sys.path.insert(0, str(root)) + try: + from models.release import load_weights + from models.vap import build_vap + except Exception as error: + raise RuntimeError(f"Cannot import MuVAP from {root}") from error + imported = Path(sys.modules[build_vap.__module__].__file__).resolve() + if root not in imported.parents: + raise RuntimeError( + f"Python resolved {build_vap.__module__} from {imported}, not {root}" + ) + payload = yaml.safe_load(config_path.read_text()) or {} + config = payload.get("vap") + if not isinstance(config, dict): + raise ValueError(f"MuVAP release has no `vap` config: {config_path}") + projection = config.get("projection_window", {}) + audio_spec = config.get("audio_encoder", "cpc") + audio_name = audio_spec if isinstance(audio_spec, str) else audio_spec.get("name") + if audio_name != "cpc": + raise ValueError("The first baseline requires the CPC VAP release") + if config.get("arch", "mono") != "mono": + raise ValueError("EgoSpeak mixed audio requires MuVAP's mono role-VAP") + if projection.get("mode") != "role_relative": + raise ValueError("Expected a role-relative VAP release") + if float(projection.get("frame_hz", -1)) != MODEL_FPS: + raise ValueError( + f"Expected a {MODEL_FPS:g} Hz VAP release, got " + f"{projection.get('frame_hz')!r}" + ) + if config.get("temporal", {}).get("causal", True) is not True: + raise ValueError("VAP temporal attention must be causal") + + self.device = torch.device(device) + if self.device.type == "cuda" and not torch.cuda.is_available(): + raise RuntimeError(f"CUDA requested but unavailable: {device}") + self.autocast_dtype = { + "float32": None, + "float16": torch.float16, + "bfloat16": torch.bfloat16, + }[dtype] + try: + model = build_vap(config) + except FileNotFoundError as error: + raise RuntimeError( + "MuVAP's CPC frontend is missing; run " + f"`python {root / 'tools' / 'fetch_frontends.py'} cpc`" + ) from error + self.model = load_weights(model, self.weights_path).to(self.device).eval() + self.model.requires_grad_(False) + self.feature_dim = int(config["temporal"]["d_model"]) + self.release_config = config + self.release_sha256 = _sha256(self.weights_path) + + @torch.inference_mode() + def encode_audio(self, waveform: np.ndarray, num_frames: int) -> torch.Tensor: + expected = num_frames * SAMPLES_PER_MODEL_FRAME + if waveform.dtype != np.float32 or waveform.shape != (expected,): + raise ValueError(f"Audio must be float32 [{expected}], got {waveform.shape}") + audio = torch.from_numpy(np.ascontiguousarray(waveform))[None, None] + audio = audio.to(self.device, dtype=torch.float32) + with torch.autocast( + device_type=self.device.type, + dtype=self.autocast_dtype, + enabled=self.autocast_dtype is not None, + ): + _, embedding = self.model(audio, return_embeddings=True) + result = embedding[0].float().cpu() + if result.shape != (num_frames, self.feature_dim): + raise RuntimeError(f"Unexpected VAP shape: {tuple(result.shape)}") + if not torch.isfinite(result).all(): + raise RuntimeError("MuVAP VAP produced NaN/Inf") + return result + + +def _chunked_vap_encode( + encoder: FrozenMuVAPVAPEncoder, + waveform: np.ndarray, + *, + num_frames: int, + output_frames: int, + context_frames: int, +) -> torch.Tensor: + """Encode every tick with a fixed, past-only VAP context.""" + + if output_frames <= 0 or context_frames < 0: + raise ValueError("VAP chunk/context frame counts are invalid") + if waveform.shape != (num_frames * SAMPLES_PER_MODEL_FRAME,): + raise ValueError("Audio/VAP high-rate lengths differ") + pieces = [] + for output_start in range(0, num_frames, output_frames): + input_start = max(0, output_start - context_frames) + output_end = min(num_frames, output_start + output_frames) + samples = waveform[ + input_start * SAMPLES_PER_MODEL_FRAME : + output_end * SAMPLES_PER_MODEL_FRAME + ] + encoded = encoder.encode_audio(samples, output_end - input_start) + pieces.append(encoded[output_start - input_start :]) + result = torch.cat(pieces) if pieces else torch.empty(0, encoder.feature_dim) + if result.shape != (num_frames, encoder.feature_dim): + raise RuntimeError("Chunked VAP output did not cover every frame once") + return result + + +def _chunked_fused_asd_encode( + encoder: FrozenMuVAPASDEncoder, + faces: np.ndarray, + audio_states: torch.Tensor, + *, + output_frames: int, + context_frames: int, +) -> torch.Tensor: + """Encode one consecutive face track into frozen ``Z_ASD`` states.""" + + if output_frames <= 0 or context_frames < 0: + raise ValueError("ASD chunk/context frame counts are invalid") + if audio_states.shape != (len(faces), encoder.audio_dim): + raise ValueError("Face/shared-audio high-rate lengths differ") + pieces = [] + for output_start in range(0, len(faces), output_frames): + input_start = max(0, output_start - context_frames) + output_end = min(len(faces), output_start + output_frames) + encoded = encoder.encode_person_fused( + faces[input_start:output_end], + audio_states[input_start:output_end], + ) + pieces.append(encoded[output_start - input_start :]) + result = torch.cat(pieces) if pieces else torch.empty(0, encoder.fused_dim) + if result.shape != (len(faces), encoder.fused_dim): + raise RuntimeError("Chunked fused ASD output did not cover every frame once") + return result + + +def extract_sample( + media: Any, + annotations: list[dict[int, list[dict[str, Any]]]], + vap: FrozenMuVAPVAPEncoder, + asd: FrozenMuVAPASDEncoder, + task_times: torch.Tensor, + args: argparse.Namespace, +) -> dict[str, Any]: + """Extract 25 Hz VAP states and ragged fused-ASD person sets.""" + + ticks_per_task = round(MODEL_FPS / OUTPUT_HZ) + if abs(ticks_per_task * OUTPUT_HZ - MODEL_FPS) > 1e-8: + raise RuntimeError("MuVAP/task grids are not integral") + num_ticks = len(task_times) * ticks_per_task + waveform = _audio_for_tick_range(media, 0, num_ticks) + vap_states = _chunked_vap_encode( + vap, + waveform, + num_frames=num_ticks, + output_frames=round(args.vap_output_seconds * MODEL_FPS), + context_frames=round(args.vap_context_seconds * MODEL_FPS), + ) + asd_audio = _chunked_audio_encode( + asd, + waveform, + num_frames=num_ticks, + output_frames=round(args.asd_output_seconds * MODEL_FPS), + context_frames=round(args.asd_context_seconds * MODEL_FPS), + ) + + rows: list[list[tuple[torch.Tensor, dict[str, Any]]]] = [ + [] for _ in range(num_ticks) + ] + required_end = len(task_times) / OUTPUT_HZ + for clip_index, clip in enumerate(media.clips): + if float(clip.dense_start) >= required_end - 1e-8: + break + start_tick, stop_tick = _model_tick_bounds(clip, required_end) + segments = _collect_clip_tracks( + media, + annotations, + clip_index, + start_tick, + stop_tick, + args, + ) + for segment in tqdm( + segments, + desc=f"fused ASD {clip.path.name}", + unit="track", + disable=args.no_progress, + leave=False, + ): + positions = torch.tensor(segment["ticks"], dtype=torch.long) + if bool(((positions < 0) | (positions >= num_ticks)).any()): + raise RuntimeError("ASD track extends beyond the VAP timeline") + states = _chunked_fused_asd_encode( + asd, + np.stack(segment["faces"]), + asd_audio[positions], + output_frames=round(args.asd_output_seconds * MODEL_FPS), + context_frames=round(args.asd_context_seconds * MODEL_FPS), + ) + for offset, tick in enumerate(segment["ticks"]): + record = segment["records"][offset] + if abs(float(record["asd_interval_end"]) - (tick + 1) / MODEL_FPS) > 1e-8: + raise RuntimeError("ASD record/tick right edges differ") + rows[tick].append((states[offset].to(torch.bfloat16), record)) + + features: list[torch.Tensor] = [] + records: list[dict[str, Any]] = [] + offsets = [0] + max_people = 0 + for observations in rows: + if observations: + features.append(torch.stack([value for value, _ in observations])) + records.extend(record for _, record in observations) + max_people = max(max_people, len(observations)) + offsets.append(offsets[-1] + len(observations)) + flat = ( + torch.cat(features) + if features + else torch.empty(0, asd.fused_dim, dtype=torch.bfloat16) + ) + task_frame_indices = torch.arange(len(task_times), dtype=torch.long) + task_frame_indices = (task_frame_indices + 1) * ticks_per_task - 1 + if len(task_times): + highrate_edges = (task_frame_indices + 1).to(torch.float64) / MODEL_FPS + if not torch.allclose(highrate_edges, task_times, atol=1e-8, rtol=0): + raise RuntimeError("25 Hz right edges do not match the 5 Hz task cutoffs") + return { + "vap": vap_states.to(torch.bfloat16), + "frame_times": torch.arange(1, num_ticks + 1, dtype=torch.float64) / MODEL_FPS, + "persons": { + "features": flat, + "frame_offsets": torch.tensor(offsets, dtype=torch.long), + "records": records, + }, + "task_frame_indices": task_frame_indices, + "max_people": max_people, + } + + +def requested_config( + args: argparse.Namespace, + vap: FrozenMuVAPVAPEncoder, + asd: FrozenMuVAPASDEncoder, +) -> dict[str, Any]: + return { + "implementation": "clean-room-paper-reconstruction-not-upstream-stage3-code", + "uses_omni_features": False, + "muvap_source_revision": vap.source_revision, + "vap_release_sha256": vap.release_sha256, + "asd_release_sha256": asd.release_sha256, + "vap_release_config": vap.release_config, + "asd_release_config": asd.release_config, + "vap_dim": vap.feature_dim, + "asd_dim": asd.fused_dim, + "model_fps": MODEL_FPS, + "task_hz": OUTPUT_HZ, + "sample_rate": SAMPLE_RATE, + "face_size": [FACE_SIZE, FACE_SIZE], + "vap_output_seconds": args.vap_output_seconds, + "vap_context_seconds": args.vap_context_seconds, + "asd_output_seconds": args.asd_output_seconds, + "asd_context_seconds": args.asd_context_seconds, + "audio_preprocessing": "causal-float-pcm-clip-only-no-future-rms", + "box_kind": args.box_kind if args.dataset == "easycom" else "ego4d-provided-tracks", + "box_coordinate_space": args.box_coordinate_space, + "crop_expansion": args.crop_expansion, + "box_protocol": "provided-boxes-not-online-detector", + "uses_speech_labels": False, + "causal_policy": "25hz-right-edge-causal-fixed-past-context", + } + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--dataset", choices=("easycom", "ego4d"), default="easycom") + parser.add_argument("--data-root", type=Path, required=True) + parser.add_argument( + "--reference-feature-root", + type=Path, + required=True, + help="Verified 5 Hz sidecars used only for timestamps/source inventory", + ) + parser.add_argument("--output-root", type=Path, required=True) + selection = parser.add_mutually_exclusive_group(required=True) + selection.add_argument("--samples", nargs="+") + selection.add_argument("--manifest", type=Path) + parser.add_argument("--split", choices=("all", "train", "test"), default="all") + parser.add_argument("--start-index", type=int, default=1) + parser.add_argument("--end-index", type=int) + parser.add_argument("--limit", type=int) + parser.add_argument("--num-shards", type=int, default=1) + parser.add_argument("--shard-index", type=int, default=0) + parser.add_argument("--ego4d-annotations", type=Path, nargs="+") + parser.add_argument( + "--ego4d-max-terminal-overrun-frames", + type=int, + default=EGO4D_MAX_TERMINAL_FRAME_OVERRUN, + ) + parser.add_argument("--muvap-root", type=Path, default=Path("third_party/MuVAP")) + parser.add_argument("--muvap-revision", default="a1ef9c8") + parser.add_argument("--vap-weights", type=Path, required=True) + parser.add_argument("--asd-weights", type=Path, required=True) + parser.add_argument("--device", default="cuda:0") + parser.add_argument( + "--dtype", + choices=("float32", "float16", "bfloat16"), + default="bfloat16", + ) + parser.add_argument("--vap-output-seconds", type=float, default=10.0) + parser.add_argument("--vap-context-seconds", type=float, default=20.0) + parser.add_argument("--asd-output-seconds", type=float, default=10.0) + parser.add_argument("--asd-context-seconds", type=float, default=2.0) + parser.add_argument("--box-kind", choices=("face", "head"), default="face") + parser.add_argument( + "--box-coordinate-space", choices=("camera", "encoded"), default=None + ) + parser.add_argument("--crop-expansion", type=float, default=1.0) + parser.add_argument( + "--decoder-backend", + choices=("decord", "torchcodec", "pynvcodec"), + default="decord", + ) + parser.add_argument("--decoder-threads", type=int, default=1) + parser.add_argument("--decoder-gpu-id", type=int, default=0) + parser.add_argument("--max-steps", type=int) + parser.add_argument("--overwrite", action="store_true") + parser.add_argument("--dry-run", action="store_true") + parser.add_argument("--no-progress", action="store_true") + args = parser.parse_args(argv) + if args.manifest is not None and args.dataset != "ego4d": + parser.error("--manifest is only supported for Ego4D") + if args.manifest is None and ( + args.split != "all" + or args.start_index != 1 + or args.end_index is not None + or args.limit is not None + ): + parser.error("--split/--start-index/--end-index/--limit require --manifest") + if args.start_index < 1 or ( + args.end_index is not None and args.end_index < args.start_index + ): + parser.error("Invalid manifest range") + if args.limit is not None and args.limit <= 0: + parser.error("--limit must be positive") + if args.num_shards <= 0 or not 0 <= args.shard_index < args.num_shards: + parser.error("Require --num-shards > 0 and a valid --shard-index") + if args.max_steps is not None and args.max_steps <= 0: + parser.error("--max-steps must be positive") + for name in ( + "vap_output_seconds", + "vap_context_seconds", + "asd_output_seconds", + "asd_context_seconds", + ): + value = getattr(args, name) + if not math.isfinite(value) or value < 0 or ( + name.endswith("output_seconds") and value == 0 + ): + parser.error(f"--{name.replace('_', '-')} is invalid") + frames = round(value * MODEL_FPS) + if not math.isclose(frames / MODEL_FPS, value, abs_tol=1e-8): + parser.error(f"--{name.replace('_', '-')} must align to the 25 Hz grid") + if name.endswith("output_seconds") and frames == 0: + parser.error(f"--{name.replace('_', '-')} is shorter than one frame") + if not math.isfinite(args.crop_expansion) or args.crop_expansion <= 0: + parser.error("--crop-expansion must be finite and positive") + if args.decoder_threads <= 0 or args.decoder_gpu_id < 0: + parser.error("decoder-threads must be positive and decoder-gpu-id nonnegative") + if args.box_coordinate_space is None: + args.box_coordinate_space = "encoded" if args.dataset == "easycom" else "camera" + args.video_fps = MODEL_FPS + return args + + +def main(argv: list[str] | None = None) -> None: + args = parse_args(argv) + names = select_samples(args) + if not names: + print(f"Shard {args.shard_index}/{args.num_shards}: no assigned samples") + return + references = { + name: resolve_feature_path(args.reference_feature_root, name, dataset=args.dataset) + for name in names + } + for name in names: + if not references[name].is_file(): + raise FileNotFoundError(references[name]) + destination = args.output_root / f"{name}.pt" + if destination.exists() and not args.overwrite: + raise FileExistsError( + f"{destination} exists; use a new root or explicit --overwrite" + ) + print(f"Shard {args.shard_index}/{args.num_shards}: {len(names)} sample(s)") + if args.dry_run: + for name in names: + reference = torch.load(references[name], map_location="cpu", weights_only=False) + times, _ = reference_alignment(reference) + steps = min(len(times), args.max_steps or len(times)) + print(f"{name}: would extract {steps * 5} MuVAP frames for {steps} targets") + return + + ego_paths: list[Path] = [] + if args.dataset == "ego4d": + ego_paths = args.ego4d_annotations or [ + args.data_root / "v1" / "annotations" / f"av_{split}.json" + for split in ("train", "val") + ] + ego_boxes = load_ego4d_boxes(ego_paths, names) if args.dataset == "ego4d" else {} + vap = FrozenMuVAPVAPEncoder( + muvap_root=args.muvap_root, + revision=args.muvap_revision, + release=args.vap_weights, + device=args.device, + dtype=args.dtype, + ) + asd = FrozenMuVAPASDEncoder( + muvap_root=args.muvap_root, + revision=args.muvap_revision, + release=args.asd_weights, + device=args.device, + dtype=args.dtype, + ) + args.output_root.mkdir(parents=True, exist_ok=True) + expected_config = requested_config(args, vap, asd) + for position, name in enumerate(names, 1): + print(f"[{position}/{len(names)}] {name}", flush=True) + reference = torch.load(references[name], map_location="cpu", weights_only=False) + times, indices = reference_alignment(reference) + full_steps = len(times) + if args.max_steps is not None: + times = times[: args.max_steps] + indices = indices[: args.max_steps] + media, annotations = build_media_and_boxes( + args, name, len(times) / OUTPUT_HZ, ego_boxes + ) + validate_source_timeline(reference, media) + output = extract_sample(media, annotations, vap, asd, times, args) + output.update( + { + "format": MUVAP_EGOSPEAK_INPUT_FORMAT, + "dataset": args.dataset, + "sample_id": name, + "times": times.to(torch.float64), + "label_indices": indices.to(torch.long), + "partial": len(times) != full_steps, + "reference_features": str(references[name].resolve()), + "direct_media": media.metadata(), + "config": expected_config, + } + ) + if args.dataset == "ego4d": + output["box_frame_audit"] = ego_boxes[name]["frame_audit"] + destination = args.output_root / f"{name}.pt" + with tempfile.NamedTemporaryFile( + dir=args.output_root, + prefix=f".{name}.", + suffix=".tmp", + delete=False, + ) as handle: + temporary = Path(handle.name) + try: + torch.save(output, temporary) + os.replace(temporary, destination) + finally: + temporary.unlink(missing_ok=True) + print( + f"Saved {destination}; MuVAP frames={len(output['vap'])}, " + f"person observations={len(output['persons']['features'])}, " + f"max_people={output['max_people']}, partial={output['partial']}", + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/infer_muvap_egospeak_stage3.py b/infer_muvap_egospeak_stage3.py new file mode 100644 index 0000000..fe1849b --- /dev/null +++ b/infer_muvap_egospeak_stage3.py @@ -0,0 +1,151 @@ +"""Run a trained MuVAP-only Stage-3 EgoSpeak head and save standard metrics.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import numpy as np +import torch +from torch.utils.data import DataLoader +from tqdm.auto import tqdm + +from easycom_head import NUM_HORIZONS, output_space_for_loss +from muvap_egospeak_stage3 import ( + MODEL_KIND, + MuVAPWindowDataset, + align_session_predictions, + build_model, + collate_stage3_windows, + evaluate_aligned, + load_muvap_session, + print_metrics, + predict_loader, +) + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--feature-root", type=Path, required=True) + parser.add_argument("--target-root", type=Path, required=True) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--sessions", nargs="+", required=True) + parser.add_argument("--batch-size", type=int, default=32) + parser.add_argument("--num-workers", type=int, default=0) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--no-progress", action="store_true") + args = parser.parse_args(argv) + if args.batch_size <= 0 or args.num_workers < 0: + parser.error("batch-size must be positive and num-workers nonnegative") + return args + + +def resolve_device(name: str) -> torch.device: + device = torch.device(name) + if device.type == "cuda" and not torch.cuda.is_available(): + raise RuntimeError(f"CUDA requested but unavailable: {name}") + return device + + +def main(argv: list[str] | None = None) -> None: + args = parse_args(argv) + if not args.checkpoint.is_file(): + raise FileNotFoundError(args.checkpoint) + if args.output_dir.exists(): + raise FileExistsError(f"Refusing to overwrite {args.output_dir}") + checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=False) + if not isinstance(checkpoint, dict) or checkpoint.get("model_kind") != MODEL_KIND: + raise ValueError(f"Not a {MODEL_KIND} checkpoint: {args.checkpoint}") + loss_name = checkpoint.get("loss") + output_space_for_loss(loss_name) + model = build_model(checkpoint.get("model_config", {})) + model.load_state_dict(checkpoint.get("model_state", {}), strict=True) + device = resolve_device(args.device) + model = model.to(device).eval() + history_seconds = model.history_frames / 25.0 + + predictions: dict[str, np.ndarray] = {} + valid: dict[str, np.ndarray] = {} + targets: dict[str, np.ndarray] = {} + sessions = [] + for raw_name in args.sessions: + session = load_muvap_session( + args.feature_root, + args.target_root, + raw_name, + expected_contract=checkpoint.get("feature_contract"), + ) + sessions.append(session.name) + dataset = MuVAPWindowDataset( + [session], + history_seconds=history_seconds, + require_full_future=False, + ) + loader = DataLoader( + dataset, + batch_size=args.batch_size, + shuffle=False, + num_workers=args.num_workers, + pin_memory=device.type == "cuda", + persistent_workers=args.num_workers > 0, + collate_fn=collate_stage3_windows, + ) + iterator = loader + if not args.no_progress: + iterator = tqdm(loader, desc=f"infer {session.name}", unit="batch") + probabilities, _, rows, names = predict_loader( + model, + iterator, + device=device, + loss_name=loss_name, + ) + expected_rows = torch.arange(len(session.targets), dtype=torch.long) + if names != [session.name] * len(session.targets) or not torch.equal(rows, expected_rows): + raise RuntimeError(f"Inference row order changed for {session.name}") + aligned, aligned_valid = align_session_predictions( + probabilities.numpy(), len(session.targets) + ) + predictions[session.name] = aligned + valid[session.name] = aligned_valid + targets[session.name] = session.targets.numpy() + print( + f"{session.name}: saved-shape={aligned.shape}, causal-valid=" + f"{aligned_valid.sum(axis=0).tolist()}" + ) + + scores = evaluate_aligned(predictions, valid, targets) + metrics = { + "dataset": "easycom", + "model_kind": MODEL_KIND, + "implementation_scope": checkpoint.get("implementation_scope"), + "sessions": sessions, + "checkpoint": str(args.checkpoint.resolve()), + "checkpoint_epoch": checkpoint.get("epoch"), + "checkpoint_best_epoch": checkpoint.get("best_epoch"), + "checkpoint_training_epochs": checkpoint.get("training_config", {}).get("epochs"), + "checkpoint_train_sessions": checkpoint.get("train_sessions"), + "checkpoint_dev_sessions": checkpoint.get("dev_sessions"), + "loss": loss_name, + "output_space": checkpoint.get("output_space"), + "layer_fusion": "muvap-stage3-frame-temporal", + "history_seconds": history_seconds, + "person_evidence": model.person_evidence, + **scores, + } + (args.output_dir / "pred").mkdir(parents=True) + (args.output_dir / "valid").mkdir(parents=True) + for name in sessions: + np.save(args.output_dir / "pred" / f"{name}.npy", predictions[name]) + np.save(args.output_dir / "valid" / f"{name}.npy", valid[name]) + (args.output_dir / "metrics.json").write_text( + json.dumps(metrics, indent=2, sort_keys=True) + "\n" + ) + print_metrics(metrics) + print(f"\nmetrics: {args.output_dir / 'metrics.json'}") + print(f"EgoSpeak-compatible predictions: {args.output_dir / 'pred'}") + + +if __name__ == "__main__": + main() diff --git a/infer_turn_head.py b/infer_turn_head.py index 7f0ba01..068e92e 100644 --- a/infer_turn_head.py +++ b/infer_turn_head.py @@ -19,6 +19,7 @@ from easycom_head import ( NUM_HORIZONS, OFFICIAL_TEST_SESSIONS, STEP_SECONDS, + TRACK_PERSON_FUSION_MODES, TurnPredictionHead, average_precision_metrics, checkpoint_head_config, @@ -149,7 +150,7 @@ def _predict_states( start, min(start + batch_size, len(states)), ) - if head.local_fusion == "track-gru-residual": + if head.local_fusion in TRACK_PERSON_FUSION_MODES: person_batch, person_mask = local_person_features.padded_history( rows, max_observations=head.person_history_frames, @@ -500,6 +501,14 @@ def main() -> None: f"{name}: saved pred {predictions.shape}, " f"causal-valid counts={valid.sum(axis=0).tolist()}" ) + if head.person_adapter_diagnostics(): + print( + "person adapter diagnostics: " + + json.dumps( + head.person_adapter_diagnostics(), + sort_keys=True, + ) + ) metrics = _evaluate(session_predictions, session_valid, session_targets) transition_metrics = ( @@ -532,10 +541,23 @@ def main() -> None: ).get("epochs"), "checkpoint_train_sessions": checkpoint_payload.get("train_sessions"), "checkpoint_dev_sessions": checkpoint_payload.get("dev_sessions"), + "initialized_pooled_checkpoint": checkpoint_payload.get( + "training_config", {} + ).get("initialize_pooled_checkpoint"), + "initialized_pooled_checkpoint_sha256": checkpoint_payload.get( + "training_config", {} + ).get("initialize_pooled_checkpoint_sha256"), + "initialized_pooled_checkpoint_epoch": checkpoint_payload.get( + "training_config", {} + ).get("initialize_pooled_checkpoint_epoch"), + "pooled_base_frozen": checkpoint_payload.get( + "training_config", {} + ).get("pooled_base_frozen"), "layer_fusion": head.layer_fusion, "local_fusion": head.local_fusion, "local_fusion_scope": LOCAL_FUSION_SCOPES[head.local_fusion], "person_interaction_gates": head.person_interaction_gates(), + "person_adapter_diagnostics": head.person_adapter_diagnostics(), "person_history_frames": head.person_history_frames, "person_track_max_gap_seconds": head.person_track_max_gap_seconds, "local_contract": local_contract, diff --git a/literature_notes/02_egospeak.md b/literature_notes/02_egospeak.md index b6d0e25..109f5bd 100644 --- a/literature_notes/02_egospeak.md +++ b/literature_notes/02_egospeak.md @@ -25,3 +25,10 @@ step（25.6 秒）、stride 4、2048D embedding、1024D hidden 和 10 个 future 匿名人物轨迹之间共享，并用零门控残差修正 matched pooled Omni head。官方实现仍应作为 单独的 sequence-GRU 强基线；逐人 track continuity 是否有效，需要 broken-ID/no-ID 控制 而不能由架构名称推断。 + +为控制逐人时序 stem 的实现差异，项目另提供 +`causal-person-adapter-egospeak-gru`。它仅采用 MiniROADA 的 +`Linear-LayerNorm-ReLU-Dropout-GRU-ReLU` 层序，仍处理逐人视觉 track，并保留 V2 的 +G+A query、人物集合注意力、动态 gate 和 pooled residual base。因此实验中必须称为 +**EgoSpeak-style per-track GRU**，不能称为 EgoSpeak/MiniROADA 复现。真正的官方 +sequence-GRU baseline 仍需单独实现全局 5 Hz 序列和 anticipation decoder。 diff --git a/literature_notes/33_muvap.md b/literature_notes/33_muvap.md index 93a5ee3..ad9b7e0 100644 --- a/literature_notes/33_muvap.md +++ b/literature_notes/33_muvap.md @@ -2,8 +2,8 @@ - 作者：Haotian Qi、Gabriel Skantze - 状态：2026 预印本；作者仓库标注 Interspeech 2026，本轮未找到正式 ISCA 条目 -- 阅读深度：重点方法、训练、事件评测及局限章节；未运行 -- 核查日期：2026-08-28 +- 阅读深度：重点方法、训练、事件评测及局限章节；VAP/ASD 源码已审计，Stage 3 未公开 +- 核查日期：2026-09-10 - 来源：[论文](https://arxiv.org/html/2606.16731)、[官方仓库](https://github.com/Haotian-Qi/MuVAP) ## 摘要与项目关系 @@ -18,4 +18,8 @@ MuVAP 不是对多人脸做普通平均。它先用 CPC 与 causal Transformer 不能把“单声道＋逐人视觉＋因果接话”作为独有贡献。我们的区别需要落在不可见 wearer、讲话对象及受召回约束的误开口上，而不是只换编码器。 -作者承认角色识别存在延迟；验证主要为第三人称双人/三人英语场景。其数据准备使用 RetinaFace/SCRFD 检测、ArcFace 特征、匈牙利匹配，并从每段完整视频随机抽取 50 帧建立全局身份锚点。这适合构建稳定标注，但会使用未来/全片信息，不能作为本项目部署时的严格在线 tracking。论文模型本身宣称 causal，不代表这套数据身份初始化也是可在线运行的。仓库当前只有 README/LICENSE，写明 code coming soon；不要称已有可直接运行的复现。 +作者承认角色识别存在延迟；验证主要为第三人称双人/三人英语场景。其数据准备使用 RetinaFace/SCRFD 检测、ArcFace 特征、匈牙利匹配，并从每段完整视频随机抽取 50 帧建立全局身份锚点。这适合构建稳定标注，但会使用未来/全片信息，不能作为本项目部署时的严格在线 tracking。论文模型本身宣称 causal，不代表这套数据身份初始化也是可在线运行的。 + +作者仓库现已公开 VAP 和 ASD 的训练代码、CPC/Mimi 前端及对应发布权重；完整 MuVAP +multiparty Stage 3 仍标注为 `to follow`，没有训练入口。因此项目中的 Stage-3 baseline +必须称为按论文描述做的 clean-room reconstruction，不能称为官方复现。 diff --git a/local_turn_features.py b/local_turn_features.py index 9feefe8..0427c4e 100644 --- a/local_turn_features.py +++ b/local_turn_features.py @@ -15,6 +15,7 @@ import torch LOCAL_FORMAT = "qwen-local-pooled-v1" ORACLE_FORMAT = "easycom-oracle-person-history-v1" +HIGHRATE_PERSON_AV_FORMAT = "causal-highrate-person-av-v2" MODALITIES = {"visual": ("visual",), "audio": ("audio",), "av": ("visual", "audio")} @@ -279,12 +280,29 @@ def _load_validated_payload( modalities: str, expected_contract: dict[str, Any] | None, oracle_analysis: bool = False, + highrate_person_av: bool = False, + allow_highrate_audio: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any], dict[str, Any], dict[str, Any]]: if modalities not in MODALITIES: raise ValueError(f"Unknown local modalities: {modalities}") times, indices = reference_alignment(reference) payload = torch.load(path, map_location="cpu", weights_only=False) - expected_format = ORACLE_FORMAT if oracle_analysis else LOCAL_FORMAT + if ( + allow_highrate_audio + and modalities == "audio" + and isinstance(payload, dict) + and payload.get("format") == HIGHRATE_PERSON_AV_FORMAT + ): + highrate_person_av = True + if oracle_analysis and highrate_person_av: + raise ValueError("Oracle and learned high-rate local formats are distinct") + expected_format = ( + ORACLE_FORMAT + if oracle_analysis + else HIGHRATE_PERSON_AV_FORMAT + if highrate_person_av + else LOCAL_FORMAT + ) if not isinstance(payload, dict) or payload.get("format") != expected_format: raise ValueError(f"Invalid local feature format: {path}") if payload.get("sample_id") != sample_id or payload.get("dataset") != dataset: @@ -308,6 +326,24 @@ def _load_validated_payload( raise ValueError( f"Oracle features lack the explicit causal/oracle contract: {path}" ) + elif highrate_person_av: + if ( + config.get("causal_policy") != "right-edge-causal-highrate-asd" + or config.get("uses_speech_labels") is not False + or config.get("model_fps") != 25.0 + or config.get("output_hz") != 5.0 + or config.get("person_feature_kind") + != "muvap-asd-causal-visual-plus-activity-logit" + or config.get("audio_feature_kind") + != "muvap-cpc-causal-audio-state" + or config.get("audio_source") != "single-shared-traditional-cpc-stream" + or config.get("person_visual_dim") != config.get("visual_dim") - 1 + or config.get("person_activity_logit_dim") != 1 + or config.get("person_tokens_duplicate_audio") is not False + ): + raise ValueError( + f"High-rate person AV features lack the causal ASD contract: {path}" + ) elif ( config.get("causal_policy") != "half-open-trailing-window" or config.get("uses_speech_labels") is not False @@ -336,6 +372,7 @@ def load_local_features( dataset=dataset, modalities=modalities, expected_contract=expected_contract, + allow_highrate_audio=True, ) chunks = [] for name in MODALITIES[modalities]: @@ -365,6 +402,7 @@ def load_frame_set_features( expected_contract: dict[str, Any] | None = None, track_max_gap_seconds: float | None = None, oracle_analysis: bool = False, + highrate_person_av: bool = False, ) -> tuple[torch.Tensor, FramePersonFeatures, dict[str, Any]]: """Load direct audio and the latest causal per-frame person set per 5 Hz row.""" @@ -376,6 +414,7 @@ def load_frame_set_features( modalities="av", expected_contract=expected_contract, oracle_analysis=oracle_analysis, + highrate_person_av=highrate_person_av, ) audio = payload.get("audio") audio_valid = payload.get("audio_valid") @@ -441,18 +480,22 @@ def load_frame_set_features( if len(frame_times) and float(frame_times[-1]) > float(times[-1]) + 1e-8: raise ValueError(f"Person frame timeline extends beyond the reference: {path}") - # Normal visual observations must be strictly earlier than the cutoff. + # Normal visual observations are instantaneous samples and must be strictly + # earlier than the cutoff. High-rate ASD states instead summarize complete + # 25 Hz intervals ending at the stored right edge, so equality with a 5 Hz + # cutoff is causal and intentional. Oracle current state has the same + # lookup geometry but remains a separate, analysis-only format. # Oracle speech state is defined and sampled at the cutoff itself, so its # explicit analysis-only format uses an inclusive current-state lookup. row_to_frame = torch.searchsorted( frame_times, times, - right=oracle_analysis, + right=oracle_analysis or highrate_person_av, ) - 1 chosen = torch.full(times.shape, float("nan"), dtype=torch.float64) has_frame = row_to_frame >= 0 chosen[has_frame] = frame_times[row_to_frame[has_frame]] - if oracle_analysis: + if oracle_analysis or highrate_person_av: if bool((chosen[has_frame] > times[has_frame]).any()): raise ValueError(f"Oracle person mapping selected future evidence: {path}") elif bool((chosen[has_frame] >= times[has_frame]).any()): @@ -482,6 +525,36 @@ def load_frame_set_features( return direct_audio, frame_sets, contract +def load_highrate_person_av_features( + path: Path, + reference: dict[str, Any], + *, + sample_id: str, + dataset: str, + expected_contract: dict[str, Any] | None = None, +) -> tuple[torch.Tensor, FramePersonFeatures, dict[str, Any]]: + """Load frozen 25 Hz causal ASD states sampled at 5 Hz right edges. + + The returned dense tensor contains the single shared causal CPC audio state + and its validity bit. The visual dimension in the contract belongs to the + ragged, variable-cardinality set of visual-track state plus ASD score; it is + not a pooled image observation consumed by the head. + """ + + direct_audio, people, contract = load_frame_set_features( + path, + reference, + sample_id=sample_id, + dataset=dataset, + expected_contract=expected_contract, + highrate_person_av=True, + ) + config = contract["config"] + if config.get("audio_source") != "single-shared-traditional-cpc-stream": + raise ValueError(f"Unexpected high-rate sidecar audio source: {path}") + return direct_audio, people, contract + + def _build_previous_track_links( records: list[Any], frame_times: torch.Tensor, diff --git a/muvap_egospeak_stage3.py b/muvap_egospeak_stage3.py new file mode 100644 index 0000000..f923b1c --- /dev/null +++ b/muvap_egospeak_stage3.py @@ -0,0 +1,671 @@ +"""Shared data/model code for the MuVAP-only EgoSpeak Stage-3 baseline.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Iterable, Sequence + +import numpy as np +import torch +from torch import nn +from torch.utils.data import Dataset + +from easycom_head import ( + CLASS_NAMES, + NUM_CLASSES, + NUM_HORIZONS, + STEP_SECONDS, + _validated_binary_activity_targets, + average_precision_metrics, + marginal_probabilities_from_logits, + normalize_sample_name, +) +MODEL_FRAME_HZ = 25.0 +MODEL_FRAMES_PER_TASK_STEP = 5 +MUVAP_EGOSPEAK_INPUT_FORMAT = "muvap-egospeak-stage3-paper-reconstruction-v1" +MODEL_KIND = "muvap-egospeak-stage3-paper-reconstruction" + + +@dataclass(frozen=True) +class MuVAPSessionFeatures: + name: str + vap: torch.Tensor + people: torch.Tensor + person_valid: torch.BoolTensor + targets: torch.Tensor + task_frame_indices: torch.LongTensor + label_indices: torch.LongTensor + times: torch.Tensor + contract: dict[str, Any] + + @property + def vap_dim(self) -> int: + return int(self.vap.shape[-1]) + + @property + def asd_dim(self) -> int: + return int(self.people.shape[-1]) + + +def _validate_exact_grid(values: torch.Tensor, expected: torch.Tensor, name: str) -> None: + if values.dtype != torch.float64 or values.shape != expected.shape: + raise ValueError(f"Invalid {name} shape/dtype") + if not torch.allclose(values, expected, atol=1e-8, rtol=0): + raise ValueError(f"{name} does not follow the required right-edge grid") + + +def load_muvap_session( + feature_root: str | Path, + target_root: str | Path, + session: str | int, + *, + expected_contract: dict[str, Any] | None = None, + allow_partial: bool = False, +) -> MuVAPSessionFeatures: + """Load one 25 Hz frozen-backbone sidecar and fail closed on alignment.""" + + name = normalize_sample_name(session) + feature_path = Path(feature_root) / f"{name}.pt" + target_path = Path(target_root) / f"{name}.npy" + if not feature_path.is_file(): + raise FileNotFoundError(feature_path) + if not target_path.is_file(): + raise FileNotFoundError(target_path) + payload = torch.load(feature_path, map_location="cpu", weights_only=False) + if not isinstance(payload, dict) or payload.get("format") != MUVAP_EGOSPEAK_INPUT_FORMAT: + raise ValueError(f"Not a MuVAP EgoSpeak Stage-3 sidecar: {feature_path}") + if payload.get("dataset") != "easycom" or payload.get("sample_id") != name: + raise ValueError(f"Dataset/sample identity mismatch: {feature_path}") + if bool(payload.get("partial", False)) and not allow_partial: + raise ValueError(f"Partial sidecar cannot be used for training: {feature_path}") + contract = { + "format": payload["format"], + "dataset": payload["dataset"], + "config": payload.get("config"), + } + if not isinstance(contract["config"], dict): + raise ValueError(f"Missing MuVAP extraction config: {feature_path}") + if contract["config"].get("uses_omni_features") is not False: + raise ValueError(f"MuVAP-only sidecar does not exclude Omni: {feature_path}") + if contract["config"].get("uses_speech_labels") is not False: + raise ValueError(f"Feature extractor consumed speech labels: {feature_path}") + if expected_contract is not None and contract != expected_contract: + raise ValueError(f"MuVAP feature contract differs across sessions: {feature_path}") + + vap = payload.get("vap") + if ( + not isinstance(vap, torch.Tensor) + or vap.ndim != 2 + or not vap.is_floating_point() + or not torch.isfinite(vap).all() + ): + raise ValueError(f"vap must be finite floating [K,D]: {feature_path}") + num_ticks = len(vap) + frame_times = payload.get("frame_times") + if not isinstance(frame_times, torch.Tensor): + raise ValueError(f"Missing frame_times: {feature_path}") + _validate_exact_grid( + frame_times, + torch.arange(1, num_ticks + 1, dtype=torch.float64) / MODEL_FRAME_HZ, + "MuVAP frame_times", + ) + + times = payload.get("times") + label_indices = payload.get("label_indices") + task_frame_indices = payload.get("task_frame_indices") + if not all(isinstance(value, torch.Tensor) for value in (times, label_indices, task_frame_indices)): + raise ValueError(f"Missing task alignment tensors: {feature_path}") + num_steps = len(times) + _validate_exact_grid( + times, + torch.arange(1, num_steps + 1, dtype=torch.float64) * STEP_SECONDS, + "task times", + ) + expected_labels = torch.arange(num_steps, dtype=torch.long) + expected_frames = (expected_labels + 1) * MODEL_FRAMES_PER_TASK_STEP - 1 + if label_indices.dtype != torch.long or not torch.equal(label_indices, expected_labels): + raise ValueError(f"label_indices must be dense and zero-based: {feature_path}") + if task_frame_indices.dtype != torch.long or not torch.equal( + task_frame_indices, expected_frames + ): + raise ValueError(f"Invalid task-to-25Hz mapping: {feature_path}") + if num_ticks != num_steps * MODEL_FRAMES_PER_TASK_STEP: + raise ValueError(f"25Hz/5Hz length mismatch: {feature_path}") + + person = payload.get("persons") + if not isinstance(person, dict): + raise ValueError(f"Missing persons block: {feature_path}") + features = person.get("features") + offsets = person.get("frame_offsets") + records = person.get("records") + if ( + not isinstance(features, torch.Tensor) + or features.ndim != 2 + or not features.is_floating_point() + or not torch.isfinite(features).all() + ): + raise ValueError(f"Invalid fused ASD features: {feature_path}") + if ( + not isinstance(offsets, torch.Tensor) + or offsets.dtype != torch.long + or offsets.shape != (num_ticks + 1,) + or int(offsets[0]) != 0 + or int(offsets[-1]) != len(features) + or bool((offsets[1:] < offsets[:-1]).any()) + ): + raise ValueError(f"Invalid person frame offsets: {feature_path}") + if not isinstance(records, list) or len(records) != len(features): + raise ValueError(f"Person records/features differ: {feature_path}") + config = contract["config"] + if config.get("vap_dim") != vap.shape[1] or config.get("asd_dim") != features.shape[1]: + raise ValueError(f"Saved feature dimensions disagree with config: {feature_path}") + if not config.get("vap_release_sha256") or not config.get("asd_release_sha256"): + raise ValueError(f"Missing frozen release identity: {feature_path}") + counts = offsets[1:] - offsets[:-1] + max_people = int(counts.max()) if len(counts) else 0 + saved_max = payload.get("max_people") + if isinstance(saved_max, bool) or not isinstance(saved_max, int) or saved_max != max_people: + raise ValueError(f"Incorrect max_people metadata: {feature_path}") + people = torch.zeros(num_ticks, max_people, features.shape[1], dtype=features.dtype) + person_valid = torch.zeros(num_ticks, max_people, dtype=torch.bool) + for tick in range(num_ticks): + start, end = int(offsets[tick]), int(offsets[tick + 1]) + count = end - start + if count: + seen: set[str] = set() + expected_edge = (tick + 1) / MODEL_FRAME_HZ + for record in records[start:end]: + if not isinstance(record, dict): + raise ValueError(f"Invalid person record at tick {tick}: {feature_path}") + person_id = str(record.get("id", "")) + if person_id in seen: + raise ValueError(f"Duplicate person ID at tick {tick}: {feature_path}") + seen.add(person_id) + if record.get("asd_model_tick") != tick: + raise ValueError(f"Person record assigned to wrong tick: {feature_path}") + edge = record.get("asd_interval_end") + source_time = record.get("source_time") + if ( + not isinstance(edge, (int, float)) + or not isinstance(source_time, (int, float)) + or abs(float(edge) - expected_edge) > 1e-8 + or float(source_time) >= expected_edge - 1e-8 + ): + raise ValueError(f"Non-causal person record at tick {tick}: {feature_path}") + people[tick, :count] = features[start:end] + person_valid[tick, :count] = True + + target_array = np.load(target_path, allow_pickle=False) + targets = torch.from_numpy(np.asarray(target_array, dtype=np.float32)) + if targets.shape != (num_steps, NUM_CLASSES): + raise ValueError( + f"Target/feature length mismatch {targets.shape} vs ({num_steps},3): {name}" + ) + _validated_binary_activity_targets(targets) + return MuVAPSessionFeatures( + name=name, + vap=vap, + people=people, + person_valid=person_valid, + targets=targets, + task_frame_indices=task_frame_indices, + label_indices=label_indices, + times=times, + contract=contract, + ) + + +@dataclass(frozen=True) +class Stage3Window: + name: str + current_row: int + vap: torch.Tensor + people: torch.Tensor + person_valid: torch.BoolTensor + history_valid: torch.BoolTensor + targets: torch.Tensor + future_valid: torch.BoolTensor + + +class MuVAPWindowDataset(Dataset[Stage3Window]): + """One fixed past-only 25 Hz window for every 5 Hz prediction cutoff.""" + + def __init__( + self, + sessions: Sequence[MuVAPSessionFeatures], + *, + history_seconds: float, + require_full_future: bool, + ) -> None: + if not sessions: + raise ValueError("At least one session is required") + history_frames = round(history_seconds * MODEL_FRAME_HZ) + if history_frames <= 0 or not math.isclose( + history_frames / MODEL_FRAME_HZ, history_seconds, abs_tol=1e-8 + ): + raise ValueError("history_seconds must be a positive 25 Hz multiple") + self.sessions = list(sessions) + self.history_frames = history_frames + self.require_full_future = require_full_future + self.examples: list[tuple[int, int]] = [] + for session_index, session in enumerate(self.sessions): + stop = len(session.targets) - NUM_HORIZONS if require_full_future else len(session.targets) + self.examples.extend((session_index, row) for row in range(max(0, stop))) + + def __len__(self) -> int: + return len(self.examples) + + def __getitem__(self, index: int) -> Stage3Window: + session_index, row = self.examples[index] + session = self.sessions[session_index] + end = int(session.task_frame_indices[row]) + 1 + start = max(0, end - self.history_frames) + observed = end - start + left = self.history_frames - observed + vap = torch.zeros(self.history_frames, session.vap_dim, dtype=session.vap.dtype) + vap[left:] = session.vap[start:end] + people = torch.zeros( + self.history_frames, + session.people.shape[1], + session.asd_dim, + dtype=session.people.dtype, + ) + person_valid = torch.zeros( + self.history_frames, session.people.shape[1], dtype=torch.bool + ) + if session.people.shape[1]: + people[left:] = session.people[start:end] + person_valid[left:] = session.person_valid[start:end] + history_valid = torch.zeros(self.history_frames, dtype=torch.bool) + history_valid[left:] = True + targets = torch.zeros(NUM_HORIZONS, NUM_CLASSES, dtype=torch.float32) + future_valid = torch.zeros(NUM_HORIZONS, dtype=torch.bool) + available = min(NUM_HORIZONS, len(session.targets) - row - 1) + if available: + targets[:available] = session.targets[row + 1 : row + 1 + available] + future_valid[:available] = True + if self.require_full_future and not bool(future_valid.all()): + raise RuntimeError("Training dataset exposed an incomplete future target") + return Stage3Window( + name=session.name, + current_row=row, + vap=vap, + people=people, + person_valid=person_valid, + history_valid=history_valid, + targets=targets, + future_valid=future_valid, + ) + + +def collate_stage3_windows(items: Sequence[Stage3Window]) -> dict[str, Any]: + if not items: + raise ValueError("Cannot collate an empty batch") + max_people = max(item.people.shape[1] for item in items) + batch_people = torch.zeros( + len(items), + items[0].people.shape[0], + max_people, + items[0].people.shape[-1], + dtype=items[0].people.dtype, + ) + batch_person_valid = torch.zeros( + len(items), items[0].people.shape[0], max_people, dtype=torch.bool + ) + for index, item in enumerate(items): + count = item.people.shape[1] + if count: + batch_people[index, :, :count] = item.people + batch_person_valid[index, :, :count] = item.person_valid + return { + "names": [item.name for item in items], + "current_rows": torch.tensor([item.current_row for item in items], dtype=torch.long), + "vap": torch.stack([item.vap for item in items]), + "people": batch_people, + "person_valid": batch_person_valid, + "history_valid": torch.stack([item.history_valid for item in items]), + "targets": torch.stack([item.targets for item in items]), + "future_valid": torch.stack([item.future_valid for item in items]), + } + + +def sinusoidal_positions(length: int, dimension: int) -> torch.Tensor: + position = torch.arange(length, dtype=torch.float32)[:, None] + frequencies = torch.exp( + torch.arange(0, dimension, 2, dtype=torch.float32) + * (-math.log(10_000.0) / dimension) + ) + result = torch.zeros(length, dimension, dtype=torch.float32) + result[:, 0::2] = torch.sin(position * frequencies) + result[:, 1::2] = torch.cos(position * frequencies[: result[:, 1::2].shape[1]]) + return result + + +class MuVAPEgoSpeakStage3(nn.Module): + """Paper-inspired MuVAP frame/temporal fusion with an EgoSpeak head. + + Published VAP and ASD states are inputs and remain outside this trainable + module. The model trains only the Stage-3 projections, same-frame person + cross-attention, causal temporal transformer and the task head. + """ + + def __init__( + self, + *, + vap_dim: int, + asd_dim: int, + model_dim: int = 256, + num_heads: int = 4, + dropout: float = 0.1, + temporal_layers: int = 1, + history_frames: int = 200, + num_classes: int = NUM_CLASSES, + person_evidence: str = "aligned", + ) -> None: + super().__init__() + if min(vap_dim, asd_dim, model_dim, num_heads, temporal_layers, history_frames) <= 0: + raise ValueError("Stage-3 dimensions must be positive") + if model_dim % num_heads: + raise ValueError("model_dim must be divisible by num_heads") + if person_evidence not in {"aligned", "none"}: + raise ValueError("person_evidence must be 'aligned' or 'none'") + self.vap_dim = vap_dim + self.asd_dim = asd_dim + self.model_dim = model_dim + self.num_heads = num_heads + self.dropout_probability = dropout + self.temporal_layers = temporal_layers + self.history_frames = history_frames + self.num_classes = num_classes + self.person_evidence = person_evidence + self.vap_projection = nn.Sequential( + nn.Linear(vap_dim, model_dim), + nn.LayerNorm(model_dim), + ) + self.asd_projection = nn.Sequential( + nn.Linear(asd_dim, model_dim), + nn.LayerNorm(model_dim), + ) + self.null_person = nn.Parameter(torch.zeros(1, 1, model_dim)) + nn.init.normal_(self.null_person, std=0.02) + self.frame_attention = nn.MultiheadAttention( + model_dim, + num_heads, + dropout=dropout, + batch_first=True, + ) + self.frame_attention_norm = nn.LayerNorm(model_dim) + self.frame_ff_norm = nn.LayerNorm(model_dim) + self.frame_ff = nn.Sequential( + nn.Linear(model_dim, model_dim * 4), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(model_dim * 4, model_dim), + ) + self.dropout = nn.Dropout(dropout) + temporal_layer = nn.TransformerEncoderLayer( + d_model=model_dim, + nhead=num_heads, + dim_feedforward=model_dim * 4, + dropout=dropout, + activation="gelu", + batch_first=True, + norm_first=True, + ) + self.temporal = nn.TransformerEncoder( + temporal_layer, + num_layers=temporal_layers, + norm=nn.LayerNorm(model_dim), + ) + self.head = nn.Linear(model_dim, NUM_HORIZONS * num_classes) + self.register_buffer( + "position_encoding", + sinusoidal_positions(history_frames, model_dim), + persistent=False, + ) + self.register_buffer( + "causal_mask", + torch.triu(torch.ones(history_frames, history_frames, dtype=torch.bool), diagonal=1), + persistent=False, + ) + + def forward( + self, + vap: torch.Tensor, + people: torch.Tensor, + person_valid: torch.BoolTensor, + history_valid: torch.BoolTensor, + *, + return_embedding: bool = False, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if vap.ndim != 3 or vap.shape[-1] != self.vap_dim: + raise ValueError(f"Expected VAP [B,H,{self.vap_dim}], got {tuple(vap.shape)}") + batch, history, _ = vap.shape + if history != self.history_frames: + raise ValueError(f"Expected {self.history_frames} history frames, got {history}") + if people.ndim != 4 or people.shape[:2] != (batch, history) or people.shape[-1] != self.asd_dim: + raise ValueError( + f"Expected people [B,H,N,{self.asd_dim}], got {tuple(people.shape)}" + ) + if person_valid.dtype != torch.bool or person_valid.shape != people.shape[:3]: + raise ValueError("person_valid must be bool [B,H,N]") + if history_valid.dtype != torch.bool or history_valid.shape != (batch, history): + raise ValueError("history_valid must be bool [B,H]") + if not bool(history_valid[:, -1].all()): + raise ValueError("Every Stage-3 window must end at an observed cutoff") + if not torch.isfinite(vap).all() or not torch.isfinite(people).all(): + raise ValueError("Stage-3 inputs contain NaN/Inf") + if self.person_evidence == "none": + # Capacity-matched VAP-only control. The same frame and temporal + # blocks remain trainable; only observed ASD tokens are hidden. + person_valid = torch.zeros_like(person_valid) + + query = self.vap_projection(vap).reshape(batch * history, 1, self.model_dim) + count = people.shape[2] + person_tokens = self.asd_projection(people).reshape( + batch * history, count, self.model_dim + ) + null = self.null_person.expand(batch * history, -1, -1) + person_tokens = torch.cat((person_tokens, null), dim=1) + key_padding = torch.cat( + ( + ~person_valid.reshape(batch * history, count), + torch.zeros(batch * history, 1, dtype=torch.bool, device=vap.device), + ), + dim=1, + ) + selected, _ = self.frame_attention( + query, + person_tokens, + person_tokens, + key_padding_mask=key_padding, + need_weights=False, + ) + frame = self.frame_attention_norm(query + self.dropout(selected)) + frame = frame + self.dropout(self.frame_ff(self.frame_ff_norm(frame))) + frame = frame.reshape(batch, history, self.model_dim) + frame = frame + self.position_encoding[None].to(frame.dtype) + + # Windows are left-padded so the final token is always the current + # 5 Hz cutoff. A plain key-padding mask leaves the earliest padded + # queries with no visible key under causality, which can yield NaNs + # (and can contaminate later layers). Give each padded query only its + # own diagonal while every observed query still masks every padded + # key. The mask is repeated because MultiheadAttention accepts one + # [H,H] matrix per batch/head pair. + temporal_mask = self.causal_mask[None].expand(batch, -1, -1).clone() + temporal_mask |= (~history_valid)[:, None, :] + padded_batch, padded_query = torch.where(~history_valid) + temporal_mask[padded_batch, padded_query, padded_query] = False + temporal_mask = temporal_mask.repeat_interleave(self.num_heads, dim=0) + temporal = self.temporal( + frame, + mask=temporal_mask, + ) + embedding = temporal[:, -1] + logits = self.head(embedding).reshape(batch, NUM_HORIZONS, self.num_classes) + if return_embedding: + return logits, embedding + return logits + + +def model_config(model: MuVAPEgoSpeakStage3) -> dict[str, Any]: + return { + "vap_dim": model.vap_dim, + "asd_dim": model.asd_dim, + "model_dim": model.model_dim, + "num_heads": model.num_heads, + "dropout": model.dropout_probability, + "temporal_layers": model.temporal_layers, + "history_frames": model.history_frames, + "num_classes": model.num_classes, + "person_evidence": model.person_evidence, + } + + +def build_model(config: dict[str, Any]) -> MuVAPEgoSpeakStage3: + required = { + "vap_dim", + "asd_dim", + "model_dim", + "num_heads", + "dropout", + "temporal_layers", + "history_frames", + "num_classes", + "person_evidence", + } + if set(config) != required: + raise ValueError(f"Invalid Stage-3 model config keys: {sorted(set(config) ^ required)}") + return MuVAPEgoSpeakStage3(**config) + + +@torch.inference_mode() +def predict_loader( + model: MuVAPEgoSpeakStage3, + loader: Iterable[dict[str, Any]], + *, + device: torch.device, + loss_name: str, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[str]]: + model.eval() + probabilities, targets, rows, names = [], [], [], [] + for batch in loader: + logits = model( + batch["vap"].to(device=device, dtype=torch.float32, non_blocking=True), + batch["people"].to(device=device, dtype=torch.float32, non_blocking=True), + batch["person_valid"].to(device=device, non_blocking=True), + batch["history_valid"].to(device=device, non_blocking=True), + ) + probabilities.append(marginal_probabilities_from_logits(logits, loss_name).cpu()) + targets.append(batch["targets"]) + rows.append(batch["current_rows"]) + names.extend(batch["names"]) + return ( + torch.cat(probabilities), + torch.cat(targets), + torch.cat(rows), + names, + ) + + +def evaluate_valid_predictions( + probabilities: torch.Tensor, + targets: torch.Tensor, +) -> dict[str, Any]: + if probabilities.shape != targets.shape or probabilities.shape[-2:] != ( + NUM_HORIZONS, + NUM_CLASSES, + ): + raise ValueError("Evaluation expects matching [N,10,3] arrays") + results = [] + for horizon in range(NUM_HORIZONS): + result = average_precision_metrics( + probabilities[:, horizon].numpy(), + targets[:, horizon].numpy(), + ) + result["horizon"] = horizon + 1 + result["seconds"] = (horizon + 1) * STEP_SECONDS + results.append(result) + return { + "horizons": results, + "mean_map": float(np.mean([result["map"] for result in results])), + } + + +def align_session_predictions( + probabilities: np.ndarray, + num_steps: int, +) -> tuple[np.ndarray, np.ndarray]: + if probabilities.shape != (num_steps, NUM_HORIZONS, NUM_CLASSES): + raise ValueError( + f"Expected [{num_steps},{NUM_HORIZONS},{NUM_CLASSES}], got " + f"{probabilities.shape}" + ) + predictions = np.zeros((num_steps, NUM_CLASSES, NUM_HORIZONS), dtype=np.float32) + valid = np.zeros((num_steps, NUM_HORIZONS), dtype=bool) + current = np.arange(num_steps) + for horizon in range(NUM_HORIZONS): + future = current + horizon + 1 + usable = future < num_steps + predictions[future[usable], :, horizon] = probabilities[usable, horizon] + valid[future[usable], horizon] = True + return predictions, valid + + +def evaluate_aligned( + session_predictions: dict[str, np.ndarray], + session_valid: dict[str, np.ndarray], + session_targets: dict[str, np.ndarray], +) -> dict[str, Any]: + protocols: dict[str, Any] = {} + for protocol in ("egospeak_full_timeline", "causal_valid_only"): + horizons = [] + for horizon in range(NUM_HORIZONS): + if protocol == "egospeak_full_timeline": + scores = np.concatenate( + [value[:, :, horizon] for value in session_predictions.values()] + ) + targets = np.concatenate(list(session_targets.values())) + else: + scores = np.concatenate( + [ + session_predictions[name][session_valid[name][:, horizon], :, horizon] + for name in session_predictions + ] + ) + targets = np.concatenate( + [ + session_targets[name][session_valid[name][:, horizon]] + for name in session_predictions + ] + ) + result = average_precision_metrics(scores, targets) + result["horizon"] = horizon + 1 + result["seconds"] = (horizon + 1) * STEP_SECONDS + horizons.append(result) + protocols[protocol] = { + "horizons": horizons, + "mean_map": float(np.mean([result["map"] for result in horizons])), + } + return protocols + + +def print_metrics(metrics: dict[str, Any]) -> None: + for protocol in ("egospeak_full_timeline", "causal_valid_only"): + print(f"\n{protocol}:") + for result in metrics[protocol]["horizons"]: + values = ", ".join( + f"{name}={result['class_ap'][name] * 100:.5f}" for name in CLASS_NAMES + ) + print( + f"Action anticipation ({result['seconds']:.2f}s) perframe mAP: " + f"{result['map'] * 100:.5f}" + ) + print(f" AP: {values}; n={result['num_steps']}") + print( + f"Action anticipation (mean) perframe mAP: " + f"{metrics[protocol]['mean_map'] * 100:.5f}" + ) diff --git a/requirements.txt b/requirements.txt index 18313e9..fc2f7cf 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,5 @@ torch +torchaudio numpy # Default for Thinker replay. The B1 local extractor also supports 4.57.6; # keep that version in existing GPU containers (see docs/local_feature_baseline.md). @@ -9,6 +10,8 @@ librosa Pillow tqdm scikit-learn +PyYAML +einops # Used by train_omni_turn_lora.py; already included in the ms-swift image. peft diff --git a/run_easycom_causal_person_adapter_v2.sh b/run_easycom_causal_person_adapter_v2.sh new file mode 100755 index 0000000..8f632b5 --- /dev/null +++ b/run_easycom_causal_person_adapter_v2.sh @@ -0,0 +1,199 @@ +#!/usr/bin/env bash +# Foreground matched EasyCom pooled AV -> causal person-track adapter ablations. +# Usage: bash run_easycom_causal_person_adapter_v2.sh +# [all|pooled|v2|egospeak-gru|test] [seed ...] +set -euo pipefail +cd "$(dirname "$0")" + +FEATURE_ROOT=${FEATURE_ROOT:-/data1/yfzx/wangx/Datasets/easycom/omni_features/open_stream_2fps_8s_last4_stack} +TARGET_ROOT=${TARGET_ROOT:-/data1/yfzx/wangx/Datasets/easycom/easycom_target_frame} +LOCAL_ROOT=${LOCAL_ROOT:-/data1/yfzx/wangx/Datasets/easycom/omni_features/local_qwen_1s_encoded_tf4576} +CHECKPOINT_ROOT=${CHECKPOINT_ROOT:-checkpoints/easycom_causal_person_adapter_v2} +RESULT_ROOT=${RESULT_ROOT:-results/easycom_causal_person_adapter_v2} +TEST_RESULT_ROOT=${TEST_RESULT_ROOT:-results/easycom_causal_person_adapter_v2_test} +EPOCHS=${EPOCHS:-4} +if [[ -z "${WARMUP_EPOCHS:-}" ]]; then + if (( EPOCHS == 1 )); then + WARMUP_EPOCHS=0 + else + WARMUP_EPOCHS=1.6 + fi +fi +LOCAL_PROJECTION_DIM=${LOCAL_PROJECTION_DIM:-256} +PERSON_ATTENTION_DIM=${PERSON_ATTENTION_DIM:-256} +PERSON_ATTENTION_HEADS=${PERSON_ATTENTION_HEADS:-4} +# Four observations at the 2 FPS local grid cover at most two causal seconds. +PERSON_HISTORY_FRAMES=${PERSON_HISTORY_FRAMES:-4} +PERSON_TRACK_MAX_GAP_SECONDS=${PERSON_TRACK_MAX_GAP_SECONDS:-1.0} +SUMMARIZE=${SUMMARIZE:-1} +read -r -a SESSION_ARGS <<< "${SESSIONS:-4 5 6 7 8 9 10 11 12}" +read -r -a DEV_SESSION_ARGS <<< "${DEV_SESSIONS:-12}" +read -r -a SCREEN_SESSION_ARGS <<< "${SCREEN_SESSIONS:-12}" +export OMP_NUM_THREADS=${OMP_NUM_THREADS:-4} +export MKL_NUM_THREADS=${MKL_NUM_THREADS:-4} +export OPENBLAS_NUM_THREADS=${OPENBLAS_NUM_THREADS:-4} + +variant=${1:-all} +if (( $# )); then shift; fi +if (( $# )); then seeds=("$@"); else seeds=(0 1 2); fi +case "$variant" in + all) variants=(pooled v2 egospeak-gru) ;; + pooled|v2|egospeak-gru) variants=("$variant") ;; + test) variants=(pooled v2 egospeak-gru) ;; + *) echo "Expected all, pooled, v2, egospeak-gru, or test" >&2; exit 2 ;; +esac + +mkdir -p "$CHECKPOINT_ROOT" "$RESULT_ROOT" + +infer_checkpoint() { + local checkpoint=$1 result=$2 + shift 2 + if [[ -e "$result/metrics.json" ]]; then + echo "Refusing to overwrite $result/metrics.json" >&2 + return 1 + fi + python infer_turn_head.py \ + --feature-root "$FEATURE_ROOT" \ + --target-root "$TARGET_ROOT" \ + --local-feature-root "$LOCAL_ROOT" \ + --checkpoint "$checkpoint" \ + --output-dir "$result" \ + --sessions "$@" \ + --expected-mode open-stream-replay \ + --device cuda:0 +} + +train_and_infer_pooled() { + local seed=$1 checkpoint result + checkpoint="$CHECKPOINT_ROOT/pooled_seed${seed}.pth" + result="$RESULT_ROOT/pooled_seed${seed}" + if [[ -e "$checkpoint" || -e "$result/metrics.json" ]]; then + echo "Refusing to overwrite $checkpoint or $result" >&2 + return 1 + fi + python train_turn_head.py \ + --feature-root "$FEATURE_ROOT" \ + --target-root "$TARGET_ROOT" \ + --local-feature-root "$LOCAL_ROOT" \ + --local-modalities av \ + --local-fusion pooled \ + --local-projection-dim "$LOCAL_PROJECTION_DIM" \ + --output "$checkpoint" \ + --sessions "${SESSION_ARGS[@]}" \ + --dev-sessions "${DEV_SESSION_ARGS[@]}" \ + --epochs "$EPOCHS" \ + --batch-size 512 \ + --learning-rate 7e-5 \ + --weight-decay 5e-5 \ + --warmup-epochs "$WARMUP_EPOCHS" \ + --warmup-factor 0.3 \ + --dropout 0.2 \ + --loss lstr-mce \ + --transition-objective none \ + --layer-fusion concat-mlp \ + --expected-mode open-stream-replay \ + --seed "$seed" \ + --device cuda:0 + infer_checkpoint "$checkpoint" "$result" "${SCREEN_SESSION_ARGS[@]}" +} + +train_and_infer_adapter() { + local seed=$1 mode=${2:-v2} base checkpoint result local_fusion + base="$CHECKPOINT_ROOT/pooled_seed${seed}.pth" + case "$mode" in + v2) + checkpoint="$CHECKPOINT_ROOT/person_adapter_v2_seed${seed}.pth" + result="$RESULT_ROOT/person_adapter_v2_seed${seed}" + local_fusion=causal-person-adapter-v2 + ;; + egospeak-gru) + checkpoint="$CHECKPOINT_ROOT/egospeak_gru_seed${seed}.pth" + result="$RESULT_ROOT/egospeak_gru_seed${seed}" + local_fusion=causal-person-adapter-egospeak-gru + ;; + *) echo "Unknown person adapter: $mode" >&2; return 2 ;; + esac + if [[ ! -s "$base" ]]; then + echo "Missing matched pooled checkpoint: $base; run pooled or all first" >&2 + return 1 + fi + if [[ -e "$checkpoint" || -e "$result/metrics.json" ]]; then + echo "Refusing to overwrite $checkpoint or $result" >&2 + return 1 + fi + python train_turn_head.py \ + --feature-root "$FEATURE_ROOT" \ + --target-root "$TARGET_ROOT" \ + --local-feature-root "$LOCAL_ROOT" \ + --local-modalities av \ + --local-fusion "$local_fusion" \ + --local-projection-dim "$LOCAL_PROJECTION_DIM" \ + --person-attention-dim "$PERSON_ATTENTION_DIM" \ + --person-attention-heads "$PERSON_ATTENTION_HEADS" \ + --person-history-frames "$PERSON_HISTORY_FRAMES" \ + --person-track-max-gap-seconds "$PERSON_TRACK_MAX_GAP_SECONDS" \ + --initialize-pooled-checkpoint "$base" \ + --output "$checkpoint" \ + --sessions "${SESSION_ARGS[@]}" \ + --dev-sessions "${DEV_SESSION_ARGS[@]}" \ + --epochs "$EPOCHS" \ + --batch-size 512 \ + --learning-rate 7e-5 \ + --weight-decay 5e-5 \ + --warmup-epochs "$WARMUP_EPOCHS" \ + --warmup-factor 0.3 \ + --dropout 0.2 \ + --loss lstr-mce \ + --transition-objective none \ + --layer-fusion concat-mlp \ + --expected-mode open-stream-replay \ + --seed "$seed" \ + --device cuda:0 + infer_checkpoint "$checkpoint" "$result" "${SCREEN_SESSION_ARGS[@]}" +} + +if [[ "$variant" == test ]]; then + mkdir -p "$TEST_RESULT_ROOT" + for selected in "${variants[@]}"; do + for seed in "${seeds[@]}"; do + if [[ "$selected" == pooled ]]; then + checkpoint="$CHECKPOINT_ROOT/pooled_seed${seed}.pth" + result="$TEST_RESULT_ROOT/pooled_seed${seed}" + elif [[ "$selected" == v2 ]]; then + checkpoint="$CHECKPOINT_ROOT/person_adapter_v2_seed${seed}.pth" + result="$TEST_RESULT_ROOT/person_adapter_v2_seed${seed}" + else + checkpoint="$CHECKPOINT_ROOT/egospeak_gru_seed${seed}.pth" + result="$TEST_RESULT_ROOT/egospeak_gru_seed${seed}" + fi + if [[ ! -s "$checkpoint" ]]; then + echo "Missing checkpoint: $checkpoint" >&2 + echo "The test mode is inference-only; train and freeze the screen first." >&2 + exit 1 + fi + echo "$(date): final test ${selected}, seed=${seed}" + infer_checkpoint "$checkpoint" "$result" 3 2 1 + done + done + summary_root=$TEST_RESULT_ROOT +else + for selected in "${variants[@]}"; do + for seed in "${seeds[@]}"; do + echo "$(date): ${selected}, seed=${seed}" + if [[ "$selected" == pooled ]]; then + train_and_infer_pooled "$seed" + else + train_and_infer_adapter "$seed" "$selected" + fi + done + done + summary_root=$RESULT_ROOT +fi + +if [[ "$SUMMARIZE" == 1 ]]; then + python summarize_local_baselines.py \ + --root "$summary_root" \ + --dataset easycom \ + --protocol causal_valid_only \ + --paired-deltas +fi diff --git a/run_easycom_highrate_person_av.sh b/run_easycom_highrate_person_av.sh new file mode 100755 index 0000000..0fa80d0 --- /dev/null +++ b/run_easycom_highrate_person_av.sh @@ -0,0 +1,159 @@ +#!/usr/bin/env bash +# Global Omni with Qwen audio, traditional audio, and high-rate person evidence. +# Usage: bash run_easycom_highrate_person_av.sh \ +# [all|global|qwen_audio|traditional_audio|highrate|test] [seed ...] +set -euo pipefail +cd "$(dirname "$0")" + +FEATURE_ROOT=${FEATURE_ROOT:-/data1/yfzx/wangx/Datasets/easycom/omni_features/open_stream_2fps_8s_last4_stack} +TARGET_ROOT=${TARGET_ROOT:-/data1/yfzx/wangx/Datasets/easycom/easycom_target_frame} +QWEN_LOCAL_ROOT=${QWEN_LOCAL_ROOT:-/data1/yfzx/wangx/Datasets/easycom/omni_features/local_qwen_1s_encoded_tf4576} +HIGHRATE_LOCAL_ROOT=${HIGHRATE_LOCAL_ROOT:-/data1/yfzx/wangx/Datasets/easycom/omni_features/local_muvap_asd_25hz_v2} +CHECKPOINT_ROOT=${CHECKPOINT_ROOT:-checkpoints/easycom_highrate_person_av} +RESULT_ROOT=${RESULT_ROOT:-results/easycom_highrate_person_av_screen} +TEST_RESULT_ROOT=${TEST_RESULT_ROOT:-results/easycom_highrate_person_av_test} +EPOCHS=${EPOCHS:-4} +WARMUP_EPOCHS=${WARMUP_EPOCHS:-1.6} +LOCAL_PROJECTION_DIM=${LOCAL_PROJECTION_DIM:-256} +PERSON_ATTENTION_DIM=${PERSON_ATTENTION_DIM:-256} +PERSON_ATTENTION_HEADS=${PERSON_ATTENTION_HEADS:-4} +SUMMARIZE=${SUMMARIZE:-1} +read -r -a TRAIN_SESSION_ARGS <<< "${TRAIN_SESSIONS:-4 5 6 7 8 9 10 11 12}" +read -r -a DEV_SESSION_ARGS <<< "${DEV_SESSIONS:-12}" +read -r -a SCREEN_SESSION_ARGS <<< "${SCREEN_SESSIONS:-12}" +export OMP_NUM_THREADS=${OMP_NUM_THREADS:-4} +export MKL_NUM_THREADS=${MKL_NUM_THREADS:-4} +export OPENBLAS_NUM_THREADS=${OPENBLAS_NUM_THREADS:-4} + +variant=${1:-all} +if (( $# )); then shift; fi +if (( $# )); then seeds=("$@"); else seeds=(0 1 2); fi +case "$variant" in + all) variants=(global qwen_audio traditional_audio highrate) ;; + global|qwen_audio|traditional_audio|highrate) variants=("$variant") ;; + test) variants=(global qwen_audio traditional_audio highrate) ;; + *) echo "Expected all, global, qwen_audio, traditional_audio, highrate, or test" >&2; exit 2 ;; +esac + +mkdir -p "$CHECKPOINT_ROOT" "$RESULT_ROOT" + +run_inference() { + local checkpoint=$1 local_root=$2 output=$3 + shift 3 + if [[ -e "$output/metrics.json" ]]; then + echo "Refusing to overwrite $output/metrics.json" >&2 + return 1 + fi + local local_args=() + if [[ -n "$local_root" ]]; then + local_args=(--local-feature-root "$local_root") + fi + python infer_turn_head.py \ + --feature-root "$FEATURE_ROOT" \ + --target-root "$TARGET_ROOT" \ + "${local_args[@]}" \ + --checkpoint "$checkpoint" \ + --output-dir "$output" \ + --sessions "$@" \ + --expected-mode open-stream-replay \ + --device cuda:0 +} + +if [[ "$variant" == test ]]; then + mkdir -p "$TEST_RESULT_ROOT" + for selected in "${variants[@]}"; do + for seed in "${seeds[@]}"; do + checkpoint="$CHECKPOINT_ROOT/${selected}_seed${seed}.pth" + [[ -s "$checkpoint" ]] || { echo "Missing checkpoint: $checkpoint" >&2; exit 1; } + if [[ "$selected" == global ]]; then + local_root= + elif [[ "$selected" == qwen_audio ]]; then + local_root=$QWEN_LOCAL_ROOT + else + local_root=$HIGHRATE_LOCAL_ROOT + fi + run_inference "$checkpoint" "$local_root" \ + "$TEST_RESULT_ROOT/${selected}_seed${seed}" 3 2 1 + done + done + summary_root=$TEST_RESULT_ROOT +else + for selected in "${variants[@]}"; do + for seed in "${seeds[@]}"; do + checkpoint="$CHECKPOINT_ROOT/${selected}_seed${seed}.pth" + result="$RESULT_ROOT/${selected}_seed${seed}" + initialization_args=() + local_args=() + if [[ -e "$checkpoint" || -e "$result/metrics.json" ]]; then + echo "Refusing to overwrite $checkpoint or $result" >&2 + exit 1 + fi + if [[ "$selected" == global ]]; then + local_root= + modalities=av + fusion=pooled + elif [[ "$selected" == qwen_audio ]]; then + local_root=$QWEN_LOCAL_ROOT + modalities=audio + fusion=pooled + elif [[ "$selected" == traditional_audio ]]; then + local_root=$HIGHRATE_LOCAL_ROOT + modalities=audio + fusion=pooled + else + local_root=$HIGHRATE_LOCAL_ROOT + modalities=av + fusion=highrate-person-av-attention + audio_checkpoint="$CHECKPOINT_ROOT/traditional_audio_seed${seed}.pth" + [[ -s "$audio_checkpoint" ]] || { + echo "Missing matched audio checkpoint: $audio_checkpoint" >&2 + echo "Run: bash run_easycom_highrate_person_av.sh traditional_audio $seed" >&2 + exit 1 + } + initialization_args=(--initialize-pooled-checkpoint "$audio_checkpoint") + fi + if [[ -n "$local_root" ]]; then + local_args=( + --local-feature-root "$local_root" + --local-modalities "$modalities" + ) + fi + echo "$(date): ${selected}, seed=${seed}" + python train_turn_head.py \ + --feature-root "$FEATURE_ROOT" \ + --target-root "$TARGET_ROOT" \ + "${local_args[@]}" \ + --local-fusion "$fusion" \ + --local-projection-dim "$LOCAL_PROJECTION_DIM" \ + --person-attention-dim "$PERSON_ATTENTION_DIM" \ + --person-attention-heads "$PERSON_ATTENTION_HEADS" \ + "${initialization_args[@]}" \ + --output "$checkpoint" \ + --sessions "${TRAIN_SESSION_ARGS[@]}" \ + --dev-sessions "${DEV_SESSION_ARGS[@]}" \ + --epochs "$EPOCHS" \ + --batch-size 512 \ + --learning-rate 7e-5 \ + --weight-decay 5e-5 \ + --warmup-epochs "$WARMUP_EPOCHS" \ + --warmup-factor 0.3 \ + --dropout 0.2 \ + --loss lstr-mce \ + --transition-objective none \ + --layer-fusion concat-mlp \ + --expected-mode open-stream-replay \ + --seed "$seed" \ + --device cuda:0 + run_inference "$checkpoint" "$local_root" "$result" "${SCREEN_SESSION_ARGS[@]}" + done + done + summary_root=$RESULT_ROOT +fi + +if [[ "$SUMMARIZE" == 1 ]]; then + python summarize_local_baselines.py \ + --root "$summary_root" \ + --dataset easycom \ + --protocol causal_valid_only \ + --paired-deltas +fi diff --git a/run_easycom_muvap_stage3.sh b/run_easycom_muvap_stage3.sh new file mode 100755 index 0000000..e3e45bd --- /dev/null +++ b/run_easycom_muvap_stage3.sh @@ -0,0 +1,144 @@ +#!/usr/bin/env bash +# MuVAP-only, paper-reconstructed Stage 3 on EgoSpeak/EasyCom targets. +# Usage: bash run_easycom_muvap_stage3.sh [all|train|screen|test] [seed ...] +set -euo pipefail +cd "$(dirname "$0")" + +MUVAP_FEATURE_ROOT=${MUVAP_FEATURE_ROOT:-/data1/yfzx/wangx/Datasets/easycom/omni_features/muvap_stage3_cpc_25hz_v1} +TARGET_ROOT=${TARGET_ROOT:-/data1/yfzx/wangx/Datasets/easycom/easycom_target_frame} +CHECKPOINT_ROOT=${CHECKPOINT_ROOT:-checkpoints/easycom_muvap_stage3_5ep} +SCREEN_ROOT=${SCREEN_ROOT:-results/easycom_muvap_stage3_screen} +TEST_ROOT=${TEST_ROOT:-results/easycom_muvap_stage3_test} +EPOCHS=${EPOCHS:-5} +BATCH_SIZE=${BATCH_SIZE:-16} +HISTORY_SECONDS=${HISTORY_SECONDS:-8} +NUM_WORKERS=${NUM_WORKERS:-0} +PERSON_EVIDENCE=${PERSON_EVIDENCE:-aligned} +export OMP_NUM_THREADS=${OMP_NUM_THREADS:-4} +export MKL_NUM_THREADS=${MKL_NUM_THREADS:-4} +export OPENBLAS_NUM_THREADS=${OPENBLAS_NUM_THREADS:-4} + +mode=${1:-all} +if (( $# )); then shift; fi +if (( $# )); then seeds=("$@"); else seeds=(0 1 2); fi +case "$mode" in + all|train|screen|test) ;; + *) echo "Expected all, train, screen, or test" >&2; exit 2 ;; +esac +case "$PERSON_EVIDENCE" in + aligned) condition=stage3 ;; + none) condition=vap_only ;; + *) echo "PERSON_EVIDENCE must be aligned or none" >&2; exit 2 ;; +esac + +mkdir -p "$CHECKPOINT_ROOT" + +write_result_readme() { + local root=$1 + mkdir -p "$root" + if [[ ! -e "$root/README.md" ]]; then + { + echo "# EasyCom MuVAP-only Stage-3 reconstruction" + echo + echo "Frozen published role-VAP and causal ASD inputs; trainable paper-reconstructed frame/temporal fusion and EgoSpeak head." + echo "Run names: stage3 uses aligned ASD person evidence; vap_only masks it in the same-capacity architecture." + echo "Split: Sessions 4--11 train, Session 12 development, Sessions 1--3 official test." + echo "Protocol: strictly causal 25 Hz input, 5 Hz cutoffs, 0.2--2.0 s horizons." + echo "See ../../docs/easycom_muvap_stage3_baseline.md for the claim boundary and commands." + } > "$root/README.md" + fi +} + +train_seed() { + local seed=$1 + local checkpoint="$CHECKPOINT_ROOT/${condition}_seed${seed}.pth" + if [[ -e "$checkpoint" || -e "${checkpoint}.json" ]]; then + echo "Refusing to overwrite $checkpoint or ${checkpoint}.json" >&2 + return 1 + fi + python train_muvap_egospeak_stage3.py \ + --feature-root "$MUVAP_FEATURE_ROOT" \ + --target-root "$TARGET_ROOT" \ + --output "$checkpoint" \ + --sessions 4 5 6 7 8 9 10 11 12 \ + --dev-sessions 12 \ + --epochs "$EPOCHS" \ + --batch-size "$BATCH_SIZE" \ + --learning-rate 1e-3 \ + --min-learning-rate 1e-4 \ + --weight-decay 1e-2 \ + --warmup-ratio 0.1 \ + --gradient-clip 1.0 \ + --history-seconds "$HISTORY_SECONDS" \ + --model-dim 256 \ + --num-heads 4 \ + --temporal-layers 1 \ + --dropout 0.1 \ + --loss lstr-mce \ + --person-evidence "$PERSON_EVIDENCE" \ + --seed "$seed" \ + --num-workers "$NUM_WORKERS" \ + --device cuda:0 +} + +infer_seed() { + local seed=$1 + local split=$2 + local checkpoint="$CHECKPOINT_ROOT/${condition}_seed${seed}.pth" + [[ -s "$checkpoint" ]] || { + echo "Missing checkpoint: $checkpoint" >&2 + return 1 + } + local root sessions + if [[ "$split" == screen ]]; then + root=$SCREEN_ROOT + sessions=(12) + else + root=$TEST_ROOT + sessions=(3 2 1) + fi + write_result_readme "$root" + local output="$root/${condition}_seed${seed}" + [[ ! -e "$output" ]] || { + echo "Refusing to overwrite $output" >&2 + return 1 + } + python infer_muvap_egospeak_stage3.py \ + --feature-root "$MUVAP_FEATURE_ROOT" \ + --target-root "$TARGET_ROOT" \ + --checkpoint "$checkpoint" \ + --output-dir "$output" \ + --sessions "${sessions[@]}" \ + --batch-size $((BATCH_SIZE * 2)) \ + --num-workers "$NUM_WORKERS" \ + --device cuda:0 +} + +for seed in "${seeds[@]}"; do + if [[ "$mode" == all || "$mode" == train ]]; then + echo "$(date): train MuVAP Stage-3 reconstruction, seed=$seed" + train_seed "$seed" + fi + if [[ "$mode" == all || "$mode" == screen ]]; then + echo "$(date): Session-12 screen, seed=$seed" + infer_seed "$seed" screen + fi + if [[ "$mode" == test ]]; then + echo "$(date): official Sessions-1--3 test, seed=$seed" + infer_seed "$seed" test + fi +done + +if [[ "$mode" == all || "$mode" == screen ]]; then + python summarize_local_baselines.py \ + --root "$SCREEN_ROOT" \ + --dataset easycom \ + --protocol causal_valid_only \ + --paired-deltas +elif [[ "$mode" == test ]]; then + python summarize_local_baselines.py \ + --root "$TEST_ROOT" \ + --dataset easycom \ + --protocol causal_valid_only \ + --paired-deltas +fi diff --git a/summarize_local_baselines.py b/summarize_local_baselines.py index af6a270..2654dc1 100644 --- a/summarize_local_baselines.py +++ b/summarize_local_baselines.py @@ -61,7 +61,15 @@ def main(): print("\nPaired deltas (percentage points; same column order; matching seeds only)") for treatment, control in (("audio", "baseline"), ("visual", "baseline"), ("av", "baseline"), ("av", "audio"), - ("track_gru", "pooled")): + ("track_gru", "pooled"), + ("person_adapter_v2", "pooled"), + ("egospeak_gru", "pooled"), + ("egospeak_gru", "person_adapter_v2"), + ("qwen_audio", "global"), + ("traditional_audio", "global"), + ("highrate", "traditional_audio"), + ("highrate", "qwen_audio"), + ("stage3", "vap_only")): left, right = by_seed.get(treatment, {}), by_seed.get(control, {}) seeds = sorted(left.keys() & right.keys()) if not seeds: diff --git a/test_local_turn_features.py b/test_local_turn_features.py index 2955caf..3de7713 100644 --- a/test_local_turn_features.py +++ b/test_local_turn_features.py @@ -31,9 +31,18 @@ from extract_turn_local_features import ( TESTED_TRANSFORMERS_VERSIONS, validate_local_transformers_version, main as extract_main, ) +from extract_highrate_person_av_features import ( + MODEL_FPS as HIGHRATE_MODEL_FPS, + _chunked_audio_encode, + _chunked_person_encode, + _model_tick_bounds, + _square_gray_face, +) from local_turn_features import ( - LOCAL_FORMAT, FramePersonFeatures, causal_frame_ticks, crop_box, easycom_boxes, - load_frame_set_features, load_local_features, load_track_features, + HIGHRATE_PERSON_AV_FORMAT, LOCAL_FORMAT, FramePersonFeatures, + causal_frame_ticks, crop_box, easycom_boxes, + load_frame_set_features, load_highrate_person_av_features, + load_local_features, load_track_features, local_head_dimensions, pooled_recent_visual, reference_alignment, ) @@ -68,6 +77,42 @@ def fixture(n=20, name="Session_4", dataset="easycom"): return reference, side, labels +def highrate_fixture(n=20, name="Session_4", dataset="easycom"): + reference, side, labels = fixture(n=n, name=name, dataset=dataset) + side["format"] = HIGHRATE_PERSON_AV_FORMAT + side["visual"] = torch.randn(n, 6) + side["visual_valid"] = torch.ones(n, dtype=torch.bool) + side["persons"] = { + "features": side["visual"].clone(), + "frame_times": reference["times"].clone(), + "frame_offsets": torch.arange(n + 1, dtype=torch.long), + "records": [ + {"id": "p", "asd_interval_end": float(reference["times"][row])} + for row in range(n) + ], + } + side["config"].update( + { + "model_fps": 25.0, + "output_hz": 5.0, + "person_feature_kind": ( + "muvap-asd-causal-visual-plus-activity-logit" + ), + "person_visual_dim": 5, + "person_activity_logit_dim": 1, + "audio_feature_kind": "muvap-cpc-causal-audio-state", + "audio_source": "single-shared-traditional-cpc-stream", + "person_tokens_duplicate_audio": False, + "asd_audio_preprocessing": ( + "causal-float-pcm-clip-only-no-rms-normalization" + ), + "asd_temporal_attention": "causal", + "causal_policy": "right-edge-causal-highrate-asd", + } + ) + return reference, side, labels + + class LocalTests(unittest.TestCase): def setUp(self): torch.set_num_threads(1) @@ -94,6 +139,110 @@ class LocalTests(unittest.TestCase): self.assertFalse(present) self.assertEqual(value.item(), 0.) + def test_highrate_person_av_right_edges_are_causal_and_inclusive(self): + reference, side, _ = highrate_fixture(n=5) + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "Session_4.pt" + torch.save(side, path) + local, people, contract = load_highrate_person_av_features( + path, + reference, + sample_id="Session_4", + dataset="easycom", + ) + packed_audio, audio_contract = load_local_features( + path, + reference, + sample_id="Session_4", + dataset="easycom", + modalities="audio", + ) + self.assertEqual(local.shape, (5, 6)) + self.assertEqual(people.row_to_frame.tolist(), [0, 1, 2, 3, 4]) + self.assertTrue(torch.equal(people.chosen_frame_times, reference["times"])) + self.assertEqual(contract["format"], HIGHRATE_PERSON_AV_FORMAT) + self.assertEqual(packed_audio.shape, (5, 6)) + self.assertEqual(audio_contract["modalities"], "audio") + values, mask = people.padded() + self.assertEqual(values.shape, (5, 1, 6)) + self.assertTrue(mask.all()) + + def test_highrate_person_av_rejects_old_or_duplicated_audio_contracts(self): + reference, side, _ = highrate_fixture(n=5) + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "Session_4.pt" + old = copy.deepcopy(side) + old["format"] = "causal-highrate-person-av-v1" + torch.save(old, path) + with self.assertRaisesRegex(ValueError, "Invalid local feature format"): + load_highrate_person_av_features( + path, + reference, + sample_id="Session_4", + dataset="easycom", + ) + + duplicated = copy.deepcopy(side) + duplicated["config"]["person_tokens_duplicate_audio"] = True + torch.save(duplicated, path) + with self.assertRaisesRegex(ValueError, "causal ASD contract"): + load_highrate_person_av_features( + path, + reference, + sample_id="Session_4", + dataset="easycom", + ) + + def test_highrate_face_geometry_and_chunked_causal_coverage(self): + image = Image.new("RGB", (40, 20), "white") + face = _square_gray_face(image) + self.assertEqual(face.shape, (112, 112)) + self.assertEqual(face.dtype, np.uint8) + self.assertEqual(int(face[0, 0]), 128) + clip = SimpleNamespace(dense_start=0.0, dense_end=1.0) + self.assertEqual(_model_tick_bounds(clip, 1.0), (0, 25)) + self.assertEqual(HIGHRATE_MODEL_FPS, 25.0) + + class FakeEncoder: + audio_dim = 2 + feature_dim = 3 + + def __init__(self): + self.audio_lengths = [] + self.person_lengths = [] + + def encode_audio(self, waveform, num_frames): + self.audio_lengths.append(num_frames) + return torch.stack( + (torch.arange(num_frames), torch.ones(num_frames)), dim=-1 + ).float() + + def encode_person(self, faces, audio_states): + self.person_lengths.append(len(faces)) + return torch.cat((audio_states, torch.ones(len(faces), 1)), dim=-1) + + encoder = FakeEncoder() + faces = np.zeros((9, 112, 112), dtype=np.uint8) + waveform = np.zeros(9 * 640, dtype=np.float32) + audio_result = _chunked_audio_encode( + encoder, + waveform, + num_frames=len(faces), + output_frames=4, + context_frames=2, + ) + person_result = _chunked_person_encode( + encoder, + faces, + audio_result, + output_frames=4, + context_frames=2, + ) + self.assertEqual(audio_result.shape, (9, 2)) + self.assertEqual(person_result.shape, (9, 3)) + self.assertEqual(encoder.audio_lengths, [4, 6, 3]) + self.assertEqual(encoder.person_lengths, [4, 6, 3]) + def test_box_formats_and_crop(self): person = {"Participant_ID": -1, "x1": 2, "y1": 3, "x2": 7, "y2": 9} faces = easycom_boxes([{"Frame_Number": 1, "Participants": [person]}], "face") @@ -505,6 +654,95 @@ class LocalTests(unittest.TestCase): restored.eval() self.assertTrue(torch.equal(restored(states, local, people, mask), logits)) + def test_highrate_person_av_starts_from_exact_audio_only_base(self): + torch.manual_seed(19) + audio = TurnPredictionHead( + hidden_size=8, + local_visual_dim=0, + local_audio_dim=5, + local_projection_dim=3, + local_fusion="pooled", + person_attention_dim=4, + person_attention_heads=2, + dropout=0, + ).eval() + highrate = TurnPredictionHead( + hidden_size=8, + local_visual_dim=6, + local_audio_dim=5, + local_projection_dim=3, + local_fusion="highrate-person-av-attention", + person_attention_dim=4, + person_attention_heads=2, + dropout=0, + ).eval() + highrate_state = highrate.state_dict() + for name, value in audio.state_dict().items(): + self.assertIn(name, highrate_state) + self.assertEqual(tuple(value.shape), tuple(highrate_state[name].shape)) + highrate_state[name] = value + highrate.load_state_dict(highrate_state) + + states = torch.randn(4, 8) + local = torch.randn(4, 6) + local[:, -1] = 1 + people = torch.randn(4, 3, 6) + mask = torch.tensor( + [[True, True, False], [True, False, False], [False, False, False], + [True, True, True]] + ) + people[~mask] = 0 + reference = audio(states, local) + actual = highrate(states, local, people, mask) + self.assertTrue(torch.equal(reference, actual)) + diagnostics = highrate.person_adapter_diagnostics() + self.assertEqual(diagnostics["residual_output_weight_norm"], 0.0) + self.assertEqual(diagnostics["delta_logit_max_abs"], 0.0) + + actual.sum().backward() + assert highrate.highrate_person_residual is not None + self.assertGreater( + float(highrate.highrate_person_residual.weight.grad.abs().sum()), + 0.0, + ) + + def test_highrate_person_av_head_uses_variable_person_sets(self): + torch.manual_seed(9) + head = TurnPredictionHead( + hidden_size=8, + local_visual_dim=6, + local_audio_dim=5, + local_projection_dim=3, + local_fusion="highrate-person-av-attention", + person_attention_dim=4, + person_attention_heads=2, + dropout=0, + ).eval() + states = torch.randn(3, 8) + local = torch.randn(3, 6) + local[:, -1] = 1 + people = torch.randn(3, 3, 6) + mask = torch.tensor( + [[True, True, False], [True, False, False], [False, False, False]] + ) + people[~mask] = 0 + logits = head(states, local, people, mask) + self.assertEqual(logits.shape, (3, 10, 3)) + permutation = torch.tensor([1, 0, 2]) + self.assertTrue( + torch.allclose( + logits, + head(states, local, people[:, permutation], mask[:, permutation]), + atol=1e-6, + rtol=1e-6, + ) + ) + changed = people.clone() + changed[~mask] = 1000 + self.assertTrue(torch.equal(logits, head(states, local, changed, mask))) + logits.square().mean().backward() + self.assertIsNotNone(head.person_attention.in_proj_weight.grad) + def test_f5_screening_audio_person_read_is_gated_and_permutation_invariant(self): torch.manual_seed(11) common = { @@ -714,6 +952,227 @@ class LocalTests(unittest.TestCase): restored.load_state_dict(track.state_dict()) self.assertTrue(torch.equal(restored(states, local, histories, mask), logits)) + def test_causal_person_adapter_v2_learns_after_exact_pooled_start(self): + common = { + "hidden_size": 8, + "num_input_layers": 4, + "layer_fusion": "concat-mlp", + "local_visual_dim": 6, + "local_audio_dim": 5, + "local_projection_dim": 3, + "dropout": 0, + } + for mode in ( + "causal-person-adapter-v2", + "causal-person-adapter-egospeak-gru", + ): + with self.subTest(mode=mode): + torch.manual_seed(23) + pooled = TurnPredictionHead( + **common, + local_fusion="pooled", + ).eval() + adapter = TurnPredictionHead( + **common, + local_fusion=mode, + person_attention_dim=4, + person_attention_heads=2, + person_history_frames=3, + ).eval() + adapter_state = adapter.state_dict() + for name, value in pooled.state_dict().items(): + self.assertIn(name, adapter_state) + self.assertEqual(adapter_state[name].shape, value.shape) + adapter_state[name] = value + adapter.load_state_dict(adapter_state) + assert adapter.person_track_residual_output is not None + self.assertEqual( + torch.count_nonzero( + adapter.person_track_residual_output.weight + ), + 0, + ) + self.assertEqual( + torch.count_nonzero( + adapter.person_track_residual_output.bias + ), + 0, + ) + frozen_parameters = {} + for name, parameter in adapter.named_parameters(): + if name == "layer_logits" or name.startswith( + ("layer_fusion_net.", "local_projections.", "net.") + ): + parameter.requires_grad_(False) + frozen_parameters[name] = parameter.detach().clone() + + states = torch.randn(3, 4, 8) + local = torch.randn(3, 13) + local[:, 6] = 1 + local[:, 12] = 1 + histories = torch.randn(3, 2, 3, 6) + mask = torch.tensor( + [ + [[True, True, True], [True, True, False]], + [[True, False, False], [False, False, False]], + [[False, False, False], [False, False, False]], + ] + ) + histories[~mask] = 0 + baseline = pooled(states, local) + initial = adapter(states, local, histories, mask) + self.assertTrue(torch.equal(initial, baseline)) + self.assertEqual( + adapter.person_adapter_diagnostics()["delta_logit_max_abs"], + 0.0, + ) + + optimizer = torch.optim.SGD(adapter.parameters(), lr=0.1) + target = torch.randn_like(initial) + (initial - target).square().mean().backward() + output_gradient = ( + adapter.person_track_residual_output.weight.grad + ) + self.assertIsNotNone(output_gradient) + self.assertGreater(int(torch.count_nonzero(output_gradient)), 0) + optimizer.step() + optimizer.zero_grad(set_to_none=True) + + trained = adapter(states, local, histories, mask) + (trained - target).square().mean().backward() + input_projection = ( + adapter.person_track_input_embedding[0].weight + if adapter.person_track_input_embedding is not None + else adapter.person_input_projection[1].weight + ) + for parameter in ( + input_projection, + adapter.person_track_gru.weight_ih_l0, + adapter.person_track_gru.weight_hh_l0, + adapter.person_attention.in_proj_weight, + adapter.person_track_local_state[1].weight, + adapter.person_track_dynamic_gate[1].weight, + ): + self.assertIsNotNone(parameter.grad) + self.assertTrue(torch.isfinite(parameter.grad).all()) + self.assertGreater(int(torch.count_nonzero(parameter.grad)), 0) + for name, expected in frozen_parameters.items(): + self.assertTrue( + torch.equal(dict(adapter.named_parameters())[name], expected), + name, + ) + + # Empty rows remain exactly on the pooled path after the + # residual opens. + self.assertTrue( + torch.equal(trained[2], pooled(states, local)[2]) + ) + permutation = torch.tensor([1, 0]) + self.assertTrue( + torch.allclose( + trained, + adapter( + states, + local, + histories[:, permutation], + mask[:, permutation], + ), + atol=1e-6, + rtol=1e-6, + ) + ) + changed_padding = histories.clone() + changed_padding[~mask] = 1000 + self.assertTrue( + torch.equal( + trained, + adapter(states, local, changed_padding, mask), + ) + ) + + config = { + **common, + "local_fusion": mode, + "person_attention_dim": 4, + "person_attention_heads": 2, + "person_history_frames": 3, + } + restored = TurnPredictionHead( + **checkpoint_head_config({"head_config": config}) + ).eval() + restored.load_state_dict(adapter.state_dict()) + self.assertTrue( + torch.equal( + restored(states, local, histories, mask), + adapter(states, local, histories, mask), + ) + ) + + def test_egospeak_style_track_gru_matches_explicit_reference(self): + torch.manual_seed(29) + head = TurnPredictionHead( + hidden_size=8, + local_visual_dim=6, + local_audio_dim=5, + local_projection_dim=3, + local_fusion="causal-person-adapter-egospeak-gru", + person_attention_dim=4, + person_attention_heads=2, + person_history_frames=3, + dropout=0, + ).eval() + embedding = head.person_track_input_embedding + self.assertIsNotNone(embedding) + self.assertEqual( + [type(module) for module in embedding], + [ + torch.nn.Linear, + torch.nn.LayerNorm, + torch.nn.ReLU, + torch.nn.Dropout, + ], + ) + self.assertEqual(head.person_track_gru.num_layers, 1) + self.assertFalse(head.person_track_gru.bidirectional) + self.assertTrue(head.person_track_gru.batch_first) + self.assertIsInstance( + head.person_track_output_activation, + torch.nn.ReLU, + ) + + histories = torch.randn(2, 2, 3, 6) + mask = torch.tensor( + [ + [[True, True, True], [True, True, False]], + [[True, False, False], [False, False, False]], + ] + ) + histories[~mask] = 1234 + actual, valid = head._encode_person_track_histories(histories, mask) + expected = torch.zeros_like(actual) + for batch in range(2): + for person in range(2): + length = int(mask[batch, person].sum()) + if not length: + continue + sequence = embedding( + histories[batch, person, :length] + )[None] + initial = sequence.new_zeros(1, 1, 4) + _, hidden = head.person_track_gru(sequence, initial) + expected[batch, person] = torch.relu(hidden[-1, 0]) + self.assertTrue(torch.equal(valid, mask.any(dim=-1))) + self.assertTrue(torch.allclose(actual, expected, atol=1e-7, rtol=0)) + self.assertGreaterEqual(float(actual.min()), 0.0) + + changed_padding = histories.clone() + changed_padding[~mask] = -9876 + changed, _ = head._encode_person_track_histories( + changed_padding, + mask, + ) + self.assertTrue(torch.equal(actual, changed)) + def test_head_shapes_gradients_missing_and_roundtrip(self): states = torch.randn(3, 4, 8) for v, a in ((0, 0), (6, 0), (0, 5), (6, 5)): @@ -983,7 +1442,7 @@ class RoundTripTests(unittest.TestCase): env = dict(os.environ, OMP_NUM_THREADS="1", MKL_NUM_THREADS="1", OPENBLAS_NUM_THREADS="1") with tempfile.TemporaryDirectory() as directory: root = Path(directory) - for folder in ("features", "local", "targets"): + for folder in ("features", "local", "highrate", "targets"): (root / folder).mkdir() train_uid = "00000000-0000-0000-0000-000000000004" test_uid = "00000000-0000-0000-0000-000000000001" @@ -994,6 +1453,11 @@ class RoundTripTests(unittest.TestCase): reference, side, targets = fixture(name=name, dataset=dataset) torch.save(reference, root / "features" / f"{name}.pt") torch.save(side, root / "local" / f"{name}.pt") + _, highrate_side, _ = highrate_fixture( + name=name, + dataset=dataset, + ) + torch.save(highrate_side, root / "highrate" / f"{name}.pt") np.save(root / "targets" / f"{name}.npy", targets) common = ["--feature-root", str(root / "features"), "--target-root", str(root / "targets"), "--device", "cpu", "--no-progress", "--expected-mode", "open-stream-replay"] @@ -1003,9 +1467,13 @@ class RoundTripTests(unittest.TestCase): "visual", "av", "track-gru", + "person-adapter-v2", + "egospeak-person-gru", "frame-set", "f5-screen", "f6-screen", + "traditional-audio", + "highrate-person-av", "ego4d", ): checkpoint = root / f"{mode}.pth" @@ -1026,9 +1494,14 @@ class RoundTripTests(unittest.TestCase): train += ["--transition-objective", "r1-supervised"] person_mode = { "track-gru": "track-gru-residual", + "person-adapter-v2": "causal-person-adapter-v2", + "egospeak-person-gru": ( + "causal-person-adapter-egospeak-gru" + ), "frame-set": "frame-set-attention", "f5-screen": "audio-person-frame-set", "f6-screen": "gated-person-state-adapter", + "highrate-person-av": "highrate-person-av-attention", }.get(mode) if person_mode is not None: train += [ @@ -1036,20 +1509,42 @@ class RoundTripTests(unittest.TestCase): "--person-attention-dim", "4", "--person-attention-heads", "2", ] - if mode == "track-gru": + if mode in ( + "track-gru", + "person-adapter-v2", + "egospeak-person-gru", + ): train += [ "--person-history-frames", "3", "--person-track-max-gap-seconds", "1", "--initialize-pooled-checkpoint", str(root / "av.pth"), + "--transition-objective", "r1-supervised", + ] + if mode == "highrate-person-av": + train += [ + "--initialize-pooled-checkpoint", + str(root / "traditional-audio.pth"), ] if mode != "baseline": modality = ( "av" - if mode in ("ego4d", "track-gru", "frame-set", "f5-screen", "f6-screen") + if mode in ( + "ego4d", "track-gru", "person-adapter-v2", + "egospeak-person-gru", + "frame-set", "f5-screen", "f6-screen", + "highrate-person-av", + ) + else "audio" + if mode == "traditional-audio" else mode ) - train += ["--local-feature-root", str(root / "local"), "--local-modalities", modality] - infer += ["--local-feature-root", str(root / "local")] + local_root = root / ( + "highrate" + if mode in ("traditional-audio", "highrate-person-av") + else "local" + ) + train += ["--local-feature-root", str(local_root), "--local-modalities", modality] + infer += ["--local-feature-root", str(local_root)] for command in (train, infer): completed = subprocess.run(command, cwd=repo, env=env, capture_output=True, text=True, timeout=60) self.assertEqual(completed.returncode, 0, f"{mode}: {completed.stdout}\n{completed.stderr}") diff --git a/test_muvap_egospeak_stage3.py b/test_muvap_egospeak_stage3.py new file mode 100644 index 0000000..bf2f97b --- /dev/null +++ b/test_muvap_egospeak_stage3.py @@ -0,0 +1,228 @@ +import tempfile +import unittest +from pathlib import Path + +import numpy as np +import torch + +from easycom_head import egospeak_multihorizon_loss +from extract_muvap_egospeak_features import ( + MUVAP_EGOSPEAK_INPUT_FORMAT, + _chunked_vap_encode, +) +from muvap_egospeak_stage3 import ( + MuVAPEgoSpeakStage3, + MuVAPWindowDataset, + align_session_predictions, + collate_stage3_windows, + load_muvap_session, +) + + +def write_fixture(root: Path, *, partial: bool = False) -> tuple[Path, Path]: + features = root / "features" + targets = root / "targets" + features.mkdir() + targets.mkdir() + steps = 14 + ticks = steps * 5 + rows = [] + offsets = [0] + records = [] + for tick in range(ticks): + if tick % 2 == 0: + rows.append(torch.full((1, 6), float(tick), dtype=torch.bfloat16)) + records.append( + { + "id": "p0", + "asd_model_tick": tick, + "asd_interval_end": (tick + 1) / 25, + "source_time": (tick + 0.5) / 25, + } + ) + offsets.append(len(rows)) + flat = torch.cat(rows) if rows else torch.empty(0, 6, dtype=torch.bfloat16) + payload = { + "format": MUVAP_EGOSPEAK_INPUT_FORMAT, + "dataset": "easycom", + "sample_id": "Session_4", + "partial": partial, + "config": { + "uses_omni_features": False, + "uses_speech_labels": False, + "vap_dim": 4, + "asd_dim": 6, + "vap_release_sha256": "vap", + "asd_release_sha256": "asd", + }, + "vap": torch.arange(ticks * 4, dtype=torch.float32).reshape(ticks, 4).to(torch.bfloat16), + "frame_times": torch.arange(1, ticks + 1, dtype=torch.float64) / 25, + "times": torch.arange(1, steps + 1, dtype=torch.float64) / 5, + "label_indices": torch.arange(steps, dtype=torch.long), + "task_frame_indices": (torch.arange(steps, dtype=torch.long) + 1) * 5 - 1, + "persons": { + "features": flat, + "frame_offsets": torch.tensor(offsets, dtype=torch.long), + "records": records, + }, + "max_people": 1, + } + torch.save(payload, features / "Session_4.pt") + target = np.zeros((steps, 3), dtype=np.float32) + target[:, 0] = 1 + target[2::3, 0] = 0 + target[2::3, 1] = 1 + target[3::5, 0] = 0 + target[3::5, 2] = 1 + np.save(targets / "Session_4.npy", target) + return features, targets + + +class FakeVAP: + feature_dim = 2 + + def encode_audio(self, waveform, num_frames): + values = torch.arange(num_frames, dtype=torch.float32) + return torch.stack((values, values + 100), dim=-1) + + +class MuVAPStage3Tests(unittest.TestCase): + def test_loader_window_alignment_and_no_future_access(self): + with tempfile.TemporaryDirectory() as directory: + feature_root, target_root = write_fixture(Path(directory)) + session = load_muvap_session(feature_root, target_root, 4) + dataset = MuVAPWindowDataset( + [session], history_seconds=0.4, require_full_future=True + ) + item = dataset[0] + self.assertEqual(tuple(item.vap.shape), (10, 4)) + self.assertEqual(int(item.history_valid.sum()), 5) + self.assertTrue(torch.equal(item.targets, session.targets[1:11])) + cutoff = int(session.task_frame_indices[0]) + 1 + original = item.vap.clone() + session.vap[cutoff:] = -123 + self.assertTrue(torch.equal(dataset[0].vap, original)) + + def test_loader_rejects_partial_feature_sidecar(self): + with tempfile.TemporaryDirectory() as directory: + feature_root, target_root = write_fixture(Path(directory), partial=True) + with self.assertRaisesRegex(ValueError, "Partial sidecar"): + load_muvap_session(feature_root, target_root, 4) + + def test_model_handles_empty_and_variable_person_sets_and_backpropagates(self): + torch.manual_seed(0) + model = MuVAPEgoSpeakStage3( + vap_dim=4, + asd_dim=6, + model_dim=8, + num_heads=2, + dropout=0.0, + temporal_layers=1, + history_frames=10, + ) + vap = torch.randn(2, 10, 4) + people = torch.randn(2, 10, 3, 6) + person_valid = torch.zeros(2, 10, 3, dtype=torch.bool) + person_valid[0, :, :2] = True + history_valid = torch.ones(2, 10, dtype=torch.bool) + logits = model(vap, people, person_valid, history_valid) + self.assertEqual(tuple(logits.shape), (2, 10, 3)) + targets = torch.zeros_like(logits) + targets[..., 0] = 1 + egospeak_multihorizon_loss(logits, targets).backward() + self.assertGreater(float(model.frame_attention.in_proj_weight.grad.abs().sum()), 0) + + def test_person_set_order_does_not_change_prediction(self): + torch.manual_seed(1) + model = MuVAPEgoSpeakStage3( + vap_dim=4, + asd_dim=6, + model_dim=8, + num_heads=2, + dropout=0.0, + temporal_layers=1, + history_frames=10, + ).eval() + vap = torch.randn(1, 10, 4) + people = torch.randn(1, 10, 2, 6) + valid = torch.ones(1, 10, 2, dtype=torch.bool) + history = torch.ones(1, 10, dtype=torch.bool) + first = model(vap, people, valid, history) + second = model(vap, people.flip(2), valid.flip(2), history) + self.assertTrue(torch.allclose(first, second, atol=1e-6, rtol=1e-6)) + + def test_left_padding_stays_finite_across_temporal_layers(self): + torch.manual_seed(3) + model = MuVAPEgoSpeakStage3( + vap_dim=4, + asd_dim=6, + model_dim=8, + num_heads=2, + dropout=0.0, + temporal_layers=2, + history_frames=10, + ).eval() + vap = torch.randn(2, 10, 4) + people = torch.randn(2, 10, 1, 6) + person_valid = torch.ones(2, 10, 1, dtype=torch.bool) + history_valid = torch.ones(2, 10, dtype=torch.bool) + history_valid[0, :7] = False + history_valid[1, :3] = False + person_valid &= history_valid[:, :, None] + logits = model(vap, people, person_valid, history_valid) + self.assertTrue(torch.isfinite(logits).all()) + + def test_vap_only_control_is_invariant_to_person_values(self): + torch.manual_seed(2) + model = MuVAPEgoSpeakStage3( + vap_dim=4, + asd_dim=6, + model_dim=8, + num_heads=2, + dropout=0.0, + temporal_layers=1, + history_frames=10, + person_evidence="none", + ).eval() + vap = torch.randn(1, 10, 4) + people = torch.randn(1, 10, 2, 6) + valid = torch.ones(1, 10, 2, dtype=torch.bool) + history = torch.ones(1, 10, dtype=torch.bool) + first = model(vap, people, valid, history) + second = model(vap, people + 1000, valid, history) + self.assertTrue(torch.equal(first, second)) + + def test_collate_pads_only_person_axis(self): + with tempfile.TemporaryDirectory() as directory: + feature_root, target_root = write_fixture(Path(directory)) + session = load_muvap_session(feature_root, target_root, 4) + dataset = MuVAPWindowDataset( + [session], history_seconds=0.4, require_full_future=True + ) + batch = collate_stage3_windows([dataset[0], dataset[1]]) + self.assertEqual(tuple(batch["vap"].shape), (2, 10, 4)) + self.assertEqual(tuple(batch["people"].shape), (2, 10, 1, 6)) + + def test_prediction_alignment_matches_egospeak_future_rows(self): + probabilities = np.zeros((12, 10, 3), dtype=np.float32) + probabilities[:, :, 1] = np.arange(12)[:, None] + predictions, valid = align_session_predictions(probabilities, 12) + self.assertEqual(predictions[1, 1, 0], 0) + self.assertEqual(predictions[5, 1, 0], 4) + self.assertEqual(predictions[5, 1, 4], 0) + self.assertFalse(valid[0, 0]) + self.assertTrue(valid[5, 4]) + + def test_chunked_vap_discards_only_past_context_prefix(self): + frames = 8 + waveform = np.zeros(frames * 640, dtype=np.float32) + result = _chunked_vap_encode( + FakeVAP(), waveform, num_frames=frames, output_frames=3, context_frames=2 + ) + self.assertEqual(tuple(result.shape), (8, 2)) + # Every chunk's kept suffix starts after its prepended past context. + self.assertTrue(torch.equal(result[:, 0], torch.tensor([0, 1, 2, 2, 3, 4, 2, 3]))) + + +if __name__ == "__main__": + unittest.main() diff --git a/third_party/EgoSpeak/.gitignore b/third_party/EgoSpeak/.gitignore new file mode 100644 index 0000000..c22eba4 --- /dev/null +++ b/third_party/EgoSpeak/.gitignore @@ -0,0 +1,6 @@ +_legacy_/ +temp/ +visualization/ +backups/ + +*.egg-info \ No newline at end of file diff --git a/third_party/EgoSpeak/.upstream-revision b/third_party/EgoSpeak/.upstream-revision new file mode 100644 index 0000000..f9a4643 --- /dev/null +++ b/third_party/EgoSpeak/.upstream-revision @@ -0,0 +1 @@ +b5bee35e796d8f2d6aa7da2bf1b07836270f8bfb diff --git a/third_party/EgoSpeak/LSTR/.gitignore b/third_party/EgoSpeak/LSTR/.gitignore new file mode 100644 index 0000000..8c953d1 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/.gitignore @@ -0,0 +1,35 @@ +*.zip +backup/ +# *.out +*-log +temp/ + +scripts/ +# Miscellaneous +*.DS_Store +*.swp +*.pyc +*.egg-info + +# PyTorch models +*.pth +checkpoints/ + +# Dataset soft links +data/THUMOS +data/TVSeries +data/HACS +data/EasyCom +data/EGO4D +data/EGO4D + +# logs +log_* +wandb/ +*.egg-info +slurm-*.out + +*.mp4 + +wandb/ +git_diff/ \ No newline at end of file diff --git a/third_party/EgoSpeak/LSTR/.gitmodules b/third_party/EgoSpeak/LSTR/.gitmodules new file mode 100644 index 0000000..9525ebe --- /dev/null +++ b/third_party/EgoSpeak/LSTR/.gitmodules @@ -0,0 +1,3 @@ +[submodule "external/rulstm"] + path = external/rulstm + url = git@github.com:fpv-iplab/rulstm.git diff --git a/third_party/EgoSpeak/LSTR/LICENSE b/third_party/EgoSpeak/LSTR/LICENSE new file mode 100644 index 0000000..67db858 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/LICENSE @@ -0,0 +1,175 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. diff --git a/third_party/EgoSpeak/LSTR/NOTICE b/third_party/EgoSpeak/LSTR/NOTICE new file mode 100644 index 0000000..739459a --- /dev/null +++ b/third_party/EgoSpeak/LSTR/NOTICE @@ -0,0 +1,9 @@ +Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + + +This codebase is based on the following repository: +https://github.com/amazon-science/long-short-term-transformer + +The codebase is modified from the original repository to fit the needs of the LSTR model. + + diff --git a/third_party/EgoSpeak/LSTR/configs/EasyCom/LSTR/lstr_aa_easycom.yaml b/third_party/EgoSpeak/LSTR/configs/EasyCom/LSTR/lstr_aa_easycom.yaml new file mode 100644 index 0000000..206dde4 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/configs/EasyCom/LSTR/lstr_aa_easycom.yaml @@ -0,0 +1,44 @@ +DATA: + DATA_NAME: 'EasyCom' +INPUT: + MODALITY: 'visual+motion' # visual+motion+object + VISUAL_FEATURE: 'rgb_kinetics_resnet50' + MOTION_FEATURE: 'audio_wav2vec2' # audio + OBJECT_FEATURE: 'audio_wav2vec2' + TARGET_PERFRAME: 'target_perframe' # 3 classes +MODEL: + MODEL_NAME: 'LSTR' + FEATURE_HEAD: + LINEAR_ENABLED: True + LINEAR_OUT_FEATURES: 1024 + LSTR: + NUM_HEADS: 16 + DIM_FEEDFORWARD: 1024 + DROPOUT: 0.2 + ACTIVATION: 'relu' + LONG_MEMORY_SECONDS: 512 + LONG_MEMORY_SAMPLE_RATE: 4 + LONG_MEMORY_USE_PE: True + WORK_MEMORY_SECONDS: 8 + WORK_MEMORY_SAMPLE_RATE: 1 + ANTICIPATION_SECONDS: 2 + ANTICIPATION_SAMPLE_RATE: 1 + ENC_MODULE: [ + [16, 1, True], [32, 2, True] + ] + DEC_MODULE: [-1, 2, True] + CRITERIONS: [['MCE', {}]] +DATA_LOADER: + BATCH_SIZE: 16 + NUM_WORKERS: 8 + PIN_MEMORY: True +SOLVER: + NUM_EPOCHS: 50 + OPTIMIZER: 'adam' + BASE_LR: 7e-05 + WEIGHT_DECAY: 5e-05 + SCHEDULER: + SCHEDULER_NAME: 'warmup_cosine' + WARMUP_FACTOR: 0.3 + WARMUP_EPOCHS: 10.0 + WARMUP_METHOD: 'linear' diff --git a/third_party/EgoSpeak/LSTR/configs/EasyCom/LSTR/lstr_aa_easycom_from_ytturn.yaml b/third_party/EgoSpeak/LSTR/configs/EasyCom/LSTR/lstr_aa_easycom_from_ytturn.yaml new file mode 100644 index 0000000..dc9b658 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/configs/EasyCom/LSTR/lstr_aa_easycom_from_ytturn.yaml @@ -0,0 +1,45 @@ +DATA: + DATA_NAME: 'EasyCom' +INPUT: + MODALITY: 'visual+motion' # visual+motion+object + VISUAL_FEATURE: 'rgb_kinetics_resnet50' + MOTION_FEATURE: 'audio_wav2vec2' # audio + OBJECT_FEATURE: 'audio_wav2vec2' + TARGET_PERFRAME: 'target_perframe' # 3 classes +MODEL: + CHECKPOINT: 'checkpoints/lstr_ytconv_pretrained_AV.pth' + MODEL_NAME: 'LSTR' + FEATURE_HEAD: + LINEAR_ENABLED: True + LINEAR_OUT_FEATURES: 1024 + LSTR: + NUM_HEADS: 16 + DIM_FEEDFORWARD: 1024 + DROPOUT: 0.2 + ACTIVATION: 'relu' + LONG_MEMORY_SECONDS: 512 + LONG_MEMORY_SAMPLE_RATE: 4 + LONG_MEMORY_USE_PE: True + WORK_MEMORY_SECONDS: 8 + WORK_MEMORY_SAMPLE_RATE: 1 + ANTICIPATION_SECONDS: 2 + ANTICIPATION_SAMPLE_RATE: 1 + ENC_MODULE: [ + [16, 1, True], [32, 2, True] + ] + DEC_MODULE: [-1, 2, True] + CRITERIONS: [['MCE', {}]] +DATA_LOADER: + BATCH_SIZE: 16 + NUM_WORKERS: 8 + PIN_MEMORY: True +SOLVER: + NUM_EPOCHS: 50 + OPTIMIZER: 'adam' + BASE_LR: 7e-05 + WEIGHT_DECAY: 5e-05 + SCHEDULER: + SCHEDULER_NAME: 'warmup_cosine' + WARMUP_FACTOR: 0.3 + WARMUP_EPOCHS: 10.0 + WARMUP_METHOD: 'linear' diff --git a/third_party/EgoSpeak/LSTR/configs/YTConv/LSTR/lstr_aa_ytconv.yaml b/third_party/EgoSpeak/LSTR/configs/YTConv/LSTR/lstr_aa_ytconv.yaml new file mode 100644 index 0000000..e5ed237 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/configs/YTConv/LSTR/lstr_aa_ytconv.yaml @@ -0,0 +1,46 @@ +DATA: + DATA_NAME: 'YTConv' + DATA_INFO: 'data/YTConv_info.json' +INPUT: + MODALITY: 'visual+motion' # visual+motion+object + VISUAL_FEATURE: 'rgb_kinetics_resnet50' + MOTION_FEATURE: 'audio_wav2vec2' # audio + OBJECT_FEATURE: 'audio_wav2vec2' + TARGET_PERFRAME: 'target_perframe' # 3 classes +MODEL: + MODEL_NAME: 'LSTR' + FEATURE_HEAD: + LINEAR_ENABLED: True + LINEAR_OUT_FEATURES: 1024 + LSTR: + NUM_HEADS: 16 + DIM_FEEDFORWARD: 1024 + DROPOUT: 0.2 + ACTIVATION: 'relu' + LONG_MEMORY_SECONDS: 512 + LONG_MEMORY_SAMPLE_RATE: 4 + LONG_MEMORY_USE_PE: True + WORK_MEMORY_SECONDS: 8 + WORK_MEMORY_SAMPLE_RATE: 1 + ANTICIPATION_SECONDS: 2 + ANTICIPATION_SAMPLE_RATE: 1 + ENC_MODULE: [ + [16, 1, True], [32, 2, True] + ] + DEC_MODULE: [-1, 2, True] + CRITERIONS: [['MCE', {}]] +DATA_LOADER: + BATCH_SIZE: 16 + NUM_WORKERS: 8 + PIN_MEMORY: True +SOLVER: + SAVE_EVERY_EPOCHS: 5 + NUM_EPOCHS: 50 + OPTIMIZER: 'adam' + BASE_LR: 7e-05 + WEIGHT_DECAY: 5e-05 + SCHEDULER: + SCHEDULER_NAME: 'warmup_cosine' + WARMUP_FACTOR: 0.3 + WARMUP_EPOCHS: 10.0 + WARMUP_METHOD: 'linear' \ No newline at end of file diff --git a/third_party/EgoSpeak/LSTR/configs/ego4d/LSTR/lstr_aa_ego4dshuffle.yaml b/third_party/EgoSpeak/LSTR/configs/ego4d/LSTR/lstr_aa_ego4dshuffle.yaml new file mode 100644 index 0000000..77e4574 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/configs/ego4d/LSTR/lstr_aa_ego4dshuffle.yaml @@ -0,0 +1,45 @@ +DATA: + DATA_NAME: 'EGO4DSHUFFLE' + # IGNORE_INDEX: 0 # calculate mAP w/o target not speaking (background) +INPUT: + MODALITY: 'visual+motion' # audio+visual + VISUAL_FEATURE: 'rgb_kinetics_resnet50' + MOTION_FEATURE: 'audio_wav2vec2' # audio + OBJECT_FEATURE: 'rgb_kinetics_resnet50' + TARGET_PERFRAME: 'target_perframe' # 3 classes +MODEL: + MODEL_NAME: 'LSTR' + FEATURE_HEAD: + LINEAR_ENABLED: True + LINEAR_OUT_FEATURES: 1024 + LSTR: + NUM_HEADS: 16 + DIM_FEEDFORWARD: 1024 + DROPOUT: 0.2 + ACTIVATION: 'relu' + LONG_MEMORY_SECONDS: 512 + LONG_MEMORY_SAMPLE_RATE: 4 + LONG_MEMORY_USE_PE: True + WORK_MEMORY_SECONDS: 8 + WORK_MEMORY_SAMPLE_RATE: 1 + ANTICIPATION_SECONDS: 2 + ANTICIPATION_SAMPLE_RATE: 1 + ENC_MODULE: [ + [16, 1, True], [32, 2, True] + ] + DEC_MODULE: [-1, 2, True] + CRITERIONS: [['MCE', {}]] +DATA_LOADER: + BATCH_SIZE: 16 + NUM_WORKERS: 8 + PIN_MEMORY: True +SOLVER: + NUM_EPOCHS: 50 + OPTIMIZER: 'adam' + BASE_LR: 7e-05 + WEIGHT_DECAY: 5e-05 + SCHEDULER: + SCHEDULER_NAME: 'warmup_cosine' + WARMUP_FACTOR: 0.3 + WARMUP_EPOCHS: 10.0 + WARMUP_METHOD: 'linear' diff --git a/third_party/EgoSpeak/LSTR/configs/ego4d/LSTR/lstr_aa_ego4dshuffle_from_ytconv.yaml b/third_party/EgoSpeak/LSTR/configs/ego4d/LSTR/lstr_aa_ego4dshuffle_from_ytconv.yaml new file mode 100644 index 0000000..7d319d5 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/configs/ego4d/LSTR/lstr_aa_ego4dshuffle_from_ytconv.yaml @@ -0,0 +1,46 @@ +DATA: + DATA_NAME: 'EGO4DSHUFFLE' + # IGNORE_INDEX: 0 # calculate mAP w/o target not speaking (background) +INPUT: + MODALITY: 'visual+motion' # audio+visual + VISUAL_FEATURE: 'rgb_kinetics_resnet50' + MOTION_FEATURE: 'audio_wav2vec2' # audio + OBJECT_FEATURE: 'rgb_kinetics_resnet50' + TARGET_PERFRAME: 'target_perframe' # 3 classes +MODEL: + CHECKPOINT: 'checkpoints/lstr_ytconv_pretrained_AV.pth' + MODEL_NAME: 'LSTR' + FEATURE_HEAD: + LINEAR_ENABLED: True + LINEAR_OUT_FEATURES: 1024 + LSTR: + NUM_HEADS: 16 + DIM_FEEDFORWARD: 1024 + DROPOUT: 0.2 + ACTIVATION: 'relu' + LONG_MEMORY_SECONDS: 512 + LONG_MEMORY_SAMPLE_RATE: 4 + LONG_MEMORY_USE_PE: True + WORK_MEMORY_SECONDS: 8 + WORK_MEMORY_SAMPLE_RATE: 1 + ANTICIPATION_SECONDS: 2 + ANTICIPATION_SAMPLE_RATE: 1 + ENC_MODULE: [ + [16, 1, True], [32, 2, True] + ] + DEC_MODULE: [-1, 2, True] + CRITERIONS: [['MCE', {}]] +DATA_LOADER: + BATCH_SIZE: 16 + NUM_WORKERS: 8 + PIN_MEMORY: True +SOLVER: + NUM_EPOCHS: 50 + OPTIMIZER: 'adam' + BASE_LR: 7e-05 + WEIGHT_DECAY: 5e-05 + SCHEDULER: + SCHEDULER_NAME: 'warmup_cosine' + WARMUP_FACTOR: 0.3 + WARMUP_EPOCHS: 10.0 + WARMUP_METHOD: 'linear' diff --git a/third_party/EgoSpeak/LSTR/data/YTConv_info.json b/third_party/EgoSpeak/LSTR/data/YTConv_info.json new file mode 100644 index 0000000..0a3269b --- /dev/null +++ b/third_party/EgoSpeak/LSTR/data/YTConv_info.json @@ -0,0 +1,431 @@ +{ + "YTConv": { + "data_root": "data/YTConv", + "class_names": [ + "background", + "camera wearer start to speak" + ], + "num_classes": 2, + "ignore_index": -100, + "metrics": "AP", + "fps": 5, + "train_session_set": [ + "m5Aoywf51L4", + "p99cIB7nHjw", + "zGgMzBfPfaI", + "O0eCtKU3FxY", + "goZQ6iOIFYo", + "9Tb-lSyhgoE", + "GBxQtGUjAMA", + "VXgt7keOC8I", + "hZ2tin2J0P4", + "Ur-P18sm4tI", + "WfYmsNBQwg8", + "AvMNaDDfrN0", + "fW7pCN3wwdo", + "xG6HBdTtS4s", + "HcDlqUeDD9Q", + "1DE3mKncHUA", + "oEMX_quSMZY", + "jmS3nWbGV5Y", + "qUJRjPRH-u4", + "aHsUgWfUcss", + "ZfdQOmJkaK0", + "4f43nYrfTnI", + "qE1FnzkNjrE", + "4NsmKhs9flk", + "3wDSJ9ABzmo", + "sY6_8QwRZsw", + "hByfAzEoeKA", + "vDEsV0AhcjM", + "XGNhX3Nu1RE", + "MkVT8K0V57A", + "yR8y-wN8Clk", + "P1dTOlO2sl8", + "6UGC_wxTeB0", + "d4Xud2Gqoxk", + "IJ1aqePculQ", + "o8Oxf5QBitQ", + "XVci_V4qxmc", + "ofDzhOS1Pv0", + "mygTXBuMRF0", + "iq6q82Jo8ZQ", + "Kzc_ekBJ4QE", + "WEeJ9B1i8HI", + "rciqpxdCc9Y", + "Rh7d2D7JVLI", + "Qkx4hSlc8jk", + "ixi7sVLFrWM", + "n96-scSyoOQ", + "YZ5xYioIQdA", + "X76QxuWm1oA", + "2I7d_LfK75Y", + "AeUWBaYFMyo", + "lsqXFgymLlI", + "Dx_0jYMWN5o", + "0raJ1-1BJ7s", + "EGX8am5hWh0", + "KjtZ_CLY1hQ", + "R4FQI2BI3hw", + "fAoqqtsNjbQ", + "l6RwyA6jUdM", + "GJW6dlagRnU", + "j_FGvmbp6aM", + "V2O15h1VrH8", + "P3c70VbDP3E", + "epAwpBdUiYU", + "6dpv_JeyL8w", + "dUKYxsvnuBM", + "L2_zn7h8sOQ", + "m09qdCkAhmY", + "7BB71tIae7A", + "6go4oViEkdc", + "MYimxshUK1c", + "v4eZXWcIRQU", + "vGI_w1Ml6S0", + "BeFndnkT00c", + "9JHRbL1jmq4", + "0rBaR5dd-Fs", + "EvhkIS_-K-w", + "pYVbLx-h5HI", + "fVciOXYDlqg", + "zWFH4N-CmDE", + "Aw4cauzfX5E", + "lxVpcQHtcWs", + "32dkF-aXgug", + "_USorgWOsUk", + "1yVMU93nmv8", + "aYkF4Ybk-l8", + "Vwt18p0BqzM", + "oOQM5Pvjut4", + "_WWSMNtNYx0", + "zzlhwg0IJ5Y", + "lJhIffkUhr8", + "PbS9nSH3mqM", + "8EGVgiFpwhg", + "5LhXpv5IFO4", + "ibD4APVzWgg", + "2nhcSQmopO4", + "gri79OfCRvo", + "lOwQN3QcRzE", + "wQyluRNtqFo", + "IZyx8zeiKB0", + "mZyNfKyzDHU", + "A9eGJWiuIKs", + "d37DZogot1w", + "7ZUSVabB6N4", + "m62GQXq-358", + "ArwPVsCZo2c", + "QK-BFkpO8s4", + "PgfX9bvmEhs", + "l1jTUhwWJYA", + "S1PGHO4evaI", + "M6uPJa5B8Ns", + "XXzz19J-1w8", + "U4uV1xmUaB0", + "skQ0OBpkKqY", + "mfW57TDpMtc", + "NAwKA75cDwE", + "eIEBI2_Y14Q", + "eWQGYpSoiYU", + "3DPIUWhYtn8", + "_elvaRJrhE0", + "PAA8IFhWJzM", + "1XDKqhq6bb0", + "ngi9gy3DzRY", + "Y0yONfuimeE", + "g9_cAgHa_4Q", + "BIG7jEw-j4w", + "TFdom0lkd-0", + "2zV2N8_DwWg", + "kTz7XB_jpB4", + "GuoglIuhNjY", + "L98U3mO38gQ", + "UfWDoRAzpC8", + "W6TakaXCIIQ", + "i6j1I6pha-M", + "issHFC8Gbio", + "4EzGoRkd6SI", + "QoS0GgF4sI4", + "A8hJg3D_TGY", + "A3N10suCRCw", + "6962tnWRsNY", + "lS4p8osXSJs", + "wbZ7ZhcZd6o", + "fsIvQvIlKwc", + "4jQSOSi2DA8", + "Vpt6tZBvXQU", + "kxOG3kpLmYw", + "Mzjlz0nq9AA", + "80cuXt2VSko", + "VvswhogSiY0", + "5Q6ky5VrOIA", + "I_XtcUpNRF8", + "QmtmTGycK9A", + "qp3x8iy7vkA", + "xGU8QDtPrQI", + "uWWUfqth_BA", + "WbM7w8r4VGQ", + "bt6Ed1qjulw", + "7vymT27scoY", + "b9AdV6oVDVg", + "nf2zqYUAOwg", + "A-LNJc-1Jc0", + "yCI-hpPiDVc", + "Uu5DME4FqIM", + "symjuH9iOTc", + "mPf2MTnTBc0", + "NDTlQ1KGvH8", + "g9vuDzFQSRU", + "B6iqnvleve0", + "9bbszAD7JL0", + "RIyMAuz7Ckg", + "mlrQHtBRMhc", + "3gaidlW7vig", + "DGdV97q9ndk", + "qthjlg0lXFg", + "gWng5YLPOow", + "iaYG-KL9XRM", + "385414AVZcA", + "7xhBnWGpI7c", + "Nh0cmSQUs-E", + "Rvtj0QHPulA", + "rAtLksYQukw", + "S9rB__zbbe0", + "SNpMVsOPfHM", + "zhF43F4kO2w", + "HOUkcInyxN0", + "x8ymkfr2Z1o", + "Ix6Khtq6O-k", + "VHs_JyPhmtE", + "zJX3byjYSMc", + "zdqX09zZNf8", + "uPwlmwNQq8Y", + "e4xNTXrk6wE", + "P2PNUWN4GMs", + "_WVHCYjHBVY", + "5Jtv1GpPvHM", + "AIZYucKOzhw", + "vbsR_VDZXh0", + "9Ew0lDHB0kM", + "ejpkPVZHuMM", + "47HlvCfubE4", + "HzPkQSIG2m0", + "e7mTRnNnOXg", + "eZHx_biEECg", + "FvP6WFKNP4Y", + "g_TJ5XFEVfY", + "8Utg2VEWsJg", + "VMdBjL-tpoo", + "zihP2Eqeqxk", + "ILVqgOlntLM", + "Z1BWQI8KPCg", + "j8TRVp-7l3Q", + "G5gJBZ-3OlY", + "o1JUlGSirSE", + "OsPdhUQ0snE", + "-5uIaW2gKWQ", + "03kiFNh-9ME", + "w5W6mgDaJJA", + "sYsgd_8AFoo", + "vLg8kVUTmWU", + "OK0J0xsuSWQ", + "d284rjoU4Cs", + "5lE6V4jrUl4", + "bS1OpC9xu48", + "ZUxpqU_J_n0", + "AUliC0c_YIQ", + "nJ58-d_viOo", + "uIy3b2kzlrM", + "vOTLfb0qjfY", + "37_7tPZjC34", + "MUwCdP5qceI", + "lYvNi_ewd8U", + "PDG7SzDBTC4", + "QHQzoHolWvI", + "YmvOEj00LB8", + "fLhkXVx7Yw0", + "Uj5nhugVgwo", + "aHbrDu2o0WY", + "yKW6vVgEPCM", + "Q0WUgyqQ-4I", + "xw3KOPJ9a3E", + "5mrTo7c3bCM", + "x0ft032W--4", + "UdQxEbrIocc", + "EHZHAFyv_Ks", + "7DojKgPRLVk", + "-1CahRD6lJ0", + "f3TuYwl7AZE", + "TgGTovCcS9Y", + "AwyrJGeIN_g", + "RbPJD4Wh-l8", + "58bFdwYkDzU", + "6gIEKGxL7zo", + "BeuzJwTuChQ", + "YebZyAzLZuc", + "9T74b8EQBVY", + "TqSVFjw1uNM", + "WUYhCmfE1a0", + "vRCVUQFUio4", + "w2t9JdMNMgQ", + "nn6-x9xXMeE", + "FDFVPzLlzFs", + "qFhvi4bH29E", + "qdpsMXLsUZM", + "JREk1QeSc9A", + "jOVEs-4r1jc", + "v4VZI9Ouzy0", + "DMg5UUvf-C4", + "26J5zRzQkOg", + "YkiJ4iJfJbw", + "_d8j8nfTEIk", + "wMNb3hQXEUE", + "1KI8ex22blk", + "eCp2Jp7MtKg", + "HNlpHBb5dC4", + "C6_pDRdIo1k", + "ZRWJJHrgP54", + "L7jfx4veKIA", + "p1a5_qltICM", + "nZ0svM8DoY0", + "miEzjCJVaZk", + "TpC2tDCxABw", + "SZmqoqCX__U", + "dZ0bYteaI2Q", + "GW-cYEzShmY", + "iUYuPhutKVw", + "97VSm6T2zMw", + "jPiQyi5jr1M", + "lrUi4nsPZ9M", + "rXLeXke-k_w", + "mEO4izUqZ7E", + "4qI1LjMXfh0", + "mVt6myGMa-0", + "VJrnhh6FqEg", + "vwCeBWbZvrE", + "c08fl8-xmVE", + "syZ3bvXIrpA", + "i3ysNaulop4", + "__t3rU0QjsE", + "jT32uQQVGvU", + "Fzvmh-R2EWk", + "QApJSKrLg6E", + "L5y3j8SdGSU", + "3t8JGblR__0", + "e1F90W40t6g", + "j6L8UIa5z28", + "Zmg2_C1FuVI", + "XIXKJXsiBFc", + "NVHB1VdWFo4", + "-x7-8f7bsl4", + "2FSxSiowkGI", + "jgFIBIeFve0", + "1BLuXK1M5CE", + "zmpx6LUT2ow", + "z-QUzJTsFro", + "0sS4v13DAK8", + "gMIBQ0mbRvM", + "9kcaEXlTB5s", + "ZYArPGMLS4k", + "tejW9GMaucA", + "z9r4bmXCWiI", + "4Zpsb46iaLc", + "52Odc8UbvC8", + "x-B6-iYDHPs", + "8GinSDc0Is8", + "5lwbp5d_w0Y", + "I4LpTQLzlHE", + "lMl6-XH7ozY", + "tOubscNbY1Q", + "7gb7AH-oWcQ", + "mDerqt9Ib-I", + "gKV3FNd9WHE", + "kH5L1E_r-3U", + "sq5TRFxVp0w", + "SrBhj7Xusk8", + "RxlxEapsz2E", + "aCD3_TziiKo", + "_HPVZOkxdkk", + "pNCjgVQ53Y8", + "iPdT5GUoGvA", + "cIrlet1oRw4", + "1dyWUMBjMwQ", + "XZQB3cq18yc", + "i75Mn_AeFBY", + "_0VwR9WPS-k", + "q1JDjgnRdGY", + "1JNEpfvRq0Y", + "2FGj6HciQZA", + "eLi59TtNwHY", + "Z3W8OST7tmA", + "F2dNYov8-xw", + "kPfyJu_CrSA", + "7I4Y2HtAwlY", + "us-andJDfWw", + "Q0Bzn1teWfo", + "pADxBfOrYeQ", + "KnA1tIR0qCE", + "lBNPfBcb6pw", + "zt9MIiBfxF8", + "FwX12hPqCws", + "6e2LXGrze2M", + "0dvclhcFjSo", + "UTkuDYOIxKo", + "4n044cOTiNg", + "VxksG15OgRY", + "4193MkWdJcQ", + "2bDX4Uzy8Wc", + "SkyBNX2auEc", + "JrSLspuJ2gk", + "GZO7V-CWSJ8", + "TYMIy4qlRzc", + "jK7q_Li7LU0", + "UiTIXDQEkPM" + ], + "test_session_set": [ + "DTGtC7FC4oI", + "CIWvk3gJ_7E", + "A1fZ6rlOwxk", + "cB8WCrmJEEA", + "gGQ_2Ank9-U", + "b91abcyMeNY", + "9MuaFeyczM0", + "rwZxhdK_UVU", + "hd_WbrqYnTE", + "4BWcqA9cFFk", + "MuOkgCkob6E", + "OdhFWA3Fgb0", + "Pe3ZabrUiUk", + "U0c_uv8DU08", + "kTqt_3YJBbs", + "eWXVVMi4bf8", + "RCD-vXumMIs", + "DuwiezJCZ3I", + "BNgEqjAiKU4", + "Cbn0t7eqaEs", + "D3M9Q0WwuFY", + "oO-Bv4iRZKI", + "bBvqtbZdSZQ", + "-mnoItCi32Y", + "yQPxKfz-bwI", + "ITVZ4cx5big", + "uY5_e0L7mjw", + "J_Pl9ftnNXE", + "Jt_pb40MRy4", + "E3hAj2niRNo", + "kl2xVSxVECQ", + "gzebuFs90Xg", + "6DQSIzJ9J-g", + "pJJlnryTYWg", + "X87qsfVLFoc", + "W7BIPWU1WHc", + "Oj506EVOiQc", + "DPEvWvcstxA", + "Pk16UT5-zeM", + "XMK1KEk5OdU", + "PpNm5llsEW8", + "sAp3rvablGk" + ] + } +} \ No newline at end of file diff --git a/third_party/EgoSpeak/LSTR/data/data_info.json b/third_party/EgoSpeak/LSTR/data/data_info.json new file mode 100644 index 0000000..11390a7 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/data/data_info.json @@ -0,0 +1,479 @@ +{ + "EasyCom": { + "data_root": "data/EasyCom", + "class_names": [ + "background", + "camera wearer speaking", + "other people speaking" + ], + "num_classes": 3, + "ignore_index": -100, + "metrics": "AP", + "fps": 5, + "train_session_set": [ + "Session_7", + "Session_11", + "Session_10", + "Session_6", + "Session_9", + "Session_4", + "Session_8", + "Session_5", + "Session_12" + ], + "test_session_set": [ + "Session_3", + "Session_2", + "Session_1" + ] + }, + "EGO4DSHUFFLE": { + "data_root": "data/EGO4D", + "class_names": [ + "background", + "camera wearer speak", + "other people speaking" + ], + "num_classes": 3, + "ignore_index": -100, + "metrics": "AP", + "fps": 5, + "train_session_set": [ + "5b8c91ca-9d16-48b7-acce-4dec221ffd71", + "0232fbc8-c9b1-405c-bd42-a9d773fd4835", + "39c74452-1957-40db-894a-b39df51e1437", + "31488ff6-05ef-4362-9786-35981b06f307", + "4decc0fa-e98e-4bf9-bbd1-d9b852b04830", + "efeaa284-2e22-44b8-bdd6-1bb4763e5b35", + "07f1486c-44a7-4c69-9532-720c81c99eb0", + "19c28cbe-21f2-4b40-a6e0-e9932b5f0605", + "1c5c0b52-1287-4df3-aa75-8ccc916b1c10", + "e92255bc-1be6-4ee9-8f8b-d1336a0e409c", + "307e8837-1fda-4cf5-bab5-d16fee20f84f", + "bb050b8a-195a-4b50-b5d0-1ed535226a39", + "ca22f118-2472-4df0-a9cd-3d70424b211b", + "08b0935e-6260-4bd6-86ca-f6fc54e388be", + "666a21c0-7469-436e-868f-64eb8022f64f", + "77e09fd6-b560-4439-addc-a1a8a3122046", + "f00bd659-56e0-41dd-8eba-c5fdb45e7232", + "c2c7f072-a02a-480d-8342-b43ffc314db9", + "ed72b048-e309-4510-b049-b0288ce4478a", + "56c5af79-f9d4-478d-96ef-6d71e0bbbdfe", + "3d8acf67-fad6-421e-841b-82b28c82f7f6", + "977ac980-f908-4c29-a2fc-dd14ec21c0bf", + "08742cb7-bb6d-47ae-af5e-3f1f306ccda6", + "f3afa5a1-8a97-4a7d-945e-1cc70931e44a", + "65172b19-a246-4dc5-98e9-1b526b00ac89", + "ee1893cb-2ff7-40e0-af8d-19518696b698", + "53144267-6da5-40d3-91f2-eaa85be367de", + "26080fb8-0588-449c-a7f2-2e0162507e02", + "c0ba997c-9ec6-4e83-ad50-a88dd677b881", + "6640a311-86c4-4de7-964c-e3a4a67a3e02", + "5f1fcf84-4001-43ec-ad88-9e06ea2f663d", + "ea5a3f9d-7b21-46d7-a17d-bcbd4955bae9", + "eb3675b8-f6a8-4ea9-82df-1080eac5b093", + "023aaf52-fbe9-4b61-af69-ed33e91f3717", + "86af705a-a326-423a-bdd9-c6734804c30c", + "55363f67-1b41-476f-a83e-7e9741786ad7", + "ce8ea1a3-65e3-4ef7-abeb-2b0bc1c84773", + "d3b7e314-f48d-4888-b496-e1eeb252cd8a", + "59000084-d286-4c83-90c8-2fbbf80d2d31", + "a47c590d-66aa-4bd1-b552-f40951ba9bf0", + "61c73721-9476-4ac2-a987-ecaf9a04dca1", + "0124376a-720d-4265-8c8c-5ff47f04d579", + "f62a8100-7ed2-492d-9709-064514d552bf", + "13d71019-9093-442b-bdb4-c79310150592", + "74e71420-9bd6-4337-9a7a-a3dcc4cab7cf", + "383466b5-cb00-4e76-9e6a-dc1b0c9b125c", + "9ef6ef95-6b97-4c78-866c-641fc7d791f5", + "ac7d3d16-d74d-4fe3-a1e0-cab0b1adf492", + "c46cbba9-2cb3-4d84-be18-71af5539ecc5", + "5eb74ea5-9b0b-484c-82b2-ee86479c5d07", + "6f80f9ae-c276-413f-accc-fad74efe8a95", + "73f99cf0-d997-4488-b6a6-993934a9fad3", + "fd50b1d6-8142-402f-88c5-7c44ca215081", + "51fee605-d077-4d78-babc-9bbcd8d898ca", + "5e3c6abc-599b-44f8-8401-d5720627bba2", + "91151b86-65b7-4854-97b0-9c5b93a251ed", + "5be9a2e2-5efc-4fc6-a65e-806e0347e19a", + "94ba75fc-8d7f-4bb4-a756-ded62a4c1761", + "c925c38c-3fe3-4eb3-9944-3fba3eaca4db", + "d97bedc8-72df-43be-a55b-4da1ae42dfd1", + "e60b110e-a9b6-42b7-8f92-082cd74c4937", + "3627996c-8188-4486-8feb-938b9b0cce2b", + "d2b674ca-9f02-46ba-820d-f0b52c76749f", + "2b90b963-e829-4787-8fcd-1a1a434ffdfd", + "ba51db0f-2bf1-4746-894a-7db0ba54e357", + "6c1bac5c-aead-40b8-9c17-be97d86b68d4", + "63483ea2-d569-45fd-82b2-f80836a8b766", + "c81eb608-6844-424e-b8b7-524ff652cf3a", + "cae5d4bf-bbf9-45dd-876b-a48579782348", + "7813b44f-04ca-4a53-868d-fe185b615b68", + "1f12c871-9b3a-4611-a2b4-9c39059052a4", + "70316529-ec02-4d26-bc68-4b436bf0030a", + "6d367aa4-56ef-4932-9d48-e7530e717ea8", + "2e651150-9338-4103-b231-2f089aa19df7", + "d2cb4045-8ef1-4288-8398-ab6d4d6a0709", + "3e4b9c29-06d2-4ade-a2be-bdec028bc68c", + "4f3c8276-99e0-4c99-9e61-5d7c06332b09", + "12634773-451f-4c82-b915-265c703ad737", + "5f615b11-2ff8-4ff1-987a-1cd61ea78dee", + "de96ee14-be42-471d-b110-aabc948260fa", + "8504aae3-4f25-4fa4-88c1-beb3fba83f13", + "d2c0f78a-978e-409f-a78f-3e40b3431909", + "0f846b5d-1cea-40b8-b18b-dc6c54544a67", + "37afa156-a068-458e-8bf7-ec6f1f2415bc", + "2d205d38-f79c-43f2-8c39-ede697065ab2", + "36e7b2ad-dddd-42fa-a427-6411ff12f9b9", + "355f8f4a-cde2-4108-821a-5ccea1d297fe", + "c37e680f-03b1-408b-9b09-1b2ae733951d", + "6958467e-388c-4f6d-a8c0-d4782166daac", + "3051e028-71f8-4f3d-a72a-010934c3566f", + "fc0f4c90-a1ee-4147-990e-de231394f484", + "c2af1389-b165-42f5-b4ee-c90402d7006b", + "1acfcabc-653a-4c03-adf7-64838d09a7b4", + "3b79017c-4d42-40fc-a1bb-4a20bc8ebca7", + "6edb7036-3369-4c03-8e3a-439ffbeb4724", + "2ac543c6-b52b-4219-b2e4-4130078c0673", + "b3ef3563-ecc0-4a15-9a7b-4feb4558d953", + "6f001eeb-fed5-43e4-bfa2-84e274cbe5cd", + "1ce2252b-31b8-4536-9902-0ba7fb039340", + "c2413391-7c1b-4fd6-8b1d-98ee7888b9f8", + "a4d53159-e2c0-4a47-9e78-af9bc9291f0e", + "66ddb7ae-90d6-4a58-9f3c-7f1d2604ecb0", + "290fc9ba-0a29-4eed-b0e0-64edc5bf54ea", + "e3a53edd-52ff-4608-a80d-77ecfbb0f961", + "50895f7f-5667-4b43-94c1-d7b4b9e238e7", + "3aea2d62-0340-4c6a-8019-3cce9117a65a", + "63ab0246-2f1c-4621-8b71-75075b30d852", + "228eb504-c1b3-43b9-818e-2f7e539a259b", + "02de9500-d574-446e-95fc-3e82c367385a", + "bd03445e-9016-423a-bb48-cd998836bbff", + "a16cb772-d0e2-4654-96d7-372c21484b90", + "0fa75089-7b28-4988-8664-a2c93542ffb0", + "fcc6af87-918f-4d0c-a0a4-8aab36891fc4", + "b9d7642e-2ced-4589-b7a6-69ad54570c4c", + "ba5b1882-c9d7-48e7-85fe-2c7b10494fac", + "92e0925c-566d-4a67-b93e-3e4b57f9b3ab", + "130a49d0-0315-4eb0-a752-a4533bd880a9", + "0fccd442-a901-4d77-aac6-dc4276cd6591", + "49c27d75-44fa-4947-899a-8d0511a10416", + "e2ffc7c9-c035-4709-97c6-3764499188a5", + "2069aea3-ac99-4b70-891e-1bfb8b15aeea", + "cdbc7718-4cb9-4e78-bbde-3d9614f8dcd1", + "32464967-b0f2-4b2e-a57c-7f774e7056a7", + "ae6466f8-1a53-4cb3-8718-96f78f391b4b", + "d66c7c66-2b5f-4e7f-bbb2-f0adae68dfcc", + "da84b37f-ec6e-414e-b52e-1231104a68e0", + "19866484-b3d8-4619-85b9-e900de7fa80d", + "c4392ca6-4706-4a96-a6b2-29e6ef544500", + "6798a65f-d79a-413b-baca-115f09caf0c4", + "b844cb2d-6428-4d5e-abac-958f2704d4c8", + "8e9c0b21-2230-4a0a-b35b-a198b90abcaa", + "3492c684-74b5-4317-9e5c-320f3ea01899", + "5beb520c-d413-4998-bd9f-7960a2baf744", + "c96fa96d-d23c-444e-bdef-481317308bb8", + "502ce8f6-bd9e-4e40-a5c9-40bdce36b87f", + "221fd24b-ba2e-40d5-a0c7-a879ece6b34f", + "fb62dc3f-a1cd-4f0e-88bc-504fe9d8585d", + "e9acbb7f-c45c-47d7-9dd8-16194319fc97", + "0d4efcc9-a336-46f9-a1db-0623b5be2869", + "84a388b6-2844-421e-a176-84d29958a76c", + "3ecc3e4d-5466-4243-acca-2feea47e09a2", + "e7b6b4ee-3f2a-41cd-bc18-c7e993cb413f", + "5998ece4-98cd-4d5b-9104-1517d13728bd", + "c9f84bf4-108c-4151-85c0-e87375c53b91", + "5eaa3bfb-cc71-43fa-88dc-d367ef929a7b", + "0fe736da-48f1-4961-bd7a-b75ce13c91b4", + "2c7796e3-b1e9-413f-a0e8-8225ac7d586e", + "69006355-3fb0-443c-826b-ea8c9d27f70b", + "8d73dd20-5349-477f-9630-218e6b08bb95", + "e7633ad6-e406-4a18-b1ba-6c834ec9c047", + "3b3a4043-3d74-4575-8aeb-511c4a879241", + "5daa132e-0b97-4d0f-be9d-b514285661fe", + "52fb1a73-2ba7-413d-8596-33615daecc73", + "9f3cb8b5-5d62-4236-be5e-8f16d59c9f50", + "efdc3dc0-275d-4073-9179-32b14b224654", + "683635b7-5a1b-467c-8d4c-c34e8bf5fdf5", + "2db7496c-aa43-4d9d-8df6-404d2d9bd810", + "9f3edf40-3dca-4086-b995-4985f6848a14", + "ef37e1e6-5be3-4e8c-9f70-c84f341f9aea", + "5df66408-7a8c-43bb-9235-d088aa3d19dc", + "453f43e9-c1d1-4749-9720-73546287b82f", + "1943bfb1-ba78-4a6c-ade7-11e44314dfb4", + "fd8cc20a-915e-45a9-93df-320ad998aee4", + "bc6cf777-fd8a-4427-9867-92a443826827", + "812f0590-4ccd-4c9e-8978-c45e6debce55", + "e6f4e36f-96b7-48f5-afd1-27cc408bdd07", + "1b6dfc9f-c873-48c6-9f25-a2d490e269cf", + "bdf1ca5a-321e-4d3c-ac9e-da2ce1f5a69c", + "604bf883-ce55-4c59-bcc1-3ceea1128211", + "679f2016-4802-4bf8-a830-d1f45fdfcd71", + "6f3f70b4-8428-4e5f-b1dd-86cd55d8e63f", + "83d8be38-ce02-41b9-b444-7291bad513d7", + "7fd21591-6a33-49cf-a32f-f8dfdf5ae697", + "a3dbec39-7126-4949-a916-eca55fee2b49", + "03cebfe9-3319-45b5-aba5-e7db10ea6fdb", + "ff9fe54b-a959-4d7e-a6fa-93e311a68168", + "172e1afa-5581-455d-9a27-590e9a1706c2", + "3ab3b82f-bfa8-4b39-a9a9-e397341ae572", + "b171dcaf-94f2-4bcd-8695-b2648bed60e1", + "063f5040-9e94-42ae-b0e5-b13413e0fe55", + "136cef9c-e17f-49f9-ae45-a17e7e5c5477", + "f906b093-281f-4872-8010-f1f3a320ee2a", + "33365c93-7f1d-4a7d-a9bc-b33010aeafbd", + "e98405ee-0a5b-4df2-bb7e-97729b8b3c11", + "bd13e435-1e38-4b22-bfd0-b620fc74cc8a", + "7dd0eab1-b83e-4bb4-ad91-7bbdf4592647", + "5f5ea963-4f59-4705-95fe-4a98c9fcd12c", + "7e6dfd31-8544-4fad-9e49-0f05516cf8cf", + "f817b176-2dce-4858-999f-fe62aac25e2e", + "83553642-1a38-4c54-8154-c4cd4fd81173", + "dfe025f3-573e-4297-b17a-30e990cfa310", + "2091513c-8f4c-4817-8c22-cdbc6303a331", + "d15eb34f-e425-4990-aa2d-2dc454fb4b20", + "0e519838-1e4a-41fd-ab30-34bf745a3339", + "c920e557-141f-4556-95ac-f494ecc0c84f", + "91e97b3c-c912-401a-ab99-da01b90682aa", + "d6adb1ae-f689-438b-ab81-8a2e37f9654e", + "cd3092e1-6022-4a11-bfc3-a4bb80487ece", + "a874ef10-022c-43e4-b4e6-e01d8fcec988", + "2e6ff051-0037-4001-90a8-de7643a96f08", + "134c4387-77d5-4ee0-911f-e50e64499b42", + "421f1c8b-5237-4735-9003-82021ba651b1", + "9ee0474c-615c-4383-a047-8432b8b24784", + "78efcc13-3718-4acb-acee-8157a54501d6", + "961a8de1-0d33-44ab-a9d5-3d92de13fde4", + "49e25cbd-1aaa-40dd-901b-bb3e524ebcc0", + "5a5b8c37-e3f6-48f2-96d3-289a043eaee4", + "394d1df6-bced-4f04-8c2a-a1ea3a04b984", + "b87c3e0f-a5af-4663-a87d-d0689d73343b", + "09194996-ff4b-4a4f-9800-e8aedc9be7f2", + "3f920762-f4a7-4bdb-86dc-05a3c3c8e875", + "5204624f-90ac-47bf-a71a-486d43f4351c", + "df24897d-a478-4fbb-816f-1d14c59da30d", + "9ba7c2a0-c86f-4967-9284-9be88ad3eacf", + "3f28dd4f-f653-4937-83b6-9f80c1284a85", + "b2735994-59ae-4478-9c5f-93f5ccf18c8b", + "5e965191-bcd1-491a-a9ce-e014e970020b", + "85718cf6-e4ba-44e3-880f-143c627f12ff", + "d0af657d-9c21-447c-a4f9-199b38deac5b", + "f8289825-2858-4dfa-b173-f4c506359158", + "8f2d26e6-239e-492f-aa1b-472f8ec85050", + "836be9c2-cee6-4fb9-8bbd-74197d213c5a", + "cd998593-e5bc-49d0-86a7-333acf9abf93", + "8deeeb88-19ed-4747-9c79-d94104bed7e0", + "6d1dc078-8b04-480e-a36f-298594ee9592", + "83464c62-929f-4bf7-9a04-43071189edde", + "1efd8b7b-65f9-4e73-bd49-d707190d6519", + "f0b69458-25df-4dda-be8b-e14383a57a50", + "88260ff4-be0b-4544-98ef-78b6a0d61ebc", + "39fb2989-1dd2-490d-b793-dce86c36b43a", + "b4f650b4-53d6-4069-a55f-43f2de7ff47b", + "15b6759d-2370-43ee-9a61-09e453db2469", + "2e4e0dac-ab21-49ff-9563-e320acd73f2f", + "986c4e43-ac2e-4983-8d56-f1dd34775133", + "12c1b77c-cbea-4c20-8b64-7154a25bbbe1", + "a3da62c3-3f46-499f-a4c6-acf93bd1c63d", + "57b6455e-2acb-4076-b8c4-1f2cdae97185", + "d6d8261e-0ee0-4002-bddf-e94fdc1075c7", + "56683d7b-7c65-4f42-b5a4-818f3cd5224c", + "9090779d-803e-4fbd-9166-d9e90ff2d36f", + "d5374572-5576-4f83-99c8-a0b12c4de160", + "ffddab3a-dbec-49b0-b114-5a1e3d805efa", + "38c02527-1b2f-458b-bad8-73a282c1d99a", + "d612edfb-e904-4de0-abce-a5fd9ecc449b", + "1261ee8c-80fc-4a3f-9dd3-29f4a918a9fe", + "8e860a91-0205-447b-a706-aa894d23d3db", + "fb07f205-868e-44c9-a704-305daaf03301", + "48f7028c-b6bf-4ca5-81e9-128147fd02c9", + "c6a6e1e4-24c4-422d-9d18-f062240ddf23", + "ba280b86-22a2-4509-b16d-b55e04680461", + "f0cb79ef-c081-4049-85ef-2623e02c9589", + "6dbfc053-7899-40d8-9827-0ccd21f3ee0a", + "50d3035c-2503-4bcf-b28b-e66c92c98f09", + "41f7665a-a692-42e0-b28d-9769c03ab99f", + "cb1f70b4-0bf6-47f1-bc98-2e9365a55c3f", + "7f566780-4ed2-41c5-9faa-06759495faab", + "0e08e1c7-3add-4e41-b28d-cd5f4413a422", + "672e0d75-df0e-4715-8c42-f1131b059a6c", + "8f15aab0-9c42-452e-b78a-47a48da1f3c1", + "edaa0b1a-d504-4852-8025-4f12b4d293ac", + "4a9ada82-9f86-495e-b17f-25b39bef2297", + "07774a7c-c5ed-4133-991f-926a316068a4", + "5a009338-0aff-4313-b02f-c48630ce1702", + "8324c0ba-7a10-42e1-90fd-afb4d3afa789", + "3db37314-adaa-4a4e-af11-c9190fa6b65b", + "b21b35d2-6d69-40a2-8c36-aa27e4a980ae", + "6daa9124-d894-42e5-8d98-2e0e6d1191ac", + "8a256d7f-2ea1-4b0e-94ea-168a1ed85678", + "08b6f5b3-4471-4cfd-a95e-4901cb1ded57", + "ea9d85e7-6a9a-451a-b019-972a6891c11b", + "f84fd67e-0797-44f4-b3ff-78da646362e2", + "ad1a9ea9-d1ab-474a-b715-1a42d892ef5d", + "ffb191f8-894a-40e8-baa5-5578aa6ba706", + "fd394f00-4049-49e5-9c07-ee3e86979003", + "9767a3ca-0a9c-480f-adfe-e4164c52c760", + "95cdb269-0088-4451-9a0d-f72349081fab", + "4ade0a82-16fd-471f-9fcc-fd2724887aef", + "fe69a78e-7773-45d1-9e0f-bacee52dac83", + "3aed7bbd-75f0-45ef-b9cc-7f5ff6f7fbbf", + "f8d71e9a-982a-437c-a968-867149c24eb7", + "8351998e-ead7-4bc1-b0a7-a7209da92b3c", + "3eb00a40-880a-4ad9-8cec-3ba0245ec472", + "602b224a-f5ac-491b-88e3-49043f362b6e", + "89125bc3-a413-44ee-a7e1-2753c511ecab", + "00792fa8-988c-4c85-8e80-73eb3ac53e80", + "094d1ab2-f34e-4a2a-ac3b-8671ce5faf65", + "4cf10990-13e5-4a1e-a913-7ec4b8c8ab02", + "a9b191de-d565-4981-8f68-02069a2ba16c", + "0eeab9d2-a55c-4708-94c3-cc1e08f9f262", + "4e78f728-06f6-43a3-ac8c-c0cb159e5b32", + "be99e9e2-82f7-4f1f-aff5-371834fde0da", + "5682f9aa-bdfb-409c-b4c2-c0105b35ee78", + "3ce15244-5293-4a5e-8c92-52ab29c6949a", + "b4b2c55b-c107-468f-bbdb-f73a0e4dca34", + "8a409551-657a-48c0-8beb-056c228ebca5", + "3c17a762-c839-492f-b361-009c1187ae4c", + "22172001-0e17-40db-bcce-825e19532297", + "5b99234b-b940-405d-9123-70218bda24b2", + "f7869fd8-9af3-493a-a627-a3e95e7c3577", + "057e9f41-e6ed-4575-906c-aca9413b35a5", + "70e88e82-5c06-44fd-959c-d4b2d40999cd", + "b95bd349-11bb-4ade-9cc6-3be8dcea06dd", + "19b880da-910e-40c7-970e-cd8718f2faf5", + "e05a2182-a97f-4d0b-87e4-66b0dd405e58", + "3072481a-04fe-4959-8410-018a41a91df1", + "fc2b2014-9dc4-4a5d-8a1d-25a6911bff7c", + "db1d3def-397e-470f-811a-85c4cc5338a9", + "ff283659-4a79-4642-8078-4366df38e1f1", + "bc8bae7f-ad9f-4f72-ae5c-b3e7ca8dd8f0", + "a3bd1be3-edfe-4f4e-89a8-dd56ce6ece08", + "2360312c-4347-476a-990f-2633b489fa7a", + "b9e32e51-8c18-420b-a582-200c79b8208d", + "0511dce9-c126-4a11-92ba-8082a4ec89bd", + "224d8b17-7eaf-4529-9b75-c6ce3078d640", + "10732476-f045-4feb-8199-e4e04685f7ca", + "5753e3fe-f1c5-426b-9f45-5a8efa6b5370", + "b69564f5-5e5e-4097-aebd-9f52f211ff75", + "97eb574d-d8bf-44a8-a6ef-1204b789db43", + "85506322-449a-45c0-b77a-48a8077b4bbd", + "ed65870d-ba08-4894-9cad-7f34e2c6f9e4", + "cf682058-962f-45c7-95ab-dbbd2f87acc0", + "65229ac2-9033-4002-af6c-492fe49bc563", + "4c85a780-8a8e-46c9-a0f8-c7e9703ef998", + "1ac6e060-a6d2-4b42-afef-599be78a533a", + "a22beaf0-377c-4b95-b472-867280a525ae", + "f5bc6a45-1591-42bf-8ed4-8286382b31bf", + "06592f50-7f29-4ca9-94a9-1b6c389d50ea", + "4ca3d9cf-4972-4e1d-aca3-28b036d6bdc5", + "076bdb81-5c75-4282-9f3a-a387624575f3", + "a5fe8210-1c37-4dd9-bd27-74051c521c3d", + "c1ca7c5c-2473-4b83-a507-7d187609111a", + "7c4079e7-4d68-4b0d-b015-5971a9866c60", + "ead54fb5-dee9-452f-8145-227ccd2405e2", + "7c776700-1d6b-4ee2-8155-b7f6855afceb", + "71a7456a-da6e-4e38-86ec-76bacaa6db5b", + "26c042d4-bf21-4d10-81f9-303d3816c114", + "c7820a4f-55c9-44eb-bf7e-b965a9d0e924", + "5b5b6b8c-9b88-4181-9ebb-14abb7783e1e", + "db12030d-7577-435f-9d48-5c1aae615c4f", + "e37c82c5-483f-4e68-8bf3-2a22dcef0b1e", + "0b4cacb1-970f-4ef0-85da-371d81f899e0", + "3949f121-af2b-4c7c-8a49-07befc08abbe", + "73c1ac5a-80a0-4e55-b143-52df5697c9ea", + "0f325614-48c9-4a6a-b7f2-731c26ddeb97", + "572117cf-a76b-44d2-804a-4901bfc545f4", + "33fd4f1f-fb01-4d3a-9b8d-b8d8024a2dab" + ], + "test_session_set": [ + "66cb6d16-f972-40e9-88c6-64b6816b900c", + "e62b75a0-f1f6-4e45-b3db-17fd7ff5756c", + "13ffc0e4-eebb-4baa-97af-5983ce46779c", + "d062849b-c60f-4369-a113-50f6ce8521fb", + "bb7f1e75-65ee-4d98-adc8-c70551a35940", + "7a6576d8-007c-4ee9-b7cc-8cc2d020e1bf", + "4ce65f3a-21ee-4f77-85c0-b3ddfe1e6381", + "783d7043-b27e-4744-a969-5f40fb40827c", + "81b74694-a2e6-4be1-b9f6-69f3aa03f50d", + "e6317c78-06f0-4c6d-9886-ca67712772ed", + "44bcb1fa-9e64-4bae-927a-4d5d53e536a6", + "218e79aa-3f56-4bd4-aa2a-14027bf7d52d", + "561fd770-f9a2-41a4-90da-beeb7d96ee7f", + "a2d94b61-57c7-4ec2-ad94-94b3585a6b64", + "764a08dc-5ac5-4c09-a281-0210cdd5a55e", + "c1057c9b-331f-458b-8914-2703511e8b13", + "3edc700d-7d14-4154-a3e7-5b0e8c3fde07", + "1492c03c-1887-4576-8889-5bc2848967de", + "2ea37ccc-e1ec-4e0d-a2a7-9707ea300b7f", + "65cad708-34b8-4158-be96-b88710288473", + "be8f4445-3958-4fc4-8cdf-b14d1d0eba73", + "3c62a0d2-10d7-4a9e-ac2c-02505663afad", + "be25331c-b1a0-4187-93fb-3b59437258c3", + "108cebaa-777d-42dd-8b1e-69ea7f443a0a", + "be3b65c8-ba74-4321-99b4-ca0e99e6fc2c", + "3bd74a94-ec3e-4138-b78f-ad76af4f07e2", + "4356f6d5-576a-4999-af40-744c242fff22", + "b0c57748-65c8-49d7-9681-ddcfcde3298b", + "12706b2d-0a0b-4511-8feb-ed8065746e50", + "f0e9a139-86f6-41bc-82f1-11283954a781", + "bb3c6009-892c-4071-9bb4-692b83ed12f3", + "90bf3f44-60ae-4430-9dc6-2365f7901fc5", + "2b5cea2f-b019-4ecf-abbc-b2346bd65a10", + "8e9f89d2-f02a-4c75-8831-6785ef6bd7f5", + "13c24739-e514-4fad-8288-6686bffa71ad", + "1828b207-9601-4521-8c0b-5d9e829fca78", + "0324c3b5-5154-4de1-9869-eb866b568e0f", + "1a746147-2dec-4a37-a51b-c94e4f1de010", + "3a20bf0a-6eb1-47c5-9bf8-8445446feb61", + "c2a8f168-06b0-4766-8865-27f08cb1bcae", + "f55f2dfd-704f-42cf-a39e-72e2560314b3", + "4ea88cb9-bff8-49e0-a855-cc1902db9a99", + "4e0053f8-566c-4952-abbe-4ea877782a0a", + "80ee109f-b097-4755-8a39-9f0b33adb170", + "933ddf59-52d2-4d4c-b653-05de44ded245", + "b432e10e-cff1-4f15-a8fe-1e2350e94b36", + "157c8e81-117a-4a15-a047-9a574181d44b", + "cfc6da55-43fc-4cd8-96d0-9de5166f4a63", + "aabeb7f9-5ada-425a-9545-c98dc443887a", + "db9978fb-fe4d-467b-9495-ab7daf3b1b55", + "6549c35b-8b1d-4577-9434-ac808e4d7057", + "e50126ee-f4c1-41a6-a2e8-79636b40b341", + "49a7e53b-aba5-4e53-b3bc-1492657fc8ae", + "7abe6d4a-f594-4a15-af08-ddd591ff500c", + "03c50078-a08d-43ad-a3c9-87916aebd865", + "a9a0ed92-fab1-4b6e-ae27-714692464641", + "3c3207e6-57b9-4be1-96fc-5b2168ae3cd3", + "895545f1-f194-4c0a-8513-87d386670f15", + "7b76be35-4e6d-42d7-9fa9-14a21fc327cf", + "e4670e4f-0b1a-450a-9778-9ee2409db554", + "34bc808a-4b98-42c5-a1ea-6f565dd4aa20", + "bec4e063-9be1-402e-a7b5-dba20730c933", + "20a3e8e5-3f43-46a5-84fa-6888bd9c5c88", + "5da4efa3-8e54-46dc-a8c2-3016fa4e12f8", + "0d998741-de68-4006-898e-7108de4167b1", + "4d144ff3-a853-4b79-9825-d7fcea248373", + "dce1d8b4-2c29-4fec-a343-77885899c434", + "dab038d7-9e08-4edf-8d03-ee7e2af5a10e", + "a1055434-9e9b-4d69-bac3-374a39f801da", + "e7086057-b5f7-4286-87c6-ab110072c5b0", + "3af65596-1aa5-4126-933b-b98ccaccedf2", + "19f880ae-9770-45a5-93f5-00047f22d5f6", + "36888bbd-1ece-4a3a-8aff-b0f6ef905752", + "f42a7e14-c3ee-4b05-a59b-8d349f8bd19b", + "d36d9514-a99b-4534-b54f-3fd98a19f0ea", + "6c4f7cb0-caa0-4a35-9071-a84c9ec2ab2f", + "666d9243-f4f9-47f5-8fba-cee8568abf41", + "2f14e742-dc58-43c8-8ebc-7244487dbba7", + "f8accf9d-4427-4458-88e9-89d51a3f5136", + "50810487-0b0f-45e0-8f37-515644189e81", + "f0445727-2729-4a76-9434-fb081a14cd9d", + "f55f890d-fd90-43f6-a11a-a1920c21890a", + "16851d4d-35e0-49b4-b3b5-28e67b564db3", + "5c8f8839-4bf9-45e0-beb0-6c17a5ead0d5", + "514d71c8-fd4e-4908-89a4-c0861e81ca8a", + "35465081-a518-4321-8384-477ef5b29b8d", + "de82a6f3-1a55-4996-ae68-35032678ff66" + ] + } +} \ No newline at end of file diff --git a/third_party/EgoSpeak/LSTR/run_LSTR_inference.sh b/third_party/EgoSpeak/LSTR/run_LSTR_inference.sh new file mode 100644 index 0000000..86a977e --- /dev/null +++ b/third_party/EgoSpeak/LSTR/run_LSTR_inference.sh @@ -0,0 +1,24 @@ +#!/bin/bash + +PATH_TO_CONFIG_FILE=configs/EasyCom/LSTR/lstr_aa_easycom.yaml +PATH_TO_CHECKPOINT_FILE=checkpoints/lstr_easycom_AV.pth + +echo "---" +echo "CUDA_VISIBLE_DEVICES: $CUDA_VISIBLE_DEVICES" +echo "Config file: $PATH_TO_CONFIG_FILE" +# echo "Checkpoint opt: $CHECKPOINT_OPT" +# echo "Epoch opt: $EPOCH_OPT" + +echo python tools/test_net.py \ + --config_file $PATH_TO_CONFIG_FILE \ + --gpu $CUDA_VISIBLE_DEVICES \ + MODEL.CHECKPOINT $PATH_TO_CHECKPOINT_FILE \ + MODEL.LSTR.INFERENCE_MODE batch \ + $@ + +python tools/test_net.py \ + --config_file $PATH_TO_CONFIG_FILE \ + --gpu $CUDA_VISIBLE_DEVICES \ + MODEL.CHECKPOINT $PATH_TO_CHECKPOINT_FILE \ + MODEL.LSTR.INFERENCE_MODE batch \ + $@ \ No newline at end of file diff --git a/third_party/EgoSpeak/LSTR/run_LSTR_train.sh b/third_party/EgoSpeak/LSTR/run_LSTR_train.sh new file mode 100644 index 0000000..710d774 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/run_LSTR_train.sh @@ -0,0 +1,25 @@ +#!/bin/bash + +PATH_TO_CONFIG_FILE=configs/EasyCom/LSTR/lstr_aa_easycom.yaml +shift # Remove the first argument so $@ contains all other arguments + +echo "Config file: $PATH_TO_CONFIG_FILE" +echo "Additional options: $@" +# EXAMPLE: +# USE_WANDB True +# INPUT.MODALITY "visual+motion+object" + +echo "---" +echo "CUDA_VISIBLE_DEVICES: $CUDA_VISIBLE_DEVICES" +echo "Config file: $PATH_TO_CONFIG_FILE" + +echo python tools/train_net.py \ + --config_file $PATH_TO_CONFIG_FILE \ + --gpu $CUDA_VISIBLE_DEVICES \ + $@ + +python tools/train_net.py \ + --config_file $PATH_TO_CONFIG_FILE \ + --gpu $CUDA_VISIBLE_DEVICES \ + $@ + diff --git a/third_party/EgoSpeak/LSTR/setup.py b/third_party/EgoSpeak/LSTR/setup.py new file mode 100644 index 0000000..d79bec1 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/setup.py @@ -0,0 +1,8 @@ +from setuptools import find_packages, setup + +setup( + name='LSTR', + version='0.1', + packages=find_packages('src') + ['configs'], + package_dir={'': 'src', 'configs': 'configs'}, +) \ No newline at end of file diff --git a/third_party/EgoSpeak/LSTR/src/__init__.py b/third_party/EgoSpeak/LSTR/src/__init__.py new file mode 100644 index 0000000..04f8b7b --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/__init__.py @@ -0,0 +1,2 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/__init__.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/__init__.py new file mode 100644 index 0000000..04f8b7b --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/__init__.py @@ -0,0 +1,2 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/config/__init__.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/config/__init__.py new file mode 100644 index 0000000..880ba93 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/config/__init__.py @@ -0,0 +1,4 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +from .defaults import get_cfg diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/config/defaults.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/config/defaults.py new file mode 100644 index 0000000..0466b69 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/config/defaults.py @@ -0,0 +1,171 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +from yacs.config import CfgNode as CN + +# ---------------------------------------------------------------------------- # +# Config Definition +# ---------------------------------------------------------------------------- # +_C = CN() + +# ---------------------------------------------------------------------------- # +# Metadata +# ---------------------------------------------------------------------------- # +_C.SEED = 0 + +# ---------------------------------------------------------------------------- # +# Model +# ---------------------------------------------------------------------------- # +_C.MODEL = CN() +_C.MODEL.MODEL_NAME = '' +_C.MODEL.CHECKPOINT = '' +_C.MODEL.EXCEPT_CLASSIFIER = False + +# ---------------------------------------------------------------------------- # +# Feature Head +# ---------------------------------------------------------------------------- # +_C.MODEL.FEATURE_HEAD = CN() +_C.MODEL.FEATURE_HEAD.LINEAR_ENABLED = True +_C.MODEL.FEATURE_HEAD.LINEAR_OUT_FEATURES = 1024 +_C.MODEL.FEATURE_HEAD.DROPOUT = 0. + +# ---------------------------------------------------------------------------- # +# Transformer Network +# ---------------------------------------------------------------------------- # +_C.MODEL.LSTR = CN() +# Hyperparameters +_C.MODEL.LSTR.NUM_HEADS = 8 +_C.MODEL.LSTR.DIM_FEEDFORWARD = 1024 +_C.MODEL.LSTR.DROPOUT = 0.2 +_C.MODEL.LSTR.DROPOUT_CLS = 0.0 +_C.MODEL.LSTR.ACTIVATION = 'relu' +# Memory choices +_C.MODEL.LSTR.AGES_MEMORY_SECONDS = 0 +_C.MODEL.LSTR.AGES_MEMORY_SAMPLE_RATE = 1 +_C.MODEL.LSTR.LONG_MEMORY_SECONDS = 0 +_C.MODEL.LSTR.LONG_MEMORY_SAMPLE_RATE = 1 +_C.MODEL.LSTR.WORK_MEMORY_SECONDS = 8 +_C.MODEL.LSTR.WORK_MEMORY_SAMPLE_RATE = 1 +_C.MODEL.LSTR.WORK_MEMORY_USE_PE = True +_C.MODEL.LSTR.LONG_MEMORY_USE_PE = False +_C.MODEL.LSTR.LONG_MEMORY_INCLUDE_WORK = False +_C.MODEL.LSTR.LONG_MEMORY_INCLUDE_WORK2 = False +# Anticipation choices +_C.MODEL.LSTR.ANTICIPATION_SECONDS = 0 +_C.MODEL.LSTR.ANTICIPATION_SAMPLE_RATE = 1 +# for EK55/100 +_C.MODEL.LSTR.V_N_CLASSIFIER = False +_C.MODEL.LSTR.LOSS_ANTICIPATE_ONLY = False +_C.MODEL.LSTR.LOSS_WEIGHT_V = 1. +_C.MODEL.LSTR.LOSS_WEIGHT_N = 1. +_C.MODEL.LSTR.FUSION_WEIGHT_V = 1. +_C.MODEL.LSTR.FUSION_WEIGHT_N = 1. +# Design choices +_C.MODEL.LSTR.ENC_ATTENTION_TYPE = 'dotproduct' +_C.MODEL.LSTR.ENC_ATTENTION_DECAY = 1.0 +_C.MODEL.LSTR.ENC_MODULE = [ + [16, 1, True], [32, 2, True] +] +_C.MODEL.LSTR.DEC_MODULE = [-1, 2, True] +_C.MODEL.LSTR.FC_NORM = False +# Inference modes +_C.MODEL.LSTR.INFERENCE_MODE = 'batch' +# Position Embedding +# _C.MODEL.LSTR.POSITION_EMBEDDING = 'default' + +# ---------------------------------------------------------------------------- # +# Criterion +# ---------------------------------------------------------------------------- # +_C.MODEL.CRITERIONS = [['MCE', {}]] + +# ---------------------------------------------------------------------------- # +# Data +# ---------------------------------------------------------------------------- # +_C.DATA = CN() +_C.DATA.DATA_INFO = 'data/data_info.json' +_C.DATA.DATA_NAME = None +_C.DATA.DATA_ROOT = None +_C.DATA.CLASS_NAMES = None +_C.DATA.NUM_CLASSES = None +_C.DATA.IGNORE_INDEX = None +_C.DATA.METRICS = None +_C.DATA.FPS = None +_C.DATA.TRAIN_SESSION_SET = None +_C.DATA.TEST_SESSION_SET = None +_C.DATA.EK_EXT_PATH = 'external/rulstm/RULSTM/data/ek55/' +_C.DATA.CLIP_MIXUP_RATE = 0. +_C.DATA.CLIP_MIXUP_SAMPLE = 'uniform' + +# ---------------------------------------------------------------------------- # +# Input +# ---------------------------------------------------------------------------- # +_C.INPUT = CN() +_C.INPUT.MODALITY = 'twostream' +# 'twostream'='visual+motion', 'threestream'='visual+motion+object' +_C.INPUT.VISUAL_FEATURE = 'rgb_anet_resnet50' +_C.INPUT.MOTION_FEATURE = 'flow_anet_resnet50' +_C.INPUT.OBJECT_FEATURE = 'flow_anet_resnet50' +_C.INPUT.TARGET_PERFRAME = 'target_perframe' + +# ---------------------------------------------------------------------------- # +# Data Loader +# ---------------------------------------------------------------------------- # +_C.DATA_LOADER = CN() +_C.DATA_LOADER.BATCH_SIZE = 32 +_C.DATA_LOADER.NUM_WORKERS = 4 +_C.DATA_LOADER.PIN_MEMORY = False + +# _C.MODEL.MINIROAD.NUM_WORKERS = 4 + +# ---------------------------------------------------------------------------- # +# Solver +# ---------------------------------------------------------------------------- # +_C.SOLVER = CN() +_C.SOLVER.START_EPOCH = 1 +_C.SOLVER.NUM_EPOCHS = 50 + +# ---------------------------------------------------------------------------- # +# Optimizer +# ---------------------------------------------------------------------------- # +_C.SOLVER.OPTIMIZER = 'adam' +_C.SOLVER.BASE_LR = 0.00005 +_C.SOLVER.WEIGHT_DECAY = 0.00005 +_C.SOLVER.MOMENTUM = 0.9 + +# _C.MODEL.MINIROAD.LR = 0.0001 +# _C.MODEL.MINIROAD.WEIGHT_DECAY = 0.05 + + +# ---------------------------------------------------------------------------- # +# Scheduler +# ---------------------------------------------------------------------------- # +_C.SOLVER.SCHEDULER = CN() +_C.SOLVER.SCHEDULER.SCHEDULER_NAME = None +_C.SOLVER.SCHEDULER.MILESTONES = [] +_C.SOLVER.SCHEDULER.GAMMA = 0.1 +_C.SOLVER.SCHEDULER.WARMUP_FACTOR = 0.3 +_C.SOLVER.SCHEDULER.WARMUP_EPOCHS = 10.0 +_C.SOLVER.SCHEDULER.WARMUP_METHOD = 'linear' + +# ---------------------------------------------------------------------------- # +# Others +# ---------------------------------------------------------------------------- # +_C.SOLVER.PHASES = ['train', 'test'] +_C.SOLVER.SAVE_EVERY_EPOCHS = 10 + +# ---------------------------------------------------------------------------- # +# Output +# ---------------------------------------------------------------------------- # +_C.OUTPUT_DIR = 'checkpoints' +_C.SESSION = '' +_C.USE_WANDB = False +_C.SAVE_HPARAM = False + +# ---------------------------------------------------------------------------- # +# Misc +# ---------------------------------------------------------------------------- # +_C.VERBOSE = False + + +def get_cfg(): + return _C.clone() diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/config/feature_sizes.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/config/feature_sizes.py new file mode 100644 index 0000000..4f3c02f --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/config/feature_sizes.py @@ -0,0 +1,20 @@ +# dimension of features extracted from different models +# used in feature_head, dataset, models + +FEATURE_SIZES = { + 'rgb_anet_resnet50': 2048, + 'flow_anet_resnet50': 2048, + 'rgb_kinetics_bninception': 1024, + 'flow_kinetics_bninception': 1024, + 'rgb_farl': 512, + 'rgb_kinetics_resnet50': 2048, + 'rgb_kinetics_resnet50_debug': 2048, + 'rgb_kinetics_resnet50_debug_resized': 2048, + 'flow_kinetics_resnet50': 2048, + 'flow_nv_kinetics_bninception': 1024, + 'rgb_kinetics_i3d': 2048, + 'flow_kinetics_i3d': 2048, + 'audio_wav2vec2': 5120, + 'audio_wav2vec2_left': 5120, + 'audio_wav2vec2_right': 5120, +} \ No newline at end of file diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/criterions/__init__.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/criterions/__init__.py new file mode 100644 index 0000000..6d5bb77 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/criterions/__init__.py @@ -0,0 +1,4 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +from .criterions import build_criterion diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/criterions/criterions.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/criterions/criterions.py new file mode 100644 index 0000000..72efcd1 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/criterions/criterions.py @@ -0,0 +1,356 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +__all__ = ['build_criterion'] + +import os.path as osp +from functools import partial + +import numpy as np +import pandas as pd +import torch +import torch.nn as nn +import torch.nn.functional as F +from rekognition_online_action_detection.utils.registry import Registry + +CRITERIONS = Registry() + + +@CRITERIONS.register('BCE') +class BinaryCrossEntropyLoss(nn.Module): + + def __init__(self, reduction='mean', ignore_index=-100): + super(BinaryCrossEntropyLoss, self).__init__() + + self.criterion = nn.BCEWithLogitsLoss(reduction=reduction) + + def forward(self, input, target): + return self.criterion(input, target) + + +@CRITERIONS.register('SCE') +class SingleCrossEntropyLoss(nn.Module): + + def __init__(self, reduction='mean', ignore_index=-100): + super(SingleCrossEntropyLoss, self).__init__() + + self.criterion = nn.CrossEntropyLoss( + reduction=reduction, ignore_index=ignore_index) + + def forward(self, input, target): + return self.criterion(input, target) + + +@CRITERIONS.register('MCE') +class MultipCrossEntropyLoss(nn.Module): + + def __init__(self, reduction='mean', ignore_index=-100): + super(MultipCrossEntropyLoss, self).__init__() + + self.reduction = reduction + self.ignore_index = ignore_index + + def forward(self, input, target): + logsoftmax = nn.LogSoftmax(dim=-1).to(input.device) + + # LSTR transform input from [16,50,3] to [800,3] + output = torch.sum(-target * logsoftmax(input), dim=1) + + if self.reduction == 'mean': + return torch.mean(output) + elif self.reduction == 'sum': + return torch.sum(output) + else: + return output + + # if self.ignore_index >= 0: + # notice_index = [i for i in range(target.shape[-1]) if i != self.ignore_index] + # output = torch.sum(-target[:, notice_index] * logsoftmax(input[:, notice_index]), dim=1) + # if (target[:, self.ignore_index] != 1).sum().item() == 0: + # return torch.mean(torch.zeros_like(output)) + # if self.reduction == 'mean': + # return torch.mean(output[target[:, self.ignore_index] != 1]) + # elif self.reduction == 'sum': + # return torch.sum(output[target[:, self.ignore_index] != 1]) + # else: + # return output[target[:, self.ignore_index] != 1] + # else: + # output = torch.sum(-target * logsoftmax(input), dim=1) + + # if self.reduction == 'mean': + # return torch.mean(output) + # elif self.reduction == 'sum': + # return torch.sum(output) + # else: + # return output + +@CRITERIONS.register('MCE_MiniROAD') +class MiniROADMultipCrossEntropyLoss(nn.Module): + + def __init__(self, reduction='mean', ignore_index=-100): + super(MiniROADMultipCrossEntropyLoss, self).__init__() + + self.reduction = reduction + self.ignore_index = ignore_index + + def forward(self, input, target): + logsoftmax = nn.LogSoftmax(dim=-1).to(input.device) + output = torch.sum(-F.normalize(target) * logsoftmax(input), dim=1) # B + if self.reduction == 'mean': + loss = torch.mean(output) + elif self.reduction == 'sum': + loss = torch.sum(output) + + return loss + +@CRITERIONS.register('MCE_ANT') +class AntMultipCrossEntropyLoss(nn.Module): + + def __init__(self, reduction='mean', ignore_index=-100): + super(AntMultipCrossEntropyLoss, self).__init__() + + self.reduction = reduction + self.ignore_index = ignore_index + + # only ant_target loss + def forward(self, input, target): + logsoftmax = nn.LogSoftmax(dim=-1).to(input.device) + import pdb; pdb.set_trace() + if self.ignore_index >= 0: + notice_index = [i for i in range(target.shape[-1]) if i != self.ignore_index] + import pdb; pdb.set_trace() + # input dim: [16, 128, 10, 3] + # target.shape: [16, 10, 3] + output = torch.sum() + + output = torch.sum(-target[:, notice_index] * logsoftmax(input[:, notice_index]), dim=1) + if (target[:, self.ignore_index] != 1).sum().item() == 0: + return torch.mean(torch.zeros_like(output)) + if self.reduction == 'mean': + return torch.mean(output[target[:, self.ignore_index] != 1]) + elif self.reduction == 'sum': + return torch.sum(output[target[:, self.ignore_index] != 1]) + else: + return output[target[:, self.ignore_index] != 1] + else: + output = torch.sum(-target * logsoftmax(input), dim=1) + + if self.reduction == 'mean': + return torch.mean(output) + elif self.reduction == 'sum': + return torch.sum(output) + else: + return output + + +@CRITERIONS.register('MCE_EQL') +class MultipCrossEntropyEqualizedLoss(nn.Module): + + def __init__(self, gamma=0.95, lambda_=3e-3, reduction='mean', ignore_index=-100, + anno_path='external/rulstm/RULSTM/data/ek55/'): + super(MultipCrossEntropyEqualizedLoss, self).__init__() + + # get label distribution + segment_list = pd.read_csv(osp.join(anno_path, 'training.csv'), + names=['id', 'video', 'start_f', 'end_f', 'verb', 'noun', 'action'], + skipinitialspace=True) + freq_info = np.zeros((max(segment_list['action']) + 1,)) + assert ignore_index == 0 + for segment in segment_list.iterrows(): + freq_info[segment[1]['action']] += 1. + freq_info = freq_info / freq_info.sum() + self.freq_info = torch.FloatTensor(freq_info) + + self.gamma = gamma + self.lambda_ = lambda_ + self.reduction = reduction + self.ignore_index = ignore_index + + def forward(self, input, target): + logsoftmax = nn.LogSoftmax(dim=1).to(input.device) + + + bg_target = target[:, self.ignore_index] + notice_index = [i for i in range(target.shape[-1]) if i != self.ignore_index] + input = input[:, notice_index] + target = target[:, notice_index] + + weight = input.new_zeros(len(notice_index)) + weight[self.freq_info < self.lambda_] = 1. + weight = weight.view(1, -1).repeat(input.shape[0], 1) + + eql_w = 1 - (torch.rand_like(target) < self.gamma) * weight * (1 - target) + input = torch.log(eql_w + 1e-8) + input + + output = torch.sum(-target * logsoftmax(input), dim=1) + if (bg_target != 1).sum().item() == 0: + return torch.mean(torch.zeros_like(output)) + if self.reduction == 'mean': + return torch.mean(output[bg_target != 1]) + elif self.reduction == 'sum': + return torch.sum(output[bg_target != 1]) + else: + return output[bg_target != 1] + + +@CRITERIONS.register('EQLv2') +class EQLv2Loss(nn.Module): + + def __init__(self, gamma=12, mu=0.8, alpha=4.0, reduction='mean', ignore_index=-100, + num_classes=3806): + super(EQLv2Loss, self).__init__() + self.num_classes = num_classes + + self.gamma = gamma + self.mu = mu + self.alpha = alpha + self.ignore_index = ignore_index + + self._pos_grad = None + self._neg_grad = None + self.pos_neg = None + + def _func(x, gamma, mu): + return 1 / (1 + torch.exp(-gamma * (x - mu))) + self.map_func = partial(_func, gamma=self.gamma, mu=self.mu) + + def forward(self, input, target): + self.n_i, self.n_c = input.size() + + pos_w, neg_w = self.get_weight(input) + + weight = pos_w * target + neg_w * (1 - target) + + cls_loss = F.binary_cross_entropy_with_logits(input, target, reduction='none') + cls_loss = torch.sum(cls_loss * weight) / self.n_i + + self.collect_grad(input.detach(), target.detach(), weight.detach()) + + return cls_loss + + def collect_grad(self, cls_score, target, weight): + prob = torch.sigmoid(cls_score) + grad = target * (prob - 1) + (1 - target) * prob + grad = torch.abs(grad) + + # do not collect grad for objectiveness branch [:-1] + pos_grad = torch.sum(grad * target * weight, dim=0)[1:] + neg_grad = torch.sum(grad * (1 - target) * weight, dim=0)[1:] + + # dist.all_reduce(pos_grad) + # dist.all_reduce(neg_grad) + + self._pos_grad += pos_grad + self._neg_grad += neg_grad + self.pos_neg = self._pos_grad / (self._neg_grad + 1e-10) + + def get_weight(self, cls_score): + if self._pos_grad is None: + self._pos_grad = cls_score.new_zeros(self.num_classes) + self._neg_grad = cls_score.new_zeros(self.num_classes) + neg_w = cls_score.new_zeros((self.n_i, self.n_c)) + pos_w = cls_score.new_zeros((self.n_i, self.n_c)) + else: + neg_w = torch.cat([cls_score.new_ones(1), self.map_func(self.pos_neg)]) + pos_w = 1 + self.alpha * (1 - neg_w) + neg_w = neg_w.view(1, -1).expand(self.n_i, self.n_c) + pos_w = pos_w.view(1, -1).expand(self.n_i, self.n_c) + return pos_w, neg_w + + +@CRITERIONS.register('MCE_SEESAW') +class MultipCrossEntropySeesawLoss(nn.Module): + + def __init__(self, reduction='mean', ignore_index=-100, + p=0.8, q=2.0, eps=1e-2, + anno_path='external/rulstm/RULSTM/data/ek55/'): + super(MultipCrossEntropySeesawLoss, self).__init__() + + # get label distribution + segment_list = pd.read_csv(osp.join(anno_path, 'training.csv'), + names=['id', 'video', 'start_f', 'end_f', 'verb', 'noun', 'action'], + skipinitialspace=True) + self.num_classes = max(segment_list['action']) + 1 + self.cum_samples = torch.zeros(self.num_classes, dtype=torch.float) + + self.p = p + self.q = q + self.eps = eps + + self.reduction = reduction + assert ignore_index == 0 + self.ignore_index = ignore_index + + def forward(self, input, target): + logsoftmax = nn.LogSoftmax(dim=1).to(input.device) + self.cum_samples = self.cum_samples.to(input.device) + + bg_target = target[:, self.ignore_index] + notice_index = [i for i in range(target.shape[-1]) if i != self.ignore_index] + input = input[:, notice_index] + target = target[:, notice_index] + + # accumulate the samples for each category + self.cum_samples += target.sum(0) + + seesaw_weights = input.new_ones(target.size()) + + if self.p > 0: + sample_ratio_matrix = self.cum_samples[None, :].clamp(min=1) / self.cum_samples[:, None].clamp(min=1) + index = (sample_ratio_matrix < 1.0).float() + sample_weights = sample_ratio_matrix.pow(self.p) * index + (1 - index) + for ni in range(input.shape[0]): + true_target = target[ni, :].nonzero(as_tuple=True)[0] + ## TODO: study mean/min/max/sum + if len(true_target) > 0: + mitigation_factor = torch.min(sample_weights[true_target, :], dim=0).values + seesaw_weights[ni, :] = seesaw_weights[ni, :] * mitigation_factor + if self.q > 0: + scores = F.softmax(input.detach(), dim=1) + self_scores = input.new_ones([input.shape[0]]) + for ni in range(input.shape[0]): + true_target = target[ni, :].nonzero(as_tuple=True)[0] + ## TODO: study mean/min/max/sum + if len(true_target) > 0: + self_scores[ni] = torch.min(scores[ni, true_target]) + score_matrix = scores / self_scores[:, None].clamp(min=self.eps) + index = (score_matrix > 1.0).float() + compensation_factor = score_matrix.pow(self.q) * index + (1 - index) + seesaw_weights = seesaw_weights * compensation_factor + + input = input + (seesaw_weights.log() * (1 - target)) + + output = torch.sum(-target * logsoftmax(input), dim=1) + if (bg_target != 1).sum().item() == 0: + return torch.mean(torch.zeros_like(output)) + if self.reduction == 'mean': + return torch.mean(output[bg_target != 1]) + elif self.reduction == 'sum': + return torch.sum(output[bg_target != 1]) + else: + return output[bg_target != 1] + + +@CRITERIONS.register('PRED_FUTURE') +class PredictFutureLoss(nn.Module): + + def __init__(self, reduction='mean', ignore_index=-100): + super(PredictFutureLoss, self).__init__() + + self.reduction = reduction + self.criterion = nn.MSELoss(reduction=reduction) + + def forward(self, pred, true): + return self.criterion(pred, true) + + +def build_criterion(cfg, device=None): + criterion = {} + for name, params in cfg.MODEL.CRITERIONS: + if name in CRITERIONS: + if 'ignore_index' not in params: + params['ignore_index'] = cfg.DATA.IGNORE_INDEX + criterion[name] = CRITERIONS[name](**params).to(device) + else: + raise RuntimeError('Unknown criterion: {}'.format(name)) + return criterion diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/datasets/__init__.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/datasets/__init__.py new file mode 100644 index 0000000..04b8ae4 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/datasets/__init__.py @@ -0,0 +1,5 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +from .datasets import build_data_loader, build_dataset +from .perframe_data_layers import * diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/datasets/datasets.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/datasets/datasets.py new file mode 100644 index 0000000..5ad5f7e --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/datasets/datasets.py @@ -0,0 +1,31 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + + +__all__ = [ + 'build_dataset', + 'build_data_loader', +] + +import torch.utils.data as data +from rekognition_online_action_detection.utils.registry import Registry + +DATA_LAYERS = Registry() + + +def build_dataset(cfg, phase, tag=''): + data_names = ['YTConv', 'EGO4D'] + data_name = next((name for name in data_names if name in cfg.DATA.DATA_NAME), cfg.DATA.DATA_NAME) + + data_layer = DATA_LAYERS[cfg.MODEL.MODEL_NAME + tag + data_name] + return data_layer(cfg, phase) + + +def build_data_loader(cfg, phase): + data_loader = data.DataLoader( + dataset=build_dataset(cfg, phase), + batch_size=cfg.DATA_LOADER.BATCH_SIZE, + shuffle=True if phase == 'train' else False, + num_workers=cfg.DATA_LOADER.NUM_WORKERS, + pin_memory=cfg.DATA_LOADER.PIN_MEMORY, + ) + return data_loader diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/datasets/perframe_data_layers.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/datasets/perframe_data_layers.py new file mode 100644 index 0000000..20e0c45 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/datasets/perframe_data_layers.py @@ -0,0 +1,453 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + + +import os.path as osp +import warnings +from bisect import bisect_right + +import numpy as np +import pandas as pd +import torch +import torch.utils.data as data +from rekognition_online_action_detection.utils.ek_utils import ( + action_to_noun_map, action_to_verb_map) + +from .datasets import DATA_LAYERS as registry + + +@registry.register('LSTRTHUMOS') +@registry.register('LSTRTVSeries') +@registry.register('LSTREK55') +@registry.register('LSTREK100') +@registry.register('LSTRYTConv') +@registry.register('LSTREGO4D') +@registry.register('LSTREasyCom') +class LSTRDataLayer(data.Dataset): + + def __init__(self, cfg, phase='train'): + self.cfg = cfg + self.data_name = cfg.DATA.DATA_NAME + self.data_root = cfg.DATA.DATA_ROOT + self.clip_mixup_rate = cfg.DATA.CLIP_MIXUP_RATE + self.clip_mixup_sample = cfg.DATA.CLIP_MIXUP_SAMPLE + self.visual_feature = cfg.INPUT.VISUAL_FEATURE + self.motion_feature = cfg.INPUT.MOTION_FEATURE + self.object_feature = cfg.INPUT.OBJECT_FEATURE + self.target_perframe = cfg.INPUT.TARGET_PERFRAME + self.sessions = getattr(cfg.DATA, phase.upper() + '_SESSION_SET') + self.long_memory_length = cfg.MODEL.LSTR.LONG_MEMORY_LENGTH + self.long_memory_sample_rate = cfg.MODEL.LSTR.LONG_MEMORY_SAMPLE_RATE + self.long_memory_num_samples = cfg.MODEL.LSTR.LONG_MEMORY_NUM_SAMPLES + self.work_memory_length = cfg.MODEL.LSTR.WORK_MEMORY_LENGTH + self.work_memory_sample_rate = cfg.MODEL.LSTR.WORK_MEMORY_SAMPLE_RATE + self.work_memory_num_samples = cfg.MODEL.LSTR.WORK_MEMORY_NUM_SAMPLES + self.anticipation_length = cfg.MODEL.LSTR.ANTICIPATION_LENGTH + self.anticipation_sample_rate = cfg.MODEL.LSTR.ANTICIPATION_SAMPLE_RATE + self.anticipation_num_samples = cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES + self.training = phase == 'train' + self._init_dataset() + + def shuffle(self): + self._init_dataset() + + def _init_dataset(self): + self.inputs = [] + + if self.data_name.startswith('EK'): + if self.data_name == 'EK55': + path_to_data = 'external/rulstm/RULSTM/data/ek55/' + elif self.data_name == 'EK100': + path_to_data = 'external/rulstm/RULSTM/data/ek100/' + else: + raise ValueError + segment_list = pd.read_csv(osp.join(path_to_data, + 'training.csv' if self.training else 'validation.csv'), + names=['id', 'video', 'start_f', 'end_f', 'verb', 'noun', 'action'], + skipinitialspace=True) + segment_list['len_f'] = segment_list['end_f'] - segment_list['start_f'] + + self.segment_by_cls = {cls: segment_list[segment_list['action'] == cls] + for cls in range(0, self.cfg.DATA.NUM_CLASSES - 1)} + for session in self.sessions: + target = np.load(osp.join(self.data_root, self.target_perframe, session + '.npy')) + if self.cfg.MODEL.LSTR.V_N_CLASSIFIER: + verb_target = np.load(osp.join(self.data_root, self.target_perframe.replace('target', 'verb'), session + '.npy')) + noun_target = np.load(osp.join(self.data_root, self.target_perframe.replace('target', 'noun'), session + '.npy')) + else: + verb_target = target + noun_target = target + segments_per_session = segment_list[segment_list['video'] == session] + for segment in segments_per_session.iterrows(): + start_tick = int(segment[1]['start_f'] / 30 * self.cfg.DATA.FPS) + end_tick = int(segment[1]['end_f'] / 30 * self.cfg.DATA.FPS) + start_tick += np.random.randint(self.anticipation_length) if self.training else 0 + start_tick = min(start_tick, end_tick) + work_end = start_tick - self.anticipation_length + work_start = work_end - self.work_memory_length + segments_before_current = segments_per_session[segments_per_session['end_f'] < segment[1]['start_f']] + if work_start < 0: + continue + self.inputs.append([ + session, work_start, work_end, + target[work_start: work_end + self.anticipation_length], + verb_target[work_start: work_end + self.anticipation_length], + noun_target[work_start: work_end + self.anticipation_length], + segments_before_current, + ]) + else: + for session in self.sessions: + target = np.load(osp.join(self.data_root, self.target_perframe, session + '.npy')) + if self.cfg.MODEL.LSTR.V_N_CLASSIFIER: + verb_target = np.load(osp.join(self.data_root, self.target_perframe.replace('target', 'verb'), session + '.npy')) + noun_target = np.load(osp.join(self.data_root, self.target_perframe.replace('target', 'noun'), session + '.npy')) + else: + verb_target = target + noun_target = target + seed = np.random.randint(self.work_memory_length + self.anticipation_length) if self.training else 0 + for work_start, work_end in zip( + range(seed, target.shape[0], self.work_memory_length + self.anticipation_length), + range(seed + self.work_memory_length, target.shape[0] - self.anticipation_length, self.work_memory_length + self.anticipation_length)): + self.inputs.append([ + session, work_start, work_end, + target[work_start: work_end + self.anticipation_length], + verb_target[work_start: work_end + self.anticipation_length], + noun_target[work_start: work_end + self.anticipation_length], + None, + ]) + + + + + + + def segment_sampler(self, start, end, num_samples): + indices = np.linspace(start, end, num_samples) + return np.sort(indices).astype(np.int32) + + def uniform_sampler(self, start, end, num_samples, sample_rate): + if start < 0: + start = (end + 1) % sample_rate + indices = np.arange(start, end + 1)[::sample_rate] + padding = num_samples - indices.shape[0] + if padding > 0: + indices = np.concatenate((np.zeros(padding), indices)) + return np.sort(indices).astype(np.int32) + + def __getitem__(self, index): + (session, work_start, work_end, target, + verb_target, noun_target, segments_before) = self.inputs[index] + visual_inputs = np.load( + osp.join(self.data_root, self.visual_feature, session + '.npy'), mmap_mode='r') + motion_inputs = np.load( + osp.join(self.data_root, self.motion_feature, session + '.npy'), mmap_mode='r') + object_inputs = np.load( + osp.join(self.data_root, self.object_feature, session + '.npy'), mmap_mode='r') + + + target = np.concatenate((target[:self.work_memory_length:self.work_memory_sample_rate], + target[self.work_memory_length::self.anticipation_sample_rate]), + axis=0) + if self.cfg.MODEL.LSTR.V_N_CLASSIFIER: + verb_target = np.concatenate((verb_target[:self.work_memory_length:self.work_memory_sample_rate], + verb_target[self.work_memory_length::self.anticipation_sample_rate]), + axis=0) + noun_target = np.concatenate((noun_target[:self.work_memory_length:self.work_memory_sample_rate], + noun_target[self.work_memory_length::self.anticipation_sample_rate]), + axis=0) + + + work_indices = np.arange(work_start, work_end).clip(0) + work_indices = work_indices[::self.work_memory_sample_rate] + work_visual_inputs = visual_inputs[work_indices] + work_motion_inputs = motion_inputs[work_indices] + work_object_inputs = object_inputs[work_indices] + + + if self.long_memory_num_samples > 0: + + long_start, long_end = work_start - self.long_memory_length, work_start - 1 + if self.training: + long_indices = self.segment_sampler( + long_start, + long_end, + self.long_memory_num_samples).clip(0) + else: + long_indices = self.uniform_sampler( + long_start, + long_end, + self.long_memory_num_samples, + self.long_memory_sample_rate).clip(0) + long_visual_inputs = visual_inputs[long_indices] + long_motion_inputs = motion_inputs[long_indices] + long_object_inputs = object_inputs[long_indices] + + if self.training and self.clip_mixup_rate > 0: + assert segments_before is not None + segments_before = segments_before[segments_before['end_f'] / 30 * self.cfg.DATA.FPS > long_start] + if self.clip_mixup_sample == 'uniform': + prob = None + elif self.clip_mixup_sample == 'by_length': + prob = (segments_before['len_f']).tolist() + prob = prob / np.sum(prob) + else: + raise ValueError + num_clip = int(len(segments_before) * self.clip_mixup_rate) + segments_to_mixup = segments_before.sample(num_clip, replace=False, + weights=prob) + for old_segment in segments_to_mixup.iterrows(): + old_start_tick = int(old_segment[1]['start_f'] / 30 * self.cfg.DATA.FPS) + old_end_tick = int(np.ceil(old_segment[1]['end_f'] / 30 * self.cfg.DATA.FPS)) + old_action = old_segment[1]['action'] + old_vid = old_segment[1]['video'] + assert old_vid == session + segment_same_class = self.segment_by_cls[old_action] + segment_same_class = segment_same_class[segment_same_class['video'] != old_vid] + valid_segments = segment_same_class[segment_same_class['len_f'] > old_segment[1]['len_f']] + if len(valid_segments) == 0: + continue + else: + sample_segment = valid_segments.sample(1) + new_vid = sample_segment['video'].values[0] + new_start_tick = int(sample_segment['start_f'] / 30 * self.cfg.DATA.FPS) + new_end_tick = int(sample_segment['end_f'] / 30 * self.cfg.DATA.FPS) + + new_visual_inputs = np.load( + osp.join(self.data_root, self.visual_feature, new_vid + '.npy'), mmap_mode='r') + new_motion_inputs = np.load( + osp.join(self.data_root, self.motion_feature, new_vid + '.npy'), mmap_mode='r') + new_object_inputs = np.load( + osp.join(self.data_root, self.object_feature, new_vid + '.npy'), mmap_mode='r') + + sel_indices = np.where((long_indices >= old_start_tick) & (long_indices <= old_end_tick)) + + shift = np.random.randint(new_end_tick - new_start_tick - len(sel_indices) + 1) + long_visual_inputs[sel_indices] = new_visual_inputs[new_start_tick + shift:new_start_tick + shift + len(sel_indices)] + long_motion_inputs[sel_indices] = new_motion_inputs[new_start_tick + shift:new_start_tick + shift + len(sel_indices)] + long_object_inputs[sel_indices] = new_object_inputs[new_start_tick + shift:new_start_tick + shift + len(sel_indices)] + + + + + + + + + + + + + + + + + memory_key_padding_mask = np.zeros(long_indices.shape[0]) + last_zero = bisect_right(long_indices, 0) - 1 + if last_zero > 0: + memory_key_padding_mask[:last_zero] = float('-inf') + else: + long_visual_inputs = None + long_motion_inputs = None + long_object_inputs = None + memory_key_padding_mask = None + + + if (long_visual_inputs is not None and + long_motion_inputs is not None and + long_object_inputs is not None): + fusion_visual_inputs = np.concatenate((long_visual_inputs, work_visual_inputs)) + fusion_motion_inputs = np.concatenate((long_motion_inputs, work_motion_inputs)) + fusion_object_inputs = np.concatenate((long_object_inputs, work_object_inputs)) + else: + fusion_visual_inputs = work_visual_inputs + fusion_motion_inputs = work_motion_inputs + fusion_object_inputs = work_object_inputs + + + fusion_visual_inputs = torch.as_tensor(fusion_visual_inputs.astype(np.float32)) + fusion_motion_inputs = torch.as_tensor(fusion_motion_inputs.astype(np.float32)) + fusion_object_inputs = torch.as_tensor(fusion_object_inputs.astype(np.float32)) + target = torch.as_tensor(target.astype(np.float32)) + if self.cfg.MODEL.LSTR.V_N_CLASSIFIER: + + + + + + + + + + + + + + + + + + + verb_target = torch.as_tensor(verb_target.astype(np.float32)) + noun_target = torch.as_tensor(noun_target.astype(np.float32)) + target = (target, verb_target, noun_target) + + assert fusion_visual_inputs.ndim == fusion_motion_inputs.ndim + if fusion_visual_inputs.ndim == 4: + + fusion_visual_inputs = fusion_visual_inputs.permute(0, 2, 3, 1) + fusion_motion_inputs = fusion_motion_inputs.permute(0, 2, 3, 1) + fusion_visual_inputs = fusion_visual_inputs.mean((1, 2)) + fusion_motion_inputs = fusion_motion_inputs.mean((1, 2)) + + + + + if memory_key_padding_mask is not None: + memory_key_padding_mask = torch.as_tensor(memory_key_padding_mask.astype(np.float32)) + return (fusion_visual_inputs, fusion_motion_inputs, + fusion_object_inputs, memory_key_padding_mask, target) + else: + return (fusion_visual_inputs, fusion_motion_inputs, + fusion_object_inputs, fusion_object_inputs, target) + + def __len__(self): + return len(self.inputs) + + +@registry.register('LSTRBatchInferenceTHUMOS') +@registry.register('LSTRBatchInferenceEGO4D') +@registry.register('LSTRBatchInferenceEasyCom') +@registry.register('LSTRBatchInferenceTVSeries') +@registry.register('LSTRBatchInferenceEK55') +@registry.register('LSTRBatchInferenceEK100') +class LSTRBatchInferenceDataLayer(data.Dataset): + + def __init__(self, cfg, phase='test'): + self.cfg = cfg + self.data_root = cfg.DATA.DATA_ROOT + self.visual_feature = cfg.INPUT.VISUAL_FEATURE + self.motion_feature = cfg.INPUT.MOTION_FEATURE + self.object_feature = cfg.INPUT.OBJECT_FEATURE + self.target_perframe = cfg.INPUT.TARGET_PERFRAME + self.sessions = getattr(cfg.DATA, phase.upper() + '_SESSION_SET') + self.long_memory_length = cfg.MODEL.LSTR.LONG_MEMORY_LENGTH + self.long_memory_sample_rate = cfg.MODEL.LSTR.LONG_MEMORY_SAMPLE_RATE + self.long_memory_num_samples = cfg.MODEL.LSTR.LONG_MEMORY_NUM_SAMPLES + self.work_memory_length = cfg.MODEL.LSTR.WORK_MEMORY_LENGTH + self.work_memory_sample_rate = cfg.MODEL.LSTR.WORK_MEMORY_SAMPLE_RATE + self.work_memory_num_samples = cfg.MODEL.LSTR.WORK_MEMORY_NUM_SAMPLES + self.anticipation_length = cfg.MODEL.LSTR.ANTICIPATION_LENGTH + self.anticipation_sample_rate = cfg.MODEL.LSTR.ANTICIPATION_SAMPLE_RATE + self.anticipation_num_samples = cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES + + if phase != 'test': + warnings.warn('originally, phase must be `test` for batch inference, got {}'.format(phase)) + + + self.inputs = [] + for session in self.sessions: + target = np.load(osp.join(self.data_root, self.target_perframe, session + '.npy')) + if self.cfg.MODEL.LSTR.V_N_CLASSIFIER: + verb_target = np.load(osp.join(self.data_root, self.target_perframe.replace('target', 'verb'), session + '.npy')) + noun_target = np.load(osp.join(self.data_root, self.target_perframe.replace('target', 'noun'), session + '.npy')) + else: + verb_target = target + noun_target = target + for work_start, work_end in zip( + range(0, target.shape[0] + 1), + range(self.work_memory_length, target.shape[0] + 1)): + self.inputs.append([ + session, work_start, work_end, + target[work_start: work_end], + verb_target[work_start: work_end], + noun_target[work_start: work_end], + target.shape[0] + ]) + + def uniform_sampler(self, start, end, num_samples, sample_rate): + if start < 0: + start = (end + 1) % sample_rate + indices = np.arange(start, end + 1)[::sample_rate] + padding = num_samples - indices.shape[0] + if padding > 0: + indices = np.concatenate((np.zeros(padding), indices)) + return np.sort(indices).astype(np.int32) + + def __getitem__(self, index): + session, work_start, work_end, target, verb_target, noun_target, num_frames = self.inputs[index] + visual_inputs = np.load( + osp.join(self.data_root, self.visual_feature, session + '.npy'), mmap_mode='r') + motion_inputs = np.load( + osp.join(self.data_root, self.motion_feature, session + '.npy'), mmap_mode='r') + object_inputs = np.load( + osp.join(self.data_root, self.object_feature, session + '.npy'), mmap_mode='r') + + + target = target[::self.work_memory_sample_rate] + verb_target = verb_target[::self.work_memory_sample_rate] + noun_target = noun_target[::self.work_memory_sample_rate] + + + work_indices = np.arange(work_start, work_end).clip(0) + work_indices = work_indices[::self.work_memory_sample_rate] + work_visual_inputs = visual_inputs[work_indices] + work_motion_inputs = motion_inputs[work_indices] + work_object_inputs = object_inputs[work_indices] + + + if self.long_memory_num_samples > 0: + + long_start, long_end = work_start - self.long_memory_length, work_start - 1 + long_indices = self.uniform_sampler( + long_start, + long_end, + self.long_memory_num_samples, + self.long_memory_sample_rate).clip(0) + long_visual_inputs = visual_inputs[long_indices] + long_motion_inputs = motion_inputs[long_indices] + long_object_inputs = object_inputs[long_indices] + + + memory_key_padding_mask = np.zeros(long_indices.shape[0]) + last_zero = bisect_right(long_indices, 0) - 1 + if last_zero > 0: + memory_key_padding_mask[:last_zero] = float('-inf') + else: + long_visual_inputs = None + long_motion_inputs = None + long_object_inputs = None + memory_key_padding_mask = None + + + if long_visual_inputs is not None and long_motion_inputs is not None: + fusion_visual_inputs = np.concatenate((long_visual_inputs, work_visual_inputs)) + fusion_motion_inputs = np.concatenate((long_motion_inputs, work_motion_inputs)) + fusion_object_inputs = np.concatenate((long_object_inputs, work_object_inputs)) + else: + fusion_visual_inputs = work_visual_inputs + fusion_motion_inputs = work_motion_inputs + fusion_object_inputs = work_object_inputs + + + fusion_visual_inputs = torch.as_tensor(fusion_visual_inputs.astype(np.float32)) + fusion_motion_inputs = torch.as_tensor(fusion_motion_inputs.astype(np.float32)) + fusion_object_inputs = torch.as_tensor(fusion_object_inputs.astype(np.float32)) + target = torch.as_tensor(target.astype(np.float32)) + + if self.cfg.MODEL.LSTR.V_N_CLASSIFIER: + verb_target = torch.as_tensor(verb_target.astype(np.float32)) + noun_target = torch.as_tensor(noun_target.astype(np.float32)) + target = (target, verb_target, noun_target) + + if memory_key_padding_mask is not None: + memory_key_padding_mask = torch.as_tensor(memory_key_padding_mask.astype(np.float32)) + return (fusion_visual_inputs, fusion_motion_inputs, + fusion_object_inputs, memory_key_padding_mask, target, + session, work_indices, num_frames) + else: + return (fusion_visual_inputs, fusion_motion_inputs, + fusion_object_inputs, target, + session, work_indices, num_frames) + + def __len__(self): + return len(self.inputs) \ No newline at end of file diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/__init__.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/__init__.py new file mode 100644 index 0000000..25ea1bb --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/__init__.py @@ -0,0 +1,6 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +from .engines import do_inference, do_train +from .lstr.lstr_inference import * +from .lstr.lstr_trainer import * diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/base_inferences/__init__.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/base_inferences/__init__.py new file mode 100644 index 0000000..04f8b7b --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/base_inferences/__init__.py @@ -0,0 +1,2 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/base_inferences/perframe_det_batch_inference.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/base_inferences/perframe_det_batch_inference.py new file mode 100644 index 0000000..62249ac --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/base_inferences/perframe_det_batch_inference.py @@ -0,0 +1,673 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import os +import os.path as osp +import pickle as pkl +import time + +import numpy as np +import pandas as pd +import torch +from rekognition_online_action_detection.datasets import build_dataset +from rekognition_online_action_detection.evaluation import compute_result +from tqdm import tqdm + + +def batch_inference(cfg, model, device, logger, data_loader): + + # Collect scores and targets + pred_scores = {} + pred_scores_verb = {} + pred_scores_noun = {} + gt_targets = {} + + start_time = time.time() + with torch.no_grad(): + pbar = tqdm(data_loader, desc='BatchInference') + for batch_idx, data in enumerate(pbar, start=1): + target = data[-4] + if cfg.MODEL.LSTR.V_N_CLASSIFIER: + target, verb_target, noun_target = target + + score = model(*[x.to(device) for x in data[:-4]]) + if cfg.MODEL.LSTR.V_N_CLASSIFIER: + score, score_verb, score_noun = score + score = score.cpu().numpy() + score_verb = score_verb.cpu().numpy() + score_noun = score_noun.cpu().numpy() + cfg.DATA.NUM_VERBS = 126 if cfg.DATA.DATA_NAME == 'EK55' else 98 + cfg.DATA.NUM_NOUNS = 353 if cfg.DATA.DATA_NAME == 'EK55' else 301 + else: + score = (score.cpu().numpy() + if cfg.DATA.DATA_NAME.startswith('EK') + else score.softmax(dim=-1).cpu().numpy()) + for bs, (session, query_indices, num_frames) in enumerate(zip(*data[-3:])): + # data[-3] is session, data[-2] is work_indices, data[-1] is num_frames + if session not in pred_scores: + if cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES > 0: + pred_scores[session] = np.zeros((num_frames, cfg.DATA.NUM_CLASSES, cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES)) + if cfg.MODEL.LSTR.V_N_CLASSIFIER: + pred_scores_verb[session] = np.zeros((num_frames, cfg.DATA.NUM_VERBS, cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES)) + pred_scores_noun[session] = np.zeros((num_frames, cfg.DATA.NUM_NOUNS, cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES)) + else: + pred_scores[session] = np.zeros((num_frames, cfg.DATA.NUM_CLASSES)) + if session not in gt_targets: + gt_targets[session] = np.zeros((num_frames, cfg.DATA.NUM_CLASSES)) + + # DEFAULT VALUES in config + # cfg.MODEL.LSTR.WORK_MEMORY_LENGTH: 40 + # cfg.MODEL.LSTR.WORK_MEMORY_SAMPLE_RATE: 1 + # cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES: 10 + # query_indices.shape: torch.Size([cfg.MODEL.LSTR.WORK_MEMORY_LENGTH]) + # query_indices at first: [0, 1, 2, 3, ..., 39] + # At first, go into this if statement + if query_indices[0] in torch.arange(0, cfg.MODEL.LSTR.WORK_MEMORY_SAMPLE_RATE): + if cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES > 0: + for t_a in range(0, cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES): + # at first, full_indices indicate the range: + # work_memory + t_a (anticiptaion timestep, from 0 to 9) + full_indices = torch.cat((query_indices, + torch.arange(query_indices[-1] + cfg.MODEL.LSTR.WORK_MEMORY_SAMPLE_RATE, + query_indices[-1] + (t_a + 1) + cfg.MODEL.LSTR.WORK_MEMORY_SAMPLE_RATE)), + dim=0) + + # then assgin the currnet logits + anticipated logits to the pred_scores + pred_scores[session][full_indices, :, t_a] = score[bs][:full_indices.shape[0]] + + if cfg.MODEL.LSTR.V_N_CLASSIFIER: + pred_scores_verb[session][full_indices, :, t_a] = score_verb[bs][:full_indices.shape[0]] + pred_scores_noun[session][full_indices, :, t_a] = score_noun[bs][:full_indices.shape[0]] + + gt_targets[session][query_indices] = target[bs][:query_indices.shape[0]] + else: + # no anticipation + pred_scores[session][query_indices] = score[bs] + gt_targets[session][query_indices] = target[bs] + else: + if cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES > 0: + for t_a in range(0, cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES): + if query_indices[-1] + t_a + 1 < num_frames: + pred_scores[session][query_indices[-1] + t_a + 1, :, t_a] = score[bs][t_a - cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES] + if cfg.MODEL.LSTR.V_N_CLASSIFIER: + pred_scores_verb[session][query_indices[-1] + t_a + 1, :, t_a] = score_verb[bs][t_a - cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES] + pred_scores_noun[session][query_indices[-1] + t_a + 1, :, t_a] = score_noun[bs][t_a - cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES] + gt_targets[session][query_indices[-1]] = target[bs][-1] + else: + pred_scores[session][query_indices[-1]] = score[bs][-1] + gt_targets[session][query_indices[-1]] = target[bs][-1] + time_taken = time.time() - start_time + # Save scores and targets + # pkl.dump({ + # 'cfg': cfg, + # 'perframe_pred_scores': pred_scores, + # 'perframe_gt_targets': gt_targets, + # }, open(osp.splitext(cfg.MODEL.CHECKPOINT)[0] + '.pkl', 'wb')) + num_frames = np.concatenate(list(gt_targets.values()), axis=0).shape[0] + logger.info(f'\nProcessed {num_frames} frames in {time_taken:.1f} seconds ({num_frames / time_taken :.1f} FPS)') + logger.info('\n=============') + # Compute results + if cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES > 0: + maps_list = [] + per_class_list = {} + + for t_a in range(0, cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES): + result = compute_result['perframe']( + cfg, + np.concatenate(list(gt_targets.values()), axis=0), + np.concatenate(list(pred_scores.values()), axis=0)[:, :, t_a], + ) + + logger.info('Action anticipation ({:.2f}s) perframe m{}: {:.5f}'.format( + (t_a + 1) / cfg.DATA.FPS * cfg.MODEL.LSTR.ANTICIPATION_SAMPLE_RATE, + cfg.DATA.METRICS, result['mean_AP'] + )) + maps_list.append(result['mean_AP']) + # per_class_AP = result['per_class_AP'] + + anticipation_time=(t_a + 1) / cfg.DATA.FPS * cfg.MODEL.LSTR.ANTICIPATION_SAMPLE_RATE + + for class_name, ap in result['per_class_AP'].items(): + if class_name not in per_class_list: + per_class_list[class_name] = [] + per_class_list[class_name].append(ap) + logger.info(f'Action anticipation ({anticipation_time:.2f}s) perframe AP of {class_name}: {ap:.5f}') + + + logger.info('Action anticipation (mean) perframe m{}: {:.5f}'.format( + cfg.DATA.METRICS, np.mean(maps_list) + )) + + for class_name, ap_list in per_class_list.items(): + logger.info(f'Action anticipation (mean) perframe AP of {class_name}: {np.mean(ap_list):.5f}') + + # EDIT: save predictions! + for session in pred_scores.keys(): + anticipation_score = pred_scores[session] + anticipation_dir = os.path.join(cfg.OUTPUT_DIR, 'pred') + if not os.path.exists(anticipation_dir): + os.makedirs(anticipation_dir) + + np.save(osp.join(anticipation_dir, session+'.npy'), anticipation_score) + + else: + result = compute_result['perframe']( + cfg, + np.concatenate(list(gt_targets.values()), axis=0), + np.concatenate(list(pred_scores.values()), axis=0), + ) + logger.info('Action detection perframe m{}: {:.5f}'.format( + cfg.DATA.METRICS, result['mean_AP'] + )) + + # segment-level evaluation on EK55/EK100 + if cfg.DATA.DATA_NAME.startswith('EK'): + try: + from external.rulstm.RULSTM.utils import ( + get_marginal_indexes, marginalize, softmax, + topk_accuracy_multiple_timesteps, + topk_recall_multiple_timesteps, tta) + except: + raise ModuleNotFoundError + if cfg.DATA.DATA_NAME == 'EK55': + path_to_data = 'external/rulstm/RULSTM/data/ek55/' + elif cfg.DATA.DATA_NAME == 'EK100': + path_to_data = 'external/rulstm/RULSTM/data/ek100/' + segment_list = pd.read_csv(osp.join(path_to_data, 'validation.csv'), + names=['id', 'video', 'start_f', 'end_f', 'verb', 'noun', 'action'], + skipinitialspace=True) + predictions, labels, ids = [], [], [] + predictions_verb, predictions_noun = [], [] + discarded_labels, discarded_ids = [], [] + for segment in segment_list.iterrows(): + start_step = int(np.floor(segment[1]['start_f'] / 30 * cfg.DATA.FPS)) + vid = segment[1]['video'].strip() + preds = (pred_scores[vid][start_step][None, 1:, ::-1]).transpose(0, 2, 1) + if cfg.MODEL.LSTR.V_N_CLASSIFIER: + preds_v = (pred_scores_verb[vid][start_step][None, 1:, ::-1]).transpose(0, 2, 1) + preds_n = (pred_scores_noun[vid][start_step][None, 1:, ::-1]).transpose(0, 2, 1) + else: + preds_v = None + preds_n = None + if start_step >= cfg.MODEL.LSTR.ANTICIPATION_LENGTH: + # pt = [] + # for t_a in range(0, cfg.MODEL.LSTR.ANTICIPATION_LENGTH): + # pt.append(pred_scores[vid][start_step - t_a, 1:, t_a][None, None, :]) + # preds = np.concatenate(pt, axis=1) + predictions.append(preds) + if preds_v is not None: + predictions_verb.append(preds_v) + if preds_n is not None: + predictions_noun.append(preds_n) + labels.append(segment[1][['verb', 'noun', 'action']].values.astype(int)) + ids.append(segment[1]['id']) + else: + discarded_labels.append(segment[1][['verb', 'noun', 'action']].values.astype(int)) + discarded_ids.append(segment[1]['id']) + + actions = pd.read_csv(osp.join(path_to_data, 'actions.csv'), index_col='id') + + action_scores = np.concatenate(predictions) + if len(predictions_verb) > 0: + verb_scores = np.concatenate(predictions_verb) + else: + verb_scores = None + if len(predictions_noun) > 0: + noun_scores = np.concatenate(predictions_noun) + else: + noun_scores = None + labels = np.array(labels) + ids = np.array(ids) + + vi = get_marginal_indexes(actions, 'verb') + ni = get_marginal_indexes(actions, 'noun') + + pkl.dump({ + 'cfg': cfg, + 'action_scores': action_scores, + 'verb_scores': verb_scores, + 'noun_scores': noun_scores, + }, open(osp.splitext(cfg.MODEL.CHECKPOINT)[0] + '.ek55_style.pkl', 'wb')) + + action_probs = softmax(action_scores.reshape(-1, action_scores.shape[-1])) + # action_probs = action_probs / action_probs.sum(-1, keepdims=True) + print(action_probs.shape) + if verb_scores is None: + verb_scores = marginalize(action_probs, vi).reshape(action_scores.shape[0], + action_scores.shape[1], + -1) + else: + verb_scores = softmax(verb_scores.reshape(-1, verb_scores.shape[-1])) + verb_scores = verb_scores.reshape(action_scores.shape[0], + action_scores.shape[1], + -1) + if noun_scores is None: + noun_scores = marginalize(action_probs, ni).reshape(action_scores.shape[0], + action_scores.shape[1], + -1) + else: + noun_scores = softmax(noun_scores.reshape(-1, noun_scores.shape[-1])) + noun_scores = noun_scores.reshape(action_scores.shape[0], + action_scores.shape[1], + -1) + + include_discarded = (cfg.DATA.DATA_NAME == 'EK100') + if include_discarded: + dlab = np.array(discarded_labels) + dislab = np.array(discarded_ids) + ids = np.concatenate([ids, dislab]) + num_disc = len(dlab) + labels = np.concatenate([labels, dlab]) + verb_scores = np.concatenate((verb_scores, np.zeros((num_disc, *verb_scores.shape[1:])))) + noun_scores = np.concatenate((noun_scores, np.zeros((num_disc, *noun_scores.shape[1:])))) + action_scores = np.concatenate((action_scores, np.zeros((num_disc, *action_scores.shape[1:])))) + + verb_labels, noun_labels, action_labels = labels[:, 0], labels[:, 1], labels[:, 2] + verb_accuracies = topk_accuracy_multiple_timesteps(verb_scores, verb_labels) + noun_accuracies = topk_accuracy_multiple_timesteps(noun_scores, noun_labels) + action_accuracies = topk_accuracy_multiple_timesteps(action_scores, action_labels) + + if cfg.DATA.DATA_NAME == 'EK55': + many_shot_verbs = pd.read_csv(osp.join(path_to_data, 'EPIC_many_shot_verbs.csv'))['verb_class'].values + many_shot_nouns = pd.read_csv(osp.join(path_to_data, 'EPIC_many_shot_nouns.csv'))['noun_class'].values + actions = pd.read_csv(osp.join(path_to_data, 'actions.csv')) + a_to_vn = {a[1]['id']: tuple(a[1][['verb', 'noun']].values) + for a in actions.iterrows()} + many_shot_actions = [a for a, (v, n) in a_to_vn.items() + if v in many_shot_verbs or n in many_shot_nouns] + + verb_recalls = topk_recall_multiple_timesteps(verb_scores, verb_labels, + k=5, classes=many_shot_verbs) + noun_recalls = topk_recall_multiple_timesteps(noun_scores, noun_labels, + k=5, classes=many_shot_nouns) + action_recalls = topk_recall_multiple_timesteps(action_scores, action_labels, + k=5, classes=many_shot_actions) + + all_accuracies = np.concatenate([verb_accuracies, noun_accuracies, action_accuracies, + verb_recalls, noun_recalls, action_recalls]) + all_accuracies = all_accuracies[[0, 1, 6, 2, 3, 7, 4, 5, 8]] + indices = [ + ('Verb', 'Top-1 Accuracy'), ('Verb', 'Top-5 Accuracy'), + ('Verb', 'Mean Top-5 Recall'), + ('Noun', 'Top-1 Accuracy'), ('Noun', 'Top-5 Accuracy'), + ('Noun', 'Mean Top-5 Recall'), + ('Action', 'Top-1 Accuracy'), ('Action', 'Top-5 Accuracy'), + ('Action', 'Mean Top-5 Recall'), + ] + else: + overall_verb_recalls = topk_recall_multiple_timesteps(verb_scores, verb_labels, k=5) + overall_noun_recalls = topk_recall_multiple_timesteps(noun_scores, noun_labels, k=5) + overall_action_recalls = topk_recall_multiple_timesteps(action_scores, action_labels, k=5) + + unseen = pd.read_csv(osp.join(path_to_data, 'validation_unseen_participants_ids.csv'), names=['id'], squeeze=True) + tail_verbs = pd.read_csv(osp.join(path_to_data, 'validation_tail_verbs_ids.csv'), names=['id'], squeeze=True) + tail_nouns = pd.read_csv(osp.join(path_to_data, 'validation_tail_nouns_ids.csv'), names=['id'], squeeze=True) + tail_actions = pd.read_csv(osp.join(path_to_data, 'validation_tail_actions_ids.csv'), names=['id'], squeeze=True) + + unseen_bool_idx = pd.Series(ids).isin(unseen).values + tail_verbs_bool_idx = pd.Series(ids).isin(tail_verbs).values + tail_nouns_bool_idx = pd.Series(ids).isin(tail_nouns).values + tail_actions_bool_idx = pd.Series(ids).isin(tail_actions).values + + tail_verb_recalls = topk_recall_multiple_timesteps( + verb_scores[tail_verbs_bool_idx], verb_labels[tail_verbs_bool_idx], k=5) + tail_noun_recalls = topk_recall_multiple_timesteps( + noun_scores[tail_nouns_bool_idx], noun_labels[tail_nouns_bool_idx], k=5) + tail_action_recalls = topk_recall_multiple_timesteps( + action_scores[tail_actions_bool_idx], action_labels[tail_actions_bool_idx], k=5) + + unseen_verb_recalls = topk_recall_multiple_timesteps( + verb_scores[unseen_bool_idx], verb_labels[unseen_bool_idx], k=5) + unseen_noun_recalls = topk_recall_multiple_timesteps( + noun_scores[unseen_bool_idx], noun_labels[unseen_bool_idx], k=5) + unseen_action_recalls = topk_recall_multiple_timesteps( + action_scores[unseen_bool_idx], action_labels[unseen_bool_idx], k=5) + + all_accuracies = np.concatenate( + [overall_verb_recalls, overall_noun_recalls, overall_action_recalls, + unseen_verb_recalls, unseen_noun_recalls, unseen_action_recalls, + tail_verb_recalls, tail_noun_recalls, tail_action_recalls] + ) #9 x 8 + + indices = [ + ('Overall Mean Top-5 Recall', 'Verb'), ('Overall Mean Top-5 Recall', 'Noun'), + ('Overall Mean Top-5 Recall', 'Action'), + ('Unseen Mean Top-5 Recall', 'Verb'), ('Unseen Mean Top-5 Recall', 'Noun'), + ('Unseen Mean Top-5 Recall', 'Action'), + ('Tail Mean Top-5 Recall', 'Verb'), ('Tail Mean Top-5 Recall', 'Noun'), + ('Tail Mean Top-5 Recall', 'Action'), + ] + + cc = np.linspace(1 / cfg.DATA.FPS * cfg.MODEL.LSTR.ANTICIPATION_LENGTH, + 1 / cfg.DATA.FPS * cfg.MODEL.LSTR.ANTICIPATION_SAMPLE_RATE, + cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES, dtype=str) + scores = pd.DataFrame(all_accuracies * 100, columns=cc, + index=pd.MultiIndex.from_tuples(indices)) + logger.info(scores) + + tta_verb = tta(verb_scores, verb_labels) + tta_noun = tta(noun_scores, noun_labels) + tta_action = tta(action_scores, action_labels) + logger.info(f'\nMean TtA(5): VERB {tta_verb:0.2f} ' + f'NOUN: {tta_noun:0.2f} ACTION: {tta_action:0.2f}') + + +def do_perframe_det_batch_inference(cfg, model, device, logger): + # Setup model to test mode + model.eval() + cfg.MODEL.CRITERIONS = [['MCE', {}]] + + data_loader = torch.utils.data.DataLoader( + dataset=build_dataset(cfg, phase='test', tag='BatchInference'), + batch_size=cfg.DATA_LOADER.BATCH_SIZE, + num_workers=cfg.DATA_LOADER.NUM_WORKERS, + pin_memory=cfg.DATA_LOADER.PIN_MEMORY, + ) + + logger.info(f"phase: test) Inference on test_dataloader") + start = time.time() + batch_inference(cfg, model, device, logger, data_loader) + end = time.time() - start + + logger.info(f"Elapsed time for function batch_inference: {end:.2f}s") + + # data_loader = torch.utils.data.DataLoader( + # dataset=build_dataset(cfg, phase='train', tag='BatchInference'), + # batch_size=cfg.DATA_LOADER.BATCH_SIZE, + # num_workers=cfg.DATA_LOADER.NUM_WORKERS, + # pin_memory=cfg.DATA_LOADER.PIN_MEMORY, + # ) + + # logger.info(f"***** phase: train") + # batch_inference(cfg, model, device, logger, data_loader) + + # # Collect scores and targets + # pred_scores = {} + # pred_scores_verb = {} + # pred_scores_noun = {} + # gt_targets = {} + + # with torch.no_grad(): + # pbar = tqdm(data_loader, desc='BatchInference') + # for batch_idx, data in enumerate(pbar, start=1): + # target = data[-4] + # if cfg.MODEL.LSTR.V_N_CLASSIFIER: + # target, verb_target, noun_target = target + + # score = model(*[x.to(device) for x in data[:-4]]) + # if cfg.MODEL.LSTR.V_N_CLASSIFIER: + # score, score_verb, score_noun = score + # score = score.cpu().numpy() + # score_verb = score_verb.cpu().numpy() + # score_noun = score_noun.cpu().numpy() + # cfg.DATA.NUM_VERBS = 126 if cfg.DATA.DATA_NAME == 'EK55' else 98 + # cfg.DATA.NUM_NOUNS = 353 if cfg.DATA.DATA_NAME == 'EK55' else 301 + # else: + # score = (score.cpu().numpy() + # if cfg.DATA.DATA_NAME.startswith('EK') + # else score.softmax(dim=-1).cpu().numpy()) + # for bs, (session, query_indices, num_frames) in enumerate(zip(*data[-3:])): + # if session not in pred_scores: + # if cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES > 0: + # pred_scores[session] = np.zeros((num_frames, cfg.DATA.NUM_CLASSES, cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES)) + # if cfg.MODEL.LSTR.V_N_CLASSIFIER: + # pred_scores_verb[session] = np.zeros((num_frames, cfg.DATA.NUM_VERBS, cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES)) + # pred_scores_noun[session] = np.zeros((num_frames, cfg.DATA.NUM_NOUNS, cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES)) + # else: + # pred_scores[session] = np.zeros((num_frames, cfg.DATA.NUM_CLASSES)) + # if session not in gt_targets: + # gt_targets[session] = np.zeros((num_frames, cfg.DATA.NUM_CLASSES)) + + # if query_indices[0] in torch.arange(0, cfg.MODEL.LSTR.WORK_MEMORY_SAMPLE_RATE): + # if cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES > 0: + # for t_a in range(0, cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES): + # full_indices = torch.cat((query_indices, + # torch.arange(query_indices[-1] + cfg.MODEL.LSTR.WORK_MEMORY_SAMPLE_RATE, + # query_indices[-1] + t_a + 1 + cfg.MODEL.LSTR.WORK_MEMORY_SAMPLE_RATE)), + # dim=0) + # pred_scores[session][full_indices, :, t_a] = score[bs][:full_indices.shape[0]] + # if cfg.MODEL.LSTR.V_N_CLASSIFIER: + # pred_scores_verb[session][full_indices, :, t_a] = score_verb[bs][:full_indices.shape[0]] + # pred_scores_noun[session][full_indices, :, t_a] = score_noun[bs][:full_indices.shape[0]] + # gt_targets[session][query_indices] = target[bs][:query_indices.shape[0]] + # else: + # pred_scores[session][query_indices] = score[bs] + # gt_targets[session][query_indices] = target[bs] + # else: + # if cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES > 0: + # for t_a in range(0, cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES): + # if query_indices[-1] + t_a + 1 < num_frames: + # pred_scores[session][query_indices[-1] + t_a + 1, :, t_a] = score[bs][t_a - cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES] + # if cfg.MODEL.LSTR.V_N_CLASSIFIER: + # pred_scores_verb[session][query_indices[-1] + t_a + 1, :, t_a] = score_verb[bs][t_a - cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES] + # pred_scores_noun[session][query_indices[-1] + t_a + 1, :, t_a] = score_noun[bs][t_a - cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES] + # gt_targets[session][query_indices[-1]] = target[bs][-1] + # else: + # pred_scores[session][query_indices[-1]] = score[bs][-1] + # gt_targets[session][query_indices[-1]] = target[bs][-1] + + # # Save scores and targets + # # pkl.dump({ + # # 'cfg': cfg, + # # 'perframe_pred_scores': pred_scores, + # # 'perframe_gt_targets': gt_targets, + # # }, open(osp.splitext(cfg.MODEL.CHECKPOINT)[0] + '.pkl', 'wb')) + + # # Compute results + # if cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES > 0: + # maps_list = [] + # for t_a in range(0, cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES): + # result = compute_result['perframe']( + # cfg, + # np.concatenate(list(gt_targets.values()), axis=0), + # np.concatenate(list(pred_scores.values()), axis=0)[:, :, t_a], + # ) + # logger.info('Action anticipation ({:.2f}s) perframe m{}: {:.5f}'.format( + # (t_a + 1) / cfg.DATA.FPS * cfg.MODEL.LSTR.ANTICIPATION_SAMPLE_RATE, + # cfg.DATA.METRICS, result['mean_AP'] + # )) + # maps_list.append(result['mean_AP']) + # logger.info('Action anticipation (mean) perframe m{}: {:.5f}'.format( + # cfg.DATA.METRICS, np.mean(maps_list) + # )) + # else: + # result = compute_result['perframe']( + # cfg, + # np.concatenate(list(gt_targets.values()), axis=0), + # np.concatenate(list(pred_scores.values()), axis=0), + # ) + # logger.info('Action detection perframe m{}: {:.5f}'.format( + # cfg.DATA.METRICS, result['mean_AP'] + # )) + + # # segment-level evaluation on EK55/EK100 + # if cfg.DATA.DATA_NAME.startswith('EK'): + # try: + # from external.rulstm.RULSTM.utils import (get_marginal_indexes, marginalize, softmax, + # topk_accuracy_multiple_timesteps, + # topk_recall_multiple_timesteps, + # tta) + # except: + # raise ModuleNotFoundError + # if cfg.DATA.DATA_NAME == 'EK55': + # path_to_data = 'external/rulstm/RULSTM/data/ek55/' + # elif cfg.DATA.DATA_NAME == 'EK100': + # path_to_data = 'external/rulstm/RULSTM/data/ek100/' + # segment_list = pd.read_csv(osp.join(path_to_data, 'validation.csv'), + # names=['id', 'video', 'start_f', 'end_f', 'verb', 'noun', 'action'], + # skipinitialspace=True) + # predictions, labels, ids = [], [], [] + # predictions_verb, predictions_noun = [], [] + # discarded_labels, discarded_ids = [], [] + # for segment in segment_list.iterrows(): + # start_step = int(np.floor(segment[1]['start_f'] / 30 * cfg.DATA.FPS)) + # vid = segment[1]['video'].strip() + # preds = (pred_scores[vid][start_step][None, 1:, ::-1]).transpose(0, 2, 1) + # if cfg.MODEL.LSTR.V_N_CLASSIFIER: + # preds_v = (pred_scores_verb[vid][start_step][None, 1:, ::-1]).transpose(0, 2, 1) + # preds_n = (pred_scores_noun[vid][start_step][None, 1:, ::-1]).transpose(0, 2, 1) + # else: + # preds_v = None + # preds_n = None + # if start_step >= cfg.MODEL.LSTR.ANTICIPATION_LENGTH: + # # pt = [] + # # for t_a in range(0, cfg.MODEL.LSTR.ANTICIPATION_LENGTH): + # # pt.append(pred_scores[vid][start_step - t_a, 1:, t_a][None, None, :]) + # # preds = np.concatenate(pt, axis=1) + # predictions.append(preds) + # if preds_v is not None: + # predictions_verb.append(preds_v) + # if preds_n is not None: + # predictions_noun.append(preds_n) + # labels.append(segment[1][['verb', 'noun', 'action']].values.astype(int)) + # ids.append(segment[1]['id']) + # else: + # discarded_labels.append(segment[1][['verb', 'noun', 'action']].values.astype(int)) + # discarded_ids.append(segment[1]['id']) + + # actions = pd.read_csv(osp.join(path_to_data, 'actions.csv'), index_col='id') + + # action_scores = np.concatenate(predictions) + # if len(predictions_verb) > 0: + # verb_scores = np.concatenate(predictions_verb) + # else: + # verb_scores = None + # if len(predictions_noun) > 0: + # noun_scores = np.concatenate(predictions_noun) + # else: + # noun_scores = None + # labels = np.array(labels) + # ids = np.array(ids) + + # vi = get_marginal_indexes(actions, 'verb') + # ni = get_marginal_indexes(actions, 'noun') + + # pkl.dump({ + # 'cfg': cfg, + # 'action_scores': action_scores, + # 'verb_scores': verb_scores, + # 'noun_scores': noun_scores, + # }, open(osp.splitext(cfg.MODEL.CHECKPOINT)[0] + '.ek55_style.pkl', 'wb')) + + # action_probs = softmax(action_scores.reshape(-1, action_scores.shape[-1])) + # # action_probs = action_probs / action_probs.sum(-1, keepdims=True) + # print(action_probs.shape) + # if verb_scores is None: + # verb_scores = marginalize(action_probs, vi).reshape(action_scores.shape[0], + # action_scores.shape[1], + # -1) + # else: + # verb_scores = softmax(verb_scores.reshape(-1, verb_scores.shape[-1])) + # verb_scores = verb_scores.reshape(action_scores.shape[0], + # action_scores.shape[1], + # -1) + # if noun_scores is None: + # noun_scores = marginalize(action_probs, ni).reshape(action_scores.shape[0], + # action_scores.shape[1], + # -1) + # else: + # noun_scores = softmax(noun_scores.reshape(-1, noun_scores.shape[-1])) + # noun_scores = noun_scores.reshape(action_scores.shape[0], + # action_scores.shape[1], + # -1) + + # include_discarded = (cfg.DATA.DATA_NAME == 'EK100') + # if include_discarded: + # dlab = np.array(discarded_labels) + # dislab = np.array(discarded_ids) + # ids = np.concatenate([ids, dislab]) + # num_disc = len(dlab) + # labels = np.concatenate([labels, dlab]) + # verb_scores = np.concatenate((verb_scores, np.zeros((num_disc, *verb_scores.shape[1:])))) + # noun_scores = np.concatenate((noun_scores, np.zeros((num_disc, *noun_scores.shape[1:])))) + # action_scores = np.concatenate((action_scores, np.zeros((num_disc, *action_scores.shape[1:])))) + + # verb_labels, noun_labels, action_labels = labels[:, 0], labels[:, 1], labels[:, 2] + # verb_accuracies = topk_accuracy_multiple_timesteps(verb_scores, verb_labels) + # noun_accuracies = topk_accuracy_multiple_timesteps(noun_scores, noun_labels) + # action_accuracies = topk_accuracy_multiple_timesteps(action_scores, action_labels) + + # if cfg.DATA.DATA_NAME == 'EK55': + # many_shot_verbs = pd.read_csv(osp.join(path_to_data, 'EPIC_many_shot_verbs.csv'))['verb_class'].values + # many_shot_nouns = pd.read_csv(osp.join(path_to_data, 'EPIC_many_shot_nouns.csv'))['noun_class'].values + # actions = pd.read_csv(osp.join(path_to_data, 'actions.csv')) + # a_to_vn = {a[1]['id']: tuple(a[1][['verb', 'noun']].values) + # for a in actions.iterrows()} + # many_shot_actions = [a for a, (v, n) in a_to_vn.items() + # if v in many_shot_verbs or n in many_shot_nouns] + + # verb_recalls = topk_recall_multiple_timesteps(verb_scores, verb_labels, + # k=5, classes=many_shot_verbs) + # noun_recalls = topk_recall_multiple_timesteps(noun_scores, noun_labels, + # k=5, classes=many_shot_nouns) + # action_recalls = topk_recall_multiple_timesteps(action_scores, action_labels, + # k=5, classes=many_shot_actions) + + # all_accuracies = np.concatenate([verb_accuracies, noun_accuracies, action_accuracies, + # verb_recalls, noun_recalls, action_recalls]) + # all_accuracies = all_accuracies[[0, 1, 6, 2, 3, 7, 4, 5, 8]] + # indices = [ + # ('Verb', 'Top-1 Accuracy'), ('Verb', 'Top-5 Accuracy'), + # ('Verb', 'Mean Top-5 Recall'), + # ('Noun', 'Top-1 Accuracy'), ('Noun', 'Top-5 Accuracy'), + # ('Noun', 'Mean Top-5 Recall'), + # ('Action', 'Top-1 Accuracy'), ('Action', 'Top-5 Accuracy'), + # ('Action', 'Mean Top-5 Recall'), + # ] + # else: + # overall_verb_recalls = topk_recall_multiple_timesteps(verb_scores, verb_labels, k=5) + # overall_noun_recalls = topk_recall_multiple_timesteps(noun_scores, noun_labels, k=5) + # overall_action_recalls = topk_recall_multiple_timesteps(action_scores, action_labels, k=5) + + # unseen = pd.read_csv(osp.join(path_to_data, 'validation_unseen_participants_ids.csv'), names=['id'], squeeze=True) + # tail_verbs = pd.read_csv(osp.join(path_to_data, 'validation_tail_verbs_ids.csv'), names=['id'], squeeze=True) + # tail_nouns = pd.read_csv(osp.join(path_to_data, 'validation_tail_nouns_ids.csv'), names=['id'], squeeze=True) + # tail_actions = pd.read_csv(osp.join(path_to_data, 'validation_tail_actions_ids.csv'), names=['id'], squeeze=True) + + # unseen_bool_idx = pd.Series(ids).isin(unseen).values + # tail_verbs_bool_idx = pd.Series(ids).isin(tail_verbs).values + # tail_nouns_bool_idx = pd.Series(ids).isin(tail_nouns).values + # tail_actions_bool_idx = pd.Series(ids).isin(tail_actions).values + + # tail_verb_recalls = topk_recall_multiple_timesteps( + # verb_scores[tail_verbs_bool_idx], verb_labels[tail_verbs_bool_idx], k=5) + # tail_noun_recalls = topk_recall_multiple_timesteps( + # noun_scores[tail_nouns_bool_idx], noun_labels[tail_nouns_bool_idx], k=5) + # tail_action_recalls = topk_recall_multiple_timesteps( + # action_scores[tail_actions_bool_idx], action_labels[tail_actions_bool_idx], k=5) + + # unseen_verb_recalls = topk_recall_multiple_timesteps( + # verb_scores[unseen_bool_idx], verb_labels[unseen_bool_idx], k=5) + # unseen_noun_recalls = topk_recall_multiple_timesteps( + # noun_scores[unseen_bool_idx], noun_labels[unseen_bool_idx], k=5) + # unseen_action_recalls = topk_recall_multiple_timesteps( + # action_scores[unseen_bool_idx], action_labels[unseen_bool_idx], k=5) + + # all_accuracies = np.concatenate( + # [overall_verb_recalls, overall_noun_recalls, overall_action_recalls, + # unseen_verb_recalls, unseen_noun_recalls, unseen_action_recalls, + # tail_verb_recalls, tail_noun_recalls, tail_action_recalls] + # ) #9 x 8 + + # indices = [ + # ('Overall Mean Top-5 Recall', 'Verb'), ('Overall Mean Top-5 Recall', 'Noun'), + # ('Overall Mean Top-5 Recall', 'Action'), + # ('Unseen Mean Top-5 Recall', 'Verb'), ('Unseen Mean Top-5 Recall', 'Noun'), + # ('Unseen Mean Top-5 Recall', 'Action'), + # ('Tail Mean Top-5 Recall', 'Verb'), ('Tail Mean Top-5 Recall', 'Noun'), + # ('Tail Mean Top-5 Recall', 'Action'), + # ] + + # cc = np.linspace(1 / cfg.DATA.FPS * cfg.MODEL.LSTR.ANTICIPATION_LENGTH, + # 1 / cfg.DATA.FPS * cfg.MODEL.LSTR.ANTICIPATION_SAMPLE_RATE, + # cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES, dtype=str) + # scores = pd.DataFrame(all_accuracies * 100, columns=cc, + # index=pd.MultiIndex.from_tuples(indices)) + # logger.info(scores) + + # tta_verb = tta(verb_scores, verb_labels) + # tta_noun = tta(noun_scores, noun_labels) + # tta_action = tta(action_scores, action_labels) + # logger.info(f'\nMean TtA(5): VERB {tta_verb:0.2f} ' + # f'NOUN: {tta_noun:0.2f} ACTION: {tta_action:0.2f}') + # f'NOUN: {tta_noun:0.2f} ACTION: {tta_action:0.2f}') diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/base_trainers/__init__.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/base_trainers/__init__.py new file mode 100644 index 0000000..04f8b7b --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/base_trainers/__init__.py @@ -0,0 +1,2 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/base_trainers/perframe_det_trainer.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/base_trainers/perframe_det_trainer.py new file mode 100644 index 0000000..971cec9 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/base_trainers/perframe_det_trainer.py @@ -0,0 +1,245 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import os +import os.path as osp +import time + +import numpy as np +import pandas as pd +import torch +import torch.nn as nn +import wandb +from rekognition_online_action_detection.evaluation import compute_result +from tqdm import tqdm + + +def log_results(phase, cfg, compute_result, gt_targets, pred_scores, losses, data_loader, epoch=None, best_mAP=None, best_epoch=None): + result = compute_result['perframe'](cfg, gt_targets, pred_scores) + per_class_AP = result['mean_AP'] + + per_class_AP_log = ' '.join(f'{class_name}: {class_AP:.5f}' for class_name, class_AP in per_class_AP.items()) + + log = [] + if epoch is not None: + log.append(f'\nEpoch {epoch:2}') + + log.append(f'{phase} det_loss: {losses["det"] / len(data_loader.dataset):.5f} det_mAP: {result["mean_AP"]:.5f} {per_class_AP_log}') + + if cfg.USE_WANDB: + wandb.log({ + f"{phase.capitalize()} mAP": result['mean_AP'], + f"{phase.capitalize()} per_class_AP": per_class_AP, + f"{phase.capitalize()} Loss": losses['det'] / len(data_loader.dataset) + }) + + if cfg.MODEL.LSTR.V_N_CLASSIFIER: + for loss_type in ['verb', 'noun']: + if loss_type in losses: + log.append(f'{phase} {loss_type}_loss: {losses[loss_type] / len(data_loader.dataset):.5f}') + + return log, result['mean_AP'] + +def do_perframe_det_train(cfg, + data_loaders, + model, + criterion, + optimizer, + scheduler, + device, + checkpointer, + logger): + # Setup model on multiple GPUs + if torch.cuda.device_count() > 1: + model = nn.DataParallel(model) + + best_mAP, best_epoch = 0, 0 + best_per_class_AP_log = None + + for epoch in range(cfg.SOLVER.START_EPOCH, cfg.SOLVER.START_EPOCH + cfg.SOLVER.NUM_EPOCHS): + # Reset + det_losses = {phase: 0.0 for phase in cfg.SOLVER.PHASES} + pred_losses = {phase: 0.0 for phase in cfg.SOLVER.PHASES} + verb_losses = {phase: 0.0 for phase in cfg.SOLVER.PHASES} + noun_losses = {phase: 0.0 for phase in cfg.SOLVER.PHASES} + det_pred_scores = [] + det_gt_targets = [] + aa_pred_scores = [] + aa_gt_targets = [] + + train_det_pred_scores = [] + train_det_gt_targets = [] + train_aa_pred_scores = [] + train_aa_gt_targets = [] + + start = time.time() + + for phase in cfg.SOLVER.PHASES: + training = phase == 'train' + model.train(training) + + with torch.set_grad_enabled(training): + pbar = tqdm(data_loaders[phase], + desc='{}ing epoch {}'.format(phase.capitalize(), epoch)) + for batch_idx, data in enumerate(pbar, start=1): + batch_size = data[0].shape[0] + + # batch_size 16 + # data[0]: visual_input, [16, 680, 2048] + # [BATCH_SIZE, 640 + 40 (LONG_MEMORY_NUM_SAMPLES + WORK_MEMORY_NUM_SAMPLES), FEATURE_SIZE] + # this if statement is for the case of EK100 dataset + if cfg.MODEL.LSTR.V_N_CLASSIFIER: + det_target, verb_target, noun_target = data[-1] + det_target = det_target.to(device) + verb_target = verb_target.to(device) + noun_target = noun_target.to(device) + else: + det_target = data[-1].to(device) + + loss_names = list(zip(*cfg.MODEL.CRITERIONS))[0] + if 'PRED_FUTURE' in list(zip(*cfg.MODEL.CRITERIONS))[0]: + det_score, feat_out, feat_ori = model(*[x.to(device) for x in data[:-1]]) + else: + det_score = model(*[x.to(device) for x in data[:-1]]) + + if cfg.MODEL.LSTR.V_N_CLASSIFIER: + det_score, verb_score, noun_score = det_score + else: + verb_score, noun_score = None, None + + if cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES > 0: + aa_score = det_score[:, -cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES:, :] + # [16, 10, 3] + aa_target = det_target[:, -cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES:, :] + + if cfg.MODEL.LSTR.LOSS_ANTICIPATE_ONLY: + det_score = det_score[:, -cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES:, :] + det_target = det_target[:, -cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES:, :] + + # det_score: [16, 50, 3], aa_score: [16, 10, 3] + # det_score can be different when aa_score + + det_score = det_score.reshape(-1, cfg.DATA.NUM_CLASSES) + det_target = det_target.reshape(-1, cfg.DATA.NUM_CLASSES) + + assert len(det_score.shape) == 2 + det_loss = criterion[loss_names[0]](det_score, det_target) + det_losses[phase] += det_loss.item() * batch_size + + if 'PRED_FUTURE' in list(zip(*cfg.MODEL.CRITERIONS))[0]: + pred_loss = criterion['PRED_FUTURE'](feat_out, feat_ori) + pred_losses[phase] += pred_loss.item() * batch_size + + if cfg.MODEL.LSTR.V_N_CLASSIFIER: + if cfg.MODEL.LSTR.LOSS_ANTICIPATE_ONLY: + verb_score = verb_score[:, -cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES:, :] + verb_target = verb_target[:, -cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES:, :] + noun_score = noun_score[:, -cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES:, :] + noun_target = noun_target[:, -cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES:, :] + verb_score = verb_score.reshape(-1, verb_score.shape[-1]) + noun_score = noun_score.reshape(-1, noun_score.shape[-1]) + verb_target = verb_target.reshape(-1, verb_target.shape[-1]) + noun_target = noun_target.reshape(-1, noun_target.shape[-1]) + verb_loss = criterion['MCE'](verb_score, verb_target) + noun_loss = criterion['MCE'](noun_score, noun_target) + verb_losses[phase] += verb_loss.item() * batch_size + noun_losses[phase] += noun_loss.item() * batch_size + + # Output log for current batch + pbar.set_postfix({ + 'lr': '{:.7f}'.format(scheduler.get_last_lr()[0]), + 'det_loss': '{:.5f}'.format(det_loss.item()), + }) + + if training: + optimizer.zero_grad() + loss = det_loss + if 'PRED_FUTURE' in list(zip(*cfg.MODEL.CRITERIONS))[0]: + loss += pred_loss + if cfg.MODEL.LSTR.V_N_CLASSIFIER: + loss += noun_loss + verb_loss + if loss.item() != 0: + loss.backward() + optimizer.step() + scheduler.step() + + if cfg.USE_WANDB: + wandb.log({"Learning Rate": scheduler.get_last_lr()[0]}) + wandb.log({"Train Loss": det_loss.item()}) + + det_score = det_score.softmax(dim=1).cpu().tolist() + det_target = det_target.cpu().tolist() + train_det_pred_scores.extend(det_score) + train_det_gt_targets.extend(det_target) + + if cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES > 0: + aa_score = aa_score.softmax(dim=1).cpu().tolist() + aa_target = aa_target.cpu().tolist() + train_aa_pred_scores.extend(aa_score) + train_aa_gt_targets.extend(aa_target) + + else: + # Prepare for evaluation + det_score = det_score.softmax(dim=1).cpu().tolist() + det_target = det_target.cpu().tolist() + det_pred_scores.extend(det_score) + det_gt_targets.extend(det_target) + + if cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES > 0: + aa_score = aa_score.softmax(dim=1).cpu().tolist() + aa_target = aa_target.cpu().tolist() + aa_pred_scores.extend(aa_score) + aa_gt_targets.extend(aa_target) + + end = time.time() + + # After the phase loop, compute and log results + for phase in cfg.SOLVER.PHASES: + if phase == 'train': + gt_targets = train_det_gt_targets + pred_scores = train_det_pred_scores + else: # 'test' + gt_targets = det_gt_targets + pred_scores = det_pred_scores + + result = compute_result['perframe'](cfg, gt_targets, pred_scores) + per_class_AP = result['per_class_AP'] + + log = [] + log.append(f'\nEpoch {epoch:2}') + log.append(f'{phase} det_loss: {det_losses[phase] / len(data_loaders[phase].dataset):.5f} det_mAP: {result["mean_AP"]:.5f}') + + per_class_AP_log = ' '.join(f'{class_name}: {class_AP:.5f}' for class_name, class_AP in per_class_AP.items()) + log.append(f'Per-class AP: {per_class_AP_log}') + + if cfg.USE_WANDB: + wandb.log({ + f"{phase.capitalize()} mAP": result['mean_AP'], + f"{phase.capitalize()} per_class_AP": per_class_AP, + f"{phase.capitalize()} Loss": det_losses[phase] / len(data_loaders[phase].dataset) + }) + + if cfg.MODEL.LSTR.V_N_CLASSIFIER: + log.append(f'{phase} verb_loss: {verb_losses[phase] / len(data_loaders[phase].dataset):.5f}') + log.append(f'{phase} noun_loss: {noun_losses[phase] / len(data_loaders[phase].dataset):.5f}') + + if phase == 'test': + if result['mean_AP'] > best_mAP: + best_mAP = result['mean_AP'] + best_epoch = epoch + best_per_class_AP_log = per_class_AP_log + checkpointer.save_best(epoch, model, optimizer) + + log.append(f'\nbest epoch: {best_epoch:2} best det_mAP: {best_mAP:.5f} best per-class AP: {best_per_class_AP_log}\n') + + logger.info(' | '.join(log)) + + log.append(f'running time: {end - start:.2f} sec') + logger.info(' | '.join(log)) + + # Save checkpoint for model and optimizer + if cfg.SOLVER.SAVE_EVERY_EPOCHS > 0 and epoch % cfg.SOLVER.SAVE_EVERY_EPOCHS == 0: + checkpointer.save(epoch, model, optimizer) + + # Shuffle dataset for next epoch + data_loaders['train'].dataset.shuffle() diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/base_trainers/perframe_det_trainer_rnn.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/base_trainers/perframe_det_trainer_rnn.py new file mode 100644 index 0000000..4ee7112 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/base_trainers/perframe_det_trainer_rnn.py @@ -0,0 +1,330 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import os +import os.path as osp +import time + +import numpy as np +import pandas as pd +import torch +import torch.nn as nn +import wandb +from rekognition_online_action_detection.evaluation import compute_result +from tqdm import tqdm + + +def log_results(phase, cfg, compute_result, gt_targets, pred_scores, losses, data_loader, epoch=None, best_mAP=None, best_epoch=None): + result = compute_result['perframe'](cfg, gt_targets, pred_scores) + per_class_AP = result['mean_AP'] + + per_class_AP_log = ' '.join(f'{class_name}: {class_AP:.5f}' for class_name, class_AP in per_class_AP.items()) + + log = [] + if epoch is not None: + log.append(f'\nEpoch {epoch:2}') + + log.append(f'{phase} det_loss: {losses["det"] / len(data_loader.dataset):.5f} det_mAP: {result["mean_AP"]:.5f} {per_class_AP_log}') + + if cfg.USE_WANDB: + wandb.log({ + f"{phase.capitalize()} mAP": result['mean_AP'], + f"{phase.capitalize()} per_class_AP": per_class_AP, + f"{phase.capitalize()} Loss": losses['det'] / len(data_loader.dataset) + }) + + if cfg.MODEL.LSTR.V_N_CLASSIFIER: + for loss_type in ['verb', 'noun']: + if loss_type in losses: + log.append(f'{phase} {loss_type}_loss: {losses[loss_type] / len(data_loader.dataset):.5f}') + + return log, result['mean_AP'] + +def do_perframe_det_train_rnn(cfg, + data_loaders, + model, + criterion, + optimizer, + scheduler, + device, + checkpointer, + logger): + # Setup model on multiple GPUs + if torch.cuda.device_count() > 1: + model = nn.DataParallel(model) + + best_mAP, best_epoch = 0, 0 + best_per_class_AP_log = None + + for epoch in range(cfg.SOLVER.START_EPOCH, cfg.SOLVER.START_EPOCH + cfg.SOLVER.NUM_EPOCHS): + # Reset + det_losses = {phase: 0.0 for phase in cfg.SOLVER.PHASES} + pred_losses = {phase: 0.0 for phase in cfg.SOLVER.PHASES} + + det_pred_scores = [] + det_gt_targets = [] + aa_pred_scores = [] + aa_gt_targets = [] + + train_det_pred_scores = [] + train_det_gt_targets = [] + train_aa_pred_scores = [] + train_aa_gt_targets = [] + + start = time.time() + + for phase in cfg.SOLVER.PHASES: + training = phase == 'train' + model.train(training) + + with torch.set_grad_enabled(training): + pbar = tqdm(data_loaders[phase], + desc='{}ing epoch {}'.format(phase.capitalize(), epoch)) + + # __getitem in LSTR: + # return (fusion_visual_inputs, fusion_motion_inputs, + # fusion_object_inputs, memory_key_padding_mask, target) + # __getitem__ in RNN: + # vid, rgb_input, flow_input, text_input, target, ant_target + for batch_idx, data in enumerate(pbar, start=1): + batch_size = data[1].shape[0] + det_target = data[-2].to(device) + aa_target = data[-1].to(device) + + loss_names = list(zip(*cfg.MODEL.CRITERIONS))[0] + assert 'PRED_FUTURE' not in list(zip(*cfg.MODEL.CRITERIONS))[0] + # if 'PRED_FUTURE' in list(zip(*cfg.MODEL.CRITERIONS))[0]: + # det_score, feat_out, feat_ori = model(*[x.to(device) for x in data[:-1]]) + # else: + # det_score = model(*[x.to(device) for x in data[1:4]]) + + out_dict = model(*[x.to(device) for x in data[1:4]]) + # out_dict['logits']: [B, window_size, num_classes] + # out_dict['anticipation_logits']: [B, window_size, ant_len, num_classes] + + # [workmem_len + ant_len, batch_size, self.d_model]: [50, 16, 1024] + + det_score = out_dict['logits'] + det_score = det_score + if cfg.MODEL.MINIROAD.ANTICIPATION_LENGTH > 0: + aa_score = out_dict['anticipation_logits'] # [B, window_size, anticipation_length, num_class] + aa_score = aa_score[:, -1, :, :].reshape(-1, cfg.DATA.NUM_CLASSES) + + aa_target = aa_target.reshape(-1, cfg.DATA.NUM_CLASSES) + det_loss = criterion[loss_names[0]](aa_score, aa_target) + anticipation_logits = out_dict['anticipation_logits'] + + else: + det_score = det_score.reshape(-1, cfg.DATA.NUM_CLASSES) + det_target = det_target.reshape(-1, cfg.DATA.NUM_CLASSES) + det_loss = criterion[loss_names[0]](det_score, det_target) + + det_loss = criterion[loss_names[0]](det_score, det_target) + + + # if cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES > 0: + # aa_score = det_score[:, -cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES:, :] + # # [16, 10, 3] + # aa_target = det_target[:, -cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES:, :] + + # if cfg.MODEL.LSTR.LOSS_ANTICIPATE_ONLY: + # det_score = det_score[:, -cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES:, :] + # det_target = det_target[:, -cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES:, :] + + # det_score: [16, 50, 3], aa_score: [16, 10, 3] + # det_score can be different when aa_score + + det_losses[phase] += det_loss.item() * batch_size + + # Output log for current batch + pbar.set_postfix({ + 'lr': '{:.7f}'.format(scheduler.get_last_lr()[0]), + 'det_loss': '{:.5f}'.format(det_loss.item()), + }) + + if training: + optimizer.zero_grad() + loss = det_loss + if 'PRED_FUTURE' in list(zip(*cfg.MODEL.CRITERIONS))[0]: + loss += pred_loss + if cfg.MODEL.LSTR.V_N_CLASSIFIER: + loss += noun_loss + verb_loss + if loss.item() != 0: + loss.backward() + optimizer.step() + scheduler.step() + + if cfg.USE_WANDB: + wandb.log({"Learning Rate": scheduler.get_last_lr()[0]}) + wandb.log({"Train Loss": det_loss.item()}) + + det_score = det_score.softmax(dim=1).cpu().tolist() + det_target = det_target.cpu().tolist() + train_det_pred_scores.extend(det_score) + train_det_gt_targets.extend(det_target) + # len(train_det_pred_scores) + 16 + + if cfg.MODEL.MINIROAD.ANTICIPATION_LENGTH > 0: + aa_score = aa_score.softmax(dim=1).cpu().tolist() + aa_target = aa_target.cpu().tolist() + train_aa_pred_scores.extend(aa_score) + train_aa_gt_targets.extend(aa_target) + + else: + # Prepare for evaluation + det_score = det_score.softmax(dim=1).cpu().tolist() + det_target = det_target.cpu().tolist() + det_pred_scores.extend(det_score) + det_gt_targets.extend(det_target) + + if cfg.MODEL.MINIROAD.ANTICIPATION_LENGTH > 0: + aa_score = aa_score.softmax(dim=1).cpu().tolist() + aa_target = aa_target.cpu().tolist() + aa_pred_scores.extend(aa_score) + aa_gt_targets.extend(aa_target) + + # print(f'len(det_score): {len(det_score)}, len(det_pred_scores): {len(det_pred_scores)}') + # print(f'len(aa_score): {len(aa_score)}, len(aa_pred_scores): {len(aa_pred_scores)}') + + end = time.time() + + # After the phase loop, compute and log results + for phase in cfg.SOLVER.PHASES: + if phase == 'train': + gt_targets = train_det_gt_targets + pred_scores = train_det_pred_scores + else: # 'test' + gt_targets = det_gt_targets + pred_scores = det_pred_scores + + result = compute_result['perframe'](cfg, gt_targets, pred_scores) + per_class_AP = result['per_class_AP'] + + log = [] + log.append(f'\nEpoch {epoch:2}') + log.append(f'{phase} det_loss: {det_losses[phase] / len(data_loaders[phase].dataset):.5f} det_mAP: {result["mean_AP"]:.5f}') + + per_class_AP_log = ' '.join(f'{class_name}: {class_AP:.5f}' for class_name, class_AP in per_class_AP.items()) + log.append(f'Per-class AP: {per_class_AP_log}') + + if cfg.USE_WANDB: + wandb.log({ + f"{phase.capitalize()} mAP": result['mean_AP'], + f"{phase.capitalize()} per_class_AP": per_class_AP, + f"{phase.capitalize()} Loss": det_losses[phase] / len(data_loaders[phase].dataset) + }) + + if cfg.MODEL.LSTR.V_N_CLASSIFIER: + log.append(f'{phase} verb_loss: {verb_losses[phase] / len(data_loaders[phase].dataset):.5f}') + log.append(f'{phase} noun_loss: {noun_losses[phase] / len(data_loaders[phase].dataset):.5f}') + + if phase == 'test': + if result['mean_AP'] > best_mAP: + best_mAP = result['mean_AP'] + best_epoch = epoch + best_per_class_AP_log = per_class_AP_log + checkpointer.save_best(epoch, model, optimizer) + + log.append(f'\nbest epoch: {best_epoch:2} best det_mAP: {best_mAP:.5f} best per-class AP: {best_per_class_AP_log}\n') + + logger.info(' | '.join(log)) + + log.append(f'running time: {end - start:.2f} sec') + logger.info(' | '.join(log)) + + # Save checkpoint for model and optimizer + if cfg.SOLVER.SAVE_EVERY_EPOCHS > 0 and epoch % cfg.SOLVER.SAVE_EVERY_EPOCHS == 0: + checkpointer.save(epoch, model, optimizer) + + # Shuffle dataset for next epoch + data_loaders['train'].dataset.shuffle() + +# @TRAINER.register("OAD") +def train_one_epoch(trainloader, model, criterion, optimizer, scaler, epoch, use_wandb = False, writer=None, scheduler=None): + epoch_loss = 0 + for it, (vid_name, rgb_input, flow_input, target) in enumerate(tqdm(trainloader, desc=f'Epoch:{epoch} Training', postfix=f'lr: {optimizer.param_groups[0]["lr"]:.7f}')): + rgb_input, flow_input, target = rgb_input.cuda(), flow_input.cuda(), target.cuda() + model.train() + if scaler != None: + with torch.cuda.amp.autocast(): + out_dict = model(rgb_input, flow_input) + loss = criterion(out_dict, target) + optimizer.zero_grad(set_to_none=True) + scaler.scale(loss).backward() + scaler.step(optimizer) + scaler.update() + else: + out_dict = model(rgb_input, flow_input) + loss = criterion(out_dict, target) + optimizer.zero_grad(set_to_none=True) + loss.backward() + optimizer.step() + + epoch_loss += loss.item() + if use_wandb == True: + wandb.log({"Train Loss": loss.item()}) + + if writer != None: + writer.add_scalar("Train Loss", loss.item(), it+epoch*len(trainloader)) + return epoch_loss + +# @TRAINER.register("OADTR") +def train_one_epoch_oadtr(trainloader, model, criterion, optimizer, scaler, epoch, use_wandb = False, writer=None, scheduler=None): + epoch_loss = 0 + for it, (vid_name, rgb_input, flow_input, enc_target, distance_target, class_h_target, dec_target) in enumerate(tqdm(trainloader, desc=f'Epoch:{epoch} Training', postfix=f'lr: {optimizer.param_groups[0]["lr"]:.7f}')): + rgb_input, flow_input, enc_target, distance_target, class_h_target = rgb_input.cuda(), flow_input.cuda(), enc_target.cuda(), distance_target.cuda(), class_h_target.cuda() + model.train() + if scaler != None: + with torch.cuda.amp.autocast(): + out_dict = model(rgb_input, flow_input) + loss = criterion(out_dict, enc_target, dec_target) + optimizer.zero_grad(set_to_none=True) + scaler.scale(loss).backward() + scaler.step(optimizer) + scaler.update() + else: + out_dict = model(rgb_input, flow_input) + loss = criterion(out_dict, enc_target, dec_target) + optimizer.zero_grad(set_to_none=True) + loss.backward() + optimizer.step() + + epoch_loss += loss.item() + if use_wandb == True: + wandb.log({"Train Loss": loss.item()}) + + if writer != None: + writer.add_scalar("Train Loss", loss.item(), it+epoch*len(trainloader)) + + return epoch_loss + + +# @TRAINER.register("ANTICIPATION") +def ant_train_one_epoch(trainloader, model, criterion, optimizer, scaler, epoch, use_wandb = False, writer=None, scheduler=None): + epoch_loss = 0 + for it, (vid_name, rgb_input, flow_input, text_input, target, ant_target) in enumerate(tqdm(trainloader, desc=f'Epoch:{epoch} Training', postfix=f'lr: {optimizer.param_groups[0]["lr"]:.7f}')): + rgb_input, flow_input, text_input, target, ant_target = rgb_input.cuda(), flow_input.cuda(), text_input.cuda(), target.cuda(), ant_target.cuda() + model.train() + if scaler != None: + with torch.cuda.amp.autocast(): + out_dict = model(rgb_input, flow_input, text_input) + loss = criterion(out_dict, target, ant_target) + optimizer.zero_grad(set_to_none=True) + scaler.scale(loss).backward() + scaler.step(optimizer) + scaler.update() + else: + out_dict = model(rgb_input, flow_input, text_input) + loss = criterion(out_dict, target, ant_target) + optimizer.zero_grad(set_to_none=True) + loss.backward() + optimizer.step() + epoch_loss += loss.item() + + if use_wandb == True: + wandb.log({"Train Loss": loss.item()}) + + if writer != None: + writer.add_scalar("Train Loss", loss.item(), it+epoch*len(trainloader)) + + return epoch_loss \ No newline at end of file diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/engines.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/engines.py new file mode 100644 index 0000000..2ac2b19 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/engines.py @@ -0,0 +1,39 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +from rekognition_online_action_detection.utils.registry import Registry + +TRAINERS = Registry() +INFERENCES = Registry() + +def do_train(cfg, + data_loaders, + model, + criterion, + optimizer, + scheduler, + device, + checkpointer, + logger): + + return TRAINERS[cfg.MODEL.MODEL_NAME]( + cfg, + data_loaders, + model, + criterion, + optimizer, + scheduler, + device, + checkpointer, + logger) + + +def do_inference(cfg, + model, + device, + logger): + return INFERENCES[cfg.MODEL.MODEL_NAME]( + cfg, + model, + device, + logger) diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/lstr/__init__.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/lstr/__init__.py new file mode 100644 index 0000000..04f8b7b --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/lstr/__init__.py @@ -0,0 +1,2 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/lstr/lstr_inference.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/lstr/lstr_inference.py new file mode 100644 index 0000000..8746c61 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/lstr/lstr_inference.py @@ -0,0 +1,417 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import os.path as osp +from bisect import bisect_right +import time +import pickle as pkl + +import torch +import numpy as np + +from rekognition_online_action_detection.evaluation import compute_result + +from ..base_inferences.perframe_det_batch_inference import do_perframe_det_batch_inference + + +from ..engines import INFERENCES as registry + + +@registry.register('LSTR') +def do_lstr_batch_inference(cfg, + model, + device, + logger): + if cfg.MODEL.LSTR.INFERENCE_MODE == 'stream': + do_lstr_stream_inference(cfg, + model, + device, + logger) + else: + do_perframe_det_batch_inference(cfg, + model, + device, + logger) + + +def do_lstr_stream_inference(cfg, model, device, logger): + # Setup model to test mode + model.eval() + + # Collect scores and targets + pred_scores = [] + gt_targets = [] + + def to_device(x, dtype=np.float32): + return torch.as_tensor(x.astype(dtype)).unsqueeze(0).to(device) + + long_memory_length = cfg.MODEL.LSTR.LONG_MEMORY_LENGTH + long_memory_sample_rate = cfg.MODEL.LSTR.LONG_MEMORY_SAMPLE_RATE + long_memory_num_samples = cfg.MODEL.LSTR.LONG_MEMORY_NUM_SAMPLES + work_memory_length = cfg.MODEL.LSTR.WORK_MEMORY_LENGTH + work_memory_sample_rate = cfg.MODEL.LSTR.WORK_MEMORY_SAMPLE_RATE + work_memory_num_samples = cfg.MODEL.LSTR.WORK_MEMORY_NUM_SAMPLES + + # if len(cfg.DATA.TEST_SESSION_SET) != 1: + # raise RuntimeError('Only support testing one video each time for stream inference, will fix later') + + pred_scores_all = {} + gt_targets_all = {} + + with torch.no_grad(): + for session_idx, session in enumerate(cfg.DATA.TEST_SESSION_SET): + model.clear_cache() + + if cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES > 0: + pred_scores = [[] for _ in range(cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES)] + else: + pred_scores = [] + gt_targets = [] + + visual_inputs = np.load(osp.join(cfg.DATA.DATA_ROOT, cfg.INPUT.VISUAL_FEATURE, session + '.npy'), mmap_mode='r') + motion_inputs = np.load(osp.join(cfg.DATA.DATA_ROOT, cfg.INPUT.MOTION_FEATURE, session + '.npy'), mmap_mode='r') + object_inputs = np.load(osp.join(cfg.DATA.DATA_ROOT, cfg.INPUT.OBJECT_FEATURE, session + '.npy'), mmap_mode='r') + target = np.load(osp.join(cfg.DATA.DATA_ROOT, cfg.INPUT.TARGET_PERFRAME, session + '.npy')) + + start_time = time.time() + + for work_start, work_end in zip(range(0, target.shape[0] + 1), + range(work_memory_length, target.shape[0] + 1)): + # Get target + # target = target[::work_memory_sample_rate] + + # Get work memory + work_indices = np.arange(work_start, work_end).clip(0) + work_indices = work_indices[::work_memory_sample_rate] + work_visual_inputs = to_device(visual_inputs[work_indices]) + work_motion_inputs = to_device(motion_inputs[work_indices]) + work_object_inputs = to_device(object_inputs[work_indices]) + + # Get long memory + PRECISE = True + if PRECISE: + long_end = work_start - long_memory_sample_rate + if long_end < 0: + long_indices = [0 for _ in range(long_memory_num_samples)] + long_indices_set = [long_indices for _ in range(long_memory_sample_rate)] + long_visual_inputs = to_device(visual_inputs[long_indices]) + long_motion_inputs = to_device(motion_inputs[long_indices]) + long_object_inputs = to_device(object_inputs[long_indices]) + else: + long_indices = long_indices_set[long_end % long_memory_sample_rate][1:] + [long_end] + long_indices_set[long_end % long_memory_sample_rate] = long_indices + long_visual_inputs = to_device(visual_inputs[[long_end]]) + long_motion_inputs = to_device(motion_inputs[[long_end]]) + long_object_inputs = to_device(object_inputs[[long_end]]) + else: + long_end = work_start - 1 + if long_end == -1: + long_indices = [0 for _ in range(long_memory_num_samples)] + long_visual_inputs = to_device(visual_inputs[long_indices]) + long_motion_inputs = to_device(motion_inputs[long_indices]) + long_object_inputs = to_device(object_inputs[long_indices]) + elif long_end % long_memory_sample_rate == 0: + long_indices = long_indices[1:] + [long_end] + long_visual_inputs = to_device(visual_inputs[[long_end]]) + long_motion_inputs = to_device(motion_inputs[[long_end]]) + long_object_inputs = to_device(object_inputs[[long_end]]) + else: + long_visual_inputs = None + long_motion_inputs = None + long_object_inputs = None + + # Get memory key padding mask + memory_key_padding_mask = np.zeros(len(long_indices)) + last_zero = bisect_right(long_indices, 0) - 1 + if last_zero > 0: + memory_key_padding_mask[:last_zero] = float('-inf') + memory_key_padding_mask = torch.as_tensor(memory_key_padding_mask.astype(np.float32)).unsqueeze(0).to(device) + + score = model.stream_inference( + long_visual_inputs, + long_motion_inputs, + long_object_inputs, + work_visual_inputs, + work_motion_inputs, + work_object_inputs, + memory_key_padding_mask, + cache_num=long_memory_sample_rate if PRECISE else 1, + cache_id=long_end % long_memory_sample_rate if PRECISE else 0)[0] + + if cfg.DATA.DATA_NAME.startswith('EK'): + score = score.cpu().numpy() + else: + score = score.softmax(dim=-1).cpu().numpy() + + if cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES > 0: + if work_start == 0: + upsample_score = score[:-cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES].repeat(work_memory_sample_rate, axis=0) + upsample_score = upsample_score[work_start:work_end-1] + anticipate_score = score[-cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES:] + # coarse but sufficient + for t_a in range(0, cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES): + combined_score = np.concatenate((upsample_score, anticipate_score[:t_a + 1]), + axis=0) + pred_scores[t_a].extend(list(combined_score)) + gt_targets.extend(list(target[:work_end-1])) + else: + for t_a in range(0, cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES): + pred_scores[t_a].append(list(score[t_a - cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES])) + gt_targets.append(list(target[work_end - 1])) + else: + if work_start == 0: + gt_targets.extend(list(target[:work_end])) + pred_scores.extend(list(score)) + else: + gt_targets.append(list(target[work_end - 1])) + pred_scores.append(list(score[-1])) + + end_time = time.time() + logger.info('Running time: {:.3f} seconds'.format(end_time - start_time)) + + if cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES > 0: + for t_a in range(0, cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES): + result = compute_result['perframe']( + cfg, + gt_targets, + pred_scores[t_a][:-1 - t_a], + ) + sec = (t_a + 1) / cfg.DATA.FPS * cfg.MODEL.LSTR.ANTICIPATION_SAMPLE_RATE + logger.info('mAP of video ({:.2f}s) {}: {:.5f}'.format(sec, session, result['mean_AP'])) + else: + result = compute_result['perframe']( + cfg, + gt_targets, + pred_scores, + ) + logger.info('mAP of video {}: {:.5f}'.format(session, result['mean_AP'])) + + + gt_targets_all[session] = np.array(gt_targets) + if cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES > 0: + for t_a in range(0, cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES): + pred_scores[t_a] = np.array(pred_scores[t_a][:- 1 - t_a]) + pred_scores_all[session] = np.stack(pred_scores, axis=0).transpose(1, 2, 0) + pred_scores_all[session] = pred_scores_all[session] + else: + pred_scores_all[session] = np.array(pred_scores) + + + # pkl.dump({ + # 'cfg': cfg, + # 'perframe_pred_scores': pred_scores_all, + # 'perframe_gt_targets': gt_targets_all, + # }, open(osp.splitext(cfg.MODEL.CHECKPOINT)[0] + '.stream.pkl', 'wb')) + + if cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES > 0: + maps_list = [] + for t_a in range(0, cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES): + result = compute_result['perframe']( + cfg, + np.concatenate(list(gt_targets_all.values()), axis=0), + np.concatenate(list(pred_scores_all.values()), axis=0)[:, :, t_a], + ) + logger.info('Action anticipation ({:.2f}s) perframe m{}: {:.5f}'.format( + (t_a + 1) / cfg.DATA.FPS * cfg.MODEL.LSTR.ANTICIPATION_SAMPLE_RATE, + cfg.DATA.METRICS, result['mean_AP'] + )) + maps_list.append(result['mean_AP']) + logger.info('Action anticipation (mean) perframe m{}: {:.5f}'.format( + cfg.DATA.METRICS, np.mean(maps_list) + )) + else: + result = compute_result['perframe']( + cfg, + np.concatenate(list(gt_targets_all.values()), axis=0), + np.concatenate(list(pred_scores_all.values()), axis=0), + ) + logger.info('Action detection perframe m{}: {:.5f}'.format( + cfg.DATA.METRICS, result['mean_AP'] + )) + + gt_targets = gt_targets_all + pred_scores = pred_scores_all + pred_scores_verb = {} + pred_scores_noun = {} + + # segment-level evaluation on EK55/EK100 + if cfg.DATA.DATA_NAME.startswith('EK'): + try: + import pandas as pd + from external.rulstm.RULSTM.utils import (get_marginal_indexes, marginalize, softmax, + topk_accuracy_multiple_timesteps, + topk_recall_multiple_timesteps, + tta) + except: + raise ModuleNotFoundError + if cfg.DATA.DATA_NAME == 'EK55': + path_to_data = 'external/rulstm/RULSTM/data/ek55/' + elif cfg.DATA.DATA_NAME == 'EK100': + path_to_data = 'external/rulstm/RULSTM/data/ek100/' + segment_list = pd.read_csv(osp.join(path_to_data, 'validation.csv'), + names=['id', 'video', 'start_f', 'end_f', 'verb', 'noun', 'action'], + skipinitialspace=True) + predictions, labels, ids = [], [], [] + predictions_verb, predictions_noun = [], [] + discarded_labels, discarded_ids = [], [] + for segment in segment_list.iterrows(): + start_step = int(np.floor(segment[1]['start_f'] / 30 * cfg.DATA.FPS)) + vid = segment[1]['video'].strip() + preds = (pred_scores[vid][start_step][None, 1:, ::-1]).transpose(0, 2, 1) + if cfg.MODEL.LSTR.V_N_CLASSIFIER: + raise NotImplementedError( + 'For TeSTra with decomposed classifier (V+N), please use batch mode, ie ``MODEL.LSTR.INFERENCE_MODE batch``' + ) + preds_v = (pred_scores_verb[vid][start_step][None, 1:, ::-1]).transpose(0, 2, 1) + preds_n = (pred_scores_noun[vid][start_step][None, 1:, ::-1]).transpose(0, 2, 1) + else: + preds_v = None + preds_n = None + if start_step >= cfg.MODEL.LSTR.ANTICIPATION_LENGTH: + predictions.append(preds) + if preds_v is not None: + predictions_verb.append(preds_v) + if preds_n is not None: + predictions_noun.append(preds_n) + labels.append(segment[1][['verb', 'noun', 'action']].values.astype(int)) + ids.append(segment[1]['id']) + else: + discarded_labels.append(segment[1][['verb', 'noun', 'action']].values.astype(int)) + discarded_ids.append(segment[1]['id']) + + actions = pd.read_csv(osp.join(path_to_data, 'actions.csv'), index_col='id') + + action_scores = np.concatenate(predictions) + if len(predictions_verb) > 0: + verb_scores = np.concatenate(predictions_verb) + else: + verb_scores = None + if len(predictions_noun) > 0: + noun_scores = np.concatenate(predictions_noun) + else: + noun_scores = None + labels = np.array(labels) + ids = np.array(ids) + + vi = get_marginal_indexes(actions, 'verb') + ni = get_marginal_indexes(actions, 'noun') + + action_probs = softmax(action_scores.reshape(-1, action_scores.shape[-1])) + if verb_scores is None: + verb_scores = marginalize(action_probs, vi).reshape(action_scores.shape[0], + action_scores.shape[1], + -1) + else: + verb_scores = softmax(verb_scores.reshape(-1, verb_scores.shape[-1])) + verb_scores = verb_scores.reshape(action_scores.shape[0], + action_scores.shape[1], + -1) + if noun_scores is None: + noun_scores = marginalize(action_probs, ni).reshape(action_scores.shape[0], + action_scores.shape[1], + -1) + else: + noun_scores = softmax(noun_scores.reshape(-1, noun_scores.shape[-1])) + noun_scores = noun_scores.reshape(action_scores.shape[0], + action_scores.shape[1], + -1) + + include_discarded = (cfg.DATA.DATA_NAME == 'EK100') + if include_discarded: + dlab = np.array(discarded_labels) + dislab = np.array(discarded_ids) + ids = np.concatenate([ids, dislab]) + num_disc = len(dlab) + labels = np.concatenate([labels, dlab]) + verb_scores = np.concatenate((verb_scores, np.zeros((num_disc, *verb_scores.shape[1:])))) + noun_scores = np.concatenate((noun_scores, np.zeros((num_disc, *noun_scores.shape[1:])))) + action_scores = np.concatenate((action_scores, np.zeros((num_disc, *action_scores.shape[1:])))) + + verb_labels, noun_labels, action_labels = labels[:, 0], labels[:, 1], labels[:, 2] + verb_accuracies = topk_accuracy_multiple_timesteps(verb_scores, verb_labels) + noun_accuracies = topk_accuracy_multiple_timesteps(noun_scores, noun_labels) + action_accuracies = topk_accuracy_multiple_timesteps(action_scores, action_labels) + + if cfg.DATA.DATA_NAME == 'EK55': + many_shot_verbs = pd.read_csv(osp.join(path_to_data, 'EPIC_many_shot_verbs.csv'))['verb_class'].values + many_shot_nouns = pd.read_csv(osp.join(path_to_data, 'EPIC_many_shot_nouns.csv'))['noun_class'].values + actions = pd.read_csv(osp.join(path_to_data, 'actions.csv')) + a_to_vn = {a[1]['id']: tuple(a[1][['verb', 'noun']].values) + for a in actions.iterrows()} + many_shot_actions = [a for a, (v, n) in a_to_vn.items() + if v in many_shot_verbs or n in many_shot_nouns] + + verb_recalls = topk_recall_multiple_timesteps(verb_scores, verb_labels, + k=5, classes=many_shot_verbs) + noun_recalls = topk_recall_multiple_timesteps(noun_scores, noun_labels, + k=5, classes=many_shot_nouns) + action_recalls = topk_recall_multiple_timesteps(action_scores, action_labels, + k=5, classes=many_shot_actions) + + all_accuracies = np.concatenate([verb_accuracies, noun_accuracies, action_accuracies, + verb_recalls, noun_recalls, action_recalls]) + all_accuracies = all_accuracies[[0, 1, 6, 2, 3, 7, 4, 5, 8]] + indices = [ + ('Verb', 'Top-1 Accuracy'), ('Verb', 'Top-5 Accuracy'), + ('Verb', 'Mean Top-5 Recall'), + ('Noun', 'Top-1 Accuracy'), ('Noun', 'Top-5 Accuracy'), + ('Noun', 'Mean Top-5 Recall'), + ('Action', 'Top-1 Accuracy'), ('Action', 'Top-5 Accuracy'), + ('Action', 'Mean Top-5 Recall'), + ] + else: + overall_verb_recalls = topk_recall_multiple_timesteps(verb_scores, verb_labels, k=5) + overall_noun_recalls = topk_recall_multiple_timesteps(noun_scores, noun_labels, k=5) + overall_action_recalls = topk_recall_multiple_timesteps(action_scores, action_labels, k=5) + + unseen = pd.read_csv(osp.join(path_to_data, 'validation_unseen_participants_ids.csv'), names=['id'], squeeze=True) + tail_verbs = pd.read_csv(osp.join(path_to_data, 'validation_tail_verbs_ids.csv'), names=['id'], squeeze=True) + tail_nouns = pd.read_csv(osp.join(path_to_data, 'validation_tail_nouns_ids.csv'), names=['id'], squeeze=True) + tail_actions = pd.read_csv(osp.join(path_to_data, 'validation_tail_actions_ids.csv'), names=['id'], squeeze=True) + + unseen_bool_idx = pd.Series(ids).isin(unseen).values + tail_verbs_bool_idx = pd.Series(ids).isin(tail_verbs).values + tail_nouns_bool_idx = pd.Series(ids).isin(tail_nouns).values + tail_actions_bool_idx = pd.Series(ids).isin(tail_actions).values + + tail_verb_recalls = topk_recall_multiple_timesteps( + verb_scores[tail_verbs_bool_idx], verb_labels[tail_verbs_bool_idx], k=5) + tail_noun_recalls = topk_recall_multiple_timesteps( + noun_scores[tail_nouns_bool_idx], noun_labels[tail_nouns_bool_idx], k=5) + tail_action_recalls = topk_recall_multiple_timesteps( + action_scores[tail_actions_bool_idx], action_labels[tail_actions_bool_idx], k=5) + + unseen_verb_recalls = topk_recall_multiple_timesteps( + verb_scores[unseen_bool_idx], verb_labels[unseen_bool_idx], k=5) + unseen_noun_recalls = topk_recall_multiple_timesteps( + noun_scores[unseen_bool_idx], noun_labels[unseen_bool_idx], k=5) + unseen_action_recalls = topk_recall_multiple_timesteps( + action_scores[unseen_bool_idx], action_labels[unseen_bool_idx], k=5) + + all_accuracies = np.concatenate( + [overall_verb_recalls, overall_noun_recalls, overall_action_recalls, + unseen_verb_recalls, unseen_noun_recalls, unseen_action_recalls, + tail_verb_recalls, tail_noun_recalls, tail_action_recalls] + ) #9 x 8 + + indices = [ + ('Overall Mean Top-5 Recall', 'Verb'), ('Overall Mean Top-5 Recall', 'Noun'), + ('Overall Mean Top-5 Recall', 'Action'), + ('Unseen Mean Top-5 Recall', 'Verb'), ('Unseen Mean Top-5 Recall', 'Noun'), + ('Unseen Mean Top-5 Recall', 'Action'), + ('Tail Mean Top-5 Recall', 'Verb'), ('Tail Mean Top-5 Recall', 'Noun'), + ('Tail Mean Top-5 Recall', 'Action'), + ] + + cc = np.linspace(1 / cfg.DATA.FPS * cfg.MODEL.LSTR.ANTICIPATION_LENGTH, + 1 / cfg.DATA.FPS * cfg.MODEL.LSTR.ANTICIPATION_SAMPLE_RATE, + cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES, dtype=str) + scores = pd.DataFrame(all_accuracies * 100, columns=cc, + index=pd.MultiIndex.from_tuples(indices)) + logger.info(scores) + + tta_verb = tta(verb_scores, verb_labels) + tta_noun = tta(noun_scores, noun_labels) + tta_action = tta(action_scores, action_labels) + logger.info(f'\nMean TtA(5): VERB {tta_verb:0.2f} ' + f'NOUN: {tta_noun:0.2f} ACTION: {tta_action:0.2f}') diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/lstr/lstr_trainer.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/lstr/lstr_trainer.py new file mode 100644 index 0000000..f3b77fd --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/engines/lstr/lstr_trainer.py @@ -0,0 +1,26 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +from ..base_trainers.perframe_det_trainer import do_perframe_det_train +from ..engines import TRAINERS as registry + + +@registry.register('LSTR') +def do_lstr_train(cfg, + data_loaders, + model, + criterion, + optimizer, + scheduler, + device, + checkpointer, + logger): + do_perframe_det_train(cfg, + data_loaders, + model, + criterion, + optimizer, + scheduler, + device, + checkpointer, + logger) diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/evaluation/__init__.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/evaluation/__init__.py new file mode 100644 index 0000000..bedd5f3 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/evaluation/__init__.py @@ -0,0 +1,4 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +from .evalution import compute_result diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/evaluation/evalution.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/evaluation/evalution.py new file mode 100644 index 0000000..f511c69 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/evaluation/evalution.py @@ -0,0 +1,42 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +from rekognition_online_action_detection.utils.registry import Registry + +from .metrics import perframe_average_precision, perstage_average_precision +from .postprocessing import postprocessing as default_pp + +compute_result = Registry() + + +@compute_result.register('perframe') +def eval_perframe(cfg, ground_truth, prediction, **kwargs): + class_names = kwargs.get('class_names', cfg.DATA.CLASS_NAMES) + ignore_index = kwargs.get('ignore_index', cfg.DATA.IGNORE_INDEX) + metrics = kwargs.get('metrics', cfg.DATA.METRICS) + postprocessing = kwargs.get('postprocessing', default_pp(cfg.DATA.DATA_NAME)) + + return perframe_average_precision( + ground_truth=ground_truth, + prediction=prediction, + class_names=class_names, + ignore_index=ignore_index, + metrics=metrics, + postprocessing=postprocessing, + ) + + +@compute_result.register('perstage') +def eval_perstage(cfg, ground_truth, prediction, **kwargs): + class_names = kwargs.get('class_names', cfg.DATA.CLASS_NAMES) + ignore_index = kwargs.get('ignore_index', cfg.DATA.IGNORE_INDEX) + metrics = kwargs.get('metrics', cfg.DATA.METRICS) + postprocessing = kwargs.get('postprocessing', default_pp(cfg.DATA.DATA_NAME)) + return perstage_average_precision( + ground_truth=ground_truth, + prediction=prediction, + class_names=class_names, + ignore_index=ignore_index, + metrics=metrics, + postprocessing=postprocessing, + ) diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/evaluation/metrics.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/evaluation/metrics.py new file mode 100644 index 0000000..7c5109f --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/evaluation/metrics.py @@ -0,0 +1,138 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +from collections import OrderedDict +from multiprocessing import Pool + +import numpy as np +from sklearn.metrics import average_precision_score + + +def calibrated_average_precision_score(y_true, y_score): + """Compute calibrated average precision (cAP), which is particularly + proposed for the TVSeries dataset. + """ + y_true_sorted = y_true[np.argsort(-y_score)] + tp = y_true_sorted.astype(float) + fp = np.abs(y_true_sorted.astype(float) - 1) + tps = np.cumsum(tp) + fps = np.cumsum(fp) + ratio = np.sum(tp == 0) / np.sum(tp) + cprec = tps / (tps + fps / (ratio + np.finfo(float).eps) + np.finfo(float).eps) + cap = np.sum(cprec[tp == 1]) / np.sum(tp) + return cap + +def perframe_average_precision(ground_truth, + prediction, + class_names, + ignore_index, + metrics, + postprocessing): + """Compute (frame-level) average precision between ground truth and + predictions data frames. + """ + result = OrderedDict() + ground_truth = np.array(ground_truth) + prediction = np.array(prediction) + + # Postprocessing + if postprocessing is not None: + ground_truth, prediction = postprocessing(ground_truth, prediction) + + # Build metrics + if metrics == 'AP': + compute_score = average_precision_score + elif metrics == 'cAP': + compute_score = calibrated_average_precision_score + else: + raise RuntimeError('Unknown metrics: {}'.format(metrics)) + + # set ignore classes including background + ignore_set = set([ignore_index]) + # enable background according to ignore_index + # ignore_set = set([ignore_index]) + + # Compute average precision + result['per_class_AP'] = OrderedDict() + + # when training, class_names becomes 2! + for idx, class_name in enumerate(class_names): + if idx not in ignore_set: + if np.any(ground_truth[:, idx]): + result['per_class_AP'][class_name] = compute_score( + ground_truth[:, idx], prediction[:, idx]) + result['mean_AP'] = np.mean(list(result['per_class_AP'].values())) + + return result + + +def get_stage_pred_scores(gt_targets, pred_scores, perc_s, perc_e): + starts = [] + ends = [] + stage_gt_targets = [] + stage_pred_scores = [] + for i in range(len(gt_targets)): + if gt_targets[i] == 0: + stage_gt_targets.append(gt_targets[i]) + stage_pred_scores.append(pred_scores[i]) + else: + if i == 0 or gt_targets[i - 1] == 0: + starts.append(i) + if i == len(gt_targets) - 1 or gt_targets[i + 1] == 0: + ends.append(i) + if len(starts) != len(ends): + raise ValueError('starts and ends cannot pair!') + + action_lens = [ends[i] - starts[i] for i in range(len(starts))] + stage_starts = [starts[i] + int(action_lens[i] * perc_s) for i in range(len(starts))] + stage_ends = [max(stage_starts[i] + 1, starts[i] + int(action_lens[i] * perc_e)) for i in range(len(starts))] + for i in range(len(starts)): + stage_gt_targets.extend(gt_targets[stage_starts[i]: stage_ends[i]]) + stage_pred_scores.extend(pred_scores[stage_starts[i]: stage_ends[i]]) + return np.array(stage_gt_targets), np.array(stage_pred_scores) + + +def perstage_average_precision(ground_truth, + prediction, + class_names, + ignore_index, + metrics, + postprocessing): + result = OrderedDict() + ground_truth = np.array(ground_truth) + prediction = np.array(prediction) + + # Postprocessing + if postprocessing is not None: + ground_truth, prediction = postprocessing(ground_truth, prediction) + + # Build metrics + if metrics == 'AP': + compute_score = average_precision_score + elif metrics == 'cAP': + compute_score = calibrated_average_precision_score + else: + raise RuntimeError('Unknown metrics: {}'.format(metrics)) + + # Ignore backgroud class + ignore_set = set([0, ignore_index]) + + # Compute average precision + for perc_s in range(10): + perc_e = perc_s + 1 + stage_name = '{:2}%_{:3}%'.format(perc_s * 10, perc_e * 10) + result[stage_name] = OrderedDict({'per_class_AP': OrderedDict()}) + for idx, class_name in enumerate(class_names): + if idx not in ignore_set: + stage_gt_targets, stage_pred_scores = get_stage_pred_scores( + (ground_truth[:, idx] == 1).astype(int), + prediction[:, idx], + perc_s / 10, + perc_e / 10, + ) + result[stage_name]['per_class_AP'][class_name] = \ + compute_score(stage_gt_targets, stage_pred_scores) + result[stage_name]['mean_AP'] = \ + np.mean(list(result[stage_name]['per_class_AP'].values())) + + return result diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/evaluation/postprocessing.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/evaluation/postprocessing.py new file mode 100644 index 0000000..65d15e7 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/evaluation/postprocessing.py @@ -0,0 +1,36 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import numpy as np + + +def postprocessing(data_name): + + def thumos_postprocessing(ground_truth, prediction, smooth=False, switch=False): + """ + We follow (Shou et al., 2017) and adopt their perframe postprocessing method on THUMOS'14 datset. + Source: https://bitbucket.org/columbiadvmm/cdc/src/master/THUMOS14/eval/PreFrameLabeling/compute_framelevel_mAP.m + """ + + # Simple temporal smoothing via NMS of 5-frames window + if smooth: + prob = np.copy(prediction) + prob1 = prob.reshape(1, prob.shape[0], prob.shape[1]) + prob2 = np.append(prob[0, :].reshape(1, -1), prob[0: -1, :], axis=0).reshape(1, prob.shape[0], prob.shape[1]) + prob3 = np.append(prob[1:, :], prob[-1, :].reshape(1, -1), axis=0).reshape(1, prob.shape[0], prob.shape[1]) + prob4 = np.append(prob[0: 2, :], prob[0: -2, :], axis=0).reshape(1, prob.shape[0], prob.shape[1]) + prob5 = np.append(prob[2:, :], prob[-2:, :], axis=0).reshape(1, prob.shape[0], prob.shape[1]) + probsmooth = np.squeeze(np.max(np.concatenate((prob1, prob2, prob3, prob4, prob5), axis=0), axis=0)) + prediction = np.copy(probsmooth) + + # Assign cliff diving (5) as diving (8) + if switch: + switch_index = np.where(prediction[:, 5] > prediction[:, 8])[0] + prediction[switch_index, 8] = prediction[switch_index, 5] + + # Remove ambiguous (21) + valid_index = np.where(ground_truth[:, 21] != 1)[0] + + return ground_truth[valid_index], prediction[valid_index] + + return {'THUMOS': thumos_postprocessing}.get(data_name, None) diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/__init__.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/__init__.py new file mode 100644 index 0000000..8c01b59 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/__init__.py @@ -0,0 +1,5 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +from .lstr import * +from .models import build_model diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/feature_head.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/feature_head.py new file mode 100644 index 0000000..47f0680 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/feature_head.py @@ -0,0 +1,95 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +__all__ = ['build_feature_head'] + +import torch +import torch.nn as nn +from rekognition_online_action_detection.config.feature_sizes import \ + FEATURE_SIZES +from rekognition_online_action_detection.utils.registry import Registry + +FEATURE_HEADS = Registry() + +@FEATURE_HEADS.register('EasyCom') +@FEATURE_HEADS.register('EGO4D') +@FEATURE_HEADS.register('EGO4DVAL') +@FEATURE_HEADS.register('THUMOS') +@FEATURE_HEADS.register('TVSeries') +@FEATURE_HEADS.register('EK55') +@FEATURE_HEADS.register('EK100') +@FEATURE_HEADS.register('YTConv') +class BaseFeatureHead(nn.Module): + def __init__(self, cfg): + super(BaseFeatureHead, self).__init__() + + if cfg.INPUT.MODALITY in ['visual', 'motion', 'object', + 'visual+motion', 'motion+object', 'visual+object', + 'visual+motion+object']: + self.with_visual = 'visual' in cfg.INPUT.MODALITY + self.with_motion = 'motion' in cfg.INPUT.MODALITY + self.with_object = 'object' in cfg.INPUT.MODALITY + else: + raise RuntimeError('Unknown modality of {}'.format(cfg.INPUT.MODALITY)) + + INPUT_VISUAL_FEATURE = cfg.INPUT.VISUAL_FEATURE + INPUT_MOTION_FEATURE = cfg.INPUT.MOTION_FEATURE + INPUT_OBJECT_FEATURE = cfg.INPUT.OBJECT_FEATURE + + if self.with_visual and self.with_motion and self.with_object: + visual_size = FEATURE_SIZES[INPUT_VISUAL_FEATURE] + motion_size = FEATURE_SIZES[INPUT_MOTION_FEATURE] + object_size = FEATURE_SIZES[INPUT_OBJECT_FEATURE] + fusion_size = visual_size + motion_size + object_size + elif self.with_visual and self.with_motion: + visual_size = FEATURE_SIZES[INPUT_VISUAL_FEATURE] + motion_size = FEATURE_SIZES[INPUT_MOTION_FEATURE] + fusion_size = visual_size + motion_size + elif self.with_visual: + fusion_size = FEATURE_SIZES[INPUT_VISUAL_FEATURE] + elif self.with_motion: + fusion_size = FEATURE_SIZES[INPUT_MOTION_FEATURE] + elif self.with_object: + fusion_size = FEATURE_SIZES[INPUT_OBJECT_FEATURE] + + self.fusion_size = fusion_size + self.d_model = fusion_size + + if cfg.MODEL.FEATURE_HEAD.LINEAR_ENABLED: + if cfg.MODEL.FEATURE_HEAD.LINEAR_OUT_FEATURES != -1: + self.d_model = cfg.MODEL.FEATURE_HEAD.LINEAR_OUT_FEATURES + self.input_linear = nn.Sequential( + nn.Linear(fusion_size, self.d_model), + nn.ReLU(inplace=True), + ) + else: + self.input_linear = nn.Identity() + if cfg.MODEL.FEATURE_HEAD.DROPOUT > 0: + self.dropout = nn.Dropout(cfg.MODEL.FEATURE_HEAD.DROPOUT) + else: + self.dropout = None + + def forward(self, visual_input, motion_input, object_input): + if self.with_visual and self.with_motion and self.with_object: + fusion_input = torch.cat((visual_input, motion_input, + object_input), dim=-1) + elif self.with_visual and self.with_motion: + fusion_input = torch.cat((visual_input, motion_input), dim=-1) + elif self.with_visual: + fusion_input = visual_input + elif self.with_motion: + fusion_input = motion_input + elif self.with_object: + fusion_input = object_input + if self.dropout is not None: + fusion_input = self.dropout(fusion_input) + fusion_input = self.input_linear(fusion_input) + return fusion_input + + +def build_feature_head(cfg): + data_names = ['YTConv', 'EGO4D'] + data_name = next((name for name in data_names if name in cfg.DATA.DATA_NAME), cfg.DATA.DATA_NAME) + + feature_head = FEATURE_HEADS[data_name] + return feature_head(cfg) diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/lstr.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/lstr.py new file mode 100644 index 0000000..8ea0ab1 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/lstr.py @@ -0,0 +1,424 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ..utils.ek_utils import action_to_noun_map, action_to_verb_map +from . import transformer as tr +from .feature_head import build_feature_head +from .models import META_ARCHITECTURES as registry +from .normalized_linear import NormalizedLinear + + +class LSTR(nn.Module): + + def __init__(self, cfg): + super(LSTR, self).__init__() + + # Build long feature heads + self.long_memory_num_samples = cfg.MODEL.LSTR.LONG_MEMORY_NUM_SAMPLES + self.long_enabled = self.long_memory_num_samples > 0 + if self.long_enabled: + self.feature_head_long = build_feature_head(cfg) + + # Build work feature head + self.work_memory_num_samples = cfg.MODEL.LSTR.WORK_MEMORY_NUM_SAMPLES + self.work_enabled = self.work_memory_num_samples > 0 + if self.work_enabled: + self.feature_head_work = build_feature_head(cfg) + + self.d_model = self.feature_head_work.d_model + self.num_heads = cfg.MODEL.LSTR.NUM_HEADS + self.dim_feedforward = cfg.MODEL.LSTR.DIM_FEEDFORWARD + self.dropout = cfg.MODEL.LSTR.DROPOUT + self.activation = cfg.MODEL.LSTR.ACTIVATION + self.num_classes = cfg.DATA.NUM_CLASSES + + self.use_v_n = cfg.MODEL.LSTR.V_N_CLASSIFIER + + self.encoder_attention_type = cfg.MODEL.LSTR.ENC_ATTENTION_TYPE + self.decay_alpha = cfg.MODEL.LSTR.ENC_ATTENTION_DECAY + + self.anticipation_num_samples = cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES + + # Build position encoding + self.pos_encoding = tr.PositionalEncoding(self.d_model, self.dropout) + + self.long_memory_use_pe = cfg.MODEL.LSTR.LONG_MEMORY_USE_PE + self.work_memory_use_pe = cfg.MODEL.LSTR.WORK_MEMORY_USE_PE + self.include_work = cfg.MODEL.LSTR.LONG_MEMORY_INCLUDE_WORK + self.include_work2 = cfg.MODEL.LSTR.LONG_MEMORY_INCLUDE_WORK2 + if self.include_work and self.include_work2: + raise ValueError('cfg.MODEL.LSTR.LONG_MEMORY_INCLUDE_WORK and' + 'cfg.MODEL.LSTR.LONG_MEMORY_INCLUDE_WORK2 should' + 'not be both True') + + # Build LSTR encoder + if self.long_enabled: + self.enc_queries = nn.ModuleList() + self.enc_modules = nn.ModuleList() + for i, param in enumerate(cfg.MODEL.LSTR.ENC_MODULE): + if param[0] != -1: + self.enc_queries.append(nn.Embedding(param[0], self.d_model)) + # in TeSTra, not like LSTR, gets atteyntion_type, decay_alpha as parameter + # however, the parameter does not change anything from LSTR + enc_layer = tr.TransformerDecoderLayer( + self.d_model, self.num_heads, self.dim_feedforward, + self.dropout, self.activation, + attention_type=self.encoder_attention_type if i == 0 else 'dotproduct', + decay_alpha=self.decay_alpha if i == 0 else 1.0) + self.enc_modules.append(tr.TransformerDecoder( + enc_layer, param[1], tr.layer_norm(self.d_model, param[2]))) + else: + self.enc_queries.append(None) + enc_layer = tr.TransformerEncoderLayer( + self.d_model, self.num_heads, self.dim_feedforward, + self.dropout, self.activation, + attention_type=self.encoder_attention_type) + self.enc_modules.append(tr.TransformerEncoder( + enc_layer, param[1], tr.layer_norm(self.d_model, param[2]))) + else: + self.register_parameter('enc_queries', None) + self.register_parameter('enc_modules', None) + + # Build LSTR decoder + if self.long_enabled: + param = cfg.MODEL.LSTR.DEC_MODULE + dec_layer = tr.TransformerDecoderLayer( + self.d_model, self.num_heads, self.dim_feedforward, + self.dropout, self.activation) + self.dec_modules = tr.TransformerDecoder( + dec_layer, param[1], tr.layer_norm(self.d_model, param[2])) + else: + param = cfg.MODEL.LSTR.DEC_MODULE + dec_layer = tr.TransformerEncoderLayer( + self.d_model, self.num_heads, self.dim_feedforward, + self.dropout, self.activation) + self.dec_modules = tr.TransformerEncoder( + dec_layer, param[1], tr.layer_norm(self.d_model, param[2])) + + # Build decode token (for anticipation) + if self.anticipation_num_samples > 0: + self.dec_query = nn.Embedding(self.anticipation_num_samples, self.d_model) + else: + self.register_parameter('dec_query', None) + + # Build classifier + if cfg.MODEL.LSTR.DROPOUT_CLS > 0: + self.dropout_cls = nn.Dropout(cfg.MODEL.LSTR.DROPOUT_CLS) + if self.use_v_n: + self.dropout_cls_v = nn.Dropout(cfg.MODEL.LSTR.DROPOUT_CLS) + self.dropout_cls_n = nn.Dropout(cfg.MODEL.LSTR.DROPOUT_CLS) + else: + self.dropout_cls = None + if cfg.MODEL.LSTR.FC_NORM: + self.classifier = NormalizedLinear(self.d_model, self.num_classes) + else: + self.classifier = nn.Linear(self.d_model, self.num_classes) + + if self.use_v_n: + a_to_v = action_to_verb_map(cfg.DATA.EK_EXT_PATH, + action_offset=True, + verb_offset=True) + a_to_n = action_to_noun_map(cfg.DATA.EK_EXT_PATH, + action_offset=True, + noun_offset=True) + num_verbs = max(set(a_to_v.values())) + 1 + num_nouns = max(set(a_to_n.values())) + 1 + print('Number of verbs (inc. 0):', num_verbs) + print('Number of nouns (inc. 0):', num_nouns) + self.classifier_verb = nn.Linear(self.d_model, num_verbs) + self.classifier_noun = nn.Linear(self.d_model, num_nouns) + + self.pred_future = 'PRED_FUTURE' in list(zip(*cfg.MODEL.CRITERIONS))[0] + + def forward(self, visual_inputs, motion_inputs, object_inputs, memory_key_padding_mask=None): + if self.long_enabled: + # Compute long memories + if self.long_memory_use_pe: + long_memories = self.pos_encoding(self.feature_head_long( + visual_inputs[:, :self.long_memory_num_samples if not self.include_work2 else None], + motion_inputs[:, :self.long_memory_num_samples if not self.include_work2 else None], + object_inputs[:, :self.long_memory_num_samples if not self.include_work2 else None], + ).transpose(0, 1)) + long_memories = long_memories.transpose(0, 1) ## transpose back + else: + long_memories = self.feature_head_long( + visual_inputs[:, :self.long_memory_num_samples if not self.include_work2 else None], + motion_inputs[:, :self.long_memory_num_samples if not self.include_work2 else None], + object_inputs[:, :self.long_memory_num_samples if not self.include_work2 else None], + ) + batch_size = long_memories.shape[0] + # if long_memories.ndim == 5: + # long_memories = self.pos_encoding_3d_spatial(long_memories) + long_memories = long_memories.view(batch_size, -1, self.d_model).transpose(0, 1) + # print(long_memories.shape) # T(*H*W), B, C + + if len(self.enc_modules) > 0: + enc_queries = [ + enc_query.weight.unsqueeze(1).repeat(1, long_memories.shape[1], 1) + if enc_query is not None else None + for enc_query in self.enc_queries + ] + + if self.include_work2: + memory_key_padding_mask = F.pad(memory_key_padding_mask, + (0, self.work_memory_num_samples), + 'constant', 0) + # Encode long memories + if enc_queries[0] is not None: + long_memories = self.enc_modules[0](enc_queries[0], long_memories, + memory_key_padding_mask=memory_key_padding_mask) + else: + long_memories = self.enc_modules[0](long_memories) + for enc_query, enc_module in zip(enc_queries[1:], self.enc_modules[1:]): + if enc_query is not None: + long_memories = enc_module(enc_query, long_memories) + else: + long_memories = enc_module(long_memories) + + # Concatenate memories + if self.long_enabled: + memory = long_memories + + if self.work_enabled: + # Compute work memories + if visual_inputs.ndim == 5: + visual_inputs_avg = visual_inputs[:, self.long_memory_num_samples:].mean((2, 3)) + motion_inputs_avg = motion_inputs[:, self.long_memory_num_samples:].mean((2, 3)) + object_inputs_avg = object_inputs[:, self.long_memory_num_samples:].mean((2, 3)) + else: + visual_inputs_avg = visual_inputs[:, self.long_memory_num_samples:] + motion_inputs_avg = motion_inputs[:, self.long_memory_num_samples:] + object_inputs_avg = object_inputs[:, self.long_memory_num_samples:] + work_memories_no_pe = self.feature_head_work( + visual_inputs_avg, + motion_inputs_avg, + object_inputs_avg + # object_inputs[:, self.long_memory_num_samples:], + ).transpose(0, 1) + if self.work_memory_use_pe: + work_memories = self.pos_encoding(work_memories_no_pe, + padding=self.long_memory_num_samples if self.long_memory_use_pe else 0) + else: + work_memories = work_memories_no_pe + if self.dec_query is not None: + anticipate_memories = self.pos_encoding( + self.dec_query.weight.unsqueeze(1).repeat(1, work_memories.shape[1], 1), + padding=(self.long_memory_num_samples + self.work_memory_num_samples + if self.long_memory_use_pe else self.work_memory_num_samples)) + work_memories = torch.cat((work_memories, anticipate_memories), dim=0) + # work_memories = self.feature_head_work( + # visual_inputs[:, self.long_memory_num_samples:], + # motion_inputs[:, self.long_memory_num_samples:], + # ).transpose(0, 1) + + + # Build mask + mask = tr.generate_square_subsequent_mask( + work_memories.shape[0]) + mask = mask.to(work_memories.device) + # Compute output + if self.long_enabled: + if self.include_work: + extended_memory = torch.cat((memory, work_memories), dim=0) + extended_mask = F.pad(mask, + (memory.shape[0], 0), + 'constant', 0) + extended_mask = extended_mask.to(extended_memory.device) + output = self.dec_modules( + work_memories, + memory=extended_memory, + tgt_mask=mask, + memory_mask=extended_mask, + ) + else: + output = self.dec_modules( + work_memories, + memory=memory, + tgt_mask=mask, + ) + else: + output = self.dec_modules( + work_memories, + src_mask=mask, + ) + + full_orig_feats = self.feature_head_work( + visual_inputs_avg, + motion_inputs_avg, + object_inputs_avg + # object_inputs[:, self.long_memory_num_samples:] + ).transpose(0, 1) + all_outputs_decoded = output + + # Compute classification score + if self.dropout_cls is not None: + output_a = self.dropout_cls(output) + else: + output_a = output + + # output_a + # [workmem_len + ant_len, batch_size, self.d_model]: [50, 16, 1024] + # score (before transpose) [len, batch_size, num_classes] + # self.classifier: nn.Linear(self.d_model, self.num_classes) + score = self.classifier(output_a) + score = score.transpose(0, 1) + + if self.use_v_n: + output_v = self.dropout_cls_v(output) + output_n = self.dropout_cls_n(output) + score_verb = self.classifier_verb(output_v).transpose(0, 1) + score_noun = self.classifier_noun(output_n).transpose(0, 1) + score = (score, score_verb, score_noun) + if self.pred_future: + return (score, + all_outputs_decoded[1: self.work_memory_num_samples + 1, ...], + full_orig_feats) + else: + return score + + +@registry.register('LSTR') +class LSTRStream(LSTR): + + def __init__(self, cfg): + super(LSTRStream, self).__init__(cfg) + + ############################ + # Cache for stream inference + ############################ + self.long_memories_cache = None + self.compressed_long_memories_cache = None + + def stream_inference(self, + long_visual_inputs, + long_motion_inputs, + long_object_inputs, + work_visual_inputs, + work_motion_inputs, + work_object_inputs, + memory_key_padding_mask=None, + cache_num=1, + cache_id=0): + assert self.long_enabled, 'Long-term memory cannot be empty for stream inference' + assert len(self.enc_modules) > 0, 'LSTR encoder cannot be disabled for stream inference' + + if ((long_visual_inputs is not None) and + (long_motion_inputs is not None) and + (long_object_inputs is not None)): + # Compute long memories + long_memories = self.feature_head_long( + long_visual_inputs if not self.include_work2 else torch.cat((long_visual_inputs, work_visual_inputs), dim=1), + long_motion_inputs if not self.include_work2 else torch.cat((long_motion_inputs, work_motion_inputs), dim=1), + long_object_inputs if not self.include_work2 else torch.cat((long_object_inputs, work_object_inputs), dim=1), + ).transpose(0, 1) + + if self.long_memories_cache is None: + self.long_memories_cache = [long_memories for _ in range(cache_num)] + else: + self.long_memories_cache[cache_id] = torch.cat(( + self.long_memories_cache[cache_id][1:], long_memories + )) + + long_memories = self.long_memories_cache[cache_id] + pos = self.pos_encoding.pe[:self.long_memory_num_samples, :] + + enc_queries = [ + enc_query.weight.unsqueeze(1).repeat(1, long_memories.shape[1], 1) + if enc_query is not None else None + for enc_query in self.enc_queries + ] + + if self.include_work2: + memory_key_padding_mask = F.pad(memory_key_padding_mask, + (0, self.work_memory_num_samples), + 'constant', 0) + # Encode long memories + long_memories = self.enc_modules[0].stream_inference(enc_queries[0], long_memories, pos, + memory_key_padding_mask=memory_key_padding_mask, + cache_num=cache_num, cache_id=cache_id) + self.compressed_long_memories_cache = long_memories + for enc_query, enc_module in zip(enc_queries[1:], self.enc_modules[1:]): + if enc_query is not None: + long_memories = enc_module(enc_query, long_memories) + else: + long_memories = enc_module(long_memories) + else: + long_memories = self.compressed_long_memories_cache + + enc_queries = [ + enc_query.weight.unsqueeze(1).repeat(1, long_memories.shape[1], 1) + if enc_query is not None else None + for enc_query in self.enc_queries + ] + + # Encode long memories + for enc_query, enc_module in zip(enc_queries[1:], self.enc_modules[1:]): + if enc_query is not None: + long_memories = enc_module(enc_query, long_memories) + else: + long_memories = enc_module(long_memories) + + # Concatenate memories + if self.long_enabled: + memory = long_memories + + if self.work_enabled: + # Compute work memories + work_memories = self.pos_encoding(self.feature_head_work( + work_visual_inputs, + work_motion_inputs, + work_object_inputs, + ).transpose(0, 1), + padding=self.long_memory_num_samples if self.long_memory_use_pe else 0) + if self.dec_query is not None: + anticipate_memories = self.pos_encoding( + self.dec_query.weight.unsqueeze(1).repeat(1, work_memories.shape[1], 1), + padding=(self.long_memory_num_samples + self.work_memory_num_samples + if self.long_memory_use_pe else self.work_memory_num_samples)) + work_memories = torch.cat((work_memories, anticipate_memories)) + + # Build mask + mask = tr.generate_square_subsequent_mask( + work_memories.shape[0]) + mask = mask.to(work_memories.device) + + # Compute output + if self.long_enabled: + if self.include_work: + extended_memory = torch.cat((memory, work_memories), dim=0) + extended_mask = F.pad(mask, + (memory.shape[0], 0), + 'constant', 0) + extended_mask = extended_mask.to(extended_memory.device) + output = self.dec_modules( + work_memories, + memory=extended_memory, + tgt_mask=mask, + memory_mask=extended_mask, + ) + else: + output = self.dec_modules( + work_memories, + memory=memory, + tgt_mask=mask, + ) + else: + output = self.dec_modules( + work_memories, + src_mask=mask, + ) + + # Compute classification score + score = self.classifier(output) + import pdb; pdb.set_trace() + + return score.transpose(0, 1) + + def clear_cache(self): + self.long_memories_cache = None + self.compressed_long_memories_cache = None + self.enc_modules[0].clear_cache() diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/models.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/models.py new file mode 100644 index 0000000..d78d3e4 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/models.py @@ -0,0 +1,15 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +__all__ = ['build_model'] + +from rekognition_online_action_detection.utils.registry import Registry + +META_ARCHITECTURES = Registry() + +def build_model(cfg, device=None): + model = META_ARCHITECTURES[cfg.MODEL.MODEL_NAME](cfg) + from .weights_init import weights_init + + model.apply(weights_init) + return model.to(device) diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/normalized_linear.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/normalized_linear.py new file mode 100644 index 0000000..b3fb2ec --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/normalized_linear.py @@ -0,0 +1,57 @@ +import math + +import torch +import torch.nn as nn +from torch import Tensor +from torch.nn.parameter import Parameter +import torch.nn.functional as F + + +class NormalizedLinear(nn.Module): + r"""Normalized linear layer which was proposed by Jiaqi Wang et. al. in + their technical report for LVIS workshop at ECCV 2020. + https://arxiv.org/pdf/2008.10032.pdf + Args: + in_features: size of each input sample + out_features: size of each output sample + bias: If set to ``False``, the layer will not learn an additive bias. + Default: ``True`` + tau: a temperature factor, set as 20 by default by following the paper. + """ + __constants__ = ['in_features', 'out_features'] + in_features: int + out_features: int + weight: Tensor + + def __init__(self, in_features: int, out_features: int, bias: bool = True, tau: float = 20) -> None: + super().__init__() + self.in_features = in_features + self.out_features = out_features + self.weight = Parameter(torch.Tensor(out_features, in_features)) + if bias: + self.bias = Parameter(torch.Tensor(out_features)) + else: + self.register_parameter('bias', None) + self.reset_parameters() + self.tau = tau + + def reset_parameters(self) -> None: + nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + if self.bias is not None: + fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight) + bound = 1 / math.sqrt(fan_in) + nn.init.uniform_(self.bias, -bound, bound) + + def forward(self, input: Tensor) -> Tensor: + + return F.linear( + input / torch.norm(input, dim=-1, keepdim=True), + self.tau * self.weight / + torch.norm(self.weight, dim=-1, keepdim=True), + self.bias + ) + + def extra_repr(self) -> str: + return 'in_features={}, out_features={}, bias={}'.format( + self.in_features, self.out_features, self.bias is not None + ) \ No newline at end of file diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/transformer/__init__.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/transformer/__init__.py new file mode 100644 index 0000000..55c1a18 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/transformer/__init__.py @@ -0,0 +1,10 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +from .position_encoding import PositionalEncoding +from .position_encoding import PositionalEncoding2D +from .transformer import MultiheadAttention +from .transformer import Transformer +from .transformer import TransformerEncoder, TransformerEncoderLayer +from .transformer import TransformerDecoder, TransformerDecoderLayer +from .utils import layer_norm, generate_square_subsequent_mask diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/transformer/multihead_attention.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/transformer/multihead_attention.py new file mode 100644 index 0000000..3946ec5 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/transformer/multihead_attention.py @@ -0,0 +1,456 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import torch +from torch._C import device +import torch.nn as nn +import torch.nn.functional as F + +import time + +class DotProductAttention(nn.Module): + + def __init__(self, dropout=0.0, alpha=1.0, num_heads=8): + super(DotProductAttention, self).__init__() + + self.dropout = dropout + self.alpha = alpha + self.num_heads = num_heads + self.decay_mat = None + + def forward(self, q, k, v, attn_mask=None): + attn_output_weights = torch.bmm(q, k.transpose(1, 2)) + + if attn_mask is not None: + attn_output_weights += attn_mask + + # attn_output_weights = F.softmax(attn_output_weights, dim=-1) + # attn_output_weights = F.dropout(attn_output_weights, + # p=self.dropout, + # training=self.training) + if self.alpha != 1.0 and self.decay_mat is None: + self.decay_mat = torch.vander(torch.tensor([self.alpha] * self.num_heads), N=v.shape[1]) + self.decay_mat = self.decay_mat.to(v.device) + if self.alpha != 1.0: + bs = attn_output_weights.shape[0] // self.num_heads + decay_mat = self.decay_mat.unsqueeze(0).repeat(bs, 1, 1) + decay_mat = decay_mat.view(-1, *decay_mat.shape[2:]) + attn_output_weights = torch.log(decay_mat[:, None, :] + 1e-10) + attn_output_weights + attn_output_weights = F.softmax(attn_output_weights, dim=-1) + attn_output_weights = F.dropout(attn_output_weights, + p=self.dropout, + training=self.training) + attn_output = torch.bmm(attn_output_weights, v) + + return attn_output + + +class DotProductAttentionStream(DotProductAttention): + + def __init__(self, dropout=0.0, alpha=1.0): + super(DotProductAttentionStream, self).__init__(dropout, alpha) + + ############################ + # Cache for stream inference + ############################ + self.k_weights_cache = None + self.k_pos_weights_cache = None + self.QK_exp_cache = None + self.recache_cnt = 0 + self.recache_freq = 99999 + self.alpha = alpha + + def clear_cache(self): + self.k_weights_cache = None + self.k_pos_weights_cache = None + self.QK_exp_cache = None + self.recache_cnt = 0 + + + def stream_inference(self, q, k, v, k_pos, v_pos, attn_mask=None, cache_num=1, cache_id=0): + ''' + if self.k_weights_cache is not None: + k_weights_new = torch.bmm(q, k[:, [-1]].transpose(1, 2)) + k_weights = torch.cat((self.k_weights_cache[:, :, 1:], k_weights_new), dim=-1) + self.k_weights_cache = k_weights + # k_pos_weights = self.k_pos_weights_cache + else: + k_weights = torch.bmm(q, k.transpose(1, 2)) + self.k_weights_cache = k_weights + # k_pos_weights = torch.bmm(q, k_pos.transpose(1, 2)) + # self.k_pos_weights_cache = k_pos_weights + # attn_output_weights = k_weights + k_pos_weights + attn_output_weights = k_weights + + if attn_mask is not None: + attn_output_weights += attn_mask + + attn_output_weights = F.softmax(attn_output_weights, dim=-1) + attn_output_weights = F.dropout(attn_output_weights, + p=self.dropout, + training=self.training) + # attn_output = torch.bmm(attn_output_weights, (v + v_pos)) + attn_output = torch.bmm(attn_output_weights, v) + return attn_output + ''' + B = 1 + num = 1 + q = q.repeat(B, 1, 1) + k = k.repeat(B, num, 1) + v = v.repeat(B, num, 1) + attn_mask = attn_mask.repeat(B, 1, num) + if self.alpha <= 1.0: + if self.QK_exp_cache is None: + QK = torch.bmm(q, k.transpose(1, 2)) # q~(H, Tq, C), k~(H, Tk, C) => O(H*Tq*Tk*C) + self.QK_max = QK.max(dim=-1, keepdim=True).values + QK_exp = torch.exp(QK - self.QK_max + attn_mask) # => O(H*Tq*Tk) + self.QK_exp_cache = [QK_exp for _ in range(cache_num)] + self.QKV_exp_cache = [torch.bmm(QK_exp, v) for _ in range(cache_num)] # => O(H*Tq*Tk*C) + self.Z = [QK_exp.sum(dim=-1, keepdim=True) for _ in range(cache_num)] # => O(H*Tq*Tk) + output = self.QKV_exp_cache[cache_id] / self.Z[cache_id] + self.recache_cnt += 1 + return output + else: + QK_push = torch.bmm(q, k[:, -num:, :].transpose(1, 2)) # => O(H*Tq*C) + QK_exp_push = torch.exp(QK_push - self.QK_max) # => O(H*Tq*Tk) + sum_push = QK_exp_push.sum(-1, keepdim=True) + QKV_exp_push = torch.bmm(QK_exp_push, v[:, -num:, :]) # => O(H*Tq*Tk) + self.QKV_exp_cache[cache_id] = self.alpha * self.QKV_exp_cache[cache_id] + QKV_exp_push # => O(H*Tq*Tk) + self.Z[cache_id] = self.alpha * self.Z[cache_id] + sum_push + + output = self.QKV_exp_cache[cache_id] / self.Z[cache_id] + self.recache_cnt += 1 + return output + + if self.QK_exp_cache is None: + QK = torch.bmm(q, k.transpose(1, 2)) # q~(H, Tq, C), k~(H, Tk, C) => O(H*Tq*Tk*C) + self.QK_max = QK.max(dim=-1, keepdim=True).values + QK_exp = torch.exp(QK - self.QK_max + attn_mask) # => O(H*Tq*Tk) + self.QK_exp_cache = [QK_exp for _ in range(cache_num)] + self.QKV_exp_cache = [torch.bmm(QK_exp, v) for _ in range(cache_num)] # => O(H*Tq*Tk*C) + + self.Z = [QK_exp.sum(dim=-1, keepdim=True) for _ in range(cache_num)] # => O(H*Tq*Tk) + self.QKV_exp_pop = [torch.bmm(self.QK_exp_cache[cache_i][:, :, :num], v[:, :num, :]) for cache_i in range(cache_num)] + self.sum_pop = [self.QK_exp_cache[cache_i][:, :, :num].sum(-1, keepdim=True) for cache_i in range(cache_num)] + + output = self.QKV_exp_cache[cache_id] / self.Z[cache_id] + self.recache_cnt += 1 + return output + elif (self.recache_cnt // cache_num) % self.recache_freq == 0: + QK = torch.bmm(q, k.transpose(1, 2)) # q~(H, Tq, C), k~(H, Tk, C) => O(H*Tq*Tk*C) + self.QK_max = QK.max(dim=-1, keepdim=True).values + QK_exp = torch.exp(QK - self.QK_max + attn_mask) # => O(H*Tq*Tk) + self.QK_exp_cache[cache_id] = QK_exp + self.QKV_exp_cache[cache_id] = torch.bmm(QK_exp, v) # => O(H*Tq*Tk*C) + + self.Z[cache_id] = QK_exp.sum(dim=-1, keepdim=True) # => O(H*Tq*Tk) + + self.QKV_exp_pop[cache_id] = torch.bmm(self.QK_exp_cache[cache_id][:, :, :num], v[:, :num, :]) + self.sum_pop[cache_id] = self.QK_exp_cache[cache_id][:, :, :num].sum(-1, keepdim=True) + output = self.QKV_exp_cache[cache_id] / self.Z[cache_id] + self.recache_cnt += 1 + return output + else: + QK_push = torch.bmm(q, k[:, -num:, :].transpose(1, 2)) # => O(H*Tq*C) + QK_exp_push = torch.exp(QK_push - self.QK_max) # => O(H*Tq*Tk) + sum_push = QK_exp_push.sum(-1, keepdim=True) + QKV_exp_push = torch.bmm(QK_exp_push, v[:, -num:, :]) # => O(H*Tq*Tk) + + # add with no decay + self.QKV_exp_cache[cache_id] = self.QKV_exp_cache[cache_id] + QKV_exp_push # => O(H*Tq*Tk) + + # add with no decay + self.Z[cache_id] = self.Z[cache_id] + sum_push # => O(H*Tq*Tk) + + self.QK_exp_cache[cache_id] = torch.cat((self.QK_exp_cache[cache_id][:, :, num:], QK_exp_push), dim=-1) + self.QKV_exp_pop[cache_id] = torch.bmm(self.QK_exp_cache[cache_id][:, :, :num], v[:, :num, :]) + self.sum_pop[cache_id] = self.QK_exp_cache[cache_id][:, :, :num].sum(-1, keepdim=True) + output = self.QKV_exp_cache[cache_id] / self.Z[cache_id] + self.recache_cnt += 1 + return output + + +class MultiheadAttention(nn.Module): + + def __init__(self, embed_dim, num_heads, dropout=0.0, bias=True, kdim=None, vdim=None, + attention_type='dotproduct', decay_alpha=1.0): + super(MultiheadAttention, self).__init__() + + self.embed_dim = embed_dim + self.num_heads = num_heads + self.kdim = kdim if kdim is not None else embed_dim + self.vdim = vdim if vdim is not None else embed_dim + self._qkv_same_embed_dim = self.kdim == embed_dim and self.vdim == embed_dim + + if self._qkv_same_embed_dim: + self.in_proj_weight = nn.Parameter(torch.empty(3 * embed_dim, embed_dim)) + else: + raise RuntimeError('Do not support q, k, v have different dimensions') + + if bias: + self.in_proj_bias = nn.Parameter(torch.empty(3 * embed_dim)) + else: + self.register_parameter('in_proj_bias', None) + + self.out_proj = nn.Linear(embed_dim, embed_dim) + + if self._qkv_same_embed_dim: + nn.init.xavier_uniform_(self.in_proj_weight) + + if self.in_proj_bias is not None: + nn.init.constant_(self.in_proj_bias, 0.) + nn.init.constant_(self.out_proj.bias, 0.) + + if attention_type == 'dotproduct': + self.attention = DotProductAttention(dropout, decay_alpha, num_heads) + elif attention_type == 'linear': + self.attention = LinearAttention(dropout) + else: + raise RuntimeError('attention_type should be [dotproduct | linear]') + + def forward(self, q, k, v, attn_mask=None, key_padding_mask=None): + tsz, bsz, embed_dim = q.shape[0], q.shape[1], q.shape[2] + + head_dim = embed_dim // self.num_heads + assert head_dim * self.num_heads == embed_dim, \ + 'embed_dim must be divisible by num_heads' + scaling = float(head_dim) ** -0.5 + + _b = self.in_proj_bias + _start = None + _end = embed_dim + _w = self.in_proj_weight[:_end, :] + if _b is not None: + _b = _b[:_end] + q = F.linear(q, _w, _b) + + _b = self.in_proj_bias + _start = embed_dim + _end = embed_dim * 2 + _w = self.in_proj_weight[_start:_end, :] + if _b is not None: + _b = _b[_start:_end] + k = F.linear(k, _w, _b) + + _b = self.in_proj_bias + _start = embed_dim * 2 + _end = None + _w = self.in_proj_weight[_start:, :] + if _b is not None: + _b = _b[_start:] + v = F.linear(v, _w, _b) + + q = q * scaling + + q = q.contiguous().view(-1, bsz * self.num_heads, head_dim).transpose(0, 1) + k = k.contiguous().view(-1, bsz * self.num_heads, head_dim).transpose(0, 1) + v = v.contiguous().view(-1, bsz * self.num_heads, head_dim).transpose(0, 1) + + if attn_mask is not None: + attn_mask = attn_mask.unsqueeze(0).repeat(bsz, 1, 1) + attn_mask = attn_mask.unsqueeze(1).repeat(1, self.num_heads, 1, 1) + attn_mask = attn_mask.reshape(-1, *attn_mask.shape[2:]) + + if key_padding_mask is not None: + if key_padding_mask.ndim == 4: + key_padding_mask = key_padding_mask.unsqueeze(1).repeat(1, tsz, 1, 1, 1) + key_padding_mask = key_padding_mask.unsqueeze(1).repeat(1, self.num_heads, 1, 1, 1, 1) + key_padding_mask = key_padding_mask.reshape(*key_padding_mask.shape[:3], -1) + else: + key_padding_mask = key_padding_mask.unsqueeze(1).repeat(1, tsz, 1) + key_padding_mask = key_padding_mask.unsqueeze(1).repeat(1, self.num_heads, 1, 1) + key_padding_mask = key_padding_mask.reshape(-1, *key_padding_mask.shape[2:]) + + if attn_mask is not None and key_padding_mask is not None: + mask = attn_mask + key_padding_mask + elif attn_mask is not None: + mask = attn_mask + elif key_padding_mask is not None: + mask = key_padding_mask + else: + mask = None + + attn_output = self.attention(q, k, v, mask) + attn_output = attn_output.transpose(0, 1).contiguous().view(tsz, bsz, + self.embed_dim) + return self.out_proj(attn_output), None + + +class MultiheadAttentionStream(MultiheadAttention): + + def __init__(self, embed_dim, num_heads, dropout=0.0, bias=True, kdim=None, vdim=None, + attention_type='dotproduct', decay_alpha=1.0): + super(MultiheadAttentionStream, self).__init__(embed_dim, num_heads, + dropout, bias, kdim, + vdim, attention_type, + decay_alpha) + + if attention_type == 'dotproduct': + self.attention = DotProductAttentionStream(dropout, decay_alpha) + elif attention_type == 'linear': + self.attention = LinearAttentionStream(dropout) + else: + raise RuntimeError('attention_type should be [dotproduct | linear]') + + ############################ + # Cache for stream inference + ############################ + self.q_cache = None + self.k_cache = None + self.v_cache = None + self.k_pos_cache = None + self.v_pos_cache = None + + def clear_cache(self): + self.q_cache = None + self.k_cache = None + self.v_cache = None + self.k_pos_cache = None + self.v_pos_cache = None + self.attention.clear_cache() + + def stream_inference(self, q, k, v, pos, attn_mask=None, key_padding_mask=None, + cache_num=1, cache_id=0): + tsz, bsz, embed_dim = q.shape[0], q.shape[1], q.shape[2] + + head_dim = embed_dim // self.num_heads + assert head_dim * self.num_heads == embed_dim, \ + 'embed_dim must be divisible by num_heads' + scaling = float(head_dim) ** -0.5 + + if self.q_cache is not None: + q = self.q_cache + else: + _b = self.in_proj_bias + _start = None + _end = embed_dim + _w = self.in_proj_weight[:_end, :] + if _b is not None: + _b = _b[:_end] + q = F.linear(q, _w, _b) + self.q_cache = q + + assert (self.k_cache is None) == (self.k_pos_cache is None) + if self.k_cache is not None: + _b = self.in_proj_bias + _start = embed_dim + _end = embed_dim * 2 + _w = self.in_proj_weight[_start:_end, :] + if _b is not None: + _b = _b[_start:_end] + k_new = F.linear(k[[-1]], _w, None) + k = torch.cat((self.k_cache[1:], k_new)) + self.k_cache = k + k_pos = self.k_pos_cache + else: + _b = self.in_proj_bias + _start = embed_dim + _end = embed_dim * 2 + _w = self.in_proj_weight[_start:_end, :] + if _b is not None: + _b = _b[_start:_end] + k = F.linear(k, _w, None) + self.k_cache = k + k_pos = F.linear(pos, _w, _b) + self.k_pos_cache = k_pos + + assert (self.v_cache is None) == (self.v_pos_cache is None) + if self.v_cache is not None: + _b = self.in_proj_bias + _start = embed_dim * 2 + _end = None + _w = self.in_proj_weight[_start:, :] + if _b is not None: + _b = _b[_start:] + v_new = F.linear(v[[-1]], _w, None) + v = torch.cat((self.v_cache[1:], v_new)) + self.v_cache = v + v_pos = self.v_pos_cache + else: + _b = self.in_proj_bias + _start = embed_dim * 2 + _end = None + _w = self.in_proj_weight[_start:, :] + if _b is not None: + _b = _b[_start:] + v = F.linear(v, _w, None) + self.v_cache = v + v_pos = F.linear(pos, _w, _b) + self.v_pos_cache = v_pos + + q = q * scaling + + q = q.contiguous().view(-1, bsz * self.num_heads, head_dim).transpose(0, 1) + k = k.contiguous().view(-1, bsz * self.num_heads, head_dim).transpose(0, 1) + v = v.contiguous().view(-1, bsz * self.num_heads, head_dim).transpose(0, 1) + k_pos = k_pos.contiguous().view(-1, bsz * self.num_heads, head_dim).transpose(0, 1) + v_pos = v_pos.contiguous().view(-1, bsz * self.num_heads, head_dim).transpose(0, 1) + + if attn_mask is not None: + attn_mask = attn_mask.unsqueeze(0).repeat(bsz, 1, 1) + attn_mask = attn_mask.unsqueeze(1).repeat(1, self.num_heads, 1, 1) + attn_mask = attn_mask.reshape(-1, *attn_mask.shape[2:]) + + if key_padding_mask is not None: + key_padding_mask = key_padding_mask.unsqueeze(1).repeat(1, tsz, 1) + key_padding_mask = key_padding_mask.unsqueeze(1).repeat(1, self.num_heads, 1, 1) + key_padding_mask = key_padding_mask.reshape(-1, *key_padding_mask.shape[2:]) + + if attn_mask is not None and key_padding_mask is not None: + mask = attn_mask + key_padding_mask + elif attn_mask is not None: + mask = attn_mask + elif key_padding_mask is not None: + mask = key_padding_mask + else: + mask = None + + attn_output = self.attention.stream_inference(q, k, v, k_pos, v_pos, mask, + cache_num=cache_num, cache_id=cache_id) + attn_output = attn_output.transpose(0, 1).contiguous().view(tsz, bsz, + self.embed_dim) + return self.out_proj(attn_output), None + + +class LinearAttention(nn.Module): + + def __init__(self, dropout=0.0, nonlinear=F.elu, + eps=1e-6): + super(LinearAttention, self).__init__() + + self.dropout = dropout + self.nonlinear = nonlinear + self.eps = eps + + def forward(self, q, k, v, attn_mask=None): + kk = self.nonlinear(k) + 1 + qq = self.nonlinear(q) + 1 + vv = self.nonlinear(v) + 1 + + kv = torch.bmm(kk.transpose(1, 2), vv) + + ks = torch.sum(kk, dim=1, keepdim=True) # (B 1 C) + Z = 1 / (torch.bmm(qq, ks.transpose(1, 2)) + self.eps) + attn_output = torch.bmm(qq, kv.transpose(1, 2)) * Z + + # if attn_mask is not None: + # print(attn_mask.shape) + # raise RuntimeError(("LinearAttention does not support arbitrary " + # "attention masks")) + + return attn_output + +class LinearAttentionStream(LinearAttention): + + def __init__(self, dropout=0.0): + super(LinearAttentionStream, self).__init__(dropout) + + ############################ + # Cache for stream inference + ############################ + self.k_weights_cache = None + self.k_pos_weights_cache = None + + def stream_inference(self, q, k, v, k_pos, v_pos, attn_mask=None): + raise NotImplementedError + diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/transformer/position_encoding.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/transformer/position_encoding.py new file mode 100644 index 0000000..ded6df3 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/transformer/position_encoding.py @@ -0,0 +1,53 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import math + +import torch +import torch.nn as nn + + +class PositionalEncoding(nn.Module): + + def __init__(self, d_model, dropout=0.1, max_len=5000): + super(PositionalEncoding, self).__init__() + + self.dropout = nn.Dropout(p=dropout) + + pe = torch.zeros(max_len, d_model) + position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) + div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) + pe[:, 0::2] = torch.sin(position * div_term) + pe[:, 1::2] = torch.cos(position * div_term) + pe = pe.unsqueeze(0).transpose(0, 1) + self.register_buffer('pe', pe) + + def forward(self, x, padding=0): + x = x + self.pe[padding: padding + x.shape[0], :] + return self.dropout(x) + + +class PositionalEncoding2D(nn.Module): + + def __init__(self, d_model, dropout=0.1, height=64, width=64): + super(PositionalEncoding2D, self).__init__() + + self.dropout = nn.Dropout(p=dropout) + + pe = torch.zeros(height, width, d_model) + position_x = torch.arange(0, width, dtype=torch.float).unsqueeze(1) + position_y = torch.arange(0, height, dtype=torch.float).unsqueeze(1) + d_model = int(d_model / 2) + div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) + pe[:, :, 0:d_model:2] = torch.sin(position_x * div_term).unsqueeze(0).repeat(height, 1, 1) + pe[:, :, 1:d_model:2] = torch.cos(position_x * div_term).unsqueeze(0).repeat(height, 1, 1) + pe[:, :, d_model::2] = torch.sin(position_y * div_term).unsqueeze(1).repeat(1, width, 1) + pe[:, :, d_model + 1::2] = torch.cos(position_y * div_term).unsqueeze(1).repeat(1, width, 1) + pe = pe.unsqueeze(0).unsqueeze(0) #.transpose(0, 1) + self.register_buffer('pe', pe) + + def forward(self, x, padding=(0,0)): + x = x + self.pe[:, :, + padding[0]: padding[0] + x.shape[2], + padding[1]: padding[1] + x.shape[3], :] + return self.dropout(x) diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/transformer/transformer.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/transformer/transformer.py new file mode 100644 index 0000000..5497e1b --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/transformer/transformer.py @@ -0,0 +1,250 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import copy + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .multihead_attention import MultiheadAttentionStream as MultiheadAttention + + +class Transformer(nn.Module): + + def __init__(self, d_model=512, nhead=8, num_encoder_layers=6, + num_decoder_layers=6, dim_feedforward=2048, dropout=0.1, + activation='relu', custom_encoder=None, custom_decoder=None): + super(Transformer, self).__init__() + + if custom_encoder is not None: + self.encoder = custom_encoder + else: + encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout, activation) + encoder_norm = nn.LayerNorm(d_model) + self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm) + + if custom_decoder is not None: + self.decoder = custom_decoder + else: + decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward, dropout, activation) + decoder_norm = nn.LayerNorm(d_model) + self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm) + + self.d_model = d_model + self.nhead = nhead + + def forward(self, src, tgt, src_mask=None, tgt_mask=None, + memory_mask=None, src_key_padding_mask=None, + tgt_key_padding_mask=None, memory_key_padding_mask=None): + if src.size(1) != tgt.size(1): + raise RuntimeError('the batch number of src and tgt must be equal') + + if src.size(2) != self.d_model or tgt.size(2) != self.d_model: + raise RuntimeError('the feature number of src and tgt must be equal to d_model') + + memory = self.encoder(src, mask=src_mask, src_key_padding_mask=src_key_padding_mask) + output = self.decoder(tgt, memory, tgt_mask=tgt_mask, memory_mask=memory_mask, + tgt_key_padding_mask=tgt_key_padding_mask, + memory_key_padding_mask=memory_key_padding_mask) + return output + + +class TransformerEncoder(nn.Module): + + def __init__(self, encoder_layer, num_layers, norm=None): + super(TransformerEncoder, self).__init__() + + self.layers = _get_clones(encoder_layer, num_layers) + self.num_layers = num_layers + self.norm = norm + + def forward(self, src, src_mask=None, src_key_padding_mask=None): + output = src + + for mod in self.layers: + output = mod(output, src_mask=src_mask, + src_key_padding_mask=src_key_padding_mask) + + if self.norm is not None: + output = self.norm(output) + + return output + + +class TransformerDecoder(nn.Module): + + def __init__(self, decoder_layer, num_layers, norm=None): + super(TransformerDecoder, self).__init__() + + self.layers = _get_clones(decoder_layer, num_layers) + self.num_layers = num_layers + self.norm = norm + + def clear_cache(self): + if len(self.layers) != 1: + raise RuntimeError('Number of layers cannot larger than 1 for stream inference') + + self.layers[0].clear_cache() + + def stream_inference(self, tgt, memory, pos, tgt_mask=None, + memory_mask=None, tgt_key_padding_mask=None, + memory_key_padding_mask=None, cache_num=1, cache_id=0): + output = tgt + + if len(self.layers) != 1: + raise RuntimeError('Number of layers cannot larger than 1 for stream inference') + + output = self.layers[0].stream_inference(output, memory, pos, + tgt_mask=tgt_mask, + memory_mask=memory_mask, + tgt_key_padding_mask=tgt_key_padding_mask, + memory_key_padding_mask=memory_key_padding_mask, + cache_num=cache_num, cache_id=cache_id) + + if self.norm is not None: + output = self.norm(output) + + return output + + def forward(self, tgt, memory, tgt_mask=None, + memory_mask=None, tgt_key_padding_mask=None, + memory_key_padding_mask=None): + output = tgt + + for mod in self.layers: + output = mod(output, memory, tgt_mask=tgt_mask, + memory_mask=memory_mask, + tgt_key_padding_mask=tgt_key_padding_mask, + memory_key_padding_mask=memory_key_padding_mask) + + if self.norm is not None: + output = self.norm(output) + + return output + + +class TransformerEncoderLayer(nn.Module): + + def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation='relu'): + super(TransformerEncoderLayer, self).__init__() + + self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout) + + # Implementation of Feedforward model + self.linear1 = nn.Linear(d_model, dim_feedforward) + self.dropout = nn.Dropout(dropout) + self.linear2 = nn.Linear(dim_feedforward, d_model) + + self.norm1 = nn.LayerNorm(d_model) + self.norm2 = nn.LayerNorm(d_model) + self.dropout1 = nn.Dropout(dropout) + self.dropout2 = nn.Dropout(dropout) + + self.activation = _get_activation_fn(activation) + + def __setstate__(self, state): + if 'activation' not in state: + state['activation'] = F.relu + super(TransformerEncoderLayer, self).__setstate__(state) + + def forward(self, src, src_mask=None, src_key_padding_mask=None): + src2 = self.self_attn(src, src, src, attn_mask=src_mask, + key_padding_mask=src_key_padding_mask)[0] + src = src + self.dropout1(src2) + src = self.norm1(src) + src2 = self.linear2(self.dropout(self.activation(self.linear1(src)))) + src = src + self.dropout2(src2) + src = self.norm2(src) + return src + + +class TransformerDecoderLayer(nn.Module): + + def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, + activation='relu', attention_type='dotproduct', + decay_alpha=1.0): + super(TransformerDecoderLayer, self).__init__() + + self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout, + attention_type=attention_type) + self.multihead_attn = MultiheadAttention(d_model, nhead, dropout=dropout, + attention_type=attention_type, + decay_alpha=decay_alpha) + + # Implementation of Feedforward model + self.linear1 = nn.Linear(d_model, dim_feedforward) + self.dropout = nn.Dropout(dropout) + self.linear2 = nn.Linear(dim_feedforward, d_model) + + self.norm1 = nn.LayerNorm(d_model) + self.norm2 = nn.LayerNorm(d_model) + self.norm3 = nn.LayerNorm(d_model) + self.dropout1 = nn.Dropout(dropout) + self.dropout2 = nn.Dropout(dropout) + self.dropout3 = nn.Dropout(dropout) + + self.activation = _get_activation_fn(activation) + + ############################ + # Cache for stream inference + ############################ + self.tgt_cache = None + + def __setstate__(self, state): + if 'activation' not in state: + state['activation'] = F.relu + super(TransformerDecoderLayer, self).__setstate__(state) + + def clear_cache(self): + self.tgt_cache = None + self.multihead_attn.clear_cache() + + def stream_inference(self, tgt, memory, pos, tgt_mask=None, memory_mask=None, + tgt_key_padding_mask=None, memory_key_padding_mask=None, + cache_num=1, cache_id=0): + if self.tgt_cache is None: + tgt2 = self.self_attn(tgt, tgt, tgt, attn_mask=tgt_mask, + key_padding_mask=tgt_key_padding_mask)[0] + tgt = tgt + self.dropout1(tgt2) + tgt = self.norm1(tgt) + self.tgt_cache = tgt + else: + tgt = self.tgt_cache + tgt2 = self.multihead_attn.stream_inference(tgt, memory, memory, pos, attn_mask=memory_mask, + key_padding_mask=memory_key_padding_mask, + cache_num=cache_num, cache_id=cache_id)[0] + tgt = tgt + self.dropout2(tgt2) + tgt = self.norm2(tgt) + tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt)))) + tgt = tgt + self.dropout3(tgt2) + tgt = self.norm3(tgt) + return tgt + + def forward(self, tgt, memory, tgt_mask=None, memory_mask=None, + tgt_key_padding_mask=None, memory_key_padding_mask=None): + tgt2 = self.self_attn(tgt, tgt, tgt, attn_mask=tgt_mask, + key_padding_mask=tgt_key_padding_mask)[0] + tgt = tgt + self.dropout1(tgt2) + tgt = self.norm1(tgt) + tgt2 = self.multihead_attn(tgt, memory, memory, attn_mask=memory_mask, + key_padding_mask=memory_key_padding_mask)[0] + tgt = tgt + self.dropout2(tgt2) + tgt = self.norm2(tgt) + tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt)))) + tgt = tgt + self.dropout3(tgt2) + tgt = self.norm3(tgt) + return tgt + + +def _get_clones(module, N): + return nn.ModuleList([copy.deepcopy(module) for i in range(N)]) + + +def _get_activation_fn(activation): + if activation == 'relu': + return F.relu + elif activation == 'gelu': + return F.gelu + + raise RuntimeError('activation should be relu/gelu, not {}'.format(activation)) diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/transformer/utils.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/transformer/utils.py new file mode 100644 index 0000000..535af91 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/transformer/utils.py @@ -0,0 +1,16 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import torch +assert torch.__version__ >= '1.6.0' +import torch.nn as nn + + +def layer_norm(d_model, condition=True): + return nn.LayerNorm(d_model) if condition else None + + +def generate_square_subsequent_mask(sz): + mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1) + mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0)) + return mask diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/weights_init.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/weights_init.py new file mode 100644 index 0000000..6ccc0fa --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/models/weights_init.py @@ -0,0 +1,33 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import math + +import torch.nn as nn + + +def weights_init(m): + if isinstance(m, nn.Linear): + nn.init.kaiming_uniform_(m.weight.data, a=math.sqrt(5)) + elif isinstance(m, nn.Conv1d): + nn.init.normal_(m.weight.data) + if m.bias is not None: + nn.init.normal_(m.bias.data) + elif isinstance(m, nn.Conv2d): + nn.init.xavier_normal_(m.weight.data) + if m.bias is not None: + nn.init.normal_(m.bias.data) + elif isinstance(m, nn.ConvTranspose1d): + nn.init.normal_(m.weight.data) + if m.bias is not None: + nn.init.normal_(m.bias.data) + elif isinstance(m, nn.ConvTranspose2d): + nn.init.xavier_normal_(m.weight.data) + if m.bias is not None: + nn.init.normal_(m.bias.data) + elif isinstance(m, nn.BatchNorm1d): + nn.init.normal_(m.weight.data, mean=1, std=0.02) + nn.init.constant_(m.bias.data, 0) + elif isinstance(m, nn.BatchNorm2d): + nn.init.normal_(m.weight.data, mean=1, std=0.02) + nn.init.constant_(m.bias.data, 0) diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/optimizers/__init__.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/optimizers/__init__.py new file mode 100644 index 0000000..8ca70d3 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/optimizers/__init__.py @@ -0,0 +1,5 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +from .optimizers import build_optimizer +from .lr_scheduler import build_scheduler diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/optimizers/lr_scheduler.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/optimizers/lr_scheduler.py new file mode 100644 index 0000000..533dd18 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/optimizers/lr_scheduler.py @@ -0,0 +1,224 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import math +from bisect import bisect_left + +import numpy as np +import torch +from torch.optim.lr_scheduler import _LRScheduler + + +def _get_warmup_factor_at_iter(warmup_method, + this_iter, + warmup_iters, + warmup_factor): + if this_iter >= warmup_iters: + return 1.0 + + if warmup_method == 'constant': + return warmup_factor + elif warmup_method == 'linear': + alpha = this_iter / warmup_iters + return warmup_factor * (1 - alpha) + alpha + else: + raise ValueError('Unknown warmup method: {}'.format(warmup_method)) + + +class MultiStepLR(_LRScheduler): + + def __init__(self, + optimizer, + milestones, + gamma=0.1, + last_epoch=-1): + if not list(milestones) == sorted(milestones): + raise ValueError( + 'Milestones should be a list of increasing integers. Got {}'.format(milestones) + ) + self.milestones = milestones + self.gamma = gamma + super().__init__(optimizer, last_epoch) + + def get_lr(self): + return [ + base_lr + * self.gamma ** bisect_left(self.milestones, self.last_epoch) + for base_lr in self.base_lrs + ] + + def _compute_values(self): + # The new interface + return self.get_lr() + + +class WarmupMultiStepLR(_LRScheduler): + + def __init__(self, + optimizer, + milestones, + gamma=0.1, + warmup_factor=0.3, + warmup_iters=500, + warmup_method='linear', + last_epoch=-1): + if not list(milestones) == sorted(milestones): + raise ValueError( + 'Milestones should be a list of increasing integers. Got {}'.format(milestones) + ) + self.milestones = milestones + self.gamma = gamma + self.warmup_factor = warmup_factor + self.warmup_iters = warmup_iters + self.warmup_method = warmup_method + super().__init__(optimizer, last_epoch) + + def get_lr(self): + warmup_factor = _get_warmup_factor_at_iter( + self.warmup_method, + self.last_epoch, + self.warmup_iters, + self.warmup_factor, + ) + import pdb; pdb.set_trace() + return [ + base_lr + * warmup_factor + * self.gamma ** bisect_left(self.milestones, self.last_epoch) + for base_lr in self.base_lrs + ] + + def _compute_values(self): + # The new interface + return self.get_lr() + + +class CosineLR(_LRScheduler): + + def __init__(self, + optimizer, + max_iters, + last_epoch=-1): + self.max_iters = max_iters + super().__init__(optimizer, last_epoch) + + def get_lr(self): + return [ + base_lr + * 0.5 + * (1.0 + math.cos(math.pi * self.last_epoch / self.max_iters)) + for base_lr in self.base_lrs + ] + + def _compute_values(self): + # The new interface + return self.get_lr() + + +class WarmupCosineLR(_LRScheduler): + + def __init__(self, + optimizer, + max_iters, + gamma=0.1, + warmup_factor=0.3, + warmup_iters=500, + warmup_method='linear', + last_epoch=-1): + self.max_iters = max_iters + self.warmup_factor = warmup_factor + self.warmup_iters = warmup_iters + self.warmup_method = warmup_method + super().__init__(optimizer, last_epoch) + + def get_lr(self): + warmup_factor = _get_warmup_factor_at_iter( + self.warmup_method, + self.last_epoch, + self.warmup_iters, + self.warmup_factor, + ) + return [ + base_lr + * warmup_factor + * 0.5 + * (1.0 + math.cos(math.pi * self.last_epoch / self.max_iters)) + for base_lr in self.base_lrs + ] + + def _compute_values(self): + # The new interface + return self.get_lr() + +class DummyScheduler: + def __init__(self, optimizer): + self.optimizer = optimizer + self.base_lrs = [group['lr'] for group in optimizer.param_groups] + + def get_last_lr(self): + return self.base_lrs + + def step(self): + pass # No operation for stepping the scheduler + + +def build_scheduler(cfg, optimizer, num_iters_per_epoch): + """Unlike the PyTorch version, our schedulers adjust the learning rate + according to iteration rather than epoch. + """ + # Set last epoch (here 'epoch' is actually 'iters') + last_epoch = (cfg.SOLVER.START_EPOCH - 1) * num_iters_per_epoch + + if cfg.SOLVER.SCHEDULER.SCHEDULER_NAME == 'none' or cfg.SOLVER.SCHEDULER.SCHEDULER_NAME is None: + scheduler = DummyScheduler(optimizer) + + elif cfg.SOLVER.SCHEDULER.SCHEDULER_NAME == 'multistep': + # Convert milestones epochs to iters + milestones = [(i - 1) * num_iters_per_epoch for i in cfg.SOLVER.SCHEDULER.MILESTONES] + + scheduler = MultiStepLR( + optimizer, + milestones=milestones, + gamma=cfg.SOLVER.SCHEDULER.GAMMA, + last_epoch=last_epoch) + elif cfg.SOLVER.SCHEDULER.SCHEDULER_NAME == 'warmup_multistep': + # Convert milestones epochs to iters + milestones = [(i - 1) * num_iters_per_epoch for i in cfg.SOLVER.SCHEDULER.MILESTONES] + + # Convert warmup epochs to iters + cfg.SOLVER.SCHEDULER.WARMUP_ITERS = int(cfg.SOLVER.SCHEDULER.WARMUP_EPOCHS * num_iters_per_epoch) + + scheduler = WarmupMultiStepLR( + optimizer, + milestones=milestones, + gamma=cfg.SOLVER.SCHEDULER.GAMMA, + warmup_factor=cfg.SOLVER.SCHEDULER.WARMUP_FACTOR, + warmup_iters=cfg.SOLVER.SCHEDULER.WARMUP_ITERS, + warmup_method=cfg.SOLVER.SCHEDULER.WARMUP_METHOD, + last_epoch=last_epoch) + elif cfg.SOLVER.SCHEDULER.SCHEDULER_NAME == 'cosine': + # Get max number of iters + max_iters = cfg.SOLVER.NUM_EPOCHS * num_iters_per_epoch + + scheduler = CosineLR( + optimizer, + max_iters=max_iters, + last_epoch=last_epoch) + elif cfg.SOLVER.SCHEDULER.SCHEDULER_NAME == 'warmup_cosine': + # Get max number of iters + max_iters = cfg.SOLVER.NUM_EPOCHS * num_iters_per_epoch + + # Convert warmup epochs to iters + cfg.SOLVER.SCHEDULER.WARMUP_ITERS = int(cfg.SOLVER.SCHEDULER.WARMUP_EPOCHS * num_iters_per_epoch) + + scheduler = WarmupCosineLR( + optimizer, + max_iters=max_iters, + gamma=cfg.SOLVER.SCHEDULER.GAMMA, + warmup_factor=cfg.SOLVER.SCHEDULER.WARMUP_FACTOR, + warmup_iters=cfg.SOLVER.SCHEDULER.WARMUP_ITERS, + warmup_method=cfg.SOLVER.SCHEDULER.WARMUP_METHOD, + last_epoch=last_epoch) + else: + raise RuntimeError('Unknown lr scheduler: {}'.format(cfg.SOLVER.SCHEDULER.SCHEDULER_NAME)) + return scheduler diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/optimizers/optimizers.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/optimizers/optimizers.py new file mode 100644 index 0000000..5e6bd9c --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/optimizers/optimizers.py @@ -0,0 +1,27 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +__all__ = ['build_optimizer'] + +import torch.optim as optim + + +def build_optimizer(cfg, model): + if cfg.SOLVER.OPTIMIZER == 'sgd': + optimizer = optim.SGD( + [{'params': model.parameters(), 'initial_lr': cfg.SOLVER.BASE_LR}], + lr=cfg.SOLVER.BASE_LR, weight_decay=cfg.SOLVER.WEIGHT_DECAY, momentum=cfg.SOLVER.MOMENTUM, + ) + elif cfg.SOLVER.OPTIMIZER == 'adam': + optimizer = optim.Adam( + [{'params': model.parameters(), 'initial_lr': cfg.SOLVER.BASE_LR}], + lr=cfg.SOLVER.BASE_LR, weight_decay=cfg.SOLVER.WEIGHT_DECAY, + ) + elif cfg.SOLVER.OPTIMIZER == 'adamw': + optimizer = optim.AdamW( + [{'params': model.parameters(), 'initial_lr': cfg.SOLVER.BASE_LR}], + lr=cfg.SOLVER.BASE_LR, weight_decay=cfg.SOLVER.WEIGHT_DECAY, + ) + else: + raise RuntimeError('Unknown optimizer: {}'.format(cfg.SOLVER.OPTIMIZER)) + return optimizer diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/utils/__init__.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/utils/__init__.py new file mode 100644 index 0000000..04f8b7b --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/utils/__init__.py @@ -0,0 +1,2 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/utils/checkpointer.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/utils/checkpointer.py new file mode 100644 index 0000000..515c843 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/utils/checkpointer.py @@ -0,0 +1,144 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +__all__ = ['setup_checkpointer'] + +import os.path as osp + +import torch + + +class Checkpointer(object): + + def __init__(self, cfg, phase): + + # Load pretrained checkpoint + + self.is_pretrained = False + self.cfg = cfg + self.checkpoint = self._load_checkpoint(cfg.MODEL.CHECKPOINT) + self.except_classifier = cfg.MODEL.EXCEPT_CLASSIFIER if phase == 'train' else False + + if self.checkpoint is not None and phase == 'train': + # further training with the same dataset: resume + model_ckpt_str = cfg['MODEL']['CHECKPOINT'].replace('_', '') + if cfg['DATA']['DATA_NAME'].lower() in model_ckpt_str.lower(): + cfg.SOLVER.START_EPOCH += self.checkpoint.get('epoch', 0) + # further training with the different dataset: finetuning on pretrained weight + else: + self.is_pretrained = True + print(f'WARNING) Finetuning on pretrained weight: {cfg.MODEL.CHECKPOINT}') + + elif self.checkpoint is None and phase != 'train': + # no checkpoint in infernece mode: error + raise RuntimeError('Cannot find checkpoint {}'.format(cfg.MODEL.CHECKPOINT)) + + self.output_dir = cfg.OUTPUT_DIR + + def load(self, model, optimizer=None): + if self.checkpoint is not None: + # try load the ckpt first, then if there is a size mismatch, do other operation + if not self.is_pretrained: # just resume + # Attempt to load the state dictionary from the checkpoint + model.load_state_dict(self.checkpoint['model_state_dict']) + # Load the optimizer state if provided + if optimizer is not None: + if 'optimizer_state_dict' in self.checkpoint: + optimizer.load_state_dict(self.checkpoint['optimizer_state_dict']) + else: + print('Optimizer state dictionary not found in checkpoint.') + + else: + # different model structure (pretrained) + print(f'WARNING) Load checkpoint separately due to possiblity of different model structure, src/rekognition_online_action_detection/utils/checkpointer.py L51') + model_dict = model.state_dict() + unnecessary_keys = [] + unmatched_keys = [] + matched_ckpt = {} + + # loaded_state_dict = self.checkpoint['model_state_dict'] + + # # Initialize a flag to track if all values match + # all_values_match = True + + # # Iterate through all keys and compare values + # for key in model_dict: + # import pdb; pdb.set_trace() + # if not torch.allclose(model_dict[key], loaded_state_dict[key]): + # print(f"Mismatch found in {key}") + # print(f"Model value: {model_dict[key]}") + # print(f"Loaded value: {loaded_state_dict[key]}") + # print(f"Difference: {model_dict[key] - loaded_state_dict[key]}") + # print("---") + # all_values_match = False + + # Filter out unnecessary keys + for k, v in self.checkpoint['model_state_dict'].items(): + if k not in model_dict: + print(f'Unnecessary key: {k} not found in model') + unnecessary_keys.append(k) + elif model_dict[k].shape != v.shape: + print(f'Unnecessary key: {k} shape {v.shape} does not match {model_dict[k].shape}') + unmatched_keys.append(k) + elif self.except_classifier is True and 'classifier' in k: + print(f'Unnecessary key: {k} is classifier') + unnecessary_keys.append(k) + else: + matched_ckpt[k] = v + + print(f'Unloaded keys: {unnecessary_keys}') + print(f'Unmathced keys: {unmatched_keys}') + # Update the existing state dict with the matched keys + # import pdb; pdb.set_trace() + model_dict.update(matched_ckpt) + model.load_state_dict(model_dict) + # for unmatched_key in unmatched_keys: + # loaded_dict_shape = self.checkpoint['model_state_dict'][unmatched_key].shape + # model_dict_shape = model_dict[unmatched_key].shape + # if len(loaded_dict_shape) == 2: + # if loaded_dict_shape[0] != model_dict_shape[0]: + # model_dict[unmatched_key][:loaded_dict_shape[0], :] = self.checkpoint['model_state_dict'][unmatched_key] + # print(f'Updated key: {unmatched_key}, model_dict[{unmatched_key}][:{loaded_dict_shape[0]}, :] = checkpoint["model_state_dict"][{unmatched_key}]') + # else: + # model_dict[unmatched_key][:, :loaded_dict_shape[1]] = self.checkpoint['model_state_dict'][unmatched_key] + # print(f'Updated key: {unmatched_key}, model_dict[{unmatched_key}][:, :{loaded_dict_shape[1]}] = checkpoint["model_state_dict"][{unmatched_key}]') + # elif len(loaded_dict_shape) == 1: + # model_dict[unmatched_key][:loaded_dict_shape[0]] = self.checkpoint['model_state_dict'][unmatched_key] + # print(f'Updated key: {unmatched_key}, model_dict[{unmatched_key}][:{loaded_dict_shape[0]}] = checkpoint["model_state_dict"][{unmatched_key}]') + # elif len(loaded_dict_shape) >= 3: + # raise ValueError(f'Unsupported shape: {loaded_dict_shape}') + # model.load_state_dict(model_dict) + + + + def save(self, epoch, model, optimizer): + torch.save({ + 'epoch': epoch, + 'model_state_dict': model.module.state_dict() if torch.cuda.device_count() > 1 else model.state_dict(), + 'optimizer_state_dict': optimizer.state_dict(), + }, osp.join(self.output_dir, f'epoch-{epoch}.pth')) + print(f'\nSave checkpoint to {osp.join(self.output_dir, f"epoch-{epoch}.pth")}') + # osp.join(self.output_dir, f'{self.cfg.INPUT.MODALITY}_epoch-{epoch}.pth')) + + def save_best(self, epoch, model, optimizer): + # save_path = osp.join(self.output_dir, f'{self.cfg.INPUT.MODALITY}_best.pth') + save_path = osp.join(self.output_dir, f'best.pth') + + torch.save({ + 'epoch': epoch, + 'model_state_dict': model.module.state_dict() if torch.cuda.device_count() > 1 else model.state_dict(), + 'optimizer_state_dict': optimizer.state_dict(), + }, save_path) + print(f'\nSave checkpoint to {save_path}') + + def _load_checkpoint(self, checkpoint): + if checkpoint is not None and osp.isfile(checkpoint): + return torch.load(checkpoint, map_location=torch.device('cpu')) + elif osp.isdir(checkpoint): + import pdb; pdb.set_trace() + + return None + + +def setup_checkpointer(cfg, phase): + return Checkpointer(cfg, phase) \ No newline at end of file diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/utils/ek_utils.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/utils/ek_utils.py new file mode 100644 index 0000000..a73670f --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/utils/ek_utils.py @@ -0,0 +1,20 @@ +import os.path as osp +import pandas as pd + + +def action_to_verb_map(path_to_data, + action_offset=False, + verb_offset=False): + actions = pd.read_csv(osp.join(path_to_data, 'actions.csv')) + a_to_v = {a[1]['id'] + action_offset: a[1]['verb'] + verb_offset + for a in actions.iterrows()} + return a_to_v + + +def action_to_noun_map(path_to_data, + action_offset=False, + noun_offset=False): + actions = pd.read_csv(osp.join(path_to_data, 'actions.csv')) + a_to_n = {a[1]['id'] + action_offset: a[1]['noun'] + noun_offset + for a in actions.iterrows()} + return a_to_n diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/utils/env.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/utils/env.py new file mode 100644 index 0000000..e2c4c8f --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/utils/env.py @@ -0,0 +1,32 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +__all__ = [ + 'setup_random_seed', + 'setup_environment', +] + +import os +import random + +import torch +import numpy as np + + +def setup_random_seed(seed): + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + torch.backends.cudnn.benchmark = False + torch.backends.cudnn.deterministic = True + + +def setup_environment(cfg): + os.environ['CUDA_VISIBLE_DEVICES'] = cfg.GPU + device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + if cfg.SEED is not None: + setup_random_seed(cfg.SEED) + return device diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/utils/logger.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/utils/logger.py new file mode 100644 index 0000000..28e6d9e --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/utils/logger.py @@ -0,0 +1,32 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +__all__ = ['setup_logger'] + +import os +import sys +import logging +import pprint + + +def setup_logger(cfg, phase, quiet=False): + logger = logging.getLogger('rekognition') + level = logging.DEBUG if cfg.VERBOSE else logging.INFO + logger.setLevel(level) + ch = logging.StreamHandler(stream=sys.stdout) + ch.setLevel(level) + logger.addHandler(ch) + + if phase == 'train': + log_file = os.path.join(cfg.OUTPUT_DIR, 'log.txt') + os.makedirs(cfg.OUTPUT_DIR, exist_ok=True) + else: + log_file = os.path.splitext(cfg.MODEL.CHECKPOINT)[0] + '.txt' + fh = logging.FileHandler(log_file) + fh.setLevel(logging.INFO) + logger.addHandler(fh) + + if not quiet: + logger.info(pprint.pformat(cfg)) + + return logger diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/utils/parser.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/utils/parser.py new file mode 100644 index 0000000..0d46c93 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/utils/parser.py @@ -0,0 +1,114 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import json +import os.path as osp + +from rekognition_online_action_detection.config.defaults import get_cfg + + +def parse_args(): + parser = argparse.ArgumentParser(description='Rekognition Face Anti-Spoofing') + parser.add_argument( + '--config_file', + default='', + type=str, + help='path to yaml config file', + ) + parser.add_argument( + '--gpu', + default='0', + type=str, + help='specify visible devices' + ) + parser.add_argument( + 'opts', + default=None, + nargs='*', + help='modify config options using the command-line', + ) + return parser.parse_args() + + +def assert_and_infer_cfg(cfg, args): + # Setup the visible devices + cfg.GPU = args.gpu + + # Infer data info + with open(cfg.DATA.DATA_INFO, 'r') as f: + data_info = json.load(f)[cfg.DATA.DATA_NAME] + + cfg.DATA.DATA_ROOT = data_info['data_root'] if cfg.DATA.DATA_ROOT is None else cfg.DATA.DATA_ROOT + cfg.DATA.CLASS_NAMES = data_info['class_names'] if cfg.DATA.CLASS_NAMES is None else cfg.DATA.CLASS_NAMES + cfg.DATA.NUM_CLASSES = data_info['num_classes'] if cfg.DATA.NUM_CLASSES is None else cfg.DATA.NUM_CLASSES + cfg.DATA.IGNORE_INDEX = data_info['ignore_index'] if cfg.DATA.IGNORE_INDEX is None else cfg.DATA.IGNORE_INDEX + cfg.DATA.METRICS = data_info['metrics'] if cfg.DATA.METRICS is None else cfg.DATA.METRICS + cfg.DATA.FPS = data_info['fps'] if cfg.DATA.FPS is None else cfg.DATA.FPS + cfg.DATA.TRAIN_SESSION_SET = data_info['train_session_set'] if cfg.DATA.TRAIN_SESSION_SET is None else cfg.DATA.TRAIN_SESSION_SET + cfg.DATA.TEST_SESSION_SET = data_info['test_session_set'] if cfg.DATA.TEST_SESSION_SET is None else cfg.DATA.TEST_SESSION_SET + + # Ignore two mis-labeled videos + if False and cfg.DATA_NAME == 'THUMOS': + cfg.DATA.TEST_SESSION_SET.remove('video_test_0000270') + cfg.DATA.TEST_SESSION_SET.remove('video_test_0001496') + + # Input assertions + if cfg.INPUT.MODALITY == 'twostream': + cfg.INPUT.MODALITY = 'visual+motion' + if cfg.INPUT.MODALITY == 'threestream': + cfg.INPUT.MODALITY = 'visual+motion+object' + assert cfg.INPUT.MODALITY in ['visual', 'motion', 'object', + 'visual+motion', 'motion+object', 'visual+object', + 'visual+motion+object'] + + # Infer memory + if cfg.MODEL.MODEL_NAME in ['LSTR']: + cfg.MODEL.LSTR.AGES_MEMORY_LENGTH = cfg.MODEL.LSTR.AGES_MEMORY_SECONDS * cfg.DATA.FPS + cfg.MODEL.LSTR.LONG_MEMORY_LENGTH = cfg.MODEL.LSTR.LONG_MEMORY_SECONDS * cfg.DATA.FPS + cfg.MODEL.LSTR.WORK_MEMORY_LENGTH = cfg.MODEL.LSTR.WORK_MEMORY_SECONDS * cfg.DATA.FPS + cfg.MODEL.LSTR.TOTAL_MEMORY_LENGTH = \ + cfg.MODEL.LSTR.AGES_MEMORY_LENGTH + \ + cfg.MODEL.LSTR.LONG_MEMORY_LENGTH + \ + cfg.MODEL.LSTR.WORK_MEMORY_LENGTH + cfg.MODEL.LSTR.ANTICIPATION_LENGTH = cfg.MODEL.LSTR.ANTICIPATION_SECONDS * cfg.DATA.FPS + assert cfg.MODEL.LSTR.AGES_MEMORY_LENGTH % cfg.MODEL.LSTR.AGES_MEMORY_SAMPLE_RATE == 0 + assert cfg.MODEL.LSTR.LONG_MEMORY_LENGTH % cfg.MODEL.LSTR.LONG_MEMORY_SAMPLE_RATE == 0 + assert cfg.MODEL.LSTR.WORK_MEMORY_LENGTH % cfg.MODEL.LSTR.WORK_MEMORY_SAMPLE_RATE == 0 + assert cfg.MODEL.LSTR.ANTICIPATION_LENGTH % cfg.MODEL.LSTR.ANTICIPATION_SAMPLE_RATE == 0 + cfg.MODEL.LSTR.AGES_MEMORY_NUM_SAMPLES = cfg.MODEL.LSTR.AGES_MEMORY_LENGTH // cfg.MODEL.LSTR.AGES_MEMORY_SAMPLE_RATE + cfg.MODEL.LSTR.LONG_MEMORY_NUM_SAMPLES = cfg.MODEL.LSTR.LONG_MEMORY_LENGTH // cfg.MODEL.LSTR.LONG_MEMORY_SAMPLE_RATE + cfg.MODEL.LSTR.WORK_MEMORY_NUM_SAMPLES = cfg.MODEL.LSTR.WORK_MEMORY_LENGTH // cfg.MODEL.LSTR.WORK_MEMORY_SAMPLE_RATE + cfg.MODEL.LSTR.TOTAL_MEMORY_NUM_SAMPLES = \ + cfg.MODEL.LSTR.AGES_MEMORY_NUM_SAMPLES + \ + cfg.MODEL.LSTR.LONG_MEMORY_NUM_SAMPLES + \ + cfg.MODEL.LSTR.WORK_MEMORY_NUM_SAMPLES + cfg.MODEL.LSTR.ANTICIPATION_NUM_SAMPLES = cfg.MODEL.LSTR.ANTICIPATION_LENGTH // cfg.MODEL.LSTR.ANTICIPATION_SAMPLE_RATE + + assert cfg.MODEL.LSTR.INFERENCE_MODE in ['batch', 'stream'] + + # Infer output dir + config_name = osp.splitext(args.config_file)[0].split('/')[1:] + cfg.OUTPUT_DIR = osp.join(cfg.OUTPUT_DIR, *config_name) + if cfg.SESSION: + cfg.OUTPUT_DIR = osp.join(cfg.OUTPUT_DIR, cfg.SESSION) + + +def load_cfg(): + args = parse_args() + cfg = get_cfg() + + for i in range(0, len(args.opts), 2): + key = args.opts[i] + value = args.opts[i+1] + if key == 'MODEL.LSTR.LONG_MEMORY_SECONDS' or key == 'MODEL.LSTR.WORK_MEMORY_SECONDS': + args.opts[i+1] = int(value) + print(f'(key, value): ({key}, {value}) is converted to integer type') + + if args.config_file is not None: + cfg.merge_from_file(args.config_file) + if args.opts is not None: + cfg.merge_from_list(args.opts) + assert_and_infer_cfg(cfg, args) + # config_file is in cfg.OUTPUT_DIR or args.config_file + return cfg diff --git a/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/utils/registry.py b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/utils/registry.py new file mode 100644 index 0000000..fd64985 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/src/rekognition_online_action_detection/utils/registry.py @@ -0,0 +1,23 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +def _register_generic(module_dict, module_name, module): + assert module_name not in module_dict + module_dict[module_name] = module + + +class Registry(dict): + + def __init__(self, *args, **kwargs): + super(Registry, self).__init__(*args, **kwargs) + + def register(self, module_name, module=None): + if module is not None: + _register_generic(self, module_name, module) + return + + def register_fn(fn): + _register_generic(self, module_name, fn) + return fn + + return register_fn diff --git a/third_party/EgoSpeak/LSTR/tools/eval/eval_fusion.py b/third_party/EgoSpeak/LSTR/tools/eval/eval_fusion.py new file mode 100644 index 0000000..4451465 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/tools/eval/eval_fusion.py @@ -0,0 +1,162 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import logging +logging.basicConfig(level=logging.INFO, format='%(message)s') +import argparse + +import numpy as np +import pickle as pkl +import pandas as pd +import os.path as osp + +from rekognition_online_action_detection.evaluation import compute_result +from external.rulstm.RULSTM.utils import (get_marginal_indexes, marginalize, softmax, + topk_accuracy_multiple_timesteps, + topk_recall_multiple_timesteps, + tta) + + +def eval_fusion(pred_scores_files, weights): + ### Comments: Per-frame scores cost too much space... + # perframe_gt_targets_fusion = {} + # perframe_pred_scores_fusion = {} + # for weight, pred_scores_file in zip(weights, pred_scores_files): + # pred_scores = pkl.load(open(pred_scores_file, 'rb')) + # cfg = pred_scores['cfg'] + # assert cfg.DATA.DATA_NAME in ['EK55', 'EK100'] + # perframe_gt_targets = pred_scores['perframe_gt_targets'] + # perframe_pred_scores = pred_scores['perframe_pred_scores'] + + # for session in perframe_pred_scores: + # if session not in perframe_pred_scores_fusion: + # perframe_pred_scores_fusion[session] = weight * perframe_pred_scores[session] + # else: + # perframe_pred_scores_fusion[session] += weight * perframe_pred_scores[session] + # perframe_gt_targets_fusion[session] = perframe_gt_targets[session] + + # # Compute results + # result = compute_result['perframe']( + # cfg, + # np.concatenate(list(perframe_gt_targets_fusion.values()), axis=0), + # np.concatenate(list(perframe_pred_scores_fusion.values()), axis=0), + # ) + # logging.info('Action detection perframe m{}: {:.5f}'.format( + # cfg.DATA.METRICS, result['mean_AP'] + # )) + + action_scores = None + for weight, pred_scores_file in zip(weights, pred_scores_files): + pred_scores = pkl.load(open(pred_scores_file, 'rb')) + cfg = pred_scores['cfg'] + assert cfg.DATA.DATA_NAME in ['EK55', 'EK100'] + if action_scores is None: + action_scores = weight * pred_scores['action_scores'] + else: + action_scores += weight * pred_scores['action_scores'] + # segment-level evaluation on EK55/EK100 + if cfg.DATA.DATA_NAME == 'EK55': + path_to_data = 'external/rulstm/RULSTM/data/ek55/' + segment_list = pd.read_csv(osp.join(path_to_data, 'validation.csv'), + names=['id', 'video', 'start_f', 'end_f', 'verb', 'noun', 'action'], + skipinitialspace=True) + predictions, labels, ids = [], [], [] + discarded_labels, discarded_ids = [], [] + for idx in segment_list['id']: + segment = segment_list[segment_list['id'] == idx] + start_step = int(np.floor(segment['start_f'].values.astype(int) / 30 * cfg.DATA.FPS)) + vid = segment['video'].values[0].strip() + # preds = (perframe_pred_scores_fusion[vid][start_step][None, 1:, ::-1]).transpose(0, 2, 1) + if start_step >= cfg.MODEL.LSTR.ANTICIPATION_LENGTH: + # predictions.append(preds) + labels.append(segment[['verb', 'noun', 'action']].values.astype(int)) + ids.append(idx) + else: + discarded_labels.append(segment[['verb', 'noun', 'action']].values.astype(int)) + discarded_ids.append(idx) + + actions = pd.read_csv(osp.join(path_to_data, 'actions.csv'), index_col='id') + + # action_scores = np.concatenate(predictions) + labels = np.concatenate(labels) + ids = np.array(ids) + + vi = get_marginal_indexes(actions, 'verb') + ni = get_marginal_indexes(actions, 'noun') + + action_probs = softmax(action_scores.reshape(-1, action_scores.shape[-1])) + # action_probs = action_scores.reshape(-1, action_scores.shape[-1]) + # action_probs = action_probs / action_probs.sum(-1, keepdims=True) + print(action_probs.shape) + verb_scores = marginalize(action_probs, vi).reshape(action_scores.shape[0], + action_scores.shape[1], + -1) + noun_scores = marginalize(action_probs, ni).reshape(action_scores.shape[0], + action_scores.shape[1], + -1) + + include_discarded = (cfg.DATA.DATA_NAME == 'EK100') + if include_discarded: + dlab = np.array(discarded_labels) + dislab = np.array(discarded_ids) + ids = np.concatenate([ids, dislab]) + num_disc = len(dlab) + labels = np.concatenate([labels, dlab]) + verb_scores = np.concatenate((verb_scores, np.zeros((num_disc, *verb_scores.shape[1:])))) + noun_scores = np.concatenate((noun_scores, np.zeros((num_disc, *noun_scores.shape[1:])))) + action_scores = np.concatenate((action_scores, np.zeros((num_disc, *action_scores.shape[1:])))) + + verb_labels, noun_labels, action_labels = labels[:, 0], labels[:, 1], labels[:, 2] + verb_accuracies = topk_accuracy_multiple_timesteps(verb_scores, verb_labels) + noun_accuracies = topk_accuracy_multiple_timesteps(noun_scores, noun_labels) + action_accuracies = topk_accuracy_multiple_timesteps(action_scores, action_labels) + + many_shot_verbs = pd.read_csv(osp.join(path_to_data, 'EPIC_many_shot_verbs.csv'))['verb_class'].values + many_shot_nouns = pd.read_csv(osp.join(path_to_data, 'EPIC_many_shot_nouns.csv'))['noun_class'].values + actions = pd.read_csv(osp.join(path_to_data, 'actions.csv')) + a_to_vn = {a[1]['id']: tuple(a[1][['verb', 'noun']].values) + for a in actions.iterrows()} + many_shot_actions = [a for a, (v, n) in a_to_vn.items() + if v in many_shot_verbs or n in many_shot_nouns] + + verb_recalls = topk_recall_multiple_timesteps(verb_scores, verb_labels, + k=5, classes=many_shot_verbs) + noun_recalls = topk_recall_multiple_timesteps(noun_scores, noun_labels, + k=5, classes=many_shot_nouns) + action_recalls = topk_recall_multiple_timesteps(action_scores, action_labels, + k=5, classes=many_shot_actions) + + all_accuracies = np.concatenate([verb_accuracies, noun_accuracies, action_accuracies, + verb_recalls, noun_recalls, action_recalls]) + all_accuracies = all_accuracies[[0, 1, 6, 2, 3, 7, 4, 5, 8]] + indices = [ + ('Verb', 'Top-1 Accuracy'), ('Verb', 'Top-5 Accuracy'), + ('Verb', 'Mean Top-5 Recall'), + ('Noun', 'Top-1 Accuracy'), ('Noun', 'Top-5 Accuracy'), + ('Noun', 'Mean Top-5 Recall'), + ('Action', 'Top-1 Accuracy'), ('Action', 'Top-5 Accuracy'), + ('Action', 'Mean Top-5 Recall'), + ] + cc = np.linspace(1 / cfg.DATA.FPS * cfg.MODEL.LSTR.ANTICIPATION_LENGTH, + 1 / cfg.DATA.FPS, cfg.MODEL.LSTR.ANTICIPATION_LENGTH, dtype=str) + scores = pd.DataFrame(all_accuracies * 100, columns=cc, + index=pd.MultiIndex.from_tuples(indices)) + print(scores) + + tta_verb = tta(verb_scores, verb_labels) + tta_noun = tta(noun_scores, noun_labels) + tta_action = tta(action_scores, action_labels) + print(f'\nMean TtA(5): VERB {tta_verb:0.2f} ' + f'NOUN: {tta_noun:0.2f} ACTION: {tta_action:0.2f}') + else: + raise NotImplementedError + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--pred_scores_files', nargs='+', required=True) + parser.add_argument('--weights', type=float, nargs='+', required=True) + args = parser.parse_args() + + assert len(args.pred_scores_files) == len(args.weights) + eval_fusion(args.pred_scores_files, args.weights) diff --git a/third_party/EgoSpeak/LSTR/tools/eval/eval_perframe.py b/third_party/EgoSpeak/LSTR/tools/eval/eval_perframe.py new file mode 100644 index 0000000..9078ecb --- /dev/null +++ b/third_party/EgoSpeak/LSTR/tools/eval/eval_perframe.py @@ -0,0 +1,36 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import logging +logging.basicConfig(level=logging.INFO, format='%(message)s') +import argparse + +import numpy as np +import pickle as pkl + +from rekognition_online_action_detection.evaluation import compute_result + + +def eval_perframe(pred_scores_file): + pred_scores = pkl.load(open(pred_scores_file, 'rb')) + cfg = pred_scores['cfg'] + perframe_gt_targets = pred_scores['perframe_gt_targets'] + perframe_pred_scores = pred_scores['perframe_pred_scores'] + + # Compute results + result = compute_result['perframe']( + cfg, + np.concatenate(list(perframe_gt_targets.values()), axis=0), + np.concatenate(list(perframe_pred_scores.values()), axis=0), + ) + logging.info('Action detection perframe m{}: {:.5f}'.format( + cfg.DATA.METRICS, result['mean_AP'] + )) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--pred_scores_file', type=str, required=True) + args = parser.parse_args() + + eval_perframe(args.pred_scores_file) diff --git a/third_party/EgoSpeak/LSTR/tools/eval/eval_perstage.py b/third_party/EgoSpeak/LSTR/tools/eval/eval_perstage.py new file mode 100644 index 0000000..0b0572b --- /dev/null +++ b/third_party/EgoSpeak/LSTR/tools/eval/eval_perstage.py @@ -0,0 +1,37 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import logging +logging.basicConfig(level=logging.INFO, format='%(message)s') +import argparse + +import numpy as np +import pickle as pkl + +from rekognition_online_action_detection.evaluation import compute_result + + +def eval_perstage(pred_scores_file): + pred_scores = pkl.load(open(pred_scores_file, 'rb')) + cfg = pred_scores['cfg'] + perframe_gt_targets = pred_scores['perframe_gt_targets'] + perframe_pred_scores = pred_scores['perframe_pred_scores'] + + # Compute results + result = compute_result['perstage']( + cfg, + np.concatenate(list(perframe_gt_targets.values()), axis=0), + np.concatenate(list(perframe_pred_scores.values()), axis=0), + ) + for stage_name in result: + logging.info('Perframe m{} of stage {}: {:.5f}'.format( + cfg.DATA.METRICS, stage_name, result[stage_name]['mean_AP'] + )) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--pred_scores_file', type=str, required=True) + args = parser.parse_args() + + eval_perstage(args.pred_scores_file) diff --git a/third_party/EgoSpeak/LSTR/tools/generate_targets_epic_kitchen.py b/third_party/EgoSpeak/LSTR/tools/generate_targets_epic_kitchen.py new file mode 100644 index 0000000..2e58222 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/tools/generate_targets_epic_kitchen.py @@ -0,0 +1,80 @@ +import argparse +import numpy as np +import pandas as pd +import os.path as osp + + +def str2sec(string): + hh, mm, ss = string.split(':') + return float(hh) * 3600 + float(mm) * 60 + float(ss) + + +def main(): + parser = argparse.ArgumentParser(description='Generate targets') + parser.add_argument('--dataset', type=str, + choices=['EK55', 'EK100'], default='EK55') + parser.add_argument('--data-path', type=str, + default='/scratch/cluster/yzhao/Datasets/EPIC-KITCHENS/') + parser.add_argument('--targets-path', type=str, + default='/scratch/cluster/yzhao/Datasets/' + 'EPIC-KITCHENS/features/target_perframe/') + parser.add_argument('--fps', type=int, default=30) + parser.add_argument('--tick', type=float, default=7.5) + parser.add_argument('--anno-path', type=str, + default='external/rulstm/RULSTM/data/ek55/') + parser.add_argument('--class-type', type=str, default='action', + choices=['verb', 'noun', 'action']) + args = parser.parse_args() + + # initialize the targets + targets = {} + if args.class_type == 'action': + num_classes = 2513 + 1 + elif args.class_type == 'verb': + actions = pd.read_csv(osp.join(args.anno_path, 'actions.csv')) + a_to_v = {a[1]['id'] + 1: a[1]['verb'] + 1 + for a in actions.iterrows()} + num_classes = max(a_to_v.values()) + 1 + elif args.class_type == 'noun': + actions = pd.read_csv(osp.join(args.anno_path, 'actions.csv')) + a_to_n = {a[1]['id'] + 1: a[1]['noun'] + 1 + for a in actions.iterrows()} + num_classes = max(a_to_n.values()) + 1 + + for subset in ['training', 'validation']: + vids = [line.strip() for line in open('{}/{}_videos.csv'.format(args.anno_path, subset))] + anno_per_video = {} + for line in open('{}/{}.csv'.format(args.anno_path, subset)): + idx, vid, f_start, f_end, c_verb, c_noun, c_verbnoun = line.strip().split(',') + vid = vid.strip() + if vid not in anno_per_video: + anno_per_video[vid] = [(int(f_start), int(f_end), int(c_verb), int(c_noun), int(c_verbnoun))] + else: + anno_per_video[vid].append((int(f_start), int(f_end), int(c_verb), int(c_noun), int(c_verbnoun))) + for vid in vids: + feature = np.load('{}/features/rgb_kinetics_resnet50/{}.npy'.format(args.data_path, vid)) + feature_len = feature.shape[0] + targets[vid] = np.zeros((feature_len, + num_classes)) + targets[vid][:, 0] = 1 + for anno in anno_per_video[vid]: + print(vid, anno) + ss = int(np.floor(anno[0] / args.tick)) + ee = min(int(np.ceil(anno[1] / args.tick)), feature_len) + if args.class_type == 'action': + ci = anno[4] + elif args.class_type == 'verb': + ci = anno[2] + elif args.class_type == 'noun': + ci = anno[3] + for si in range(ss, ee): + targets[vid][si, ci + 1] = 1 + targets[vid][si, 0] = 0 + + print('Dumping npy files..') + for vid in targets: + np.save(args.targets_path + '{}.npy'.format(vid), targets[vid]) + + +if __name__ == '__main__': + main() diff --git a/third_party/EgoSpeak/LSTR/tools/test_net.py b/third_party/EgoSpeak/LSTR/tools/test_net.py new file mode 100644 index 0000000..29a1815 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/tools/test_net.py @@ -0,0 +1,45 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import os + +from rekognition_online_action_detection.engines import do_inference +from rekognition_online_action_detection.models import build_model +from rekognition_online_action_detection.utils.checkpointer import \ + setup_checkpointer +from rekognition_online_action_detection.utils.env import setup_environment +from rekognition_online_action_detection.utils.logger import setup_logger +from rekognition_online_action_detection.utils.parser import load_cfg +from utils import validate_cfg + + +def main(cfg): + validate_cfg(cfg) + config_name = os.path.basename(cfg.OUTPUT_DIR) + + print(f"cfg.MODEL.CHECKPOINT: {cfg.MODEL.CHECKPOINT}") + print(f"cfg.OUTPUT_DIR: {cfg.OUTPUT_DIR}") + + # Setup configurations + device = setup_environment(cfg) + checkpointer = setup_checkpointer(cfg, phase='test') + logger = setup_logger(cfg, phase='test') + + print(cfg) + print(f"DATA.IGNORE_INDEX: {cfg.DATA.IGNORE_INDEX}") + # Build model + model = build_model(cfg, device) + + # Load pretrained model + checkpointer.load(model) + + do_inference( + cfg, + model, + device, + logger, + ) + + +if __name__ == '__main__': + main(load_cfg()) diff --git a/third_party/EgoSpeak/LSTR/tools/train_net.py b/third_party/EgoSpeak/LSTR/tools/train_net.py new file mode 100644 index 0000000..12ce01b --- /dev/null +++ b/third_party/EgoSpeak/LSTR/tools/train_net.py @@ -0,0 +1,67 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import os +import time + +import wandb +from rekognition_online_action_detection.criterions import build_criterion +from rekognition_online_action_detection.datasets import build_data_loader +from rekognition_online_action_detection.engines import do_train +from rekognition_online_action_detection.models import build_model +from rekognition_online_action_detection.optimizers import (build_optimizer, + build_scheduler) +from rekognition_online_action_detection.utils.checkpointer import \ + setup_checkpointer +from rekognition_online_action_detection.utils.env import setup_environment +from rekognition_online_action_detection.utils.logger import setup_logger +from rekognition_online_action_detection.utils.parser import load_cfg +from utils import parse_config + + +def main(cfg): + parse_config(cfg) + # Setup configurations + device = setup_environment(cfg) + checkpointer = setup_checkpointer(cfg, phase="train") + logger = setup_logger(cfg, phase="train") + + print(cfg) + print(f"cfg['SOLVER']['NUM_EPOCHS']: {cfg['SOLVER']['NUM_EPOCHS']}") + print(f"cfg['MODEL']['CHECKPOINT']: {cfg['MODEL']['CHECKPOINT']}") + + # Build data loaders + data_loaders = {phase: build_data_loader(cfg, phase) for phase in cfg.SOLVER.PHASES} + + # Build model + model = build_model(cfg, device) + + # Build criterion + criterion = build_criterion(cfg, device) + + # Build optimizer + optimizer = build_optimizer(cfg, model) + + # Load pretrained model and optimizer + checkpointer.load(model, optimizer) + + # Build scheduler + scheduler = build_scheduler(cfg, optimizer, len(data_loaders["train"])) + + do_train( + cfg, + data_loaders, + model, + criterion, + optimizer, + scheduler, + device, + checkpointer, + logger, + ) + + wandb.finish() if cfg.USE_WANDB else None + + +if __name__ == "__main__": + main(load_cfg()) diff --git a/third_party/EgoSpeak/LSTR/tools/utils.py b/third_party/EgoSpeak/LSTR/tools/utils.py new file mode 100644 index 0000000..87efcd2 --- /dev/null +++ b/third_party/EgoSpeak/LSTR/tools/utils.py @@ -0,0 +1,97 @@ +import os + +import wandb + + +def parse_config(cfg): + modality = get_modality(cfg) + + cfg_filename = os.path.basename(cfg.OUTPUT_DIR) + + os.environ["WANDB__SERVICE_WAIT"] = "3000" + cfg_filename_abbr = cfg_filename.split('_1x_')[-1] + wandb_postfix = generate_wandb_postfix(cfg, modality) + scheduler_name = cfg.SOLVER.SCHEDULER.SCHEDULER_NAME + learning_rate = cfg.SOLVER.BASE_LR + weight_decay = cfg.SOLVER.WEIGHT_DECAY + + cfg_filename = cfg_filename.replace('long_512', f'long_{cfg.MODEL.LSTR.LONG_MEMORY_SECONDS}') + cfg_filename = cfg_filename.replace('work_8', f'work_{cfg.MODEL.LSTR.WORK_MEMORY_SECONDS}') + + cfg.OUTPUT_DIR = os.path.join(os.path.dirname(cfg.OUTPUT_DIR), cfg_filename) + + if cfg.SAVE_HPARAM: + # save different dir containing hparam + hparam_info = '_'.join([f"lr_{learning_rate}", f"wd_{weight_decay}", f"schd_{scheduler_name}"]) + cfg.OUTPUT_DIR = os.path.join(cfg.OUTPUT_DIR, hparam_info, modality) + print(f"HPARAM: {hparam_info}") + print(f"cfg.OUTPUT_DIR: {cfg.OUTPUT_DIR}") + + wandb_name = cfg_filename_abbr + '_' + wandb_postfix + + if cfg.SAVE_HPARAM: + wandb_name += '_' + '_'.join([f"lr_{learning_rate}", f"wd_{weight_decay}", f"schd_{scheduler_name}"]) + + if cfg.USE_WANDB: + print("*** init wandb ***") + print(f"wandb_name: {wandb_name}") + +def get_modality(cfg): + modality = '' + if 'motion' in cfg.INPUT.MODALITY: + modality += 'A' + if 'visual' in cfg.INPUT.MODALITY: + modality += 'V' + if 'object' in cfg.INPUT.MODALITY: + modality += 'T' + elif cfg['INPUT']['MODALITY'] == 'twostream': + modality = 'AV' + + if modality == '': + raise ValueError(f"Unknown modality: {cfg['INPUT']['MODALITY']}") + + return modality + + +def generate_wandb_postfix(cfg, modality): + visual_feature_type = '_'.join(cfg['INPUT']['VISUAL_FEATURE'].split('_')[1:]) + audio_feature_type = '_'.join(cfg['INPUT']['MOTION_FEATURE'].split('_')[1:]) + text_feature_type = '_'.join(cfg['INPUT']['OBJECT_FEATURE'].split('_')[1:]) + + if cfg.MODEL.LSTR.LONG_MEMORY_USE_PE == True: + model_name = 'LSTR_TESTRA' + else: + model_name = 'TESTRA' + + if '_aa' in cfg.OUTPUT_DIR: + model_name += '_aa' + + name_parts = [ + model_name, + cfg['DATA']['DATA_NAME'], + modality, + # visual_feature_type, + # audio_feature_type, + # text_feature_type, + # f"epoch_{cfg['SOLVER']['NUM_EPOCHS']}" + ] + if cfg['MODEL']['CHECKPOINT']: + name_parts.append(f"ckpt_{os.path.basename(cfg['MODEL']['CHECKPOINT']).split('.')[0]}") + return '_'.join(name_parts) + +def get_wandb_name(cfg, cfg_filename, modality): + cfg_filename_abbr = cfg_filename.split('_1x_')[-1] + wandb_postfix = generate_wandb_postfix(cfg, modality) + scheduler_name = cfg.SOLVER.SCHEDULER.SCHEDULER_NAME + learning_rate = cfg.SOLVER.BASE_LR + weight_decay = cfg.SOLVER.WEIGHT_DECAY + wandb_name = cfg_filename_abbr + '_' + wandb_postfix + '_' + '_'.join([f"lr_{learning_rate}", f"wd_{weight_decay}", f"schd_{scheduler_name}"]) + return wandb_name + +def validate_cfg(cfg): + if cfg.DATA.NUM_CLASSES != len(cfg.DATA.CLASS_NAMES): + print(f"len(cfg.DATA.CLASS_NAMES): {len(cfg.DATA.CLASS_NAMES)} is not equal to cfg.DATA.NUM_CLASSES: {cfg.DATA.NUM_CLASSES}") + cfg.DATA.CLASS_NAMES = cfg.DATA.CLASS_NAMES[:cfg.DATA.NUM_CLASSES] + print(f"Adjusted cfg.DATA.CLASS_NAMES: {cfg.DATA.CLASS_NAMES}") + + \ No newline at end of file diff --git a/third_party/EgoSpeak/MiniROAD/.gitignore b/third_party/EgoSpeak/MiniROAD/.gitignore new file mode 100644 index 0000000..1aa7129 --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/.gitignore @@ -0,0 +1,27 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] + +# PyTorch models +*.pth + +# Miscellaneous +*.DS_Store +*.swp +*.pyc +*.pkl + +*.out + +dataset/ +# Dataset soft links +data/ + +# Output files +output/ +output_*/ + +checkpoints/ + +notebooks/ +wandb/ \ No newline at end of file diff --git a/third_party/EgoSpeak/MiniROAD/configs/EasyCom/mamba_aa_easycom.yaml b/third_party/EgoSpeak/MiniROAD/configs/EasyCom/mamba_aa_easycom.yaml new file mode 100755 index 0000000..eda5be4 --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/configs/EasyCom/mamba_aa_easycom.yaml @@ -0,0 +1,32 @@ +model: 'MambaA' +data_name: 'EasyCom_ANTICIPATION' +task: 'ANTICIPATION' +loss: 'ANTICIPATION' # NONUNIFORM, originally +metric: 'AP' # AP or cAP +optimizer: 'AdamW' +device: 'cuda' +feature_pretrained: 'resnet50' +root_path: './data/EasyCom' +rgb_type: 'rgb_kinetics_resnet50' +flow_type: 'audio_wav2vec2' +annotation_type: 'target_perframe' +video_list_path: 'data_info/data_info.json' +output_path: './output_cfg' +window_size: 128 +batch_size: 64 +test_batch_size: 1 +# loss_offset: 4 # loss weight a_t when L - {loss_offset} <= t <= L +num_epoch: 30 +lr: 0.0001 +weight_decay: 0.05 +num_workers: 4 +dropout: 0.20 +num_classes: 3 +embedding_dim: 2048 +hidden_dim: 1024 +num_layers: 1 +stride: 4 +anticipation_length: 10 +actionness: False +ignore_index: -100 +seed: 20 diff --git a/third_party/EgoSpeak/MiniROAD/configs/EasyCom/miniroad_aa_easycom.yaml b/third_party/EgoSpeak/MiniROAD/configs/EasyCom/miniroad_aa_easycom.yaml new file mode 100755 index 0000000..2e7a9d0 --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/configs/EasyCom/miniroad_aa_easycom.yaml @@ -0,0 +1,32 @@ +model: 'MiniROADA' +data_name: 'EasyCom_ANTICIPATION' +task: 'ANTICIPATION' +loss: 'ANTICIPATION' # NONUNIFORM, originally +metric: 'AP' # AP or cAP +optimizer: 'AdamW' +device: 'cuda' +feature_pretrained: 'resnet50' +root_path: './data/EasyCom' +rgb_type: 'rgb_kinetics_resnet50' +flow_type: 'audio_wav2vec2' +annotation_type: 'target_perframe' +video_list_path: 'data_info/data_info.json' +output_path: './output_cfg' +window_size: 128 +batch_size: 64 +test_batch_size: 1 +# loss_offset: 4 # loss weight a_t when L - {loss_offset} <= t <= L +num_epoch: 30 +lr: 0.0001 +weight_decay: 0.05 +num_workers: 4 +dropout: 0.20 +num_classes: 3 # [background, target speaker speaking, other person speaking] +embedding_dim: 2048 +hidden_dim: 1024 +num_layers: 1 +stride: 4 +anticipation_length: 10 +actionness: False +ignore_index: -100 +seed: 20 \ No newline at end of file diff --git a/third_party/EgoSpeak/MiniROAD/configs/YTConv/mamba_aa_ytconv.yaml b/third_party/EgoSpeak/MiniROAD/configs/YTConv/mamba_aa_ytconv.yaml new file mode 100755 index 0000000..4df5bfb --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/configs/YTConv/mamba_aa_ytconv.yaml @@ -0,0 +1,32 @@ +model: 'MambaA' +data_name: 'YTConv_ANTICIPATION' +data_info_name: 'shard_0' +task: 'ANTICIPATION' +loss: 'ANTICIPATION' # NONUNIFORM, originally +metric: 'AP' # AP or cAP +optimizer: 'AdamW' +device: 'cuda' +feature_pretrained: 'resnet50' +root_path: './data/YTConv/shard_0' +rgb_type: 'rgb_kinetics_resnet50' +flow_type: 'audio_wav2vec2' # or 'none'? +annotation_type: 'target_perframe/case1' +video_list_path: 'data_info/YTConv_info.json' +output_path: './output_cfg' +window_size: 128 +batch_size: 32 +test_batch_size: 1 +# loss_offset: 4 # loss weight a_t when L - {loss_offset} <= t <= L +num_epoch: 30 +lr: 0.0001 +weight_decay: 0.05 +num_workers: 4 +dropout: 0.20 +num_classes: 2 # for YTConv pretraining, there is no +embedding_dim: 2048 +hidden_dim: 1024 +num_layers: 1 +stride: 4 +anticipation_length: 10 +actionness: False +ignore_index: -100 \ No newline at end of file diff --git a/third_party/EgoSpeak/MiniROAD/configs/YTConv/miniroad_aa_ytconv.yaml b/third_party/EgoSpeak/MiniROAD/configs/YTConv/miniroad_aa_ytconv.yaml new file mode 100644 index 0000000..bcbb0a0 --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/configs/YTConv/miniroad_aa_ytconv.yaml @@ -0,0 +1,32 @@ +model: 'MiniROADA' +data_name: 'YTConv_ANTICIPATION' +data_info_name: 'shard_0' +task: 'ANTICIPATION' +loss: 'ANTICIPATION' # NONUNIFORM, originally +metric: 'AP' # AP or cAP +optimizer: 'AdamW' +device: 'cuda' +feature_pretrained: 'resnet50' +root_path: './data/YTConv/shard_0' +rgb_type: 'rgb_kinetics_resnet50' +flow_type: 'audio_wav2vec2' +annotation_type: 'target_perframe' +video_list_path: 'data_info/YTConv_info.json' +output_path: './output_cfg' +window_size: 128 +batch_size: 64 +test_batch_size: 1 +# loss_offset: 4 # loss weight a_t when L - {loss_offset} <= t <= L +num_epoch: 30 +lr: 0.0001 +weight_decay: 0.05 +num_workers: 4 +dropout: 0.20 +num_classes: 2 # for pretraining, there is no target speaker +embedding_dim: 2048 +hidden_dim: 1024 +num_layers: 1 +stride: 4 +anticipation_length: 10 +actionness: False +ignore_index: -100 \ No newline at end of file diff --git a/third_party/EgoSpeak/MiniROAD/configs/__init__.py b/third_party/EgoSpeak/MiniROAD/configs/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/third_party/EgoSpeak/MiniROAD/configs/ego4d/mamba_aa_ego4dshuffle.yaml b/third_party/EgoSpeak/MiniROAD/configs/ego4d/mamba_aa_ego4dshuffle.yaml new file mode 100755 index 0000000..22e083d --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/configs/ego4d/mamba_aa_ego4dshuffle.yaml @@ -0,0 +1,36 @@ +model: 'MambaA' +data_name: 'EGO4D_ANTICIPATION' +data_info_name: 'EGO4DSHUFFLE' +task: 'ANTICIPATION' +loss: 'ANTICIPATION' # NONUNIFORM, originally +metric: 'AP' # AP or cAP +optimizer: 'AdamW' +device: 'cuda' +feature_pretrained: 'resnet50' +root_path: './data/EGO4D' +rgb_type: 'rgb_kinetics_resnet50' +flow_type: 'audio_wav2vec2' +annotation_type: 'target_perframe' +video_list_path: 'data_info/data_info.json' +output_path: './output_cfg' +window_size: 128 +batch_size: 64 +test_batch_size: 1 +# loss_offset: 4 # loss weight a_t when L - {loss_offset} <= t <= L +num_epoch: 30 +lr: 0.0001 +weight_decay: 0.05 +num_workers: 4 +dropout: 0.20 +num_classes: 3 # including background +embedding_dim: 2048 +hidden_dim: 1024 +d_state: 16 +d_conv: 4 +expand: 2 +num_layers: 1 +stride: 4 +anticipation_length: 10 +actionness: False +ignore_index: -100 +seed: 20 \ No newline at end of file diff --git a/third_party/EgoSpeak/MiniROAD/configs/ego4d/miniroad_aa_ego4dshuffle.yaml b/third_party/EgoSpeak/MiniROAD/configs/ego4d/miniroad_aa_ego4dshuffle.yaml new file mode 100755 index 0000000..b171251 --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/configs/ego4d/miniroad_aa_ego4dshuffle.yaml @@ -0,0 +1,32 @@ +model: 'MiniROADA' +data_name: 'EGO4D_ANTICIPATION' +data_info_name: 'EGO4DSHUFFLE' +task: 'ANTICIPATION' +loss: 'ANTICIPATION' # NONUNIFORM, originally +metric: 'AP' # AP or cAP +optimizer: 'AdamW' +device: 'cuda' +feature_pretrained: 'resnet50' +root_path: './data/EGO4D' +rgb_type: 'rgb_kinetics_resnet50' +flow_type: 'audio_wav2vec2' # or 'none'? +annotation_type: 'target_perframe' +video_list_path: 'data_info/data_info.json' +output_path: './output_cfg' +window_size: 128 +batch_size: 64 +test_batch_size: 1 +# loss_offset: 4 # loss weight a_t when L - {loss_offset} <= t <= L +num_epoch: 30 +lr: 0.0001 +weight_decay: 0.05 +num_workers: 4 +dropout: 0.20 +num_classes: 3 +embedding_dim: 2048 +hidden_dim: 1024 +num_layers: 1 +stride: 4 +anticipation_length: 10 +actionness: False +ignore_index: -100 \ No newline at end of file diff --git a/third_party/EgoSpeak/MiniROAD/configs/feature_sizes.py b/third_party/EgoSpeak/MiniROAD/configs/feature_sizes.py new file mode 100644 index 0000000..abe8bc4 --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/configs/feature_sizes.py @@ -0,0 +1,26 @@ +# dimension of features extracted from different models +# used in feature_head, dataset, models + +FEATURE_SIZES = { + 'rgb_anet_resnet50': 2048, + 'flow_anet_resnet50': 2048, + 'rgb_kinetics_bninception': 1024, + 'flow_kinetics_bninception': 1024, + 'rgb_farl': 512, + 'rgb_kinetics_resnet50': 2048, + 'rgb_kinetics_resnet50_debug': 2048, + 'rgb_kinetics_resnet50_debug_resized': 2048, + 'flow_kinetics_resnet50': 2048, + 'flow_nv_kinetics_bninception': 1024, + 'rgb_kinetics_i3d': 2048, + 'flow_kinetics_i3d': 2048, + 'audio_wav2vec2': 5120, + 'audio_wav2vec2_left': 5120, + 'audio_wav2vec2_right': 5120, + 'text_word2vec': 300, + 'text_fasttext_crawl': 300, + 'text_gpt2_last_subword_window_size_0': 768, + 'text_gpt2_last_subword_window_size_20': 768, + 'text_gpt2_last_subword_window_size_50': 768, + 'text_gpt2_last_subword_window_size_100': 768, +} \ No newline at end of file diff --git a/third_party/EgoSpeak/MiniROAD/criterions/__init__.py b/third_party/EgoSpeak/MiniROAD/criterions/__init__.py new file mode 100755 index 0000000..5fe67eb --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/criterions/__init__.py @@ -0,0 +1,2 @@ +from .loss import OadLoss +from .loss_builder import build_criterion \ No newline at end of file diff --git a/third_party/EgoSpeak/MiniROAD/criterions/loss.py b/third_party/EgoSpeak/MiniROAD/criterions/loss.py new file mode 100755 index 0000000..5399a99 --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/criterions/loss.py @@ -0,0 +1,205 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from criterions.loss_builder import CRITERIONS + + +@CRITERIONS.register('NONUNIFORM') +class OadLoss(nn.Module): + + def __init__(self, cfg, reduction='mean'): + super(OadLoss, self).__init__() + self.reduction = reduction + self.num_classes = cfg['num_classes'] + self.loss_offset = cfg['loss_offset'] if 'loss_offset' in cfg else 0 + assert self.loss_offset==0, f'loss_offset is not 0: {self.loss_offset}' + self.loss = self.end_loss + + def end_loss(self, out_dict, target): + # logits: (B, seq, K) target: (B, seq, K) + assert self.loss_offset < out_dict['logits'].shape[1], f'loss_offset is longer than the sequence loss_offset: {self.loss_offset}, logits.shape: {out_dict["logits"].shape}' + logits = out_dict['logits'] + logits = logits[:,-1 - self.loss_offset,:].contiguous() + target = target[:,-1 - self.loss_offset,:].contiguous() + ce_loss = self.mlce_loss(logits, target) + return ce_loss + + def mlce_loss(self, logits, target): + ''' + multi label cross entropy loss. + logits: (B, K) target: (B, K) + ''' + logsoftmax = nn.LogSoftmax(dim=-1).to(logits.device) + output = torch.sum(-F.normalize(target) * logsoftmax(logits), dim=1) # B + if self.reduction == 'mean': + loss = torch.mean(output) + elif self.reduction == 'sum': + loss = torch.sum(output) + return loss + + def forward(self, out_dict, target): + return self.loss(out_dict, target) + + +@CRITERIONS.register('UNIFORM') +class OadUniformLoss(nn.Module): + + def __init__(self, cfg, reduction='mean'): + super(OadUniformLoss, self).__init__() + self.reduction = reduction + self.num_classes = cfg['num_classes'] + # self.loss_offset = cfg['loss_offset'] if 'loss_offset' in cfg else 0 + self.loss = self.end_loss + + def end_loss(self, out_dict, target): + # logits: (B, seq, K) target: (B, seq, K) + logits = out_dict['logits'] + # logits = logits[:,-1 - self.loss_offset,:].contiguous() + # target = target[:,-1 - self.loss_offset,:].contiguous() + ce_loss = self.mlce_loss(logits, target) + return ce_loss + + def mlce_loss(self, logits, target): + ''' + multi label cross entropy loss. + logits: (B, K) target: (B, K) + ''' + logsoftmax = nn.LogSoftmax(dim=-1).to(logits.device) + output = torch.sum(-target * logsoftmax(logits), dim=1) # B + # output = torch.sum(-F.normalize(target) * logsoftmax(logits), dim=1) # B + if self.reduction == 'mean': + loss = torch.mean(output) + elif self.reduction == 'sum': + loss = torch.sum(output) + return loss + + def forward(self, out_dict, target): + return self.loss(out_dict, target) + +@CRITERIONS.register('OADTR') +class OadTRLoss(nn.Module): + + def __init__(self, cfg, reduction='mean'): + super(OadTRLoss, self).__init__() + self.reduction = reduction + self.num_classes = cfg['num_classes'] + # self.loss_offset = cfg['loss_offset'] if 'loss_offset' in cfg else 0 + self.loss = self.end_loss + + def end_loss(self, out_dict, target): + # logits: (B, seq, K) target: (B, seq, K) + assert self.loss_offset < out_dict['logits'].shape[1], f'loss_offset is longer than the sequence loss_offset: {self.loss_offset}, logits.shape: {out_dict["logits"].shape}' + logits = out_dict['logits'] + # logits = logits[:,-1 - self.loss_offset,:].contiguous() + # target = target[:,-1 - self.loss_offset,:].contiguous() + ce_loss = self.mlce_loss(logits, target) + return ce_loss + + def mlce_loss(self, logits, target): + ''' + multi label cross entropy loss. + logits: (B, K) target: (B, K) + ''' + logsoftmax = nn.LogSoftmax(dim=-1).to(logits.device) + output = torch.sum(-F.normalize(target) * logsoftmax(logits), dim=1) # B + if self.reduction == 'mean': + loss = torch.mean(output) + elif self.reduction == 'sum': + loss = torch.sum(output) + return loss + + def forward(self, out_dict, enc_target, dec_target): + return self.loss(out_dict, enc_target, dec_target) + +@CRITERIONS.register('ANTICIPATION') +class OadAntLoss(nn.Module): + + # reduction was sum, not mean, what if we change to sum? + def __init__(self, cfg, reduction='mean'): + super(OadAntLoss, self).__init__() + self.reduction = reduction + self.loss = self.anticipation_loss + self.num_classes = cfg['num_classes'] + + def anticipation_loss(self, out_dict, target, ant_target): + # anticipation_logits: (B, seq, ANTICIPATION_LENGTH, num_class) + # ant_target: (B, ANTICIPATION_LENGTH, num_class) + anticipation_logits = out_dict['anticipation_logits'] + # [B, window_size, anticipation_length, num_class] + pred_anticipation_logits = anticipation_logits[:,-1,:,:].contiguous().view(-1, self.num_classes) + + anticipation_logit_targets = ant_target.view(-1, self.num_classes) + ant_loss = self.mlce_loss(pred_anticipation_logits, anticipation_logit_targets) + + return ant_loss + + def ce_loss(self, out_dict, target): + # logits: (B, seq, K) target: (B, seq, K) + logits = out_dict['logits'] + logits = logits[:,-1,:].contiguous() + target = target[:,-1,:].contiguous() + ce_loss = self.mlce_loss(logits, target) + return ce_loss + + def mlce_loss(self, logits, target): + ''' + multi label cross entropy loss. + logits: (B, K) target: (B, K) + ''' + logsoftmax = nn.LogSoftmax(dim=-1).to(logits.device) + output = torch.sum(-F.normalize(target) * logsoftmax(logits), dim=1) # B + if self.reduction == 'mean': + loss = torch.mean(output) + elif self.reduction == 'sum': + loss = torch.sum(output) + + return loss + + def forward(self, out_dict, target, ant_target): + return self.loss(out_dict, target, ant_target) + + + +@CRITERIONS.register('UNIFORM_ANTICIPATION') +class OadUniformAntLoss(nn.Module): + + # reduction was sum, not mean, what if we change to sum? + def __init__(self, cfg, reduction='mean'): + super(OadUniformAntLoss, self).__init__() + self.reduction = reduction + self.loss = self.anticipation_loss + self.num_classes = cfg['num_classes'] + + def anticipation_loss(self, out_dict, target, ant_target): + anticipation_logits = out_dict['anticipation_logits'] + pred_anticipation_logits = anticipation_logits[:,-1,:,:].contiguous().view(-1, self.num_classes) + anticipation_logit_targets = ant_target.view(-1, self.num_classes) + ant_loss = self.mlce_loss(pred_anticipation_logits, anticipation_logit_targets) + return ant_loss + + def ce_loss(self, out_dict, target): + # logits: (B, seq, K) target: (B, seq, K) + logits = out_dict['logits'] + logits = logits[:,-1,:].contiguous() + target = target[:,-1,:].contiguous() + ce_loss = self.mlce_loss(logits, target) + return ce_loss + + def mlce_loss(self, logits, target): + ''' + multi label cross entropy loss. + logits: (B, K) target: (B, K) + ''' + logsoftmax = nn.LogSoftmax(dim=-1).to(logits.device) + output = torch.sum(-target * logsoftmax(logits), dim=1) # B + + # output = torch.sum(-F.normalize(target) * logsoftmax(logits), dim=1) # B + if self.reduction == 'mean': + loss = torch.mean(output) + elif self.reduction == 'sum': + loss = torch.sum(output) + + return loss + + def forward(self, out_dict, target, ant_target): + return self.loss(out_dict, target, ant_target) diff --git a/third_party/EgoSpeak/MiniROAD/criterions/loss_builder.py b/third_party/EgoSpeak/MiniROAD/criterions/loss_builder.py new file mode 100755 index 0000000..826931b --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/criterions/loss_builder.py @@ -0,0 +1,11 @@ +__all__ = [ + 'build_criterion' +] + +from utils import Registry + +CRITERIONS = Registry() + +def build_criterion(cfg, device=None): + criterion = CRITERIONS[cfg["loss"]](cfg) + return criterion.to(device) \ No newline at end of file diff --git a/third_party/EgoSpeak/MiniROAD/data_info/YTConv_info.json b/third_party/EgoSpeak/MiniROAD/data_info/YTConv_info.json new file mode 100644 index 0000000..726091a --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/data_info/YTConv_info.json @@ -0,0 +1,426 @@ +{ + "shard_0": { + "class_index": [ + "background", + "camera wearer start to speak" + ], + "train_session_set": [ + "m5Aoywf51L4", + "p99cIB7nHjw", + "zGgMzBfPfaI", + "O0eCtKU3FxY", + "goZQ6iOIFYo", + "9Tb-lSyhgoE", + "GBxQtGUjAMA", + "VXgt7keOC8I", + "hZ2tin2J0P4", + "Ur-P18sm4tI", + "WfYmsNBQwg8", + "AvMNaDDfrN0", + "fW7pCN3wwdo", + "xG6HBdTtS4s", + "HcDlqUeDD9Q", + "1DE3mKncHUA", + "oEMX_quSMZY", + "jmS3nWbGV5Y", + "qUJRjPRH-u4", + "aHsUgWfUcss", + "ZfdQOmJkaK0", + "4f43nYrfTnI", + "qE1FnzkNjrE", + "4NsmKhs9flk", + "3wDSJ9ABzmo", + "sY6_8QwRZsw", + "hByfAzEoeKA", + "vDEsV0AhcjM", + "XGNhX3Nu1RE", + "MkVT8K0V57A", + "yR8y-wN8Clk", + "P1dTOlO2sl8", + "6UGC_wxTeB0", + "d4Xud2Gqoxk", + "IJ1aqePculQ", + "o8Oxf5QBitQ", + "XVci_V4qxmc", + "ofDzhOS1Pv0", + "mygTXBuMRF0", + "iq6q82Jo8ZQ", + "Kzc_ekBJ4QE", + "WEeJ9B1i8HI", + "rciqpxdCc9Y", + "Rh7d2D7JVLI", + "Qkx4hSlc8jk", + "ixi7sVLFrWM", + "n96-scSyoOQ", + "YZ5xYioIQdA", + "X76QxuWm1oA", + "2I7d_LfK75Y", + "AeUWBaYFMyo", + "lsqXFgymLlI", + "Dx_0jYMWN5o", + "0raJ1-1BJ7s", + "EGX8am5hWh0", + "KjtZ_CLY1hQ", + "R4FQI2BI3hw", + "fAoqqtsNjbQ", + "l6RwyA6jUdM", + "GJW6dlagRnU", + "j_FGvmbp6aM", + "V2O15h1VrH8", + "P3c70VbDP3E", + "epAwpBdUiYU", + "6dpv_JeyL8w", + "dUKYxsvnuBM", + "L2_zn7h8sOQ", + "m09qdCkAhmY", + "7BB71tIae7A", + "6go4oViEkdc", + "MYimxshUK1c", + "v4eZXWcIRQU", + "vGI_w1Ml6S0", + "BeFndnkT00c", + "9JHRbL1jmq4", + "0rBaR5dd-Fs", + "EvhkIS_-K-w", + "pYVbLx-h5HI", + "fVciOXYDlqg", + "zWFH4N-CmDE", + "Aw4cauzfX5E", + "lxVpcQHtcWs", + "32dkF-aXgug", + "_USorgWOsUk", + "1yVMU93nmv8", + "aYkF4Ybk-l8", + "Vwt18p0BqzM", + "oOQM5Pvjut4", + "_WWSMNtNYx0", + "zzlhwg0IJ5Y", + "lJhIffkUhr8", + "PbS9nSH3mqM", + "8EGVgiFpwhg", + "5LhXpv5IFO4", + "ibD4APVzWgg", + "2nhcSQmopO4", + "gri79OfCRvo", + "lOwQN3QcRzE", + "wQyluRNtqFo", + "IZyx8zeiKB0", + "mZyNfKyzDHU", + "A9eGJWiuIKs", + "d37DZogot1w", + "7ZUSVabB6N4", + "m62GQXq-358", + "ArwPVsCZo2c", + "QK-BFkpO8s4", + "PgfX9bvmEhs", + "l1jTUhwWJYA", + "S1PGHO4evaI", + "M6uPJa5B8Ns", + "XXzz19J-1w8", + "U4uV1xmUaB0", + "skQ0OBpkKqY", + "mfW57TDpMtc", + "NAwKA75cDwE", + "eIEBI2_Y14Q", + "eWQGYpSoiYU", + "3DPIUWhYtn8", + "_elvaRJrhE0", + "PAA8IFhWJzM", + "1XDKqhq6bb0", + "ngi9gy3DzRY", + "Y0yONfuimeE", + "g9_cAgHa_4Q", + "BIG7jEw-j4w", + "TFdom0lkd-0", + "2zV2N8_DwWg", + "kTz7XB_jpB4", + "GuoglIuhNjY", + "L98U3mO38gQ", + "UfWDoRAzpC8", + "W6TakaXCIIQ", + "i6j1I6pha-M", + "issHFC8Gbio", + "4EzGoRkd6SI", + "QoS0GgF4sI4", + "A8hJg3D_TGY", + "A3N10suCRCw", + "6962tnWRsNY", + "lS4p8osXSJs", + "wbZ7ZhcZd6o", + "fsIvQvIlKwc", + "4jQSOSi2DA8", + "Vpt6tZBvXQU", + "kxOG3kpLmYw", + "Mzjlz0nq9AA", + "80cuXt2VSko", + "VvswhogSiY0", + "5Q6ky5VrOIA", + "I_XtcUpNRF8", + "QmtmTGycK9A", + "qp3x8iy7vkA", + "xGU8QDtPrQI", + "uWWUfqth_BA", + "WbM7w8r4VGQ", + "bt6Ed1qjulw", + "7vymT27scoY", + "b9AdV6oVDVg", + "nf2zqYUAOwg", + "A-LNJc-1Jc0", + "yCI-hpPiDVc", + "Uu5DME4FqIM", + "symjuH9iOTc", + "mPf2MTnTBc0", + "NDTlQ1KGvH8", + "g9vuDzFQSRU", + "B6iqnvleve0", + "9bbszAD7JL0", + "RIyMAuz7Ckg", + "mlrQHtBRMhc", + "3gaidlW7vig", + "DGdV97q9ndk", + "qthjlg0lXFg", + "gWng5YLPOow", + "iaYG-KL9XRM", + "385414AVZcA", + "7xhBnWGpI7c", + "Nh0cmSQUs-E", + "Rvtj0QHPulA", + "rAtLksYQukw", + "S9rB__zbbe0", + "SNpMVsOPfHM", + "zhF43F4kO2w", + "HOUkcInyxN0", + "x8ymkfr2Z1o", + "Ix6Khtq6O-k", + "VHs_JyPhmtE", + "zJX3byjYSMc", + "zdqX09zZNf8", + "uPwlmwNQq8Y", + "e4xNTXrk6wE", + "P2PNUWN4GMs", + "_WVHCYjHBVY", + "5Jtv1GpPvHM", + "AIZYucKOzhw", + "vbsR_VDZXh0", + "9Ew0lDHB0kM", + "ejpkPVZHuMM", + "47HlvCfubE4", + "HzPkQSIG2m0", + "e7mTRnNnOXg", + "eZHx_biEECg", + "FvP6WFKNP4Y", + "g_TJ5XFEVfY", + "8Utg2VEWsJg", + "VMdBjL-tpoo", + "zihP2Eqeqxk", + "ILVqgOlntLM", + "Z1BWQI8KPCg", + "j8TRVp-7l3Q", + "G5gJBZ-3OlY", + "o1JUlGSirSE", + "OsPdhUQ0snE", + "-5uIaW2gKWQ", + "03kiFNh-9ME", + "w5W6mgDaJJA", + "sYsgd_8AFoo", + "vLg8kVUTmWU", + "OK0J0xsuSWQ", + "d284rjoU4Cs", + "5lE6V4jrUl4", + "bS1OpC9xu48", + "ZUxpqU_J_n0", + "AUliC0c_YIQ", + "nJ58-d_viOo", + "uIy3b2kzlrM", + "vOTLfb0qjfY", + "37_7tPZjC34", + "MUwCdP5qceI", + "lYvNi_ewd8U", + "PDG7SzDBTC4", + "QHQzoHolWvI", + "YmvOEj00LB8", + "fLhkXVx7Yw0", + "Uj5nhugVgwo", + "aHbrDu2o0WY", + "yKW6vVgEPCM", + "Q0WUgyqQ-4I", + "xw3KOPJ9a3E", + "5mrTo7c3bCM", + "x0ft032W--4", + "UdQxEbrIocc", + "EHZHAFyv_Ks", + "7DojKgPRLVk", + "-1CahRD6lJ0", + "f3TuYwl7AZE", + "TgGTovCcS9Y", + "AwyrJGeIN_g", + "RbPJD4Wh-l8", + "58bFdwYkDzU", + "6gIEKGxL7zo", + "BeuzJwTuChQ", + "YebZyAzLZuc", + "9T74b8EQBVY", + "TqSVFjw1uNM", + "WUYhCmfE1a0", + "vRCVUQFUio4", + "w2t9JdMNMgQ", + "nn6-x9xXMeE", + "FDFVPzLlzFs", + "qFhvi4bH29E", + "qdpsMXLsUZM", + "JREk1QeSc9A", + "jOVEs-4r1jc", + "v4VZI9Ouzy0", + "DMg5UUvf-C4", + "26J5zRzQkOg", + "YkiJ4iJfJbw", + "_d8j8nfTEIk", + "wMNb3hQXEUE", + "1KI8ex22blk", + "eCp2Jp7MtKg", + "HNlpHBb5dC4", + "C6_pDRdIo1k", + "ZRWJJHrgP54", + "L7jfx4veKIA", + "p1a5_qltICM", + "nZ0svM8DoY0", + "miEzjCJVaZk", + "TpC2tDCxABw", + "SZmqoqCX__U", + "dZ0bYteaI2Q", + "GW-cYEzShmY", + "iUYuPhutKVw", + "97VSm6T2zMw", + "jPiQyi5jr1M", + "lrUi4nsPZ9M", + "rXLeXke-k_w", + "mEO4izUqZ7E", + "4qI1LjMXfh0", + "mVt6myGMa-0", + "VJrnhh6FqEg", + "vwCeBWbZvrE", + "c08fl8-xmVE", + "syZ3bvXIrpA", + "i3ysNaulop4", + "__t3rU0QjsE", + "jT32uQQVGvU", + "Fzvmh-R2EWk", + "QApJSKrLg6E", + "L5y3j8SdGSU", + "3t8JGblR__0", + "e1F90W40t6g", + "j6L8UIa5z28", + "Zmg2_C1FuVI", + "XIXKJXsiBFc", + "NVHB1VdWFo4", + "-x7-8f7bsl4", + "2FSxSiowkGI", + "jgFIBIeFve0", + "1BLuXK1M5CE", + "zmpx6LUT2ow", + "z-QUzJTsFro", + "0sS4v13DAK8", + "gMIBQ0mbRvM", + "9kcaEXlTB5s", + "ZYArPGMLS4k", + "tejW9GMaucA", + "z9r4bmXCWiI", + "4Zpsb46iaLc", + "52Odc8UbvC8", + "x-B6-iYDHPs", + "8GinSDc0Is8", + "5lwbp5d_w0Y", + "I4LpTQLzlHE", + "lMl6-XH7ozY", + "tOubscNbY1Q", + "7gb7AH-oWcQ", + "mDerqt9Ib-I", + "gKV3FNd9WHE", + "kH5L1E_r-3U", + "sq5TRFxVp0w", + "SrBhj7Xusk8", + "RxlxEapsz2E", + "aCD3_TziiKo", + "_HPVZOkxdkk", + "pNCjgVQ53Y8", + "iPdT5GUoGvA", + "cIrlet1oRw4", + "1dyWUMBjMwQ", + "XZQB3cq18yc", + "i75Mn_AeFBY", + "_0VwR9WPS-k", + "q1JDjgnRdGY", + "1JNEpfvRq0Y", + "2FGj6HciQZA", + "eLi59TtNwHY", + "Z3W8OST7tmA", + "F2dNYov8-xw", + "kPfyJu_CrSA", + "7I4Y2HtAwlY", + "us-andJDfWw", + "Q0Bzn1teWfo", + "pADxBfOrYeQ", + "KnA1tIR0qCE", + "lBNPfBcb6pw", + "zt9MIiBfxF8", + "FwX12hPqCws", + "6e2LXGrze2M", + "0dvclhcFjSo", + "UTkuDYOIxKo", + "4n044cOTiNg", + "VxksG15OgRY", + "4193MkWdJcQ", + "2bDX4Uzy8Wc", + "SkyBNX2auEc", + "JrSLspuJ2gk", + "GZO7V-CWSJ8", + "TYMIy4qlRzc", + "jK7q_Li7LU0", + "UiTIXDQEkPM" + ], + "test_session_set": [ + "DTGtC7FC4oI", + "CIWvk3gJ_7E", + "A1fZ6rlOwxk", + "cB8WCrmJEEA", + "gGQ_2Ank9-U", + "b91abcyMeNY", + "9MuaFeyczM0", + "rwZxhdK_UVU", + "hd_WbrqYnTE", + "4BWcqA9cFFk", + "MuOkgCkob6E", + "OdhFWA3Fgb0", + "Pe3ZabrUiUk", + "U0c_uv8DU08", + "kTqt_3YJBbs", + "eWXVVMi4bf8", + "RCD-vXumMIs", + "DuwiezJCZ3I", + "BNgEqjAiKU4", + "Cbn0t7eqaEs", + "D3M9Q0WwuFY", + "oO-Bv4iRZKI", + "bBvqtbZdSZQ", + "-mnoItCi32Y", + "yQPxKfz-bwI", + "ITVZ4cx5big", + "uY5_e0L7mjw", + "J_Pl9ftnNXE", + "Jt_pb40MRy4", + "E3hAj2niRNo", + "kl2xVSxVECQ", + "gzebuFs90Xg", + "6DQSIzJ9J-g", + "pJJlnryTYWg", + "X87qsfVLFoc", + "W7BIPWU1WHc", + "Oj506EVOiQc", + "DPEvWvcstxA", + "Pk16UT5-zeM", + "XMK1KEk5OdU", + "PpNm5llsEW8", + "sAp3rvablGk" + ] + } +} \ No newline at end of file diff --git a/third_party/EgoSpeak/MiniROAD/data_info/data_info.json b/third_party/EgoSpeak/MiniROAD/data_info/data_info.json new file mode 100644 index 0000000..1f077f2 --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/data_info/data_info.json @@ -0,0 +1,1898 @@ +{ + "EGO4D": { + "class_index": [ + "background", + "camera wearer speak", + "other people speaking" + ], + "train_session_set": [ + "fc2b2014-9dc4-4a5d-8a1d-25a6911bff7c", + "cf682058-962f-45c7-95ab-dbbd2f87acc0", + "8e9c0b21-2230-4a0a-b35b-a198b90abcaa", + "3db37314-adaa-4a4e-af11-c9190fa6b65b", + "ba51db0f-2bf1-4746-894a-7db0ba54e357", + "c1ca7c5c-2473-4b83-a507-7d187609111a", + "8504aae3-4f25-4fa4-88c1-beb3fba83f13", + "65172b19-a246-4dc5-98e9-1b526b00ac89", + "49c27d75-44fa-4947-899a-8d0511a10416", + "19b880da-910e-40c7-970e-cd8718f2faf5", + "8a256d7f-2ea1-4b0e-94ea-168a1ed85678", + "de96ee14-be42-471d-b110-aabc948260fa", + "421f1c8b-5237-4735-9003-82021ba651b1", + "f0b69458-25df-4dda-be8b-e14383a57a50", + "83d8be38-ce02-41b9-b444-7291bad513d7", + "cd998593-e5bc-49d0-86a7-333acf9abf93", + "90bf3f44-60ae-4430-9dc6-2365f7901fc5", + "683635b7-5a1b-467c-8d4c-c34e8bf5fdf5", + "3492c684-74b5-4317-9e5c-320f3ea01899", + "49e25cbd-1aaa-40dd-901b-bb3e524ebcc0", + "9f3edf40-3dca-4086-b995-4985f6848a14", + "55363f67-1b41-476f-a83e-7e9741786ad7", + "fd50b1d6-8142-402f-88c5-7c44ca215081", + "36888bbd-1ece-4a3a-8aff-b0f6ef905752", + "33fd4f1f-fb01-4d3a-9b8d-b8d8024a2dab", + "8deeeb88-19ed-4747-9c79-d94104bed7e0", + "6958467e-388c-4f6d-a8c0-d4782166daac", + "85718cf6-e4ba-44e3-880f-143c627f12ff", + "8e9f89d2-f02a-4c75-8831-6785ef6bd7f5", + "5b99234b-b940-405d-9123-70218bda24b2", + "5a009338-0aff-4313-b02f-c48630ce1702", + "130a49d0-0315-4eb0-a752-a4533bd880a9", + "db12030d-7577-435f-9d48-5c1aae615c4f", + "b0c57748-65c8-49d7-9681-ddcfcde3298b", + "4a9ada82-9f86-495e-b17f-25b39bef2297", + "f42a7e14-c3ee-4b05-a59b-8d349f8bd19b", + "6549c35b-8b1d-4577-9434-ac808e4d7057", + "d2b674ca-9f02-46ba-820d-f0b52c76749f", + "02de9500-d574-446e-95fc-3e82c367385a", + "ff9fe54b-a959-4d7e-a6fa-93e311a68168", + "5c8f8839-4bf9-45e0-beb0-6c17a5ead0d5", + "023aaf52-fbe9-4b61-af69-ed33e91f3717", + "c920e557-141f-4556-95ac-f494ecc0c84f", + "b4b2c55b-c107-468f-bbdb-f73a0e4dca34", + "2069aea3-ac99-4b70-891e-1bfb8b15aeea", + "39c74452-1957-40db-894a-b39df51e1437", + "66cb6d16-f972-40e9-88c6-64b6816b900c", + "2db7496c-aa43-4d9d-8df6-404d2d9bd810", + "56683d7b-7c65-4f42-b5a4-818f3cd5224c", + "b2735994-59ae-4478-9c5f-93f5ccf18c8b", + "5e3c6abc-599b-44f8-8401-d5720627bba2", + "4d144ff3-a853-4b79-9825-d7fcea248373", + "3ecc3e4d-5466-4243-acca-2feea47e09a2", + "4decc0fa-e98e-4bf9-bbd1-d9b852b04830", + "5753e3fe-f1c5-426b-9f45-5a8efa6b5370", + "7fd21591-6a33-49cf-a32f-f8dfdf5ae697", + "572117cf-a76b-44d2-804a-4901bfc545f4", + "f5bc6a45-1591-42bf-8ed4-8286382b31bf", + "a5fe8210-1c37-4dd9-bd27-74051c521c3d", + "77e09fd6-b560-4439-addc-a1a8a3122046", + "91e97b3c-c912-401a-ab99-da01b90682aa", + "ffb191f8-894a-40e8-baa5-5578aa6ba706", + "d612edfb-e904-4de0-abce-a5fd9ecc449b", + "35465081-a518-4321-8384-477ef5b29b8d", + "4e78f728-06f6-43a3-ac8c-c0cb159e5b32", + "0eeab9d2-a55c-4708-94c3-cc1e08f9f262", + "e7b6b4ee-3f2a-41cd-bc18-c7e993cb413f", + "e98405ee-0a5b-4df2-bb7e-97729b8b3c11", + "19866484-b3d8-4619-85b9-e900de7fa80d", + "a16cb772-d0e2-4654-96d7-372c21484b90", + "d062849b-c60f-4369-a113-50f6ce8521fb", + "764a08dc-5ac5-4c09-a281-0210cdd5a55e", + "ee1893cb-2ff7-40e0-af8d-19518696b698", + "3edc700d-7d14-4154-a3e7-5b0e8c3fde07", + "c46cbba9-2cb3-4d84-be18-71af5539ecc5", + "a3dbec39-7126-4949-a916-eca55fee2b49", + "502ce8f6-bd9e-4e40-a5c9-40bdce36b87f", + "3051e028-71f8-4f3d-a72a-010934c3566f", + "0f325614-48c9-4a6a-b7f2-731c26ddeb97", + "a47c590d-66aa-4bd1-b552-f40951ba9bf0", + "4c85a780-8a8e-46c9-a0f8-c7e9703ef998", + "3949f121-af2b-4c7c-8a49-07befc08abbe", + "bc8bae7f-ad9f-4f72-ae5c-b3e7ca8dd8f0", + "e4670e4f-0b1a-450a-9778-9ee2409db554", + "094d1ab2-f34e-4a2a-ac3b-8671ce5faf65", + "290fc9ba-0a29-4eed-b0e0-64edc5bf54ea", + "f8289825-2858-4dfa-b173-f4c506359158", + "cfc6da55-43fc-4cd8-96d0-9de5166f4a63", + "355f8f4a-cde2-4108-821a-5ccea1d297fe", + "3627996c-8188-4486-8feb-938b9b0cce2b", + "986c4e43-ac2e-4983-8d56-f1dd34775133", + "97eb574d-d8bf-44a8-a6ef-1204b789db43", + "6c1bac5c-aead-40b8-9c17-be97d86b68d4", + "ed72b048-e309-4510-b049-b0288ce4478a", + "73f99cf0-d997-4488-b6a6-993934a9fad3", + "c81eb608-6844-424e-b8b7-524ff652cf3a", + "2d205d38-f79c-43f2-8c39-ede697065ab2", + "5f5ea963-4f59-4705-95fe-4a98c9fcd12c", + "6c4f7cb0-caa0-4a35-9071-a84c9ec2ab2f", + "bd13e435-1e38-4b22-bfd0-b620fc74cc8a", + "86af705a-a326-423a-bdd9-c6734804c30c", + "b9e32e51-8c18-420b-a582-200c79b8208d", + "f8d71e9a-982a-437c-a968-867149c24eb7", + "836be9c2-cee6-4fb9-8bbd-74197d213c5a", + "9090779d-803e-4fbd-9166-d9e90ff2d36f", + "15b6759d-2370-43ee-9a61-09e453db2469", + "0fa75089-7b28-4988-8664-a2c93542ffb0", + "3c3207e6-57b9-4be1-96fc-5b2168ae3cd3", + "4356f6d5-576a-4999-af40-744c242fff22", + "977ac980-f908-4c29-a2fc-dd14ec21c0bf", + "12c1b77c-cbea-4c20-8b64-7154a25bbbe1", + "cae5d4bf-bbf9-45dd-876b-a48579782348", + "0324c3b5-5154-4de1-9869-eb866b568e0f", + "961a8de1-0d33-44ab-a9d5-3d92de13fde4", + "df24897d-a478-4fbb-816f-1d14c59da30d", + "f84fd67e-0797-44f4-b3ff-78da646362e2", + "3af65596-1aa5-4126-933b-b98ccaccedf2", + "2b5cea2f-b019-4ecf-abbc-b2346bd65a10", + "394d1df6-bced-4f04-8c2a-a1ea3a04b984", + "dce1d8b4-2c29-4fec-a343-77885899c434", + "59000084-d286-4c83-90c8-2fbbf80d2d31", + "0e08e1c7-3add-4e41-b28d-cd5f4413a422", + "f7869fd8-9af3-493a-a627-a3e95e7c3577", + "22172001-0e17-40db-bcce-825e19532297", + "41f7665a-a692-42e0-b28d-9769c03ab99f", + "933ddf59-52d2-4d4c-b653-05de44ded245", + "a2d94b61-57c7-4ec2-ad94-94b3585a6b64", + "5b8c91ca-9d16-48b7-acce-4dec221ffd71", + "ead54fb5-dee9-452f-8145-227ccd2405e2", + "1f12c871-9b3a-4611-a2b4-9c39059052a4", + "ef37e1e6-5be3-4e8c-9f70-c84f341f9aea", + "c4392ca6-4706-4a96-a6b2-29e6ef544500", + "ac7d3d16-d74d-4fe3-a1e0-cab0b1adf492", + "c9f84bf4-108c-4151-85c0-e87375c53b91", + "44bcb1fa-9e64-4bae-927a-4d5d53e536a6", + "c1057c9b-331f-458b-8914-2703511e8b13", + "03c50078-a08d-43ad-a3c9-87916aebd865", + "12706b2d-0a0b-4511-8feb-ed8065746e50", + "f55f2dfd-704f-42cf-a39e-72e2560314b3", + "5beb520c-d413-4998-bd9f-7960a2baf744", + "78efcc13-3718-4acb-acee-8157a54501d6", + "e60b110e-a9b6-42b7-8f92-082cd74c4937", + "5da4efa3-8e54-46dc-a8c2-3016fa4e12f8", + "b844cb2d-6428-4d5e-abac-958f2704d4c8", + "5eaa3bfb-cc71-43fa-88dc-d367ef929a7b", + "e7633ad6-e406-4a18-b1ba-6c834ec9c047", + "fb07f205-868e-44c9-a704-305daaf03301", + "8e860a91-0205-447b-a706-aa894d23d3db", + "a1055434-9e9b-4d69-bac3-374a39f801da", + "812f0590-4ccd-4c9e-8978-c45e6debce55", + "f55f890d-fd90-43f6-a11a-a1920c21890a", + "63ab0246-2f1c-4621-8b71-75075b30d852", + "10732476-f045-4feb-8199-e4e04685f7ca", + "6f001eeb-fed5-43e4-bfa2-84e274cbe5cd", + "f62a8100-7ed2-492d-9709-064514d552bf", + "9f3cb8b5-5d62-4236-be5e-8f16d59c9f50", + "92e0925c-566d-4a67-b93e-3e4b57f9b3ab", + "057e9f41-e6ed-4575-906c-aca9413b35a5", + "94ba75fc-8d7f-4bb4-a756-ded62a4c1761", + "3aed7bbd-75f0-45ef-b9cc-7f5ff6f7fbbf", + "ea5a3f9d-7b21-46d7-a17d-bcbd4955bae9", + "bb3c6009-892c-4071-9bb4-692b83ed12f3", + "5f1fcf84-4001-43ec-ad88-9e06ea2f663d", + "2e651150-9338-4103-b231-2f089aa19df7", + "7813b44f-04ca-4a53-868d-fe185b615b68", + "3072481a-04fe-4959-8410-018a41a91df1", + "13c24739-e514-4fad-8288-6686bffa71ad", + "83464c62-929f-4bf7-9a04-43071189edde", + "a3da62c3-3f46-499f-a4c6-acf93bd1c63d", + "8f15aab0-9c42-452e-b78a-47a48da1f3c1", + "f3afa5a1-8a97-4a7d-945e-1cc70931e44a", + "cb1f70b4-0bf6-47f1-bc98-2e9365a55c3f", + "1efd8b7b-65f9-4e73-bd49-d707190d6519", + "08b6f5b3-4471-4cfd-a95e-4901cb1ded57", + "06592f50-7f29-4ca9-94a9-1b6c389d50ea", + "13d71019-9093-442b-bdb4-c79310150592", + "37afa156-a068-458e-8bf7-ec6f1f2415bc", + "383466b5-cb00-4e76-9e6a-dc1b0c9b125c", + "063f5040-9e94-42ae-b0e5-b13413e0fe55", + "84a388b6-2844-421e-a176-84d29958a76c", + "a4d53159-e2c0-4a47-9e78-af9bc9291f0e", + "895545f1-f194-4c0a-8513-87d386670f15", + "19c28cbe-21f2-4b40-a6e0-e9932b5f0605", + "e9acbb7f-c45c-47d7-9dd8-16194319fc97", + "6d367aa4-56ef-4932-9d48-e7530e717ea8", + "2f14e742-dc58-43c8-8ebc-7244487dbba7", + "34bc808a-4b98-42c5-a1ea-6f565dd4aa20", + "d2cb4045-8ef1-4288-8398-ab6d4d6a0709", + "604bf883-ce55-4c59-bcc1-3ceea1128211", + "7b76be35-4e6d-42d7-9fa9-14a21fc327cf", + "5f615b11-2ff8-4ff1-987a-1cd61ea78dee", + "a3bd1be3-edfe-4f4e-89a8-dd56ce6ece08", + "d6adb1ae-f689-438b-ab81-8a2e37f9654e", + "514d71c8-fd4e-4908-89a4-c0861e81ca8a", + "0d998741-de68-4006-898e-7108de4167b1", + "c7820a4f-55c9-44eb-bf7e-b965a9d0e924", + "172e1afa-5581-455d-9a27-590e9a1706c2", + "5b5b6b8c-9b88-4181-9ebb-14abb7783e1e", + "c0ba997c-9ec6-4e83-ad50-a88dd677b881", + "679f2016-4802-4bf8-a830-d1f45fdfcd71", + "38c02527-1b2f-458b-bad8-73a282c1d99a", + "e6f4e36f-96b7-48f5-afd1-27cc408bdd07", + "65cad708-34b8-4158-be96-b88710288473", + "9ef6ef95-6b97-4c78-866c-641fc7d791f5", + "4ce65f3a-21ee-4f77-85c0-b3ddfe1e6381", + "f0e9a139-86f6-41bc-82f1-11283954a781", + "d5374572-5576-4f83-99c8-a0b12c4de160", + "c925c38c-3fe3-4eb3-9944-3fba3eaca4db", + "efeaa284-2e22-44b8-bdd6-1bb4763e5b35", + "6640a311-86c4-4de7-964c-e3a4a67a3e02", + "136cef9c-e17f-49f9-ae45-a17e7e5c5477", + "12634773-451f-4c82-b915-265c703ad737", + "fc0f4c90-a1ee-4147-990e-de231394f484", + "50895f7f-5667-4b43-94c1-d7b4b9e238e7", + "03cebfe9-3319-45b5-aba5-e7db10ea6fdb", + "0124376a-720d-4265-8c8c-5ff47f04d579", + "0b4cacb1-970f-4ef0-85da-371d81f899e0", + "da84b37f-ec6e-414e-b52e-1231104a68e0", + "70316529-ec02-4d26-bc68-4b436bf0030a", + "81b74694-a2e6-4be1-b9f6-69f3aa03f50d", + "eb3675b8-f6a8-4ea9-82df-1080eac5b093", + "39fb2989-1dd2-490d-b793-dce86c36b43a", + "1ac6e060-a6d2-4b42-afef-599be78a533a", + "9ee0474c-615c-4383-a047-8432b8b24784", + "95cdb269-0088-4451-9a0d-f72349081fab", + "48f7028c-b6bf-4ca5-81e9-128147fd02c9", + "f00bd659-56e0-41dd-8eba-c5fdb45e7232", + "1ce2252b-31b8-4536-9902-0ba7fb039340", + "8f2d26e6-239e-492f-aa1b-472f8ec85050", + "5daa132e-0b97-4d0f-be9d-b514285661fe", + "cdbc7718-4cb9-4e78-bbde-3d9614f8dcd1", + "fd394f00-4049-49e5-9c07-ee3e86979003", + "20a3e8e5-3f43-46a5-84fa-6888bd9c5c88", + "bd03445e-9016-423a-bb48-cd998836bbff", + "1492c03c-1887-4576-8889-5bc2848967de", + "5eb74ea5-9b0b-484c-82b2-ee86479c5d07", + "4e0053f8-566c-4952-abbe-4ea877782a0a", + "57b6455e-2acb-4076-b8c4-1f2cdae97185", + "3f28dd4f-f653-4937-83b6-9f80c1284a85", + "71a7456a-da6e-4e38-86ec-76bacaa6db5b", + "8d73dd20-5349-477f-9630-218e6b08bb95", + "0fccd442-a901-4d77-aac6-dc4276cd6591", + "5df66408-7a8c-43bb-9235-d088aa3d19dc", + "a9b191de-d565-4981-8f68-02069a2ba16c", + "ffddab3a-dbec-49b0-b114-5a1e3d805efa", + "61c73721-9476-4ac2-a987-ecaf9a04dca1", + "b9d7642e-2ced-4589-b7a6-69ad54570c4c", + "5998ece4-98cd-4d5b-9104-1517d13728bd", + "3ce15244-5293-4a5e-8c92-52ab29c6949a", + "3e4b9c29-06d2-4ade-a2be-bdec028bc68c", + "e05a2182-a97f-4d0b-87e4-66b0dd405e58", + "2e6ff051-0037-4001-90a8-de7643a96f08", + "e62b75a0-f1f6-4e45-b3db-17fd7ff5756c", + "d15eb34f-e425-4990-aa2d-2dc454fb4b20", + "1b6dfc9f-c873-48c6-9f25-a2d490e269cf", + "be8f4445-3958-4fc4-8cdf-b14d1d0eba73", + "85506322-449a-45c0-b77a-48a8077b4bbd", + "ba5b1882-c9d7-48e7-85fe-2c7b10494fac", + "3ab3b82f-bfa8-4b39-a9a9-e397341ae572", + "b69564f5-5e5e-4097-aebd-9f52f211ff75", + "13ffc0e4-eebb-4baa-97af-5983ce46779c", + "c96fa96d-d23c-444e-bdef-481317308bb8", + "3b3a4043-3d74-4575-8aeb-511c4a879241", + "672e0d75-df0e-4715-8c42-f1131b059a6c", + "9767a3ca-0a9c-480f-adfe-e4164c52c760", + "e2ffc7c9-c035-4709-97c6-3764499188a5", + "07774a7c-c5ed-4133-991f-926a316068a4", + "69006355-3fb0-443c-826b-ea8c9d27f70b", + "c2c7f072-a02a-480d-8342-b43ffc314db9", + "d2c0f78a-978e-409f-a78f-3e40b3431909", + "7c776700-1d6b-4ee2-8155-b7f6855afceb", + "db1d3def-397e-470f-811a-85c4cc5338a9", + "be99e9e2-82f7-4f1f-aff5-371834fde0da", + "157c8e81-117a-4a15-a047-9a574181d44b", + "de82a6f3-1a55-4996-ae68-35032678ff66", + "218e79aa-3f56-4bd4-aa2a-14027bf7d52d", + "ff283659-4a79-4642-8078-4366df38e1f1", + "d0af657d-9c21-447c-a4f9-199b38deac5b", + "2ac543c6-b52b-4219-b2e4-4130078c0673", + "83553642-1a38-4c54-8154-c4cd4fd81173", + "b3ef3563-ecc0-4a15-9a7b-4feb4558d953", + "4cf10990-13e5-4a1e-a913-7ec4b8c8ab02", + "16851d4d-35e0-49b4-b3b5-28e67b564db3", + "7c4079e7-4d68-4b0d-b015-5971a9866c60", + "1943bfb1-ba78-4a6c-ade7-11e44314dfb4", + "1a746147-2dec-4a37-a51b-c94e4f1de010", + "2e4e0dac-ab21-49ff-9563-e320acd73f2f", + "561fd770-f9a2-41a4-90da-beeb7d96ee7f", + "5be9a2e2-5efc-4fc6-a65e-806e0347e19a", + "d66c7c66-2b5f-4e7f-bbb2-f0adae68dfcc", + "3bd74a94-ec3e-4138-b78f-ad76af4f07e2", + "5204624f-90ac-47bf-a71a-486d43f4351c", + "0e519838-1e4a-41fd-ab30-34bf745a3339", + "c37e680f-03b1-408b-9b09-1b2ae733951d", + "453f43e9-c1d1-4749-9720-73546287b82f", + "666d9243-f4f9-47f5-8fba-cee8568abf41", + "2091513c-8f4c-4817-8c22-cdbc6303a331", + "307e8837-1fda-4cf5-bab5-d16fee20f84f", + "783d7043-b27e-4744-a969-5f40fb40827c", + "e6317c78-06f0-4c6d-9886-ca67712772ed", + "ba280b86-22a2-4509-b16d-b55e04680461", + "51fee605-d077-4d78-babc-9bbcd8d898ca", + "63483ea2-d569-45fd-82b2-f80836a8b766", + "221fd24b-ba2e-40d5-a0c7-a879ece6b34f", + "6d1dc078-8b04-480e-a36f-298594ee9592", + "e37c82c5-483f-4e68-8bf3-2a22dcef0b1e", + "d3b7e314-f48d-4888-b496-e1eeb252cd8a", + "8351998e-ead7-4bc1-b0a7-a7209da92b3c", + "53144267-6da5-40d3-91f2-eaa85be367de", + "4ade0a82-16fd-471f-9fcc-fd2724887aef", + "228eb504-c1b3-43b9-818e-2f7e539a259b", + "db9978fb-fe4d-467b-9495-ab7daf3b1b55", + "fb62dc3f-a1cd-4f0e-88bc-504fe9d8585d", + "32464967-b0f2-4b2e-a57c-7f774e7056a7", + "efdc3dc0-275d-4073-9179-32b14b224654", + "50d3035c-2503-4bcf-b28b-e66c92c98f09", + "f8accf9d-4427-4458-88e9-89d51a3f5136", + "ca22f118-2472-4df0-a9cd-3d70424b211b", + "1261ee8c-80fc-4a3f-9dd3-29f4a918a9fe", + "666a21c0-7469-436e-868f-64eb8022f64f", + "e3a53edd-52ff-4608-a80d-77ecfbb0f961", + "50810487-0b0f-45e0-8f37-515644189e81", + "bdf1ca5a-321e-4d3c-ac9e-da2ce1f5a69c", + "f817b176-2dce-4858-999f-fe62aac25e2e", + "31488ff6-05ef-4362-9786-35981b06f307", + "bec4e063-9be1-402e-a7b5-dba20730c933", + "3d8acf67-fad6-421e-841b-82b28c82f7f6", + "c2a8f168-06b0-4766-8865-27f08cb1bcae", + "a874ef10-022c-43e4-b4e6-e01d8fcec988", + "2ea37ccc-e1ec-4e0d-a2a7-9707ea300b7f", + "4f3c8276-99e0-4c99-9e61-5d7c06332b09", + "e50126ee-f4c1-41a6-a2e8-79636b40b341", + "fd8cc20a-915e-45a9-93df-320ad998aee4", + "89125bc3-a413-44ee-a7e1-2753c511ecab", + "c6a6e1e4-24c4-422d-9d18-f062240ddf23", + "3c62a0d2-10d7-4a9e-ac2c-02505663afad", + "52fb1a73-2ba7-413d-8596-33615daecc73", + "91151b86-65b7-4854-97b0-9c5b93a251ed", + "19f880ae-9770-45a5-93f5-00047f22d5f6", + "26c042d4-bf21-4d10-81f9-303d3816c114", + "602b224a-f5ac-491b-88e3-49043f362b6e", + "36e7b2ad-dddd-42fa-a427-6411ff12f9b9", + "3f920762-f4a7-4bdb-86dc-05a3c3c8e875", + "a9a0ed92-fab1-4b6e-ae27-714692464641", + "fcc6af87-918f-4d0c-a0a4-8aab36891fc4", + "d6d8261e-0ee0-4002-bddf-e94fdc1075c7", + "09194996-ff4b-4a4f-9800-e8aedc9be7f2", + "dfe025f3-573e-4297-b17a-30e990cfa310", + "c2af1389-b165-42f5-b4ee-c90402d7006b", + "49a7e53b-aba5-4e53-b3bc-1492657fc8ae", + "1828b207-9601-4521-8c0b-5d9e829fca78", + "70e88e82-5c06-44fd-959c-d4b2d40999cd", + "3a20bf0a-6eb1-47c5-9bf8-8445446feb61", + "bc6cf777-fd8a-4427-9867-92a443826827", + "74e71420-9bd6-4337-9a7a-a3dcc4cab7cf", + "2c7796e3-b1e9-413f-a0e8-8225ac7d586e", + "ed65870d-ba08-4894-9cad-7f34e2c6f9e4", + "26080fb8-0588-449c-a7f2-2e0162507e02", + "0f846b5d-1cea-40b8-b18b-dc6c54544a67", + "7abe6d4a-f594-4a15-af08-ddd591ff500c", + "d36d9514-a99b-4534-b54f-3fd98a19f0ea", + "00792fa8-988c-4c85-8e80-73eb3ac53e80", + "f906b093-281f-4872-8010-f1f3a320ee2a", + "ce8ea1a3-65e3-4ef7-abeb-2b0bc1c84773", + "e92255bc-1be6-4ee9-8f8b-d1336a0e409c", + "6798a65f-d79a-413b-baca-115f09caf0c4", + "b171dcaf-94f2-4bcd-8695-b2648bed60e1", + "be3b65c8-ba74-4321-99b4-ca0e99e6fc2c", + "3eb00a40-880a-4ad9-8cec-3ba0245ec472", + "ae6466f8-1a53-4cb3-8718-96f78f391b4b", + "7dd0eab1-b83e-4bb4-ad91-7bbdf4592647", + "3c17a762-c839-492f-b361-009c1187ae4c", + "be25331c-b1a0-4187-93fb-3b59437258c3", + "b432e10e-cff1-4f15-a8fe-1e2350e94b36", + "9ba7c2a0-c86f-4967-9284-9be88ad3eacf", + "1acfcabc-653a-4c03-adf7-64838d09a7b4", + "108cebaa-777d-42dd-8b1e-69ea7f443a0a", + "edaa0b1a-d504-4852-8025-4f12b4d293ac", + "6daa9124-d894-42e5-8d98-2e0e6d1191ac", + "0232fbc8-c9b1-405c-bd42-a9d773fd4835", + "134c4387-77d5-4ee0-911f-e50e64499b42", + "b21b35d2-6d69-40a2-8c36-aa27e4a980ae", + "bb7f1e75-65ee-4d98-adc8-c70551a35940" + ], + "test_session_set": [ + "c2413391-7c1b-4fd6-8b1d-98ee7888b9f8", + "fe69a78e-7773-45d1-9e0f-bacee52dac83", + "3b79017c-4d42-40fc-a1bb-4a20bc8ebca7", + "6dbfc053-7899-40d8-9827-0ccd21f3ee0a", + "7e6dfd31-8544-4fad-9e49-0f05516cf8cf", + "56c5af79-f9d4-478d-96ef-6d71e0bbbdfe", + "3aea2d62-0340-4c6a-8019-3cce9117a65a", + "d97bedc8-72df-43be-a55b-4da1ae42dfd1", + "f0cb79ef-c081-4049-85ef-2623e02c9589", + "08b0935e-6260-4bd6-86ca-f6fc54e388be", + "6f80f9ae-c276-413f-accc-fad74efe8a95", + "b95bd349-11bb-4ade-9cc6-3be8dcea06dd", + "e7086057-b5f7-4286-87c6-ab110072c5b0", + "8324c0ba-7a10-42e1-90fd-afb4d3afa789", + "2b90b963-e829-4787-8fcd-1a1a434ffdfd", + "66ddb7ae-90d6-4a58-9f3c-7f1d2604ecb0", + "ad1a9ea9-d1ab-474a-b715-1a42d892ef5d", + "4ca3d9cf-4972-4e1d-aca3-28b036d6bdc5", + "2360312c-4347-476a-990f-2633b489fa7a", + "0fe736da-48f1-4961-bd7a-b75ce13c91b4", + "1c5c0b52-1287-4df3-aa75-8ccc916b1c10", + "07f1486c-44a7-4c69-9532-720c81c99eb0", + "80ee109f-b097-4755-8a39-9f0b33adb170", + "88260ff4-be0b-4544-98ef-78b6a0d61ebc", + "aabeb7f9-5ada-425a-9545-c98dc443887a", + "65229ac2-9033-4002-af6c-492fe49bc563", + "7f566780-4ed2-41c5-9faa-06759495faab", + "4ea88cb9-bff8-49e0-a855-cc1902db9a99", + "bb050b8a-195a-4b50-b5d0-1ed535226a39", + "b4f650b4-53d6-4069-a55f-43f2de7ff47b", + "8a409551-657a-48c0-8beb-056c228ebca5", + "6f3f70b4-8428-4e5f-b1dd-86cd55d8e63f", + "0511dce9-c126-4a11-92ba-8082a4ec89bd", + "33365c93-7f1d-4a7d-a9bc-b33010aeafbd", + "ea9d85e7-6a9a-451a-b019-972a6891c11b", + "cd3092e1-6022-4a11-bfc3-a4bb80487ece", + "5682f9aa-bdfb-409c-b4c2-c0105b35ee78", + "b87c3e0f-a5af-4663-a87d-d0689d73343b", + "6edb7036-3369-4c03-8e3a-439ffbeb4724", + "0d4efcc9-a336-46f9-a1db-0623b5be2869", + "7a6576d8-007c-4ee9-b7cc-8cc2d020e1bf", + "f0445727-2729-4a76-9434-fb081a14cd9d", + "73c1ac5a-80a0-4e55-b143-52df5697c9ea", + "224d8b17-7eaf-4529-9b75-c6ce3078d640", + "a22beaf0-377c-4b95-b472-867280a525ae", + "5a5b8c37-e3f6-48f2-96d3-289a043eaee4", + "dab038d7-9e08-4edf-8d03-ee7e2af5a10e", + "08742cb7-bb6d-47ae-af5e-3f1f306ccda6", + "5e965191-bcd1-491a-a9ce-e014e970020b", + "076bdb81-5c75-4282-9f3a-a387624575f3" + ] + }, + "EasyCom": { + "class_index": [ + "background", + "camera wearer speaking", + "other people speaking" + ], + "train_session_set": [ + "Session_7", + "Session_11", + "Session_10", + "Session_6", + "Session_9", + "Session_4", + "Session_8", + "Session_5", + "Session_12" + ], + "test_session_set": [ + "Session_3", + "Session_2", + "Session_1" + ] + }, + "EasyCom_inference": { + "class_index": [ + "background", + "camera wearer speaking", + "other people speaking" + ], + "train_session_set": [ + "Session_7", + "Session_11", + "Session_10", + "Session_6", + "Session_9", + "Session_4", + "Session_8", + "Session_5", + "Session_12" + ], + "test_session_set": [ + "Session_12", + "Session_11", + "Session_10", + "Session_9", + "Session_8", + "Session_7", + "Session_6", + "Session_5", + "Session_4", + "Session_3", + "Session_2", + "Session_1" + ] + }, + "TVSeries": { + "class_index": [ + "Background", + "Pick something up", + "Point", + "Drink", + "Stand up", + "Run", + "Sit down", + "Read", + "Smoke", + "Drive car", + "Open door", + "Give something", + "Use computer", + "Write", + "Go down stairway", + "Close door", + "Throw something", + "Go up stairway", + "Get in/out of car", + "Hang up phone", + "Eat", + "Answer phone", + "Dress up", + "Clap", + "Undress", + "Kiss", + "Fall/trip", + "Wave", + "Pour", + "Punch", + "Fire weapon" + ], + "train_session_set": [ + "24_ep1", + "24_ep2", + "24_ep3", + "Breaking_Bad_ep1", + "Breaking_Bad_ep2", + "How_I_Met_Your_Mother_ep1", + "How_I_Met_Your_Mother_ep2", + "How_I_Met_Your_Mother_ep3", + "How_I_Met_Your_Mother_ep4", + "How_I_Met_Your_Mother_ep5", + "How_I_Met_Your_Mother_ep6", + "Mad_Men_ep1", + "Mad_Men_ep2", + "Modern_Family_ep1", + "Modern_Family_ep2", + "Modern_Family_ep3", + "Modern_Family_ep4", + "Modern_Family_ep6", + "Sons_of_Anarchy_ep1", + "Sons_of_Anarchy_ep2" + ], + "test_session_set": [ + "24_ep4", + "Breaking_Bad_ep3", + "Mad_Men_ep3", + "How_I_Met_Your_Mother_ep7", + "How_I_Met_Your_Mother_ep8", + "Modern_Family_ep5", + "Sons_of_Anarchy_ep3" + ] + }, + "THUMOS": { + "class_index": [ + "Background", + "BaseballPitch", + "BasketballDunk", + "Billiards", + "CleanAndJerk", + "CliffDiving", + "CricketBowling", + "CricketShot", + "Diving", + "FrisbeeCatch", + "GolfSwing", + "HammerThrow", + "HighJump", + "JavelinThrow", + "LongJump", + "PoleVault", + "Shotput", + "SoccerPenalty", + "TennisSwing", + "ThrowDiscus", + "VolleyballSpiking", + "Ambiguous" + ], + "train_session_set": [ + "video_validation_0000690", + "video_validation_0000288", + "video_validation_0000289", + "video_validation_0000416", + "video_validation_0000282", + "video_validation_0000283", + "video_validation_0000281", + "video_validation_0000286", + "video_validation_0000287", + "video_validation_0000284", + "video_validation_0000285", + "video_validation_0000202", + "video_validation_0000203", + "video_validation_0000201", + "video_validation_0000206", + "video_validation_0000207", + "video_validation_0000204", + "video_validation_0000205", + "video_validation_0000790", + "video_validation_0000208", + "video_validation_0000209", + "video_validation_0000420", + "video_validation_0000364", + "video_validation_0000853", + "video_validation_0000950", + "video_validation_0000937", + "video_validation_0000367", + "video_validation_0000290", + "video_validation_0000210", + "video_validation_0000059", + "video_validation_0000058", + "video_validation_0000057", + "video_validation_0000056", + "video_validation_0000055", + "video_validation_0000054", + "video_validation_0000053", + "video_validation_0000052", + "video_validation_0000051", + "video_validation_0000933", + "video_validation_0000949", + "video_validation_0000948", + "video_validation_0000945", + "video_validation_0000944", + "video_validation_0000947", + "video_validation_0000946", + "video_validation_0000941", + "video_validation_0000940", + "video_validation_0000190", + "video_validation_0000942", + "video_validation_0000261", + "video_validation_0000262", + "video_validation_0000263", + "video_validation_0000264", + "video_validation_0000265", + "video_validation_0000266", + "video_validation_0000267", + "video_validation_0000268", + "video_validation_0000269", + "video_validation_0000989", + "video_validation_0000060", + "video_validation_0000370", + "video_validation_0000938", + "video_validation_0000935", + "video_validation_0000668", + "video_validation_0000669", + "video_validation_0000664", + "video_validation_0000665", + "video_validation_0000932", + "video_validation_0000667", + "video_validation_0000934", + "video_validation_0000661", + "video_validation_0000662", + "video_validation_0000663", + "video_validation_0000181", + "video_validation_0000180", + "video_validation_0000183", + "video_validation_0000182", + "video_validation_0000185", + "video_validation_0000184", + "video_validation_0000187", + "video_validation_0000186", + "video_validation_0000189", + "video_validation_0000188", + "video_validation_0000936", + "video_validation_0000270", + "video_validation_0000854", + "video_validation_0000178", + "video_validation_0000179", + "video_validation_0000174", + "video_validation_0000175", + "video_validation_0000176", + "video_validation_0000177", + "video_validation_0000170", + "video_validation_0000171", + "video_validation_0000172", + "video_validation_0000173", + "video_validation_0000670", + "video_validation_0000419", + "video_validation_0000943", + "video_validation_0000485", + "video_validation_0000369", + "video_validation_0000368", + "video_validation_0000318", + "video_validation_0000319", + "video_validation_0000415", + "video_validation_0000414", + "video_validation_0000413", + "video_validation_0000412", + "video_validation_0000411", + "video_validation_0000311", + "video_validation_0000312", + "video_validation_0000313", + "video_validation_0000314", + "video_validation_0000315", + "video_validation_0000316", + "video_validation_0000317", + "video_validation_0000418", + "video_validation_0000365", + "video_validation_0000482", + "video_validation_0000169", + "video_validation_0000168", + "video_validation_0000167", + "video_validation_0000166", + "video_validation_0000165", + "video_validation_0000164", + "video_validation_0000163", + "video_validation_0000162", + "video_validation_0000161", + "video_validation_0000160", + "video_validation_0000857", + "video_validation_0000856", + "video_validation_0000855", + "video_validation_0000366", + "video_validation_0000488", + "video_validation_0000489", + "video_validation_0000851", + "video_validation_0000484", + "video_validation_0000361", + "video_validation_0000486", + "video_validation_0000487", + "video_validation_0000481", + "video_validation_0000910", + "video_validation_0000483", + "video_validation_0000363", + "video_validation_0000990", + "video_validation_0000939", + "video_validation_0000362", + "video_validation_0000987", + "video_validation_0000859", + "video_validation_0000787", + "video_validation_0000786", + "video_validation_0000785", + "video_validation_0000784", + "video_validation_0000783", + "video_validation_0000782", + "video_validation_0000781", + "video_validation_0000981", + "video_validation_0000983", + "video_validation_0000982", + "video_validation_0000985", + "video_validation_0000984", + "video_validation_0000417", + "video_validation_0000788", + "video_validation_0000152", + "video_validation_0000153", + "video_validation_0000151", + "video_validation_0000156", + "video_validation_0000157", + "video_validation_0000154", + "video_validation_0000155", + "video_validation_0000158", + "video_validation_0000159", + "video_validation_0000901", + "video_validation_0000903", + "video_validation_0000902", + "video_validation_0000905", + "video_validation_0000904", + "video_validation_0000907", + "video_validation_0000906", + "video_validation_0000909", + "video_validation_0000908", + "video_validation_0000490", + "video_validation_0000860", + "video_validation_0000858", + "video_validation_0000988", + "video_validation_0000320", + "video_validation_0000688", + "video_validation_0000689", + "video_validation_0000686", + "video_validation_0000687", + "video_validation_0000684", + "video_validation_0000685", + "video_validation_0000682", + "video_validation_0000683", + "video_validation_0000681", + "video_validation_0000789", + "video_validation_0000986", + "video_validation_0000931", + "video_validation_0000852", + "video_validation_0000666" + ], + "test_session_set": [ + "video_test_0000292", + "video_test_0001078", + "video_test_0000896", + "video_test_0000897", + "video_test_0000950", + "video_test_0001159", + "video_test_0001079", + "video_test_0000807", + "video_test_0000179", + "video_test_0000173", + "video_test_0001072", + "video_test_0001075", + "video_test_0000767", + "video_test_0001076", + "video_test_0000007", + "video_test_0000006", + "video_test_0000556", + "video_test_0001307", + "video_test_0001153", + "video_test_0000718", + "video_test_0000716", + "video_test_0001309", + "video_test_0000714", + "video_test_0000558", + "video_test_0001267", + "video_test_0000367", + "video_test_0001324", + "video_test_0000085", + "video_test_0000887", + "video_test_0001281", + "video_test_0000882", + "video_test_0000671", + "video_test_0000964", + "video_test_0001164", + "video_test_0001114", + "video_test_0000771", + "video_test_0001163", + "video_test_0001118", + "video_test_0001201", + "video_test_0001040", + "video_test_0001207", + "video_test_0000723", + "video_test_0000569", + "video_test_0000672", + "video_test_0000673", + "video_test_0000278", + "video_test_0001162", + "video_test_0000405", + "video_test_0000073", + "video_test_0000560", + "video_test_0001276", + "video_test_0000270", + "video_test_0000273", + "video_test_0000374", + "video_test_0000372", + "video_test_0001168", + "video_test_0000379", + "video_test_0001446", + "video_test_0001447", + "video_test_0001098", + "video_test_0000873", + "video_test_0000039", + "video_test_0000442", + "video_test_0001219", + "video_test_0000762", + "video_test_0000611", + "video_test_0000617", + "video_test_0000615", + "video_test_0001270", + "video_test_0000740", + "video_test_0000293", + "video_test_0000504", + "video_test_0000505", + "video_test_0000665", + "video_test_0000664", + "video_test_0000577", + "video_test_0000814", + "video_test_0001369", + "video_test_0001194", + "video_test_0001195", + "video_test_0001512", + "video_test_0001235", + "video_test_0001459", + "video_test_0000691", + "video_test_0000765", + "video_test_0001452", + "video_test_0000188", + "video_test_0000591", + "video_test_0001268", + "video_test_0000593", + "video_test_0000864", + "video_test_0000601", + "video_test_0001135", + "video_test_0000004", + "video_test_0000903", + "video_test_0000285", + "video_test_0001174", + "video_test_0000046", + "video_test_0000045", + "video_test_0001223", + "video_test_0001358", + "video_test_0001134", + "video_test_0000698", + "video_test_0000461", + "video_test_0001182", + "video_test_0000450", + "video_test_0000602", + "video_test_0001229", + "video_test_0000989", + "video_test_0000357", + "video_test_0001039", + "video_test_0000355", + "video_test_0000353", + "video_test_0001508", + "video_test_0000981", + "video_test_0000242", + "video_test_0000854", + "video_test_0001484", + "video_test_0000635", + "video_test_0001129", + "video_test_0001339", + "video_test_0001483", + "video_test_0001123", + "video_test_0001127", + "video_test_0000689", + "video_test_0000756", + "video_test_0001431", + "video_test_0000129", + "video_test_0001433", + "video_test_0001343", + "video_test_0000324", + "video_test_0001064", + "video_test_0001531", + "video_test_0001532", + "video_test_0000413", + "video_test_0000991", + "video_test_0001255", + "video_test_0000464", + "video_test_0001202", + "video_test_0001080", + "video_test_0001081", + "video_test_0000847", + "video_test_0000028", + "video_test_0000844", + "video_test_0000622", + "video_test_0000026", + "video_test_0001325", + "video_test_0001496", + "video_test_0001495", + "video_test_0000624", + "video_test_0000724", + "video_test_0001409", + "video_test_0000131", + "video_test_0000448", + "video_test_0000444", + "video_test_0000443", + "video_test_0001038", + "video_test_0000238", + "video_test_0001527", + "video_test_0001522", + "video_test_0000051", + "video_test_0001058", + "video_test_0001391", + "video_test_0000429", + "video_test_0000426", + "video_test_0000785", + "video_test_0000786", + "video_test_0001314", + "video_test_0000392", + "video_test_0000423", + "video_test_0001146", + "video_test_0001313", + "video_test_0001008", + "video_test_0001247", + "video_test_0000737", + "video_test_0001319", + "video_test_0000308", + "video_test_0000730", + "video_test_0000058", + "video_test_0000538", + "video_test_0001556", + "video_test_0000113", + "video_test_0000626", + "video_test_0000839", + "video_test_0000220", + "video_test_0001389", + "video_test_0000437", + "video_test_0000940", + "video_test_0000211", + "video_test_0000946", + "video_test_0001558", + "video_test_0000796", + "video_test_0000062", + "video_test_0000793", + "video_test_0000987", + "video_test_0001066", + "video_test_0000412", + "video_test_0000798", + "video_test_0001549", + "video_test_0000011", + "video_test_0001257", + "video_test_0000541", + "video_test_0000701", + "video_test_0000250", + "video_test_0000254", + "video_test_0000549", + "video_test_0001209", + "video_test_0001463", + "video_test_0001460", + "video_test_0000319", + "video_test_0001468", + "video_test_0000846", + "video_test_0001292" + ] + }, + "EGO4DVAL": { + "class_index": [ + "background", + "camera wearer speak", + "other people speaking" + ], + "train_session_set": [ + "5b8c91ca-9d16-48b7-acce-4dec221ffd71", + "0232fbc8-c9b1-405c-bd42-a9d773fd4835", + "39c74452-1957-40db-894a-b39df51e1437", + "31488ff6-05ef-4362-9786-35981b06f307", + "4decc0fa-e98e-4bf9-bbd1-d9b852b04830", + "efeaa284-2e22-44b8-bdd6-1bb4763e5b35", + "07f1486c-44a7-4c69-9532-720c81c99eb0", + "19c28cbe-21f2-4b40-a6e0-e9932b5f0605", + "1c5c0b52-1287-4df3-aa75-8ccc916b1c10", + "e92255bc-1be6-4ee9-8f8b-d1336a0e409c", + "307e8837-1fda-4cf5-bab5-d16fee20f84f", + "bb050b8a-195a-4b50-b5d0-1ed535226a39", + "ca22f118-2472-4df0-a9cd-3d70424b211b", + "08b0935e-6260-4bd6-86ca-f6fc54e388be", + "666a21c0-7469-436e-868f-64eb8022f64f", + "77e09fd6-b560-4439-addc-a1a8a3122046", + "f00bd659-56e0-41dd-8eba-c5fdb45e7232", + "c2c7f072-a02a-480d-8342-b43ffc314db9", + "ed72b048-e309-4510-b049-b0288ce4478a", + "56c5af79-f9d4-478d-96ef-6d71e0bbbdfe", + "3d8acf67-fad6-421e-841b-82b28c82f7f6", + "977ac980-f908-4c29-a2fc-dd14ec21c0bf", + "08742cb7-bb6d-47ae-af5e-3f1f306ccda6", + "f3afa5a1-8a97-4a7d-945e-1cc70931e44a", + "65172b19-a246-4dc5-98e9-1b526b00ac89", + "ee1893cb-2ff7-40e0-af8d-19518696b698", + "53144267-6da5-40d3-91f2-eaa85be367de", + "26080fb8-0588-449c-a7f2-2e0162507e02", + "c0ba997c-9ec6-4e83-ad50-a88dd677b881", + "6640a311-86c4-4de7-964c-e3a4a67a3e02", + "5f1fcf84-4001-43ec-ad88-9e06ea2f663d", + "ea5a3f9d-7b21-46d7-a17d-bcbd4955bae9", + "eb3675b8-f6a8-4ea9-82df-1080eac5b093", + "023aaf52-fbe9-4b61-af69-ed33e91f3717", + "86af705a-a326-423a-bdd9-c6734804c30c", + "55363f67-1b41-476f-a83e-7e9741786ad7", + "ce8ea1a3-65e3-4ef7-abeb-2b0bc1c84773", + "d3b7e314-f48d-4888-b496-e1eeb252cd8a", + "59000084-d286-4c83-90c8-2fbbf80d2d31", + "a47c590d-66aa-4bd1-b552-f40951ba9bf0", + "61c73721-9476-4ac2-a987-ecaf9a04dca1", + "0124376a-720d-4265-8c8c-5ff47f04d579", + "f62a8100-7ed2-492d-9709-064514d552bf", + "13d71019-9093-442b-bdb4-c79310150592", + "74e71420-9bd6-4337-9a7a-a3dcc4cab7cf", + "383466b5-cb00-4e76-9e6a-dc1b0c9b125c", + "9ef6ef95-6b97-4c78-866c-641fc7d791f5", + "ac7d3d16-d74d-4fe3-a1e0-cab0b1adf492", + "c46cbba9-2cb3-4d84-be18-71af5539ecc5", + "5eb74ea5-9b0b-484c-82b2-ee86479c5d07", + "6f80f9ae-c276-413f-accc-fad74efe8a95", + "73f99cf0-d997-4488-b6a6-993934a9fad3", + "fd50b1d6-8142-402f-88c5-7c44ca215081", + "51fee605-d077-4d78-babc-9bbcd8d898ca", + "5e3c6abc-599b-44f8-8401-d5720627bba2", + "91151b86-65b7-4854-97b0-9c5b93a251ed", + "5be9a2e2-5efc-4fc6-a65e-806e0347e19a", + "94ba75fc-8d7f-4bb4-a756-ded62a4c1761", + "c925c38c-3fe3-4eb3-9944-3fba3eaca4db", + "d97bedc8-72df-43be-a55b-4da1ae42dfd1", + "e60b110e-a9b6-42b7-8f92-082cd74c4937", + "3627996c-8188-4486-8feb-938b9b0cce2b", + "d2b674ca-9f02-46ba-820d-f0b52c76749f", + "2b90b963-e829-4787-8fcd-1a1a434ffdfd", + "ba51db0f-2bf1-4746-894a-7db0ba54e357", + "6c1bac5c-aead-40b8-9c17-be97d86b68d4", + "63483ea2-d569-45fd-82b2-f80836a8b766", + "c81eb608-6844-424e-b8b7-524ff652cf3a", + "cae5d4bf-bbf9-45dd-876b-a48579782348", + "7813b44f-04ca-4a53-868d-fe185b615b68", + "1f12c871-9b3a-4611-a2b4-9c39059052a4", + "70316529-ec02-4d26-bc68-4b436bf0030a", + "6d367aa4-56ef-4932-9d48-e7530e717ea8", + "2e651150-9338-4103-b231-2f089aa19df7", + "d2cb4045-8ef1-4288-8398-ab6d4d6a0709", + "3e4b9c29-06d2-4ade-a2be-bdec028bc68c", + "4f3c8276-99e0-4c99-9e61-5d7c06332b09", + "12634773-451f-4c82-b915-265c703ad737", + "5f615b11-2ff8-4ff1-987a-1cd61ea78dee", + "de96ee14-be42-471d-b110-aabc948260fa", + "8504aae3-4f25-4fa4-88c1-beb3fba83f13", + "d2c0f78a-978e-409f-a78f-3e40b3431909", + "0f846b5d-1cea-40b8-b18b-dc6c54544a67", + "37afa156-a068-458e-8bf7-ec6f1f2415bc", + "2d205d38-f79c-43f2-8c39-ede697065ab2", + "36e7b2ad-dddd-42fa-a427-6411ff12f9b9", + "355f8f4a-cde2-4108-821a-5ccea1d297fe", + "c37e680f-03b1-408b-9b09-1b2ae733951d", + "6958467e-388c-4f6d-a8c0-d4782166daac", + "3051e028-71f8-4f3d-a72a-010934c3566f", + "fc0f4c90-a1ee-4147-990e-de231394f484", + "c2af1389-b165-42f5-b4ee-c90402d7006b", + "1acfcabc-653a-4c03-adf7-64838d09a7b4", + "3b79017c-4d42-40fc-a1bb-4a20bc8ebca7", + "6edb7036-3369-4c03-8e3a-439ffbeb4724", + "2ac543c6-b52b-4219-b2e4-4130078c0673", + "b3ef3563-ecc0-4a15-9a7b-4feb4558d953", + "6f001eeb-fed5-43e4-bfa2-84e274cbe5cd", + "1ce2252b-31b8-4536-9902-0ba7fb039340", + "c2413391-7c1b-4fd6-8b1d-98ee7888b9f8", + "a4d53159-e2c0-4a47-9e78-af9bc9291f0e", + "66ddb7ae-90d6-4a58-9f3c-7f1d2604ecb0", + "290fc9ba-0a29-4eed-b0e0-64edc5bf54ea", + "e3a53edd-52ff-4608-a80d-77ecfbb0f961", + "50895f7f-5667-4b43-94c1-d7b4b9e238e7", + "3aea2d62-0340-4c6a-8019-3cce9117a65a", + "63ab0246-2f1c-4621-8b71-75075b30d852", + "228eb504-c1b3-43b9-818e-2f7e539a259b", + "02de9500-d574-446e-95fc-3e82c367385a", + "bd03445e-9016-423a-bb48-cd998836bbff", + "a16cb772-d0e2-4654-96d7-372c21484b90", + "0fa75089-7b28-4988-8664-a2c93542ffb0", + "fcc6af87-918f-4d0c-a0a4-8aab36891fc4", + "b9d7642e-2ced-4589-b7a6-69ad54570c4c", + "ba5b1882-c9d7-48e7-85fe-2c7b10494fac", + "92e0925c-566d-4a67-b93e-3e4b57f9b3ab", + "130a49d0-0315-4eb0-a752-a4533bd880a9", + "0fccd442-a901-4d77-aac6-dc4276cd6591", + "49c27d75-44fa-4947-899a-8d0511a10416", + "e2ffc7c9-c035-4709-97c6-3764499188a5", + "2069aea3-ac99-4b70-891e-1bfb8b15aeea", + "cdbc7718-4cb9-4e78-bbde-3d9614f8dcd1", + "32464967-b0f2-4b2e-a57c-7f774e7056a7", + "ae6466f8-1a53-4cb3-8718-96f78f391b4b", + "d66c7c66-2b5f-4e7f-bbb2-f0adae68dfcc", + "da84b37f-ec6e-414e-b52e-1231104a68e0", + "19866484-b3d8-4619-85b9-e900de7fa80d", + "c4392ca6-4706-4a96-a6b2-29e6ef544500", + "6798a65f-d79a-413b-baca-115f09caf0c4", + "b844cb2d-6428-4d5e-abac-958f2704d4c8", + "8e9c0b21-2230-4a0a-b35b-a198b90abcaa", + "3492c684-74b5-4317-9e5c-320f3ea01899", + "5beb520c-d413-4998-bd9f-7960a2baf744", + "c96fa96d-d23c-444e-bdef-481317308bb8", + "502ce8f6-bd9e-4e40-a5c9-40bdce36b87f", + "221fd24b-ba2e-40d5-a0c7-a879ece6b34f", + "fb62dc3f-a1cd-4f0e-88bc-504fe9d8585d", + "e9acbb7f-c45c-47d7-9dd8-16194319fc97", + "0d4efcc9-a336-46f9-a1db-0623b5be2869", + "84a388b6-2844-421e-a176-84d29958a76c", + "3ecc3e4d-5466-4243-acca-2feea47e09a2", + "e7b6b4ee-3f2a-41cd-bc18-c7e993cb413f", + "5998ece4-98cd-4d5b-9104-1517d13728bd", + "c9f84bf4-108c-4151-85c0-e87375c53b91", + "5eaa3bfb-cc71-43fa-88dc-d367ef929a7b", + "0fe736da-48f1-4961-bd7a-b75ce13c91b4", + "2c7796e3-b1e9-413f-a0e8-8225ac7d586e", + "69006355-3fb0-443c-826b-ea8c9d27f70b", + "8d73dd20-5349-477f-9630-218e6b08bb95", + "e7633ad6-e406-4a18-b1ba-6c834ec9c047", + "3b3a4043-3d74-4575-8aeb-511c4a879241", + "5daa132e-0b97-4d0f-be9d-b514285661fe", + "52fb1a73-2ba7-413d-8596-33615daecc73", + "9f3cb8b5-5d62-4236-be5e-8f16d59c9f50", + "efdc3dc0-275d-4073-9179-32b14b224654", + "683635b7-5a1b-467c-8d4c-c34e8bf5fdf5", + "2db7496c-aa43-4d9d-8df6-404d2d9bd810", + "9f3edf40-3dca-4086-b995-4985f6848a14", + "ef37e1e6-5be3-4e8c-9f70-c84f341f9aea", + "5df66408-7a8c-43bb-9235-d088aa3d19dc", + "453f43e9-c1d1-4749-9720-73546287b82f", + "1943bfb1-ba78-4a6c-ade7-11e44314dfb4", + "fd8cc20a-915e-45a9-93df-320ad998aee4", + "bc6cf777-fd8a-4427-9867-92a443826827", + "812f0590-4ccd-4c9e-8978-c45e6debce55", + "e6f4e36f-96b7-48f5-afd1-27cc408bdd07", + "1b6dfc9f-c873-48c6-9f25-a2d490e269cf", + "bdf1ca5a-321e-4d3c-ac9e-da2ce1f5a69c", + "604bf883-ce55-4c59-bcc1-3ceea1128211", + "679f2016-4802-4bf8-a830-d1f45fdfcd71", + "6f3f70b4-8428-4e5f-b1dd-86cd55d8e63f", + "83d8be38-ce02-41b9-b444-7291bad513d7", + "7fd21591-6a33-49cf-a32f-f8dfdf5ae697", + "a3dbec39-7126-4949-a916-eca55fee2b49", + "03cebfe9-3319-45b5-aba5-e7db10ea6fdb", + "ff9fe54b-a959-4d7e-a6fa-93e311a68168", + "172e1afa-5581-455d-9a27-590e9a1706c2", + "3ab3b82f-bfa8-4b39-a9a9-e397341ae572", + "b171dcaf-94f2-4bcd-8695-b2648bed60e1", + "063f5040-9e94-42ae-b0e5-b13413e0fe55", + "136cef9c-e17f-49f9-ae45-a17e7e5c5477", + "f906b093-281f-4872-8010-f1f3a320ee2a", + "33365c93-7f1d-4a7d-a9bc-b33010aeafbd", + "e98405ee-0a5b-4df2-bb7e-97729b8b3c11", + "bd13e435-1e38-4b22-bfd0-b620fc74cc8a", + "7dd0eab1-b83e-4bb4-ad91-7bbdf4592647", + "5f5ea963-4f59-4705-95fe-4a98c9fcd12c", + "7e6dfd31-8544-4fad-9e49-0f05516cf8cf", + "f817b176-2dce-4858-999f-fe62aac25e2e", + "83553642-1a38-4c54-8154-c4cd4fd81173", + "dfe025f3-573e-4297-b17a-30e990cfa310", + "2091513c-8f4c-4817-8c22-cdbc6303a331", + "d15eb34f-e425-4990-aa2d-2dc454fb4b20", + "0e519838-1e4a-41fd-ab30-34bf745a3339", + "c920e557-141f-4556-95ac-f494ecc0c84f", + "91e97b3c-c912-401a-ab99-da01b90682aa", + "d6adb1ae-f689-438b-ab81-8a2e37f9654e", + "cd3092e1-6022-4a11-bfc3-a4bb80487ece", + "a874ef10-022c-43e4-b4e6-e01d8fcec988", + "2e6ff051-0037-4001-90a8-de7643a96f08", + "134c4387-77d5-4ee0-911f-e50e64499b42", + "421f1c8b-5237-4735-9003-82021ba651b1", + "9ee0474c-615c-4383-a047-8432b8b24784", + "78efcc13-3718-4acb-acee-8157a54501d6", + "961a8de1-0d33-44ab-a9d5-3d92de13fde4", + "49e25cbd-1aaa-40dd-901b-bb3e524ebcc0", + "5a5b8c37-e3f6-48f2-96d3-289a043eaee4", + "394d1df6-bced-4f04-8c2a-a1ea3a04b984", + "b87c3e0f-a5af-4663-a87d-d0689d73343b", + "09194996-ff4b-4a4f-9800-e8aedc9be7f2", + "3f920762-f4a7-4bdb-86dc-05a3c3c8e875", + "5204624f-90ac-47bf-a71a-486d43f4351c", + "df24897d-a478-4fbb-816f-1d14c59da30d", + "9ba7c2a0-c86f-4967-9284-9be88ad3eacf", + "3f28dd4f-f653-4937-83b6-9f80c1284a85", + "b2735994-59ae-4478-9c5f-93f5ccf18c8b", + "5e965191-bcd1-491a-a9ce-e014e970020b", + "85718cf6-e4ba-44e3-880f-143c627f12ff", + "d0af657d-9c21-447c-a4f9-199b38deac5b", + "f8289825-2858-4dfa-b173-f4c506359158", + "8f2d26e6-239e-492f-aa1b-472f8ec85050", + "836be9c2-cee6-4fb9-8bbd-74197d213c5a", + "cd998593-e5bc-49d0-86a7-333acf9abf93", + "8deeeb88-19ed-4747-9c79-d94104bed7e0", + "6d1dc078-8b04-480e-a36f-298594ee9592", + "83464c62-929f-4bf7-9a04-43071189edde", + "1efd8b7b-65f9-4e73-bd49-d707190d6519", + "f0b69458-25df-4dda-be8b-e14383a57a50", + "88260ff4-be0b-4544-98ef-78b6a0d61ebc", + "39fb2989-1dd2-490d-b793-dce86c36b43a", + "b4f650b4-53d6-4069-a55f-43f2de7ff47b", + "15b6759d-2370-43ee-9a61-09e453db2469", + "2e4e0dac-ab21-49ff-9563-e320acd73f2f", + "986c4e43-ac2e-4983-8d56-f1dd34775133", + "12c1b77c-cbea-4c20-8b64-7154a25bbbe1", + "a3da62c3-3f46-499f-a4c6-acf93bd1c63d", + "57b6455e-2acb-4076-b8c4-1f2cdae97185", + "d6d8261e-0ee0-4002-bddf-e94fdc1075c7", + "56683d7b-7c65-4f42-b5a4-818f3cd5224c", + "9090779d-803e-4fbd-9166-d9e90ff2d36f", + "d5374572-5576-4f83-99c8-a0b12c4de160", + "ffddab3a-dbec-49b0-b114-5a1e3d805efa", + "38c02527-1b2f-458b-bad8-73a282c1d99a", + "d612edfb-e904-4de0-abce-a5fd9ecc449b", + "1261ee8c-80fc-4a3f-9dd3-29f4a918a9fe", + "8e860a91-0205-447b-a706-aa894d23d3db", + "fb07f205-868e-44c9-a704-305daaf03301", + "48f7028c-b6bf-4ca5-81e9-128147fd02c9", + "c6a6e1e4-24c4-422d-9d18-f062240ddf23", + "ba280b86-22a2-4509-b16d-b55e04680461", + "f0cb79ef-c081-4049-85ef-2623e02c9589", + "6dbfc053-7899-40d8-9827-0ccd21f3ee0a", + "50d3035c-2503-4bcf-b28b-e66c92c98f09", + "41f7665a-a692-42e0-b28d-9769c03ab99f", + "cb1f70b4-0bf6-47f1-bc98-2e9365a55c3f", + "7f566780-4ed2-41c5-9faa-06759495faab", + "0e08e1c7-3add-4e41-b28d-cd5f4413a422", + "672e0d75-df0e-4715-8c42-f1131b059a6c", + "8f15aab0-9c42-452e-b78a-47a48da1f3c1", + "edaa0b1a-d504-4852-8025-4f12b4d293ac", + "4a9ada82-9f86-495e-b17f-25b39bef2297", + "07774a7c-c5ed-4133-991f-926a316068a4", + "5a009338-0aff-4313-b02f-c48630ce1702", + "8324c0ba-7a10-42e1-90fd-afb4d3afa789", + "3db37314-adaa-4a4e-af11-c9190fa6b65b", + "b21b35d2-6d69-40a2-8c36-aa27e4a980ae", + "6daa9124-d894-42e5-8d98-2e0e6d1191ac", + "8a256d7f-2ea1-4b0e-94ea-168a1ed85678", + "08b6f5b3-4471-4cfd-a95e-4901cb1ded57", + "ea9d85e7-6a9a-451a-b019-972a6891c11b", + "f84fd67e-0797-44f4-b3ff-78da646362e2", + "ad1a9ea9-d1ab-474a-b715-1a42d892ef5d", + "ffb191f8-894a-40e8-baa5-5578aa6ba706", + "fd394f00-4049-49e5-9c07-ee3e86979003", + "9767a3ca-0a9c-480f-adfe-e4164c52c760", + "95cdb269-0088-4451-9a0d-f72349081fab", + "4ade0a82-16fd-471f-9fcc-fd2724887aef", + "fe69a78e-7773-45d1-9e0f-bacee52dac83", + "3aed7bbd-75f0-45ef-b9cc-7f5ff6f7fbbf", + "f8d71e9a-982a-437c-a968-867149c24eb7", + "8351998e-ead7-4bc1-b0a7-a7209da92b3c", + "3eb00a40-880a-4ad9-8cec-3ba0245ec472", + "602b224a-f5ac-491b-88e3-49043f362b6e", + "89125bc3-a413-44ee-a7e1-2753c511ecab", + "00792fa8-988c-4c85-8e80-73eb3ac53e80", + "094d1ab2-f34e-4a2a-ac3b-8671ce5faf65", + "4cf10990-13e5-4a1e-a913-7ec4b8c8ab02", + "a9b191de-d565-4981-8f68-02069a2ba16c", + "0eeab9d2-a55c-4708-94c3-cc1e08f9f262", + "4e78f728-06f6-43a3-ac8c-c0cb159e5b32", + "be99e9e2-82f7-4f1f-aff5-371834fde0da", + "5682f9aa-bdfb-409c-b4c2-c0105b35ee78", + "3ce15244-5293-4a5e-8c92-52ab29c6949a", + "b4b2c55b-c107-468f-bbdb-f73a0e4dca34", + "8a409551-657a-48c0-8beb-056c228ebca5", + "3c17a762-c839-492f-b361-009c1187ae4c", + "22172001-0e17-40db-bcce-825e19532297", + "5b99234b-b940-405d-9123-70218bda24b2", + "f7869fd8-9af3-493a-a627-a3e95e7c3577", + "057e9f41-e6ed-4575-906c-aca9413b35a5", + "70e88e82-5c06-44fd-959c-d4b2d40999cd", + "b95bd349-11bb-4ade-9cc6-3be8dcea06dd", + "19b880da-910e-40c7-970e-cd8718f2faf5", + "e05a2182-a97f-4d0b-87e4-66b0dd405e58", + "3072481a-04fe-4959-8410-018a41a91df1", + "fc2b2014-9dc4-4a5d-8a1d-25a6911bff7c", + "db1d3def-397e-470f-811a-85c4cc5338a9", + "ff283659-4a79-4642-8078-4366df38e1f1", + "bc8bae7f-ad9f-4f72-ae5c-b3e7ca8dd8f0", + "a3bd1be3-edfe-4f4e-89a8-dd56ce6ece08", + "2360312c-4347-476a-990f-2633b489fa7a", + "b9e32e51-8c18-420b-a582-200c79b8208d", + "0511dce9-c126-4a11-92ba-8082a4ec89bd", + "224d8b17-7eaf-4529-9b75-c6ce3078d640", + "10732476-f045-4feb-8199-e4e04685f7ca", + "5753e3fe-f1c5-426b-9f45-5a8efa6b5370", + "b69564f5-5e5e-4097-aebd-9f52f211ff75", + "97eb574d-d8bf-44a8-a6ef-1204b789db43", + "85506322-449a-45c0-b77a-48a8077b4bbd", + "ed65870d-ba08-4894-9cad-7f34e2c6f9e4", + "cf682058-962f-45c7-95ab-dbbd2f87acc0", + "65229ac2-9033-4002-af6c-492fe49bc563", + "4c85a780-8a8e-46c9-a0f8-c7e9703ef998", + "1ac6e060-a6d2-4b42-afef-599be78a533a", + "a22beaf0-377c-4b95-b472-867280a525ae", + "f5bc6a45-1591-42bf-8ed4-8286382b31bf", + "06592f50-7f29-4ca9-94a9-1b6c389d50ea", + "4ca3d9cf-4972-4e1d-aca3-28b036d6bdc5", + "076bdb81-5c75-4282-9f3a-a387624575f3", + "a5fe8210-1c37-4dd9-bd27-74051c521c3d", + "c1ca7c5c-2473-4b83-a507-7d187609111a", + "7c4079e7-4d68-4b0d-b015-5971a9866c60", + "ead54fb5-dee9-452f-8145-227ccd2405e2", + "7c776700-1d6b-4ee2-8155-b7f6855afceb", + "71a7456a-da6e-4e38-86ec-76bacaa6db5b", + "26c042d4-bf21-4d10-81f9-303d3816c114", + "c7820a4f-55c9-44eb-bf7e-b965a9d0e924", + "5b5b6b8c-9b88-4181-9ebb-14abb7783e1e", + "db12030d-7577-435f-9d48-5c1aae615c4f", + "e37c82c5-483f-4e68-8bf3-2a22dcef0b1e", + "0b4cacb1-970f-4ef0-85da-371d81f899e0", + "3949f121-af2b-4c7c-8a49-07befc08abbe", + "73c1ac5a-80a0-4e55-b143-52df5697c9ea", + "0f325614-48c9-4a6a-b7f2-731c26ddeb97", + "572117cf-a76b-44d2-804a-4901bfc545f4", + "33fd4f1f-fb01-4d3a-9b8d-b8d8024a2dab" + ], + "test_session_set": [ + "66cb6d16-f972-40e9-88c6-64b6816b900c", + "e62b75a0-f1f6-4e45-b3db-17fd7ff5756c", + "13ffc0e4-eebb-4baa-97af-5983ce46779c", + "d062849b-c60f-4369-a113-50f6ce8521fb", + "bb7f1e75-65ee-4d98-adc8-c70551a35940", + "7a6576d8-007c-4ee9-b7cc-8cc2d020e1bf", + "4ce65f3a-21ee-4f77-85c0-b3ddfe1e6381", + "783d7043-b27e-4744-a969-5f40fb40827c", + "81b74694-a2e6-4be1-b9f6-69f3aa03f50d", + "e6317c78-06f0-4c6d-9886-ca67712772ed", + "44bcb1fa-9e64-4bae-927a-4d5d53e536a6", + "218e79aa-3f56-4bd4-aa2a-14027bf7d52d", + "561fd770-f9a2-41a4-90da-beeb7d96ee7f", + "a2d94b61-57c7-4ec2-ad94-94b3585a6b64", + "764a08dc-5ac5-4c09-a281-0210cdd5a55e", + "c1057c9b-331f-458b-8914-2703511e8b13", + "3edc700d-7d14-4154-a3e7-5b0e8c3fde07", + "1492c03c-1887-4576-8889-5bc2848967de", + "2ea37ccc-e1ec-4e0d-a2a7-9707ea300b7f", + "65cad708-34b8-4158-be96-b88710288473", + "be8f4445-3958-4fc4-8cdf-b14d1d0eba73", + "3c62a0d2-10d7-4a9e-ac2c-02505663afad", + "be25331c-b1a0-4187-93fb-3b59437258c3", + "108cebaa-777d-42dd-8b1e-69ea7f443a0a", + "be3b65c8-ba74-4321-99b4-ca0e99e6fc2c", + "3bd74a94-ec3e-4138-b78f-ad76af4f07e2", + "4356f6d5-576a-4999-af40-744c242fff22", + "b0c57748-65c8-49d7-9681-ddcfcde3298b", + "12706b2d-0a0b-4511-8feb-ed8065746e50", + "f0e9a139-86f6-41bc-82f1-11283954a781", + "bb3c6009-892c-4071-9bb4-692b83ed12f3", + "90bf3f44-60ae-4430-9dc6-2365f7901fc5", + "2b5cea2f-b019-4ecf-abbc-b2346bd65a10", + "8e9f89d2-f02a-4c75-8831-6785ef6bd7f5", + "13c24739-e514-4fad-8288-6686bffa71ad", + "1828b207-9601-4521-8c0b-5d9e829fca78", + "0324c3b5-5154-4de1-9869-eb866b568e0f", + "1a746147-2dec-4a37-a51b-c94e4f1de010", + "3a20bf0a-6eb1-47c5-9bf8-8445446feb61", + "c2a8f168-06b0-4766-8865-27f08cb1bcae", + "f55f2dfd-704f-42cf-a39e-72e2560314b3", + "4ea88cb9-bff8-49e0-a855-cc1902db9a99", + "4e0053f8-566c-4952-abbe-4ea877782a0a", + "80ee109f-b097-4755-8a39-9f0b33adb170", + "933ddf59-52d2-4d4c-b653-05de44ded245", + "b432e10e-cff1-4f15-a8fe-1e2350e94b36", + "157c8e81-117a-4a15-a047-9a574181d44b", + "cfc6da55-43fc-4cd8-96d0-9de5166f4a63", + "aabeb7f9-5ada-425a-9545-c98dc443887a", + "db9978fb-fe4d-467b-9495-ab7daf3b1b55", + "6549c35b-8b1d-4577-9434-ac808e4d7057", + "e50126ee-f4c1-41a6-a2e8-79636b40b341", + "49a7e53b-aba5-4e53-b3bc-1492657fc8ae", + "7abe6d4a-f594-4a15-af08-ddd591ff500c", + "03c50078-a08d-43ad-a3c9-87916aebd865", + "a9a0ed92-fab1-4b6e-ae27-714692464641", + "3c3207e6-57b9-4be1-96fc-5b2168ae3cd3", + "895545f1-f194-4c0a-8513-87d386670f15", + "7b76be35-4e6d-42d7-9fa9-14a21fc327cf", + "e4670e4f-0b1a-450a-9778-9ee2409db554", + "34bc808a-4b98-42c5-a1ea-6f565dd4aa20", + "bec4e063-9be1-402e-a7b5-dba20730c933", + "20a3e8e5-3f43-46a5-84fa-6888bd9c5c88", + "5da4efa3-8e54-46dc-a8c2-3016fa4e12f8", + "0d998741-de68-4006-898e-7108de4167b1", + "4d144ff3-a853-4b79-9825-d7fcea248373", + "dce1d8b4-2c29-4fec-a343-77885899c434", + "dab038d7-9e08-4edf-8d03-ee7e2af5a10e", + "a1055434-9e9b-4d69-bac3-374a39f801da", + "e7086057-b5f7-4286-87c6-ab110072c5b0", + "3af65596-1aa5-4126-933b-b98ccaccedf2", + "19f880ae-9770-45a5-93f5-00047f22d5f6", + "36888bbd-1ece-4a3a-8aff-b0f6ef905752", + "f42a7e14-c3ee-4b05-a59b-8d349f8bd19b", + "d36d9514-a99b-4534-b54f-3fd98a19f0ea", + "6c4f7cb0-caa0-4a35-9071-a84c9ec2ab2f", + "666d9243-f4f9-47f5-8fba-cee8568abf41", + "2f14e742-dc58-43c8-8ebc-7244487dbba7", + "f8accf9d-4427-4458-88e9-89d51a3f5136", + "50810487-0b0f-45e0-8f37-515644189e81", + "f0445727-2729-4a76-9434-fb081a14cd9d", + "f55f890d-fd90-43f6-a11a-a1920c21890a", + "16851d4d-35e0-49b4-b3b5-28e67b564db3", + "5c8f8839-4bf9-45e0-beb0-6c17a5ead0d5", + "514d71c8-fd4e-4908-89a4-c0861e81ca8a", + "35465081-a518-4321-8384-477ef5b29b8d", + "de82a6f3-1a55-4996-ae68-35032678ff66" + ] + }, + "EGO4DSHUFFLE": { + "class_index": [ + "background", + "camera wearer speak", + "other people speaking" + ], + "train_session_set": [ + "5b8c91ca-9d16-48b7-acce-4dec221ffd71", + "0232fbc8-c9b1-405c-bd42-a9d773fd4835", + "39c74452-1957-40db-894a-b39df51e1437", + "31488ff6-05ef-4362-9786-35981b06f307", + "4decc0fa-e98e-4bf9-bbd1-d9b852b04830", + "efeaa284-2e22-44b8-bdd6-1bb4763e5b35", + "07f1486c-44a7-4c69-9532-720c81c99eb0", + "19c28cbe-21f2-4b40-a6e0-e9932b5f0605", + "1c5c0b52-1287-4df3-aa75-8ccc916b1c10", + "e92255bc-1be6-4ee9-8f8b-d1336a0e409c", + "307e8837-1fda-4cf5-bab5-d16fee20f84f", + "bb050b8a-195a-4b50-b5d0-1ed535226a39", + "ca22f118-2472-4df0-a9cd-3d70424b211b", + "08b0935e-6260-4bd6-86ca-f6fc54e388be", + "666a21c0-7469-436e-868f-64eb8022f64f", + "77e09fd6-b560-4439-addc-a1a8a3122046", + "f00bd659-56e0-41dd-8eba-c5fdb45e7232", + "c2c7f072-a02a-480d-8342-b43ffc314db9", + "ed72b048-e309-4510-b049-b0288ce4478a", + "56c5af79-f9d4-478d-96ef-6d71e0bbbdfe", + "3d8acf67-fad6-421e-841b-82b28c82f7f6", + "977ac980-f908-4c29-a2fc-dd14ec21c0bf", + "08742cb7-bb6d-47ae-af5e-3f1f306ccda6", + "f3afa5a1-8a97-4a7d-945e-1cc70931e44a", + "65172b19-a246-4dc5-98e9-1b526b00ac89", + "ee1893cb-2ff7-40e0-af8d-19518696b698", + "53144267-6da5-40d3-91f2-eaa85be367de", + "26080fb8-0588-449c-a7f2-2e0162507e02", + "c0ba997c-9ec6-4e83-ad50-a88dd677b881", + "6640a311-86c4-4de7-964c-e3a4a67a3e02", + "5f1fcf84-4001-43ec-ad88-9e06ea2f663d", + "ea5a3f9d-7b21-46d7-a17d-bcbd4955bae9", + "eb3675b8-f6a8-4ea9-82df-1080eac5b093", + "023aaf52-fbe9-4b61-af69-ed33e91f3717", + "86af705a-a326-423a-bdd9-c6734804c30c", + "55363f67-1b41-476f-a83e-7e9741786ad7", + "ce8ea1a3-65e3-4ef7-abeb-2b0bc1c84773", + "d3b7e314-f48d-4888-b496-e1eeb252cd8a", + "59000084-d286-4c83-90c8-2fbbf80d2d31", + "a47c590d-66aa-4bd1-b552-f40951ba9bf0", + "61c73721-9476-4ac2-a987-ecaf9a04dca1", + "0124376a-720d-4265-8c8c-5ff47f04d579", + "f62a8100-7ed2-492d-9709-064514d552bf", + "13d71019-9093-442b-bdb4-c79310150592", + "74e71420-9bd6-4337-9a7a-a3dcc4cab7cf", + "383466b5-cb00-4e76-9e6a-dc1b0c9b125c", + "9ef6ef95-6b97-4c78-866c-641fc7d791f5", + "ac7d3d16-d74d-4fe3-a1e0-cab0b1adf492", + "c46cbba9-2cb3-4d84-be18-71af5539ecc5", + "5eb74ea5-9b0b-484c-82b2-ee86479c5d07", + "6f80f9ae-c276-413f-accc-fad74efe8a95", + "73f99cf0-d997-4488-b6a6-993934a9fad3", + "fd50b1d6-8142-402f-88c5-7c44ca215081", + "51fee605-d077-4d78-babc-9bbcd8d898ca", + "5e3c6abc-599b-44f8-8401-d5720627bba2", + "91151b86-65b7-4854-97b0-9c5b93a251ed", + "5be9a2e2-5efc-4fc6-a65e-806e0347e19a", + "94ba75fc-8d7f-4bb4-a756-ded62a4c1761", + "c925c38c-3fe3-4eb3-9944-3fba3eaca4db", + "d97bedc8-72df-43be-a55b-4da1ae42dfd1", + "e60b110e-a9b6-42b7-8f92-082cd74c4937", + "3627996c-8188-4486-8feb-938b9b0cce2b", + "d2b674ca-9f02-46ba-820d-f0b52c76749f", + "2b90b963-e829-4787-8fcd-1a1a434ffdfd", + "ba51db0f-2bf1-4746-894a-7db0ba54e357", + "6c1bac5c-aead-40b8-9c17-be97d86b68d4", + "63483ea2-d569-45fd-82b2-f80836a8b766", + "c81eb608-6844-424e-b8b7-524ff652cf3a", + "cae5d4bf-bbf9-45dd-876b-a48579782348", + "7813b44f-04ca-4a53-868d-fe185b615b68", + "1f12c871-9b3a-4611-a2b4-9c39059052a4", + "70316529-ec02-4d26-bc68-4b436bf0030a", + "6d367aa4-56ef-4932-9d48-e7530e717ea8", + "2e651150-9338-4103-b231-2f089aa19df7", + "d2cb4045-8ef1-4288-8398-ab6d4d6a0709", + "3e4b9c29-06d2-4ade-a2be-bdec028bc68c", + "4f3c8276-99e0-4c99-9e61-5d7c06332b09", + "12634773-451f-4c82-b915-265c703ad737", + "5f615b11-2ff8-4ff1-987a-1cd61ea78dee", + "de96ee14-be42-471d-b110-aabc948260fa", + "8504aae3-4f25-4fa4-88c1-beb3fba83f13", + "d2c0f78a-978e-409f-a78f-3e40b3431909", + "0f846b5d-1cea-40b8-b18b-dc6c54544a67", + "37afa156-a068-458e-8bf7-ec6f1f2415bc", + "2d205d38-f79c-43f2-8c39-ede697065ab2", + "36e7b2ad-dddd-42fa-a427-6411ff12f9b9", + "355f8f4a-cde2-4108-821a-5ccea1d297fe", + "c37e680f-03b1-408b-9b09-1b2ae733951d", + "6958467e-388c-4f6d-a8c0-d4782166daac", + "3051e028-71f8-4f3d-a72a-010934c3566f", + "fc0f4c90-a1ee-4147-990e-de231394f484", + "c2af1389-b165-42f5-b4ee-c90402d7006b", + "1acfcabc-653a-4c03-adf7-64838d09a7b4", + "3b79017c-4d42-40fc-a1bb-4a20bc8ebca7", + "6edb7036-3369-4c03-8e3a-439ffbeb4724", + "2ac543c6-b52b-4219-b2e4-4130078c0673", + "b3ef3563-ecc0-4a15-9a7b-4feb4558d953", + "6f001eeb-fed5-43e4-bfa2-84e274cbe5cd", + "1ce2252b-31b8-4536-9902-0ba7fb039340", + "c2413391-7c1b-4fd6-8b1d-98ee7888b9f8", + "a4d53159-e2c0-4a47-9e78-af9bc9291f0e", + "66ddb7ae-90d6-4a58-9f3c-7f1d2604ecb0", + "290fc9ba-0a29-4eed-b0e0-64edc5bf54ea", + "e3a53edd-52ff-4608-a80d-77ecfbb0f961", + "50895f7f-5667-4b43-94c1-d7b4b9e238e7", + "3aea2d62-0340-4c6a-8019-3cce9117a65a", + "63ab0246-2f1c-4621-8b71-75075b30d852", + "228eb504-c1b3-43b9-818e-2f7e539a259b", + "02de9500-d574-446e-95fc-3e82c367385a", + "bd03445e-9016-423a-bb48-cd998836bbff", + "a16cb772-d0e2-4654-96d7-372c21484b90", + "0fa75089-7b28-4988-8664-a2c93542ffb0", + "fcc6af87-918f-4d0c-a0a4-8aab36891fc4", + "b9d7642e-2ced-4589-b7a6-69ad54570c4c", + "ba5b1882-c9d7-48e7-85fe-2c7b10494fac", + "92e0925c-566d-4a67-b93e-3e4b57f9b3ab", + "130a49d0-0315-4eb0-a752-a4533bd880a9", + "0fccd442-a901-4d77-aac6-dc4276cd6591", + "49c27d75-44fa-4947-899a-8d0511a10416", + "e2ffc7c9-c035-4709-97c6-3764499188a5", + "2069aea3-ac99-4b70-891e-1bfb8b15aeea", + "cdbc7718-4cb9-4e78-bbde-3d9614f8dcd1", + "32464967-b0f2-4b2e-a57c-7f774e7056a7", + "ae6466f8-1a53-4cb3-8718-96f78f391b4b", + "d66c7c66-2b5f-4e7f-bbb2-f0adae68dfcc", + "da84b37f-ec6e-414e-b52e-1231104a68e0", + "19866484-b3d8-4619-85b9-e900de7fa80d", + "c4392ca6-4706-4a96-a6b2-29e6ef544500", + "6798a65f-d79a-413b-baca-115f09caf0c4", + "b844cb2d-6428-4d5e-abac-958f2704d4c8", + "8e9c0b21-2230-4a0a-b35b-a198b90abcaa", + "3492c684-74b5-4317-9e5c-320f3ea01899", + "5beb520c-d413-4998-bd9f-7960a2baf744", + "c96fa96d-d23c-444e-bdef-481317308bb8", + "502ce8f6-bd9e-4e40-a5c9-40bdce36b87f", + "221fd24b-ba2e-40d5-a0c7-a879ece6b34f", + "fb62dc3f-a1cd-4f0e-88bc-504fe9d8585d", + "e9acbb7f-c45c-47d7-9dd8-16194319fc97", + "0d4efcc9-a336-46f9-a1db-0623b5be2869", + "84a388b6-2844-421e-a176-84d29958a76c", + "3ecc3e4d-5466-4243-acca-2feea47e09a2", + "e7b6b4ee-3f2a-41cd-bc18-c7e993cb413f", + "5998ece4-98cd-4d5b-9104-1517d13728bd", + "c9f84bf4-108c-4151-85c0-e87375c53b91", + "5eaa3bfb-cc71-43fa-88dc-d367ef929a7b", + "0fe736da-48f1-4961-bd7a-b75ce13c91b4", + "2c7796e3-b1e9-413f-a0e8-8225ac7d586e", + "69006355-3fb0-443c-826b-ea8c9d27f70b", + "8d73dd20-5349-477f-9630-218e6b08bb95", + "e7633ad6-e406-4a18-b1ba-6c834ec9c047", + "3b3a4043-3d74-4575-8aeb-511c4a879241", + "5daa132e-0b97-4d0f-be9d-b514285661fe", + "52fb1a73-2ba7-413d-8596-33615daecc73", + "9f3cb8b5-5d62-4236-be5e-8f16d59c9f50", + "efdc3dc0-275d-4073-9179-32b14b224654", + "683635b7-5a1b-467c-8d4c-c34e8bf5fdf5", + "2db7496c-aa43-4d9d-8df6-404d2d9bd810", + "9f3edf40-3dca-4086-b995-4985f6848a14", + "ef37e1e6-5be3-4e8c-9f70-c84f341f9aea", + "5df66408-7a8c-43bb-9235-d088aa3d19dc", + "453f43e9-c1d1-4749-9720-73546287b82f", + "1943bfb1-ba78-4a6c-ade7-11e44314dfb4", + "fd8cc20a-915e-45a9-93df-320ad998aee4", + "bc6cf777-fd8a-4427-9867-92a443826827", + "812f0590-4ccd-4c9e-8978-c45e6debce55", + "e6f4e36f-96b7-48f5-afd1-27cc408bdd07", + "1b6dfc9f-c873-48c6-9f25-a2d490e269cf", + "bdf1ca5a-321e-4d3c-ac9e-da2ce1f5a69c", + "604bf883-ce55-4c59-bcc1-3ceea1128211", + "679f2016-4802-4bf8-a830-d1f45fdfcd71", + "6f3f70b4-8428-4e5f-b1dd-86cd55d8e63f", + "83d8be38-ce02-41b9-b444-7291bad513d7", + "7fd21591-6a33-49cf-a32f-f8dfdf5ae697", + "a3dbec39-7126-4949-a916-eca55fee2b49", + "03cebfe9-3319-45b5-aba5-e7db10ea6fdb", + "ff9fe54b-a959-4d7e-a6fa-93e311a68168", + "172e1afa-5581-455d-9a27-590e9a1706c2", + "3ab3b82f-bfa8-4b39-a9a9-e397341ae572", + "b171dcaf-94f2-4bcd-8695-b2648bed60e1", + "063f5040-9e94-42ae-b0e5-b13413e0fe55", + "136cef9c-e17f-49f9-ae45-a17e7e5c5477", + "f906b093-281f-4872-8010-f1f3a320ee2a", + "33365c93-7f1d-4a7d-a9bc-b33010aeafbd", + "e98405ee-0a5b-4df2-bb7e-97729b8b3c11", + "bd13e435-1e38-4b22-bfd0-b620fc74cc8a", + "7dd0eab1-b83e-4bb4-ad91-7bbdf4592647", + "5f5ea963-4f59-4705-95fe-4a98c9fcd12c", + "7e6dfd31-8544-4fad-9e49-0f05516cf8cf", + "f817b176-2dce-4858-999f-fe62aac25e2e", + "83553642-1a38-4c54-8154-c4cd4fd81173", + "dfe025f3-573e-4297-b17a-30e990cfa310", + "2091513c-8f4c-4817-8c22-cdbc6303a331", + "d15eb34f-e425-4990-aa2d-2dc454fb4b20", + "0e519838-1e4a-41fd-ab30-34bf745a3339", + "c920e557-141f-4556-95ac-f494ecc0c84f", + "91e97b3c-c912-401a-ab99-da01b90682aa", + "d6adb1ae-f689-438b-ab81-8a2e37f9654e", + "cd3092e1-6022-4a11-bfc3-a4bb80487ece", + "a874ef10-022c-43e4-b4e6-e01d8fcec988", + "2e6ff051-0037-4001-90a8-de7643a96f08", + "134c4387-77d5-4ee0-911f-e50e64499b42", + "421f1c8b-5237-4735-9003-82021ba651b1", + "9ee0474c-615c-4383-a047-8432b8b24784", + "78efcc13-3718-4acb-acee-8157a54501d6", + "961a8de1-0d33-44ab-a9d5-3d92de13fde4", + "49e25cbd-1aaa-40dd-901b-bb3e524ebcc0", + "5a5b8c37-e3f6-48f2-96d3-289a043eaee4", + "394d1df6-bced-4f04-8c2a-a1ea3a04b984", + "b87c3e0f-a5af-4663-a87d-d0689d73343b", + "09194996-ff4b-4a4f-9800-e8aedc9be7f2", + "3f920762-f4a7-4bdb-86dc-05a3c3c8e875", + "5204624f-90ac-47bf-a71a-486d43f4351c", + "df24897d-a478-4fbb-816f-1d14c59da30d", + "9ba7c2a0-c86f-4967-9284-9be88ad3eacf", + "3f28dd4f-f653-4937-83b6-9f80c1284a85", + "b2735994-59ae-4478-9c5f-93f5ccf18c8b", + "5e965191-bcd1-491a-a9ce-e014e970020b", + "85718cf6-e4ba-44e3-880f-143c627f12ff", + "d0af657d-9c21-447c-a4f9-199b38deac5b", + "f8289825-2858-4dfa-b173-f4c506359158", + "8f2d26e6-239e-492f-aa1b-472f8ec85050", + "836be9c2-cee6-4fb9-8bbd-74197d213c5a", + "cd998593-e5bc-49d0-86a7-333acf9abf93", + "8deeeb88-19ed-4747-9c79-d94104bed7e0", + "6d1dc078-8b04-480e-a36f-298594ee9592", + "83464c62-929f-4bf7-9a04-43071189edde", + "1efd8b7b-65f9-4e73-bd49-d707190d6519", + "f0b69458-25df-4dda-be8b-e14383a57a50", + "88260ff4-be0b-4544-98ef-78b6a0d61ebc", + "39fb2989-1dd2-490d-b793-dce86c36b43a", + "b4f650b4-53d6-4069-a55f-43f2de7ff47b", + "15b6759d-2370-43ee-9a61-09e453db2469", + "2e4e0dac-ab21-49ff-9563-e320acd73f2f", + "986c4e43-ac2e-4983-8d56-f1dd34775133", + "12c1b77c-cbea-4c20-8b64-7154a25bbbe1", + "a3da62c3-3f46-499f-a4c6-acf93bd1c63d", + "57b6455e-2acb-4076-b8c4-1f2cdae97185", + "d6d8261e-0ee0-4002-bddf-e94fdc1075c7", + "56683d7b-7c65-4f42-b5a4-818f3cd5224c", + "9090779d-803e-4fbd-9166-d9e90ff2d36f", + "d5374572-5576-4f83-99c8-a0b12c4de160", + "ffddab3a-dbec-49b0-b114-5a1e3d805efa", + "38c02527-1b2f-458b-bad8-73a282c1d99a", + "d612edfb-e904-4de0-abce-a5fd9ecc449b", + "1261ee8c-80fc-4a3f-9dd3-29f4a918a9fe", + "8e860a91-0205-447b-a706-aa894d23d3db", + "fb07f205-868e-44c9-a704-305daaf03301", + "48f7028c-b6bf-4ca5-81e9-128147fd02c9", + "c6a6e1e4-24c4-422d-9d18-f062240ddf23", + "ba280b86-22a2-4509-b16d-b55e04680461", + "f0cb79ef-c081-4049-85ef-2623e02c9589", + "6dbfc053-7899-40d8-9827-0ccd21f3ee0a", + "50d3035c-2503-4bcf-b28b-e66c92c98f09", + "41f7665a-a692-42e0-b28d-9769c03ab99f", + "cb1f70b4-0bf6-47f1-bc98-2e9365a55c3f", + "7f566780-4ed2-41c5-9faa-06759495faab", + "0e08e1c7-3add-4e41-b28d-cd5f4413a422", + "672e0d75-df0e-4715-8c42-f1131b059a6c", + "8f15aab0-9c42-452e-b78a-47a48da1f3c1", + "edaa0b1a-d504-4852-8025-4f12b4d293ac", + "4a9ada82-9f86-495e-b17f-25b39bef2297", + "07774a7c-c5ed-4133-991f-926a316068a4", + "5a009338-0aff-4313-b02f-c48630ce1702", + "8324c0ba-7a10-42e1-90fd-afb4d3afa789", + "3db37314-adaa-4a4e-af11-c9190fa6b65b", + "b21b35d2-6d69-40a2-8c36-aa27e4a980ae", + "6daa9124-d894-42e5-8d98-2e0e6d1191ac", + "8a256d7f-2ea1-4b0e-94ea-168a1ed85678", + "08b6f5b3-4471-4cfd-a95e-4901cb1ded57", + "ea9d85e7-6a9a-451a-b019-972a6891c11b", + "f84fd67e-0797-44f4-b3ff-78da646362e2", + "ad1a9ea9-d1ab-474a-b715-1a42d892ef5d", + "ffb191f8-894a-40e8-baa5-5578aa6ba706", + "fd394f00-4049-49e5-9c07-ee3e86979003", + "9767a3ca-0a9c-480f-adfe-e4164c52c760", + "95cdb269-0088-4451-9a0d-f72349081fab", + "4ade0a82-16fd-471f-9fcc-fd2724887aef", + "fe69a78e-7773-45d1-9e0f-bacee52dac83", + "3aed7bbd-75f0-45ef-b9cc-7f5ff6f7fbbf", + "f8d71e9a-982a-437c-a968-867149c24eb7", + "8351998e-ead7-4bc1-b0a7-a7209da92b3c", + "3eb00a40-880a-4ad9-8cec-3ba0245ec472", + "602b224a-f5ac-491b-88e3-49043f362b6e", + "89125bc3-a413-44ee-a7e1-2753c511ecab", + "00792fa8-988c-4c85-8e80-73eb3ac53e80", + "094d1ab2-f34e-4a2a-ac3b-8671ce5faf65", + "4cf10990-13e5-4a1e-a913-7ec4b8c8ab02", + "a9b191de-d565-4981-8f68-02069a2ba16c", + "0eeab9d2-a55c-4708-94c3-cc1e08f9f262", + "4e78f728-06f6-43a3-ac8c-c0cb159e5b32", + "be99e9e2-82f7-4f1f-aff5-371834fde0da", + "5682f9aa-bdfb-409c-b4c2-c0105b35ee78", + "3ce15244-5293-4a5e-8c92-52ab29c6949a", + "b4b2c55b-c107-468f-bbdb-f73a0e4dca34", + "8a409551-657a-48c0-8beb-056c228ebca5", + "3c17a762-c839-492f-b361-009c1187ae4c", + "22172001-0e17-40db-bcce-825e19532297", + "5b99234b-b940-405d-9123-70218bda24b2", + "f7869fd8-9af3-493a-a627-a3e95e7c3577", + "057e9f41-e6ed-4575-906c-aca9413b35a5", + "70e88e82-5c06-44fd-959c-d4b2d40999cd", + "b95bd349-11bb-4ade-9cc6-3be8dcea06dd", + "19b880da-910e-40c7-970e-cd8718f2faf5", + "e05a2182-a97f-4d0b-87e4-66b0dd405e58", + "3072481a-04fe-4959-8410-018a41a91df1", + "fc2b2014-9dc4-4a5d-8a1d-25a6911bff7c", + "db1d3def-397e-470f-811a-85c4cc5338a9", + "ff283659-4a79-4642-8078-4366df38e1f1", + "bc8bae7f-ad9f-4f72-ae5c-b3e7ca8dd8f0", + "a3bd1be3-edfe-4f4e-89a8-dd56ce6ece08", + "2360312c-4347-476a-990f-2633b489fa7a", + "b9e32e51-8c18-420b-a582-200c79b8208d", + "0511dce9-c126-4a11-92ba-8082a4ec89bd", + "224d8b17-7eaf-4529-9b75-c6ce3078d640", + "10732476-f045-4feb-8199-e4e04685f7ca", + "5753e3fe-f1c5-426b-9f45-5a8efa6b5370", + "b69564f5-5e5e-4097-aebd-9f52f211ff75", + "97eb574d-d8bf-44a8-a6ef-1204b789db43", + "85506322-449a-45c0-b77a-48a8077b4bbd", + "ed65870d-ba08-4894-9cad-7f34e2c6f9e4", + "cf682058-962f-45c7-95ab-dbbd2f87acc0", + "65229ac2-9033-4002-af6c-492fe49bc563", + "4c85a780-8a8e-46c9-a0f8-c7e9703ef998", + "1ac6e060-a6d2-4b42-afef-599be78a533a", + "a22beaf0-377c-4b95-b472-867280a525ae", + "f5bc6a45-1591-42bf-8ed4-8286382b31bf", + "06592f50-7f29-4ca9-94a9-1b6c389d50ea", + "4ca3d9cf-4972-4e1d-aca3-28b036d6bdc5", + "076bdb81-5c75-4282-9f3a-a387624575f3", + "a5fe8210-1c37-4dd9-bd27-74051c521c3d", + "c1ca7c5c-2473-4b83-a507-7d187609111a", + "7c4079e7-4d68-4b0d-b015-5971a9866c60", + "ead54fb5-dee9-452f-8145-227ccd2405e2", + "7c776700-1d6b-4ee2-8155-b7f6855afceb", + "71a7456a-da6e-4e38-86ec-76bacaa6db5b", + "26c042d4-bf21-4d10-81f9-303d3816c114", + "c7820a4f-55c9-44eb-bf7e-b965a9d0e924", + "5b5b6b8c-9b88-4181-9ebb-14abb7783e1e", + "db12030d-7577-435f-9d48-5c1aae615c4f", + "e37c82c5-483f-4e68-8bf3-2a22dcef0b1e", + "0b4cacb1-970f-4ef0-85da-371d81f899e0", + "3949f121-af2b-4c7c-8a49-07befc08abbe", + "73c1ac5a-80a0-4e55-b143-52df5697c9ea", + "0f325614-48c9-4a6a-b7f2-731c26ddeb97", + "572117cf-a76b-44d2-804a-4901bfc545f4", + "33fd4f1f-fb01-4d3a-9b8d-b8d8024a2dab" + ], + "test_session_set": [ + "66cb6d16-f972-40e9-88c6-64b6816b900c", + "e62b75a0-f1f6-4e45-b3db-17fd7ff5756c", + "13ffc0e4-eebb-4baa-97af-5983ce46779c", + "d062849b-c60f-4369-a113-50f6ce8521fb", + "bb7f1e75-65ee-4d98-adc8-c70551a35940", + "7a6576d8-007c-4ee9-b7cc-8cc2d020e1bf", + "4ce65f3a-21ee-4f77-85c0-b3ddfe1e6381", + "783d7043-b27e-4744-a969-5f40fb40827c", + "81b74694-a2e6-4be1-b9f6-69f3aa03f50d", + "e6317c78-06f0-4c6d-9886-ca67712772ed", + "44bcb1fa-9e64-4bae-927a-4d5d53e536a6", + "218e79aa-3f56-4bd4-aa2a-14027bf7d52d", + "561fd770-f9a2-41a4-90da-beeb7d96ee7f", + "a2d94b61-57c7-4ec2-ad94-94b3585a6b64", + "764a08dc-5ac5-4c09-a281-0210cdd5a55e", + "c1057c9b-331f-458b-8914-2703511e8b13", + "3edc700d-7d14-4154-a3e7-5b0e8c3fde07", + "1492c03c-1887-4576-8889-5bc2848967de", + "2ea37ccc-e1ec-4e0d-a2a7-9707ea300b7f", + "65cad708-34b8-4158-be96-b88710288473", + "be8f4445-3958-4fc4-8cdf-b14d1d0eba73", + "3c62a0d2-10d7-4a9e-ac2c-02505663afad", + "be25331c-b1a0-4187-93fb-3b59437258c3", + "108cebaa-777d-42dd-8b1e-69ea7f443a0a", + "be3b65c8-ba74-4321-99b4-ca0e99e6fc2c", + "3bd74a94-ec3e-4138-b78f-ad76af4f07e2", + "4356f6d5-576a-4999-af40-744c242fff22", + "b0c57748-65c8-49d7-9681-ddcfcde3298b", + "12706b2d-0a0b-4511-8feb-ed8065746e50", + "f0e9a139-86f6-41bc-82f1-11283954a781", + "bb3c6009-892c-4071-9bb4-692b83ed12f3", + "90bf3f44-60ae-4430-9dc6-2365f7901fc5", + "2b5cea2f-b019-4ecf-abbc-b2346bd65a10", + "8e9f89d2-f02a-4c75-8831-6785ef6bd7f5", + "13c24739-e514-4fad-8288-6686bffa71ad", + "1828b207-9601-4521-8c0b-5d9e829fca78", + "0324c3b5-5154-4de1-9869-eb866b568e0f", + "1a746147-2dec-4a37-a51b-c94e4f1de010", + "3a20bf0a-6eb1-47c5-9bf8-8445446feb61", + "c2a8f168-06b0-4766-8865-27f08cb1bcae", + "f55f2dfd-704f-42cf-a39e-72e2560314b3", + "4ea88cb9-bff8-49e0-a855-cc1902db9a99", + "4e0053f8-566c-4952-abbe-4ea877782a0a", + "80ee109f-b097-4755-8a39-9f0b33adb170", + "933ddf59-52d2-4d4c-b653-05de44ded245", + "b432e10e-cff1-4f15-a8fe-1e2350e94b36", + "157c8e81-117a-4a15-a047-9a574181d44b", + "cfc6da55-43fc-4cd8-96d0-9de5166f4a63", + "aabeb7f9-5ada-425a-9545-c98dc443887a", + "db9978fb-fe4d-467b-9495-ab7daf3b1b55", + "6549c35b-8b1d-4577-9434-ac808e4d7057", + "e50126ee-f4c1-41a6-a2e8-79636b40b341", + "49a7e53b-aba5-4e53-b3bc-1492657fc8ae", + "7abe6d4a-f594-4a15-af08-ddd591ff500c", + "03c50078-a08d-43ad-a3c9-87916aebd865", + "a9a0ed92-fab1-4b6e-ae27-714692464641", + "3c3207e6-57b9-4be1-96fc-5b2168ae3cd3", + "895545f1-f194-4c0a-8513-87d386670f15", + "7b76be35-4e6d-42d7-9fa9-14a21fc327cf", + "e4670e4f-0b1a-450a-9778-9ee2409db554", + "34bc808a-4b98-42c5-a1ea-6f565dd4aa20", + "bec4e063-9be1-402e-a7b5-dba20730c933", + "20a3e8e5-3f43-46a5-84fa-6888bd9c5c88", + "5da4efa3-8e54-46dc-a8c2-3016fa4e12f8", + "0d998741-de68-4006-898e-7108de4167b1", + "4d144ff3-a853-4b79-9825-d7fcea248373", + "dce1d8b4-2c29-4fec-a343-77885899c434", + "dab038d7-9e08-4edf-8d03-ee7e2af5a10e", + "a1055434-9e9b-4d69-bac3-374a39f801da", + "e7086057-b5f7-4286-87c6-ab110072c5b0", + "3af65596-1aa5-4126-933b-b98ccaccedf2", + "19f880ae-9770-45a5-93f5-00047f22d5f6", + "36888bbd-1ece-4a3a-8aff-b0f6ef905752", + "f42a7e14-c3ee-4b05-a59b-8d349f8bd19b", + "d36d9514-a99b-4534-b54f-3fd98a19f0ea", + "6c4f7cb0-caa0-4a35-9071-a84c9ec2ab2f", + "666d9243-f4f9-47f5-8fba-cee8568abf41", + "2f14e742-dc58-43c8-8ebc-7244487dbba7", + "f8accf9d-4427-4458-88e9-89d51a3f5136", + "50810487-0b0f-45e0-8f37-515644189e81", + "f0445727-2729-4a76-9434-fb081a14cd9d", + "f55f890d-fd90-43f6-a11a-a1920c21890a", + "16851d4d-35e0-49b4-b3b5-28e67b564db3", + "5c8f8839-4bf9-45e0-beb0-6c17a5ead0d5", + "514d71c8-fd4e-4908-89a4-c0861e81ca8a", + "35465081-a518-4321-8384-477ef5b29b8d", + "de82a6f3-1a55-4996-ae68-35032678ff66" + ] + } +} \ No newline at end of file diff --git a/third_party/EgoSpeak/MiniROAD/datasets/__init__.py b/third_party/EgoSpeak/MiniROAD/datasets/__init__.py new file mode 100755 index 0000000..caa902b --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/datasets/__init__.py @@ -0,0 +1,2 @@ +from .dataset_builder import build_dataset, build_data_loader +from .dataset import THUMOSDataset, FINEACTIONDataset \ No newline at end of file diff --git a/third_party/EgoSpeak/MiniROAD/datasets/dataset.py b/third_party/EgoSpeak/MiniROAD/datasets/dataset.py new file mode 100755 index 0000000..bc16bb0 --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/datasets/dataset.py @@ -0,0 +1,314 @@ +import gc +import json +import os.path as osp +import sys + +import numpy as np +import torch +import torch.utils.data as data +from configs.feature_sizes import FEATURE_SIZES +from datasets.dataset_builder import DATA_LAYERS +from utils.util import parse_flow_type, parse_rgb_type + + +@DATA_LAYERS.register("THUMOS") +@DATA_LAYERS.register("TVSERIES") +@DATA_LAYERS.register("EasyCom") +@DATA_LAYERS.register("EGO4D") +@DATA_LAYERS.register("IEMOCAP") +@DATA_LAYERS.register("IEMOCAP_DEBUG") +@DATA_LAYERS.register("YTConv") +class THUMOSDataset(data.Dataset): + + def __init__(self, cfg, mode='train'): + self.root_path = cfg['root_path'] + self.mode = mode + self.training = mode == 'train' + self.window_size = cfg['window_size'] + self.enc_steps = cfg['window_size'] if 'window_size' in cfg else None + self.dec_steps = cfg['query_num'] if 'query_num' in cfg else None + + self.stride = cfg['stride'] + data_name = cfg['data_name'] + self.model = cfg['model'] + self.vids = json.load(open(cfg['video_list_path']))[data_name][mode + '_session_set'] + self.num_classes = cfg['num_classes'] + self.inputs = [] + + self._load_features(cfg) + self._init_features() + + def _load_features(self, cfg): + self.annotation_type = cfg['annotation_type'] + self.rgb_type = cfg['rgb_type'] + self.flow_type = cfg['flow_type'] + self.target_all = {} + self.rgb_inputs = {} + self.flow_inputs = {} + dummy_target = np.zeros((self.window_size-1, self.num_classes)) + + rgb_type = parse_rgb_type(cfg['rgb_type']) + flow_type = parse_flow_type(cfg['flow_type']) + + dummy_rgb = np.zeros((self.window_size-1, FEATURE_SIZES[rgb_type])) + dummy_flow = np.zeros((self.window_size-1, FEATURE_SIZES[flow_type])) + for vid in self.vids: + target = np.load(osp.join(self.root_path, self.annotation_type, vid + '.npy')) + rgb = np.load(osp.join(self.root_path, self.rgb_type, vid + '.npy')) + flow = np.load(osp.join(self.root_path, self.flow_type, vid + '.npy')) + + if self.training: + self.target_all[vid] = np.concatenate((dummy_target, target), axis=0) + self.rgb_inputs[vid] = np.concatenate((dummy_rgb, rgb), axis=0) + self.flow_inputs[vid] = np.concatenate((dummy_flow, flow), axis=0) + else: + self.target_all[vid] = target + self.rgb_inputs[vid] = rgb + self.flow_inputs[vid] = flow + + def _init_features(self): + del self.inputs + gc.collect() + self.inputs = [] + if self.model == 'MiniROAD': + for vid in self.vids: + target = self.target_all[vid] + + + if self.training: + seed = np.random.randint(self.stride) + for start, end in zip(range(seed, target.shape[0], self.stride), + range(seed + self.window_size, target.shape[0]+1, self.stride)): + self.inputs.append([ + vid, start, end, target[start:end] + ]) + + else: + start = 0 + end = target.shape[0] + self.inputs.append([ + vid, start, end, target[start:end] + ]) + + elif self.model == 'OADTR': + for vid in self.vids: + target = self.target_all[vid] + + + seed = np.random.randint(self.stride) if self.training else 0 + + for start, end in zip( + range(seed, target.shape[0], 1), + range(seed + self.window_size, target.shape[0]+1, 1) + ): + enc_target = target[start:end] + dec_target = target[end:end + self.dec_steps] + distance_target, class_h_target = self.get_distance_target(target[start:end]) + + self.inputs.append([ + vid, start, end, enc_target, distance_target, class_h_target, dec_target + ]) + + else: + for vid in self.vids: + target = self.target_all[vid] + seed = np.random.randint(self.stride) if self.training else 0 + for start, end in zip( + range(seed, target.shape[0], 1), + range(seed + self.window_size, target.shape[0]+1, 1) + ): + self.inputs.append([ + vid, start, end, target[start:end] + ]) + + def get_distance_target(self, target_vector): + target_matrix = np.zeros(self.enc_steps - 1) + target_argmax = target_vector[self.enc_steps-1].argmax() + for i in range(self.enc_steps - 1): + if target_vector[i].argmax() == target_argmax: + target_matrix[i] = 1. + return target_matrix, target_vector[self.enc_steps-1] + + def __getitem__(self, index): + if self.model == 'MiniROAD' or self.model == 'Transformer': + vid, start, end, target = self.inputs[index] + rgb_input = self.rgb_inputs[vid][start:end] + flow_input = self.flow_inputs[vid][start:end] + rgb_input = torch.tensor(rgb_input.astype(np.float32)) + flow_input = torch.tensor(flow_input.astype(np.float32)) + target = torch.tensor(target.astype(np.float32)) + return vid, rgb_input, flow_input, target + + elif self.model == 'OADTR': + vid, start, end, enc_target, distance_target, class_h_target, dec_target = self.inputs[index] + rgb_input = self.rgb_inputs[vid][start:end] + flow_input = self.flow_inputs[vid][start:end] + rgb_input = torch.tensor(rgb_input.astype(np.float32)) + flow_input = torch.tensor(flow_input.astype(np.float32)) + enc_target = torch.tensor(enc_target.astype(np.float32)) + distance_target = torch.tensor(distance_target.astype(np.float32)) + class_h_target = torch.tensor(class_h_target.astype(np.float32)) + dec_target = torch.tensor(dec_target.astype(np.float32)) + return vid, rgb_input, flow_input, enc_target, distance_target, class_h_target, dec_target + + def __len__(self): + return len(self.inputs) + + +@DATA_LAYERS.register("THUMOS_ANTICIPATION") +@DATA_LAYERS.register("TVSERIES_ANTICIPATION") +@DATA_LAYERS.register("IEMOCAP_ANTICIPATION") +@DATA_LAYERS.register("EasyCom_ANTICIPATION") +@DATA_LAYERS.register("YTConv_ANTICIPATION") +@DATA_LAYERS.register("EGO4D_ANTICIPATION") +class THUMOSDataset(data.Dataset): + + def __init__(self, cfg, mode='train'): + self.cfg = cfg + self.root_path = cfg['root_path'] + self.mode = mode + self.training = mode == 'train' + self.window_size = cfg['window_size'] + self.stride = cfg['stride'] + self.anticipation_length = cfg['anticipation_length'] + data_name = cfg["data_name"].split('_')[0] + if 'data_info_name' in cfg.keys() and cfg['data_info_name'] is not None: + data_name = cfg['data_info_name'] + + self.vids = json.load(open(cfg['video_list_path']))[data_name][mode + '_session_set'] + print(f'data_name: {data_name}') + print(f'{mode} session set length: {len(self.vids)}') + self.num_classes = cfg['num_classes'] + self.inputs = [] + + self._load_features(cfg) + self._init_features() + + def _load_features(self, cfg): + self.annotation_type = cfg['annotation_type'] + self.rgb_type = cfg['rgb_type'] + self.flow_type = cfg['flow_type'] + + self.target_all = {} + self.rgb_inputs = {} + self.flow_inputs = {} + + dummy_target = np.zeros((self.window_size-1, self.num_classes)) + dummy_rgb = np.zeros((self.window_size-1, FEATURE_SIZES[cfg['rgb_type']])) + dummy_flow = np.zeros((self.window_size-1, FEATURE_SIZES[cfg['flow_type']])) + + for vid in self.vids: + target = np.load(osp.join(self.root_path, self.annotation_type, vid + '.npy')) + rgb = np.load(osp.join(self.root_path, self.rgb_type, vid + '.npy')) + flow = np.load(osp.join(self.root_path, self.flow_type, vid + '.npy')) + + if self.training: + self.target_all[vid] = np.concatenate((dummy_target, target), axis=0) + self.rgb_inputs[vid] = np.concatenate((dummy_rgb, rgb), axis=0) + self.flow_inputs[vid] = np.concatenate((dummy_flow, flow), axis=0) + else: + self.target_all[vid] = target + self.rgb_inputs[vid] = rgb + self.flow_inputs[vid] = flow + + def _init_features(self): + del self.inputs + gc.collect() + self.inputs = [] + + for vid in self.vids: + target = self.target_all[vid] + if self.training: + seed = np.random.randint(self.stride) + for start, end in zip(range(seed, target.shape[0], self.stride), + range(seed + self.window_size, target.shape[0]-self.anticipation_length, self.stride)): + self.inputs.append([ + vid, start, end, target[start:end], target[end:end+self.anticipation_length] + ]) + else: + start = 0 + end = target.shape[0] - self.anticipation_length + ant_target = [] + + for s in range(1, target.shape[0] - self.anticipation_length + 1): + ant_target.append(target[s:s + self.anticipation_length]) + + self.inputs.append([ + vid, start, end, target[start:end], np.array(ant_target) + ]) + + + + def __getitem__(self, index): + vid, start, end, target, ant_target = self.inputs[index] + + + rgb_input = self.rgb_inputs[vid][start:end] + flow_input = self.flow_inputs[vid][start:end] + + rgb_input = torch.tensor(rgb_input.astype(np.float32)) + flow_input = torch.tensor(flow_input.astype(np.float32)) + + target = torch.tensor(target.astype(np.float32)) + ant_target = torch.tensor(ant_target.astype(np.float32)) + + return vid, rgb_input, flow_input, target, ant_target + + def __len__(self): + return len(self.inputs) + + +@DATA_LAYERS.register("FINEACTION") +class FINEACTIONDataset(data.Dataset): + + def __init__(self, cfg, mode='train'): + self.root_path = cfg['root_path'] + self.mode = mode + self.training = mode == 'train' + self.window_size = cfg['window_size'] + self.stride = cfg['stride'] + data_name = cfg['data_name'] + self.vids = json.load(open(cfg['video_list_path']))[data_name][mode + '_session_set'] + self.num_classes = cfg['num_classes'] + self.inputs = [] + self._load_features(cfg) + self._init_features() + + def _load_features(self, cfg): + self.annotation_type = cfg['annotation_type'] + self.rgb_type = cfg['rgb_type'] + self.flow_type = cfg['flow_type'] + + def _init_features(self, seed=0): + + del self.inputs + gc.collect() + self.inputs = [] + for vid in self.vids: + target = np.load(osp.join(self.root_path, self.annotation_type, vid + '.npy')) + if self.training: + seed = np.random.randint(self.stride) + for start, end in zip(range(seed, target.shape[0], self.stride), + range(seed + self.window_size, target.shape[0]+1, self.stride)): + self.inputs.append([ + vid, start, end + ]) + else: + start = 0 + end = target.shape[0] + self.inputs.append([ + vid, start, end + ]) + + def __getitem__(self, index): + vid, start, end = self.inputs[index] + rgb_input = np.load(osp.join(self.root_path, self.rgb_type, vid + '.npy'), mmap_mode='r')[start:end] + flow_input = np.load(osp.join(self.root_path, self.flow_type, vid + '.npy'), mmap_mode='r')[start:end] + target = np.load(osp.join(self.root_path, self.annotation_type, vid + '.npy'), mmap_mode='r')[start:end] + rgb_input = torch.tensor(rgb_input.astype(np.float32)) + flow_input = torch.tensor(flow_input.astype(np.float32)) + target = torch.tensor(target.astype(np.float32)) + return rgb_input, flow_input, target + + def __len__(self): + return len(self.inputs) \ No newline at end of file diff --git a/third_party/EgoSpeak/MiniROAD/datasets/dataset_builder.py b/third_party/EgoSpeak/MiniROAD/datasets/dataset_builder.py new file mode 100755 index 0000000..a1079e5 --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/datasets/dataset_builder.py @@ -0,0 +1,32 @@ +__all__ = [ + 'build_dataset', + 'build_data_loader', +] + +import torch.utils.data as data +from utils import Registry + +DATA_LAYERS = Registry() + +def build_dataset(cfg): + data_layer = DATA_LAYERS[f'{cfg["data_name"]}'] + return data_layer + +def build_data_loader(cfg, mode): + data_layer = build_dataset(cfg) + data_loader = data.DataLoader( + dataset=data_layer(cfg, mode), + batch_size=cfg["batch_size"] if mode == 'train' else cfg["test_batch_size"], + shuffle=True if mode == 'train' else False, + num_workers=cfg["num_workers"], + pin_memory=False, + ) + + + + + + + + + return data_loader diff --git a/third_party/EgoSpeak/MiniROAD/main.py b/third_party/EgoSpeak/MiniROAD/main.py new file mode 100755 index 0000000..45b0519 --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/main.py @@ -0,0 +1,294 @@ +import argparse +import os +import os.path as osp + +import torch +import wandb +import yaml +from criterions import build_criterion +from model import build_model +from torch.utils.tensorboard import SummaryWriter +from trainer import build_eval, build_trainer +from utils import * +from utils import get_logger + +from datasets import build_data_loader + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--config", type=str) + parser.add_argument("--eval", action="store_true") + parser.add_argument("--ckpt", type=str) + parser.add_argument("--amp", action="store_true") + parser.add_argument("--tensorboard", action="store_true") + parser.add_argument("--wandb", action="store_true") + parser.add_argument("--lr_scheduler", action="store_true") + parser.add_argument("--no_rgb", action="store_true") + parser.add_argument("--no_flow", action="store_true") + parser.add_argument( + "--save-as-cfg", + action="store_true", + help="Save outputs with config name in the output directory", + ) + parser.add_argument("--debug", action="store_true") + parser.add_argument( + "--save-preds", + action="store_true", + help="Save predictions in the output directory", + ) + # parser.add_argument('--model', type=str, default=None) + args = parser.parse_args() + + # combine argparse and yaml + opt = yaml.load(open(args.config), Loader=yaml.FullLoader) + opt.update(vars(args)) + cfg = opt + + if "checkpoint" in cfg and cfg["checkpoint"] is not None and args.ckpt is None: + args.ckpt = cfg["checkpoint"] + print(f"checkpoint is given in the config, set args.ckpt to {args.ckpt}") + + seed = 20 if "seed" not in cfg else cfg["seed"] + set_seed(seed) + + device = "cuda" if torch.cuda.is_available() else "cpu" + + annotation_type = os.path.basename(cfg["annotation_type"]) + + identifier = os.path.basename(cfg["config"])[:-5] + identifier += "_" + if cfg["no_flow"] == False: + identifier += "A" + if cfg["no_rgb"] == False: + identifier += "V" + + if cfg["save_as_cfg"] == True: + output_dir = osp.join("output_cfg", identifier) + result_path = create_outdir(output_dir, overwrite=True) + else: + identifier = get_identifier(cfg) + annotation_type = os.path.basename(cfg["annotation_type"]) + identifier += f"_{annotation_type}" + result_path = create_outdir(osp.join(cfg["output_path"], identifier)) + + logger = get_logger(result_path) + logger.info(f"Identifier: {identifier}") + logger.info(f"result_path: {result_path}") + logger.info(cfg) + + testloader = build_data_loader(cfg, mode="test") + print("testloader loaded") + model = build_model(cfg, device) + evaluate = build_eval(cfg) + + trainloader = build_data_loader(cfg, mode="train") + print("trainloader loaded") + + # evaluation + if args.eval == True and cfg["ckpt"] == None: + print("Automatically find the best checkpoint") + ckpt_root = os.path.join("output_cfg", identifier, "ckpts") + ckpt_dirs = os.listdir(ckpt_root) + + ckpt_path = os.path.join(ckpt_root, "best.pth") + if os.path.exists(ckpt_path): + args.ckpt = ckpt_path + print(f"args.ckpt set to {args.ckpt}") + + else: + ckpt_path = os.path.join(ckpt_root, "pred", "best.pth") + if os.path.exists(ckpt_path): + args.ckpt = ckpt_path + print(f"args.ckpt set to {args.ckpt}") + + print("Warninig) no ckpt_path found, there can be multiple ckpts") + print(f"ckpt_dirs: {ckpt_dirs}") + pred_dir_number = 0 + ckpt_dirs.sort(reverse=True) + + if args.ckpt == None: + for ckpt_dir in ckpt_dirs: + ckpt_path = os.path.join(ckpt_root, ckpt_dir, "best.pth") + + if os.path.exists(ckpt_path): + args.ckpt = ckpt_path + print(f"Found ckpt, args.ckpt was None, set to {args.ckpt}") + break + + if args.ckpt == None: + raise ValueError("No checkpoint found, args.ckpt is None") + + # load checkpoint if args.ckpt is given + if args.eval == True: + if args.ckpt == None: + raise ValueError("Please specify the checkpoint path") + checkpoint = torch.load(args.ckpt) + model.load_state_dict(checkpoint) + + elif args.eval == False and args.ckpt != None: + checkpoint = torch.load(args.ckpt) + print(f"Load checkpoint from {args.ckpt}") + model_dict = model.state_dict() + + unnecessary_keys = [] + matched_ckpt = {} + # Filter out unnecessary keys + for k, v in checkpoint.items(): + if not (k in model_dict and model_dict[k].shape == checkpoint[k].shape): + unnecessary_keys.append(k) + elif ( + "except_classifier" in cfg + and cfg["except_classifier"] == True + and "classification" in k + ): + print(f"Unnecessary key: {k} is classifier") + unnecessary_keys.append(k) + else: + matched_ckpt[k] = v + + # Overwrite entries in the existing state dict + model_dict.update(matched_ckpt) + + # Load the new state dict + model.load_state_dict(model_dict) + + print(f"Unnecessary keys: {unnecessary_keys}") + + if cfg["wandb"] == True: + # args.tensorboard = True + print("wandb init") + os.environ["WANDB__SERVICE_WAIT"] = "3000" + wandb.init( + project="egospeak", + tags=[cfg["model"], cfg["data_name"], annotation_type], + config=cfg, + name=os.path.basename(result_path), + ) + + if args.eval == True or args.debug == True: + if args.eval == True and args.ckpt == None: + raise ValueError("Please specify the checkpoint path") + + logger.info(f'Dataset: {cfg["data_name"]}, Model: {cfg["model"]}') + logger.info(f"Result path:{result_path}") + + for phase in ["test"]: + if phase == "train": + raise NotImplementedError("evaluation on train split is not implemented") + dataloader = trainloader + + elif phase == "test": + dataloader = testloader + + mAP, mcAP, preds = evaluate( + model, dataloader, logger, result_path=result_path + ) + + pred_path = osp.join(result_path, "eval") + os.makedirs(pred_path, exist_ok=True) + + if args.save_preds == True: + print(f"pred_path: {pred_path}") + for k, v in preds.items(): + np.save(osp.join(pred_path, f"{k}.npy"), v) + + logger.info(f'{phase}) {cfg["task"]} result: {mAP*100:.2f} mAP') + logger.info(f'{phase}) {cfg["task"]} result: {mcAP*100:.2f} mcAP') + exit() + + criterion = build_criterion(cfg, device) + train_one_epoch = build_trainer(cfg) + optim = torch.optim.AdamW if cfg["optimizer"] == "AdamW" else torch.optim.Adam + optimizer = optim( + [{"params": model.parameters(), "initial_lr": cfg["lr"]}], + lr=cfg["lr"], + weight_decay=cfg["weight_decay"], + ) + + scheduler = ( + build_lr_scheduler(cfg, optimizer, len(trainloader)) + if args.lr_scheduler + else None + ) + writer = SummaryWriter(osp.join(result_path, "runs")) if args.tensorboard else None + scaler = torch.cuda.amp.GradScaler() if args.amp else None + total_params = sum(p.numel() for p in model.parameters()) + + logger.info(f'Dataset: {cfg["data_name"]}, Model: {cfg["model"]}') + logger.info( + f'lr:{cfg["lr"]} | Weight Decay:{cfg["weight_decay"]} | Window Size:{cfg["window_size"]} | Batch Size:{cfg["batch_size"]}' + ) + logger.info( + f'Total epoch:{cfg["num_epoch"]} | Total Params:{total_params/1e6:.1f} M | Optimizer: {cfg["optimizer"]}' + ) + logger.info(f"Output Path:{result_path}") + + best_mAP, best_epoch = 0, 0 + best_cAP, best_cAP_epoch = 0, 0 + + pred_save_path = osp.join(result_path, "ckpts", "pred") + print(f"pred_save_path: {pred_save_path}") + + # if you want to overwrite, change the code + os.makedirs(pred_save_path, exist_ok=True) + + for epoch in range(1, cfg["num_epoch"] + 1): + epoch_loss = train_one_epoch( + trainloader, + model, + criterion, + optimizer, + scaler, + epoch, + cfg["wandb"], + writer, + scheduler=scheduler, + ) + trainloader.dataset._init_features() + + mAP, cAP, preds = evaluate( + model, testloader, logger, criterion, cfg["wandb"], epoch, writer + ) + + if mAP > best_mAP: + best_mAP = mAP + best_epoch = epoch + torch.save(model.state_dict(), osp.join(pred_save_path, "best.pth")) + + if args.save_preds == True: + for k, v in preds.items(): + np.save(osp.join(pred_save_path, f"{k}.npy"), v) + + if cAP > best_cAP: + best_cAP = cAP + best_cAP_epoch = epoch + torch.save(model.state_dict(), osp.join(pred_save_path, "best_cAP.pth")) + + logger.info( + f'Epoch {epoch} mAP: {mAP*100:.2f} | Best mAP: {best_mAP*100:.2f} at epoch {best_epoch}, iter {epoch*cfg["batch_size"]*len(trainloader)} | train_loss: {epoch_loss/len(trainloader):.4f}, lr: {optimizer.param_groups[0]["lr"]:.7f}' + ) + logger.info( + f'Epoch {epoch} cAP: {cAP*100:.2f} | Best cAP: {best_cAP*100:.2f} at epoch {best_epoch}, iter {epoch*cfg["batch_size"]*len(trainloader)} | train_loss: {epoch_loss/len(trainloader):.4f}, lr: {optimizer.param_groups[0]["lr"]:.7f}' + ) + + + if cfg["wandb"] == True: + logger.info("Wandb logging") + wandb.log( + { + "Epoch": epoch, + "Train Loss (Epoch)": epoch_loss / len(trainloader), + "Learning Rate": optimizer.param_groups[0]["lr"], + "anticipation average mAP": mAP, + "anticipation average cAP": cAP, + } + ) + + # os.rename(osp.join(result_path, 'ckpts', 'best.pth'), osp.join(result_path, 'ckpts', f'best_{best_mAP*100:.2f}.pth')) + # save the training result and config as txt file + with open(osp.join(result_path, "result.txt"), "w") as f: + f.write(f"Best mAP: {best_mAP*100:.2f} at epoch {best_epoch}\n") + f.write(f"Best cAP: {best_cAP*100:.2f} at epoch {best_cAP_epoch}\n") + f.write(f"Config: {cfg}") + + wandb.finish() if args.wandb else None diff --git a/third_party/EgoSpeak/MiniROAD/model/__init__.py b/third_party/EgoSpeak/MiniROAD/model/__init__.py new file mode 100755 index 0000000..feef232 --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/model/__init__.py @@ -0,0 +1,4 @@ +from .mamba import OADMamba +from .model_builder import build_model +from .rnn import MROAD +from .transformer_models import ViTEnc diff --git a/third_party/EgoSpeak/MiniROAD/model/mamba/__init__.py b/third_party/EgoSpeak/MiniROAD/model/mamba/__init__.py new file mode 100644 index 0000000..62f284b --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/model/mamba/__init__.py @@ -0,0 +1 @@ +from .mamba import OADMamba, OADMambaA diff --git a/third_party/EgoSpeak/MiniROAD/model/mamba/mamba.py b/third_party/EgoSpeak/MiniROAD/model/mamba/mamba.py new file mode 100644 index 0000000..6cc8f39 --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/model/mamba/mamba.py @@ -0,0 +1,172 @@ +import sys + +import torch +import torch.nn as nn +import torch.nn.functional as F +from configs.feature_sizes import FEATURE_SIZES +from mamba_ssm import Mamba +# from mambapy.mamba import Mamba, MambaConfig +from model.model_builder import META_ARCHITECTURES +from utils.util import parse_flow_type, parse_rgb_type + + +@META_ARCHITECTURES.register("Mamba") +class OADMamba(nn.Module): + + def __init__(self, cfg): + super(OADMamba, self).__init__() + self.use_flow = not cfg['no_flow'] + self.use_rgb = not cfg['no_rgb'] + + self.input_dim = 0 + + rgb_type = parse_rgb_type(cfg['rgb_type']) + flow_type = parse_flow_type(cfg['flow_type']) + + if self.use_rgb: + self.input_dim += FEATURE_SIZES[rgb_type] + if self.use_flow: + self.input_dim += FEATURE_SIZES[flow_type] + + self.hidden_dim = cfg['hidden_dim'] + self.num_layers = cfg['num_layers'] + + # raise NotImplementedError("The following code is not implemented yet") + self.out_dim = cfg['num_classes'] + self.window_size = cfg['window_size'] + + self.relu = nn.ReLU() + + self.embedding_dim = cfg['embedding_dim'] + self.d_state = cfg['d_state'] + self.d_conv = cfg['d_conv'] + self.expand = cfg['expand'] + + # self.mamba_config = MambaConfig(d_model=self.embedding_dim, n_layers=self.num_layers) + # self.mamba = Mamba(self.mamba_config) + self.mamba = Mamba( + d_model=self.embedding_dim, # Model dimension d_model + d_state=self.d_state, # SSM state expansion factor + d_conv=self.d_conv, # Local convolution width + expand=self.expand, # Block expansion factor + ) + + # self.gru = nn.GRU(self.embedding_dim, self.hidden_dim, self.num_layers, batch_first=True) + self.layer1 = nn.Sequential( + nn.Linear(self.input_dim, self.embedding_dim), + nn.LayerNorm(self.embedding_dim), + nn.ReLU(), + nn.Dropout(p=cfg['dropout']), + ) + self.f_classification = nn.Sequential( + nn.Linear(self.hidden_dim, self.out_dim) + ) + # self.h0 = torch.nn.Parameter(torch.zeros(1, 1, self.hidden_dim)) + self.h0 = torch.zeros(self.num_layers, 1, self.hidden_dim) + + def forward(self, rgb_input, flow_input): + x = torch.empty(0, device='cuda:0') + if self.use_rgb: + x = torch.cat((x, rgb_input), dim=2) # Concatenating along the channel dimension + if self.use_flow: + x = torch.cat((x, flow_input), dim=2) + + x = self.layer1(x) + ht = self.mamba(x) + + logits = self.f_classification(ht) + out_dict = {} + if self.training: + out_dict['logits'] = logits + else: + pred_scores = F.softmax(logits, dim=-1) + out_dict['logits'] = pred_scores + return out_dict + +@META_ARCHITECTURES.register("MambaA") +class OADMambaA(nn.Module): + + def __init__(self, cfg): + super(OADMambaA, self).__init__() + self.use_flow = not cfg['no_flow'] + self.use_rgb = not cfg['no_rgb'] + self.input_dim = 0 + if self.use_rgb: + self.input_dim += FEATURE_SIZES[cfg['rgb_type']] + if self.use_flow: + self.input_dim += FEATURE_SIZES[cfg['flow_type']] + + self.hidden_dim = cfg['hidden_dim'] + self.num_layers = cfg['num_layers'] + self.anticipation_length = cfg["anticipation_length"] + self.out_dim = cfg['num_classes'] + + self.embedding_dim = cfg['embedding_dim'] + self.d_state = cfg['d_state'] if 'd_state' in cfg else 16 + self.d_conv = cfg['d_conv'] if 'd_conv' in cfg else 4 + self.expand = cfg['expand'] if 'expand' in cfg else 2 + + self.layer1 = nn.Sequential( + nn.Linear(self.input_dim, self.embedding_dim), + nn.LayerNorm(self.embedding_dim), + nn.ReLU(), + nn.Dropout(p=cfg['dropout']) + ) + self.actionness = cfg['actionness'] + if self.actionness: + self.f_actionness = nn.Sequential( + nn.Linear(self.hidden_dim, 1), + ) + self.relu = nn.ReLU() + # self.mamba_config = MambaConfig(d_model=self.embedding_dim, n_layers=self.num_layers) + # self.mamba = Mamba(self.mamba_config) + self.mamba = Mamba( + d_model=self.embedding_dim, # Model dimension d_model + d_state=self.d_state, # SSM state expansion factor + d_conv=self.d_conv, # Local convolution width + expand=self.expand, # Block expansion facto + ) + + # self.gru = nn.GRU(self.embedding_dim, self.hidden_dim, self.num_layers, batch_first=True) + self.f_classification = nn.Sequential( + nn.Linear(self.embedding_dim, self.out_dim) + ) + self.anticipation_layer = nn.Sequential( + nn.Linear(self.embedding_dim, self.anticipation_length*self.embedding_dim), + ) + self.cfg = cfg + + def forward(self, rgb_input, flow_input): + x = torch.empty(0, device='cuda:0') + if self.use_rgb: + # [64, 128, 2048] + x = torch.cat((x, rgb_input), dim=2) # Concatenating along the channel dimension + if self.use_flow: + # [64, 128, 5120] + x = torch.cat((x, flow_input), dim=2) + + # audio channel is too huge to directly concat... + # to c + + # B, L, D: mamba + B, S, feat_dim = x.shape # batch, length, dim + # + x = self.layer1(x) + ht = self.mamba(x) + # h0 = torch.zeros(1, B, self.hidden_dim).to(x.device) + # ht, _ = self.gru(x, h0) + logits = self.f_classification(self.relu(ht)) + anticipation_ht = self.anticipation_layer(self.relu(ht)).view(B, S, self.anticipation_length, self.embedding_dim) + anticipation_logits = self.f_classification(self.relu(anticipation_ht)) + + out_dict = {} + if self.training: + out_dict['logits'] = logits + out_dict['anticipation_logits'] = anticipation_logits + else: + pred_scores = F.softmax(logits, dim=-1) + pred_anticipation_scores = F.softmax(anticipation_logits, dim=-1) + out_dict['logits'] = pred_scores + out_dict['anticipation_logits'] = pred_anticipation_scores + + return out_dict \ No newline at end of file diff --git a/third_party/EgoSpeak/MiniROAD/model/model_builder.py b/third_party/EgoSpeak/MiniROAD/model/model_builder.py new file mode 100755 index 0000000..9350a84 --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/model/model_builder.py @@ -0,0 +1,9 @@ +__all__ = ['build_model'] + +from utils import Registry + +META_ARCHITECTURES = Registry() + +def build_model(cfg, device=None): + model = META_ARCHITECTURES[cfg["model"]](cfg) + return model.to(device) \ No newline at end of file diff --git a/third_party/EgoSpeak/MiniROAD/model/rnn/__init__.py b/third_party/EgoSpeak/MiniROAD/model/rnn/__init__.py new file mode 100755 index 0000000..3c98c5e --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/model/rnn/__init__.py @@ -0,0 +1 @@ +from .rnn import MROAD, MROADA \ No newline at end of file diff --git a/third_party/EgoSpeak/MiniROAD/model/rnn/rnn.py b/third_party/EgoSpeak/MiniROAD/model/rnn/rnn.py new file mode 100755 index 0000000..1588302 --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/model/rnn/rnn.py @@ -0,0 +1,153 @@ +import sys + +import torch +import torch.nn as nn +import torch.nn.functional as F +from configs.feature_sizes import FEATURE_SIZES +from model.model_builder import META_ARCHITECTURES +from utils.util import parse_flow_type, parse_rgb_type + + +@META_ARCHITECTURES.register("MiniROAD") +class MROAD(nn.Module): + + def __init__(self, cfg): + super(MROAD, self).__init__() + self.use_flow = not cfg['no_flow'] + self.use_rgb = not cfg['no_rgb'] + + self.input_dim = 0 + + rgb_type = parse_rgb_type(cfg['rgb_type']) + flow_type = parse_flow_type(cfg['flow_type']) + + if self.use_rgb: + self.input_dim += FEATURE_SIZES[rgb_type] + if self.use_flow: + self.input_dim += FEATURE_SIZES[flow_type] + + self.hidden_dim = cfg['hidden_dim'] + self.num_layers = cfg['num_layers'] + self.out_dim = cfg['num_classes'] + self.window_size = cfg['window_size'] + + self.relu = nn.ReLU() + self.embedding_dim = cfg['embedding_dim'] + self.gru = nn.GRU(self.embedding_dim, self.hidden_dim, self.num_layers, batch_first=True) + self.layer1 = nn.Sequential( + nn.Linear(self.input_dim, self.embedding_dim), + nn.LayerNorm(self.embedding_dim), + nn.ReLU(), + nn.Dropout(p=cfg['dropout']), + ) + self.f_classification = nn.Sequential( + nn.Linear(self.hidden_dim, self.out_dim) + ) + # self.h0 = torch.nn.Parameter(torch.zeros(1, 1, self.hidden_dim)) + self.h0 = torch.zeros(self.num_layers, 1, self.hidden_dim) + + def forward(self, rgb_input, flow_input): + x = torch.empty(0, device='cuda:0') + if self.use_rgb: + x = torch.cat((x, rgb_input), dim=2) # Concatenating along the channel dimension + if self.use_flow: + x = torch.cat((x, flow_input), dim=2) + + # if self.use_rgb and self.use_flow: + # x = torch.cat((rgb_input, flow_input), 2) + # elif self.use_rgb: + # x = rgb_input + # elif self.use_flow: + # x = flow_input + + x = self.layer1(x) + B, _, _ = x.shape + h0 = self.h0.expand(-1, B, -1).to(x.device) + ht, _ = self.gru(x, h0) + ht = self.relu(ht) + # ht = self.relu(ht + x) + logits = self.f_classification(ht) + out_dict = {} + if self.training: + out_dict['logits'] = logits + else: + pred_scores = F.softmax(logits, dim=-1) + out_dict['logits'] = pred_scores + return out_dict + +@META_ARCHITECTURES.register("MiniROADA") +class MROADA(nn.Module): + + def __init__(self, cfg): + super(MROADA, self).__init__() + self.use_flow = not cfg['no_flow'] + self.use_rgb = not cfg['no_rgb'] + self.input_dim = 0 + if self.use_rgb: + self.input_dim += FEATURE_SIZES[cfg['rgb_type']] + if self.use_flow: + self.input_dim += FEATURE_SIZES[cfg['flow_type']] + + self.embedding_dim = cfg['embedding_dim'] + self.hidden_dim = cfg['hidden_dim'] + self.num_layers = cfg['num_layers'] + self.anticipation_length = cfg["anticipation_length"] + self.out_dim = cfg['num_classes'] + + self.layer1 = nn.Sequential( + nn.Linear(self.input_dim, self.embedding_dim), + nn.LayerNorm(self.embedding_dim), + nn.ReLU(), + nn.Dropout(p=cfg['dropout']) + ) + + self.actionness = cfg['actionness'] + if self.actionness: + self.f_actionness = nn.Sequential( + nn.Linear(self.hidden_dim, 1), + ) + self.relu = nn.ReLU() + self.gru = nn.GRU(self.embedding_dim, self.hidden_dim, self.num_layers, batch_first=True) + self.f_classification = nn.Sequential( + nn.Linear(self.hidden_dim, self.out_dim) + ) + self.anticipation_layer = nn.Sequential( + nn.Linear(self.hidden_dim, self.anticipation_length*self.hidden_dim), + ) + self.cfg = cfg + + def forward(self, rgb_input, flow_input): + x = torch.empty(0, device='cuda:0') + if self.use_rgb: + # [64, 128, 2048] + x = torch.cat((x, rgb_input), dim=2) # Concatenating along the channel dimension + if self.use_flow: + # [64, 128, 5120] + x = torch.cat((x, flow_input), dim=2) + + # audio channel is too huge to directly concat... + # to c + + B, S, feat_dim = x.shape # batch, length, dim + # + x = self.layer1(x) + h0 = torch.zeros(1, B, self.hidden_dim).to(x.device) + ht, _ = self.gru(x, h0) + ht = self.relu(ht) + + logits = self.f_classification(self.relu(ht)) + anticipation_ht = self.anticipation_layer(self.relu(ht)).view(B, S, self.anticipation_length, self.hidden_dim) + anticipation_logits = self.f_classification(self.relu(anticipation_ht)) + + + out_dict = {} + if self.training: + out_dict['logits'] = logits + out_dict['anticipation_logits'] = anticipation_logits + else: + pred_scores = F.softmax(logits, dim=-1) + pred_anticipation_scores = F.softmax(anticipation_logits, dim=-1) + out_dict['logits'] = pred_scores + out_dict['anticipation_logits'] = pred_anticipation_scores + + return out_dict \ No newline at end of file diff --git a/third_party/EgoSpeak/MiniROAD/model/transformer_models/Attention.py b/third_party/EgoSpeak/MiniROAD/model/transformer_models/Attention.py new file mode 100755 index 0000000..f0473e4 --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/model/transformer_models/Attention.py @@ -0,0 +1,155 @@ +import math +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class SelfAttention(nn.Module): + def __init__( + self, dim, heads=8, qkv_bias=False, qk_scale=None, dropout_rate=0.0 + ): + super().__init__() + self.num_heads = heads + head_dim = dim // heads + self.scale = qk_scale or head_dim ** -0.5 + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.attn_drop = nn.Dropout(dropout_rate) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(dropout_rate) + + def forward(self, x): + B, N, C = x.shape + qkv = ( + self.qkv(x) + .reshape(B, N, 3, self.num_heads, C // self.num_heads) + .permute(2, 0, 3, 1, 4) + ) + q, k, v = ( + qkv[0], + qkv[1], + qkv[2], + ) # make torchscript happy (cannot use tensor as tuple) + + attn = (q @ k.transpose(-2, -1)) * self.scale + attn = attn.softmax(dim=-1) + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + + +class AxialAttention(nn.Module): + def __init__( + self, + in_planes, + out_planes, + groups=8, + kernel_size=56, + stride=1, + bias=False, + width=False, + ): + assert (in_planes % groups == 0) and (out_planes % groups == 0) + super(AxialAttention, self).__init__() + self.in_planes = in_planes + self.out_planes = out_planes + self.groups = groups + self.group_planes = out_planes // groups + self.kernel_size = kernel_size + self.stride = stride + self.bias = bias + self.width = width + + # Multi-head self attention + self.qkv_transform = nn.Conv1d( + in_planes, + out_planes * 2, + kernel_size=1, + stride=1, + padding=0, + bias=False, + ) + self.bn_qkv = nn.BatchNorm1d(out_planes * 2) + self.bn_similarity = nn.BatchNorm2d(groups * 3) + self.bn_output = nn.BatchNorm1d(out_planes * 2) + + # Position embedding + self.relative = nn.Parameter( + torch.randn(self.group_planes * 2, kernel_size * 2 - 1), + requires_grad=True, + ) + query_index = torch.arange(kernel_size).unsqueeze(0) + key_index = torch.arange(kernel_size).unsqueeze(1) + relative_index = key_index - query_index + kernel_size - 1 + self.register_buffer('flatten_index', relative_index.view(-1)) + if stride > 1: + self.pooling = nn.AvgPool2d(stride, stride=stride) + + self.reset_parameters() + + def forward(self, x): + if self.width: + x = x.permute(0, 2, 1, 3) + else: + x = x.permute(0, 3, 1, 2) # N, W, C, H + N, W, C, H = x.shape + x = x.contiguous().view(N * W, C, H) + + # Transformations + qkv = self.bn_qkv(self.qkv_transform(x)) + q, k, v = torch.split( + qkv.reshape(N * W, self.groups, self.group_planes * 2, H), + [self.group_planes // 2, self.group_planes // 2, self.group_planes], + dim=2, + ) + + # Calculate position embedding + all_embeddings = torch.index_select( + self.relative, 1, self.flatten_index + ).view(self.group_planes * 2, self.kernel_size, self.kernel_size) + q_embedding, k_embedding, v_embedding = torch.split( + all_embeddings, + [self.group_planes // 2, self.group_planes // 2, self.group_planes], + dim=0, + ) + qr = torch.einsum('bgci,cij->bgij', q, q_embedding) + kr = torch.einsum('bgci,cij->bgij', k, k_embedding).transpose(2, 3) + qk = torch.einsum('bgci, bgcj->bgij', q, k) + stacked_similarity = torch.cat([qk, qr, kr], dim=1) + stacked_similarity = ( + self.bn_similarity(stacked_similarity) + .view(N * W, 3, self.groups, H, H) + .sum(dim=1) + ) + + # (N, groups, H, H, W) + similarity = F.softmax(stacked_similarity, dim=3) + sv = torch.einsum('bgij,bgcj->bgci', similarity, v) + sve = torch.einsum('bgij,cij->bgci', similarity, v_embedding) + stacked_output = torch.cat([sv, sve], dim=-1).view( + N * W, self.out_planes * 2, H + ) + output = ( + self.bn_output(stacked_output) + .view(N, W, self.out_planes, 2, H) + .sum(dim=-2) + ) + + if self.width: + output = output.permute(0, 2, 1, 3) + else: + output = output.permute(0, 2, 3, 1) + + if self.stride > 1: + output = self.pooling(output) + + return output + + def reset_parameters(self): + self.qkv_transform.weight.data.normal_( + 0, math.sqrt(1.0 / self.in_planes) + ) + nn.init.normal_(self.relative, 0.0, math.sqrt(1.0 / self.group_planes)) diff --git a/third_party/EgoSpeak/MiniROAD/model/transformer_models/AxialNet.py b/third_party/EgoSpeak/MiniROAD/model/transformer_models/AxialNet.py new file mode 100755 index 0000000..d2238c6 --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/model/transformer_models/AxialNet.py @@ -0,0 +1,244 @@ +import torch +import torch.nn as nn +from .Attention import AxialAttention + + +def conv1x1(in_planes, out_planes, stride=1): + """1x1 convolution""" + return nn.Conv2d( + in_planes, out_planes, kernel_size=1, stride=stride, bias=False + ) + + +class AxialBlock(nn.Module): + expansion = 2 + + def __init__( + self, + inplanes, + planes, + stride=1, + downsample=None, + groups=1, + base_width=64, + dilation=1, + norm_layer=None, + kernel_size=56, + ): + super(AxialBlock, self).__init__() + if norm_layer is None: + norm_layer = nn.BatchNorm2d + width = int(planes * (base_width / 64.0)) + # Both self.conv2 and self.downsample layers downsample the input when stride != 1 + self.conv_down = conv1x1(inplanes, width) + self.bn1 = norm_layer(width) + self.hight_block = AxialAttention( + width, width, groups=groups, kernel_size=kernel_size + ) + self.width_block = AxialAttention( + width, + width, + groups=groups, + kernel_size=kernel_size, + stride=stride, + width=True, + ) + self.conv_up = conv1x1(width, planes * self.expansion) + self.bn2 = norm_layer(planes * self.expansion) + self.relu = nn.ReLU(inplace=True) + self.downsample = downsample + self.stride = stride + + def forward(self, x): + identity = x + + out = self.conv_down(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.hight_block(out) + out = self.width_block(out) + out = self.relu(out) + + out = self.conv_up(out) + out = self.bn2(out) + + if self.downsample is not None: + identity = self.downsample(x) + + out += identity + out = self.relu(out) + + return out + + +class AxialAttentionNet(nn.Module): + def __init__( + self, + layers, + num_classes=1000, + zero_init_residual=True, + groups=8, + width_per_group=64, + replace_stride_with_dilation=None, + norm_layer=None, + s=0.5, + ): + super(AxialAttentionNet, self).__init__() + block = AxialBlock + if norm_layer is None: + norm_layer = nn.BatchNorm2d + self._norm_layer = norm_layer + + self.inplanes = int(64 * s) + self.dilation = 1 + if replace_stride_with_dilation is None: + # each element in the tuple indicates if we should replace + # the 2x2 stride with a dilated convolution instead + replace_stride_with_dilation = [False, False, False] + if len(replace_stride_with_dilation) != 3: + raise ValueError( + "replace_stride_with_dilation should be None " + "or a 3-element tuple, got {}".format( + replace_stride_with_dilation + ) + ) + self.groups = groups + self.base_width = width_per_group + self.conv1 = nn.Conv2d( + 3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False + ) + self.bn1 = norm_layer(self.inplanes) + self.relu = nn.ReLU(inplace=True) + self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) + self.layer1 = self._make_layer( + block, int(128 * s), layers[0], kernel_size=56 + ) + self.layer2 = self._make_layer( + block, + int(256 * s), + layers[1], + stride=2, + kernel_size=56, + dilate=replace_stride_with_dilation[0], + ) + self.layer3 = self._make_layer( + block, + int(512 * s), + layers[2], + stride=2, + kernel_size=28, + dilate=replace_stride_with_dilation[1], + ) + self.layer4 = self._make_layer( + block, + int(1024 * s), + layers[3], + stride=2, + kernel_size=14, + dilate=replace_stride_with_dilation[2], + ) + self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) + self.fc = nn.Linear(int(1024 * block.expansion * s), num_classes) + + for m in self.modules(): + if isinstance(m, (nn.Conv2d, nn.Conv1d)): + if isinstance(m, nn.Conv1d): + pass + else: + nn.init.kaiming_normal_( + m.weight, mode='fan_out', nonlinearity='relu' + ) + elif isinstance(m, (nn.BatchNorm2d, nn.BatchNorm1d, nn.GroupNorm)): + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + + # Zero-initialize the last BN in each residual branch, + # so that the residual branch starts with zeros, and each residual block behaves like an identity. + # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677 + if zero_init_residual: + for m in self.modules(): + if isinstance(m, AxialBlock): + nn.init.constant_(m.bn2.weight, 0) + + def _make_layer( + self, block, planes, blocks, kernel_size=56, stride=1, dilate=False + ): + norm_layer = self._norm_layer + downsample = None + previous_dilation = self.dilation + if dilate: + self.dilation *= stride + stride = 1 + if stride != 1 or self.inplanes != planes * block.expansion: + downsample = nn.Sequential( + conv1x1(self.inplanes, planes * block.expansion, stride), + norm_layer(planes * block.expansion), + ) + + layers = [] + layers.append( + block( + self.inplanes, + planes, + stride, + downsample, + groups=self.groups, + base_width=self.base_width, + dilation=previous_dilation, + norm_layer=norm_layer, + kernel_size=kernel_size, + ) + ) + self.inplanes = planes * block.expansion + if stride != 1: + kernel_size = kernel_size // 2 + + for _ in range(1, blocks): + layers.append( + block( + self.inplanes, + planes, + groups=self.groups, + base_width=self.base_width, + dilation=self.dilation, + norm_layer=norm_layer, + kernel_size=kernel_size, + ) + ) + + return nn.Sequential(*layers) + + def _forward_impl(self, x, include_conv5=False, include_top=False): + # See note [TorchScript super()] + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + x = self.maxpool(x) + + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + if include_conv5: + x = self.layer4(x) + + if include_top: + x = self.avgpool(x) + x = torch.flatten(x, 1) + x = self.fc(x) + + return x + + def forward(self, x, include_conv5=False, include_top=False): + return self._forward_impl(x, include_conv5, include_top) + + +if __name__ == "__main__": + model = AxialAttentionNet([1, 2, 4, 1], s=0.5) + # model = AxialAttentionNet([3, 4, 6, 3], s=0.25) + # print(model) + + x = torch.randn([8, 3, 224, 224]) + print(model(x).shape) + print(model(x, include_conv5=True).shape) + print(model(x, include_conv5=True, include_top=True).shape) diff --git a/third_party/EgoSpeak/MiniROAD/model/transformer_models/BiT.py b/third_party/EgoSpeak/MiniROAD/model/transformer_models/BiT.py new file mode 100755 index 0000000..d551156 --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/model/transformer_models/BiT.py @@ -0,0 +1,138 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from collections import OrderedDict + + +class StdConv2d(nn.Conv2d): + def forward(self, x): + w = self.weight + v, m = torch.var_mean(w, dim=[1, 2, 3], keepdim=True, unbiased=False) + w = (w - m) / torch.sqrt(v + 1e-10) + return F.conv2d( + x, + w, + self.bias, + self.stride, + self.padding, + self.dilation, + self.groups, + ) + + +def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1): + """3x3 convolution with padding""" + return StdConv2d( + in_planes, + out_planes, + kernel_size=3, + stride=stride, + padding=dilation, + groups=groups, + bias=False, + dilation=dilation, + ) + + +def conv1x1(in_planes, out_planes, stride=1): + """1x1 convolution""" + return StdConv2d( + in_planes, out_planes, kernel_size=1, stride=stride, bias=False + ) + + +class PreActBottleneck(nn.Module): + """Pre-activation (v2) bottleneck block. + Follows the implementation of "Identity Mappings in Deep Residual Networks": + https://github.com/KaimingHe/resnet-1k-layers/blob/master/resnet-pre-act.lua + Except it puts the stride on 3x3 conv when available. + """ + + def __init__(self, in_planes, out_planes=None, mid_planes=None, stride=1): + super(PreActBottleneck, self).__init__() + out_planes = out_planes or in_planes + mid_planes = mid_planes or out_planes // 4 + + self.gn1 = nn.GroupNorm(32, in_planes) + self.conv1 = conv1x1(in_planes, mid_planes) + self.gn2 = nn.GroupNorm(32, mid_planes) + self.conv2 = conv3x3(mid_planes, mid_planes, stride) + self.gn3 = nn.GroupNorm(32, mid_planes) + self.conv3 = conv1x1(mid_planes, out_planes) + self.relu = nn.ReLU(inplace=True) + + if stride != 1 or in_planes != out_planes: + # Projection also with pre-activation according to paper. + self.downsample = conv1x1(in_planes, out_planes, stride) + + def forward(self, x): + out = self.relu(self.gn1(x)) + + # Residual branch + residual = x + if hasattr(self, 'downsample'): + residual = self.downsample(out) + + # Unit's branch + out = self.conv1(out) + out = self.conv2(self.relu(self.gn2(out))) + out = self.conv3(self.relu(self.gn3(out))) + + return out + residual + + +class ResNetV2Model(nn.Module): + """Implementation of Pre-activation (v2) ResNet mode.""" + + def __init__(self, block_units, width_factor, head_size=21843): + super(ResNetV2Model, self).__init__() + wf = width_factor # shortcut 'cause we'll use it a lot. + + # The following will be unreadable if we split lines. + # pylint: disable=line-too-long + self.conv1 = nn.Sequential(OrderedDict([ + ('conv', StdConv2d(3, 64*wf, kernel_size=7, stride=2, padding=3, bias=False)), + ('pad', nn.ConstantPad2d(1, 0)), + ('pool', nn.MaxPool2d(kernel_size=3, stride=2, padding=0)), + ])) + + self.conv2 = nn.Sequential(OrderedDict( + [('unit01', PreActBottleneck(in_planes=64*wf, out_planes=256*wf, mid_planes=64*wf))] + + [(f'unit{i:02d}', PreActBottleneck(in_planes=256*wf, out_planes=256*wf, mid_planes=64*wf)) for i in range(2, block_units[0] + 1)], + )) + self.conv3 = nn.Sequential(OrderedDict( + [('unit01', PreActBottleneck(in_planes=256*wf, out_planes=512*wf, mid_planes=128*wf, stride=2))] + + [(f'unit{i:02d}', PreActBottleneck(in_planes=512*wf, out_planes=512*wf, mid_planes=128*wf)) for i in range(2, block_units[1] + 1)], + )) + self.conv4 = nn.Sequential(OrderedDict( + [('unit01', PreActBottleneck(in_planes=512*wf, out_planes=1024*wf, mid_planes=256*wf, stride=2))] + + [(f'unit{i:02d}', PreActBottleneck(in_planes=1024*wf, out_planes=1024*wf, mid_planes=256*wf)) for i in range(2, block_units[2] + 1)], + )) + self.conv5 = nn.Sequential(OrderedDict( + [('unit01', PreActBottleneck(in_planes=1024*wf, out_planes=2048*wf, mid_planes=512*wf, stride=2))] + + [(f'unit{i:02d}', PreActBottleneck(in_planes=2048*wf, out_planes=2048*wf, mid_planes=512*wf)) for i in range(2, block_units[3] + 1)], + )) + # pylint: enable=line-too-long + + self.head = nn.Sequential(OrderedDict([ + ('gn', nn.GroupNorm(32, 2048*wf)), + ('relu', nn.ReLU(inplace=True)), + ('avg', nn.AdaptiveAvgPool2d(output_size=1)), + ('conv', nn.Conv2d(2048*wf, head_size, kernel_size=1, bias=True)), + ])) + + def forward(self, x, include_conv5=False, include_top=False): + x = self.conv1(x) + x = self.conv2(x) + x = self.conv3(x) + x = self.conv4(x) + if include_conv5: + x = self.conv5(x) + if include_top: + x = self.head(x) + + if include_top and include_conv5: + assert x.shape[-2:] == (1,1,) + return x[..., 0, 0] + + return x diff --git a/third_party/EgoSpeak/MiniROAD/model/transformer_models/HybridViT.py b/third_party/EgoSpeak/MiniROAD/model/transformer_models/HybridViT.py new file mode 100755 index 0000000..6434b0d --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/model/transformer_models/HybridViT.py @@ -0,0 +1,223 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from .BiT import ResNetV2Model +from .AxialNet import AxialAttentionNet +from .Transformer import TransformerModel +from .PositionalEncoding import ( + FixedPositionalEncoding, + LearnedPositionalEncoding, +) + + +class HybridVisionTransformer(nn.Module): + def __init__( + self, + img_dim, + out_dim, + num_channels, + embedding_dim, + num_heads, + num_layers, + hidden_dim, + include_conv5, + dropout_rate, + positional_encoding_type, + backbone=None, + ): + super(HybridVisionTransformer, self).__init__() + + assert embedding_dim % num_heads == 0 + + self.embedding_dim = embedding_dim + self.num_heads = num_heads + self.out_dim = out_dim + self.num_channels = num_channels + self.include_conv5 = include_conv5 + self.backbone = backbone + + self.backbone_model, self.flatten_dim = self.configure_backbone() + + self.projection_encoding = nn.Linear(self.flatten_dim, embedding_dim) + self.cls_token = nn.Parameter(torch.randn(1, 1, embedding_dim)) + + self.decoder_dim = int(img_dim / 16.0) ** 2 + if self.include_conv5: + self.decoder_dim = int(img_dim / 32.0) ** 2 + + self.decoder_dim += 1 # for the cls token + + if positional_encoding_type == "learned": + self.position_encoding = LearnedPositionalEncoding( + self.decoder_dim, self.embedding_dim, self.decoder_dim + ) + elif positional_encoding_type == "fixed": + self.position_encoding = FixedPositionalEncoding( + self.embedding_dim, + ) + + self.transformer = TransformerModel( + embedding_dim, num_layers, num_heads, hidden_dim + ) + self.mlp_head = nn.Sequential( + nn.Linear(embedding_dim, hidden_dim), + nn.GELU(), + nn.Linear(hidden_dim, out_dim), + ) + self.to_cls_token = nn.Identity() + + def forward(self, x): + # apply bit backbone + x = self.backbone_model(x, include_conv5=self.include_conv5) + x = x.view(x.size(0), -1, self.flatten_dim) + + x = self.projection_encoding(x) + cls_tokens = self.cls_token.expand(x.shape[0], -1, -1) + x = torch.cat((cls_tokens, x), dim=1) + x = self.position_encoding(x) + + # apply transformer + x = self.transformer(x) + x = self.to_cls_token(x[:, 0]) + x = self.mlp_head(x) + x = F.log_softmax(x, dim=-1) + + return x + + def configure_backbone(self): + raise NotImplementedError("Method to be called in child class!!") + + +class ResNetHybridViT(HybridVisionTransformer): + def __init__( + self, + img_dim, + out_dim, + num_channels, + embedding_dim, + num_heads, + num_layers, + hidden_dim, + include_conv5=False, + dropout_rate=0.1, + positional_encoding_type="learned", + backbone='r50x1', + ): + super(ResNetHybridViT, self).__init__( + img_dim=img_dim, + out_dim=out_dim, + num_channels=num_channels, + embedding_dim=embedding_dim, + num_heads=num_heads, + num_layers=num_layers, + hidden_dim=hidden_dim, + include_conv5=include_conv5, + dropout_rate=dropout_rate, + positional_encoding_type=positional_encoding_type, + backbone=backbone, + ) + + def configure_backbone(self): + """ + Current support offered for all BiT models + KNOWN_MODELS in https://github.com/google-research/big_transfer/blob/master/bit_pytorch/models.py + + expects model name of style 'r{depth}x{width}' + where depth in [50, 101, 152] + where width in [1,2,3,4] + """ + backbone = self.backbone + out_dim = self.out_dim + + splits = backbone.split('x') + model_name = splits[0] + width_factor = int(splits[1]) + + if model_name in ['r50', 'r101'] and width_factor in [2, 4]: + return ValueError( + "Invalid Configuration of models -- expect 50x1, 50x3, 101x1, 101x3" + ) + elif model_name == 'r152' and width_factor in [1, 3]: + return ValueError( + "Invalid Configuration of models -- expect 152x2, 152x4" + ) + + block_units_dict = { + 'r50': [3, 4, 6, 3], + 'r101': [3, 4, 23, 3], + 'r152': [3, 8, 36, 3], + } + block_units = block_units_dict.get(model_name, [3, 4, 6, 3]) + model = ResNetV2Model(block_units, width_factor, head_size=out_dim) + + if self.num_channels == 3: + flatten_dim = 1024 * width_factor + if self.include_conv5: + flatten_dim *= 2 + + return model, flatten_dim + + +class AxialNetHybridViT(HybridVisionTransformer): + def __init__( + self, + img_dim, + out_dim, + num_channels, + embedding_dim, + num_heads, + num_layers, + hidden_dim, + include_conv5=False, + dropout_rate=0.1, + positional_encoding_type="learned", + backbone='a50m', + ): + super(AxialNetHybridViT, self).__init__( + img_dim=img_dim, + out_dim=out_dim, + num_channels=num_channels, + embedding_dim=embedding_dim, + num_heads=num_heads, + num_layers=num_layers, + hidden_dim=hidden_dim, + include_conv5=include_conv5, + dropout_rate=dropout_rate, + positional_encoding_type=positional_encoding_type, + backbone=backbone, + ) + + def configure_backbone(self): + """ + Current support offered for all BiT models + models from https://github.com/csrhddlam/axial-deeplab/blob/master/lib/models/axialnet.py + + expects model name of style 'a{depth}{width}' + where depth in [26, 50, 101] + where width in [s, m, l] + """ + backbone = self.backbone + out_dim = self.out_dim + + model_name = backbone[:3] + width = backbone[-1] + + block_units_dict = { + 'a26': [1, 2, 4, 1], + 'a50': [3, 4, 6, 3], + 'a101': [3, 4, 23, 3], + } + block_units = block_units_dict.get(model_name, [3, 4, 6, 3]) + + scale_factor_dict = {'s': 0.5, 'm': 0.75, 'l': 1.0} + scale_factor = scale_factor_dict.get(width, 0.75) + model = AxialAttentionNet( + block_units, s=scale_factor, num_classes=out_dim + ) + + if self.num_channels == 3: + flatten_dim = int(512 * float(scale_factor / 0.5)) + if self.include_conv5: + flatten_dim *= 2 + + return model, flatten_dim diff --git a/third_party/EgoSpeak/MiniROAD/model/transformer_models/OADTR_ViT_origin.py b/third_party/EgoSpeak/MiniROAD/model/transformer_models/OADTR_ViT_origin.py new file mode 100644 index 0000000..9ca4c40 --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/model/transformer_models/OADTR_ViT_origin.py @@ -0,0 +1,331 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from .decoder import Decoder, DecoderLayer +from .attn import FullAttention, ProbAttention, AttentionLayer +from .Transformer import TransformerModel +from ipdb import set_trace +from .PositionalEncoding import ( + FixedPositionalEncoding, + LearnedPositionalEncoding, +) + +__all__ = ['ViT_B16', 'ViT_B32', 'ViT_L16', 'ViT_L32', 'ViT_H14'] + + +class VisionTransformer_v3(nn.Module): + def __init__( + self, + args, + img_dim, + patch_dim, + out_dim, + embedding_dim, + num_heads, + num_layers, + hidden_dim, + dropout_rate=0.0, + attn_dropout_rate=0.0, + use_representation=True, + conv_patch_representation=False, + positional_encoding_type="learned", with_camera=True, with_motion=True, num_channels=3072, + ): + super(VisionTransformer_v3, self).__init__() + + assert embedding_dim % num_heads == 0 + assert img_dim % patch_dim == 0 + self.with_camera = with_camera + self.with_motion = with_motion + self.embedding_dim = embedding_dim + self.num_heads = num_heads + self.patch_dim = patch_dim + # num_channels = img_dim + self.num_channels = num_channels + self.dropout_rate = dropout_rate + self.attn_dropout_rate = attn_dropout_rate + self.conv_patch_representation = conv_patch_representation + + # self.num_patches = int((img_dim // patch_dim) ** 2) + self.num_patches = int(img_dim // patch_dim) + self.seq_length = self.num_patches + 1 + self.flatten_dim = patch_dim * patch_dim * num_channels + self.cls_token = nn.Parameter(torch.zeros(1, 1, embedding_dim)) + + self.linear_encoding = nn.Linear(self.flatten_dim, embedding_dim) + if positional_encoding_type == "learned": + self.position_encoding = LearnedPositionalEncoding( + self.seq_length, self.embedding_dim, self.seq_length + ) + elif positional_encoding_type == "fixed": + self.position_encoding = FixedPositionalEncoding( + self.embedding_dim, + ) + print('position encoding :', positional_encoding_type) + + self.pe_dropout = nn.Dropout(p=self.dropout_rate) + + self.encoder = TransformerModel( + embedding_dim, + num_layers, + num_heads, + hidden_dim, + self.dropout_rate, + self.attn_dropout_rate, + ) + self.pre_head_ln = nn.LayerNorm(embedding_dim) + + d_model = args.decoder_embedding_dim + use_representation = False # False + if use_representation: + self.mlp_head = nn.Sequential( + nn.Linear(embedding_dim + d_model, hidden_dim//2), + # nn.Tanh(), + nn.ReLU(), + nn.Linear(hidden_dim//2, out_dim), + ) + else: + self.mlp_head = nn.Linear(embedding_dim + d_model, out_dim) + + if self.conv_patch_representation: + # self.conv_x = nn.Conv2d( + # self.num_channels, + # self.embedding_dim, + # kernel_size=(self.patch_dim, self.patch_dim), + # stride=(self.patch_dim, self.patch_dim), + # padding=self._get_padding( + # 'VALID', (self.patch_dim, self.patch_dim), + # ), + # ) + self.conv_x = nn.Conv1d( + self.num_channels, + self.embedding_dim, + kernel_size=self.patch_dim, + stride=self.patch_dim, + padding=self._get_padding( + 'VALID', (self.patch_dim), + ), + ) + else: + self.conv_x = None + + self.to_cls_token = nn.Identity() + + # Decoder + factor = 1 # 5 + dropout = args.decoder_attn_dropout_rate + # d_model = args.decoder_embedding_dim + n_heads = args.decoder_num_heads + d_layers = args.decoder_layers + d_ff = args.decoder_embedding_dim_out # args.decoder_embedding_dim_out or 4*args.decoder_embedding_dim None + activation = 'gelu' # 'gelu' + self.decoder = Decoder( + [ + DecoderLayer( + AttentionLayer(FullAttention(True, factor, attention_dropout=dropout), # True + d_model, n_heads), # ProbAttention FullAttention + AttentionLayer(FullAttention(False, factor, attention_dropout=dropout), # False + d_model, n_heads), + d_model, + d_ff, + dropout=dropout, + activation=activation, + ) + for l in range(d_layers) + ], + norm_layer=torch.nn.LayerNorm(d_model) + ) + self.decoder_cls_token = nn.Parameter(torch.zeros(1, args.query_num, d_model)) + if positional_encoding_type == "learned": + self.decoder_position_encoding = LearnedPositionalEncoding( + args.query_num, self.embedding_dim, args.query_num + ) + elif positional_encoding_type == "fixed": + self.decoder_position_encoding = FixedPositionalEncoding( + self.embedding_dim, + ) + print('position decoding :', positional_encoding_type) + self.classifier = nn.Linear(d_model, out_dim) + self.after_dropout = nn.Dropout(p=self.dropout_rate) + # self.merge_fc = nn.Linear(d_model, 1) + # self.merge_sigmoid = nn.Sigmoid() + + def forward(self, sequence_input_rgb, sequence_input_flow): + if self.with_camera and self.with_motion: + x = torch.cat((sequence_input_rgb, sequence_input_flow), 2) + elif self.with_camera: + x = sequence_input_rgb + elif self.with_motion: + x = sequence_input_flow + + x = self.linear_encoding(x) + cls_tokens = self.cls_token.expand(x.shape[0], -1, -1) + # x = torch.cat((cls_tokens, x), dim=1) + x = torch.cat((x, cls_tokens), dim=1) + x = self.position_encoding(x) + x = self.pe_dropout(x) # not delete + + # apply transformer + x = self.encoder(x) + x = self.pre_head_ln(x) # [128, 33, 1024] + # x = self.after_dropout(x) # add + # decoder + decoder_cls_token = self.decoder_cls_token.expand(x.shape[0], -1, -1) + # decoder_cls_token = self.after_dropout(decoder_cls_token) # add + # decoder_cls_token = self.decoder_position_encoding(decoder_cls_token) # [128, 8, 1024] + dec = self.decoder(decoder_cls_token, x) # [128, 8, 1024] + dec = self.after_dropout(dec) # add + # merge_atte = self.merge_sigmoid(self.merge_fc(dec)) # [128, 8, 1] + # dec_for_token = (merge_atte*dec).sum(dim=1) # [128, 1024] + # dec_for_token = (merge_atte*dec).sum(dim=1)/(merge_atte.sum(dim=-2) + 0.0001) + dec_for_token = dec.mean(dim=1) + # dec_for_token = dec.max(dim=1)[0] + dec_cls_out = self.classifier(dec) + # set_trace() + # x = self.to_cls_token(x[:, 0]) + x = torch.cat((self.to_cls_token(x[:, -1]), dec_for_token), dim=1) + x = self.mlp_head(x) + # x = F.log_softmax(x, dim=-1) + + return x, dec_cls_out + + def _get_padding(self, padding_type, kernel_size): + assert padding_type in ['SAME', 'VALID'] + if padding_type == 'SAME': + _list = [(k - 1) // 2 for k in kernel_size] + return tuple(_list) + return tuple(0 for _ in kernel_size) + + +def ViT_B16(dataset='imagenet'): + if dataset == 'imagenet': + img_dim = 224 + out_dim = 1000 + patch_dim = 16 + elif 'cifar' in dataset: + img_dim = 32 + out_dim = 10 + patch_dim = 4 + + return VisionTransformer( + img_dim=img_dim, + patch_dim=patch_dim, + out_dim=out_dim, + num_channels=3, + embedding_dim=768, + num_heads=12, + num_layers=12, + hidden_dim=3072, + dropout_rate=0.1, + attn_dropout_rate=0.0, + use_representation=False, + conv_patch_representation=False, + positional_encoding_type="learned", + ) + + +def ViT_B32(dataset='imagenet'): + if dataset == 'imagenet': + img_dim = 224 + out_dim = 1000 + patch_dim = 32 + elif 'cifar' in dataset: + img_dim = 32 + out_dim = 10 + patch_dim = 4 + + return VisionTransformer( + img_dim=img_dim, + patch_dim=patch_dim, + out_dim=out_dim, + num_channels=3, + embedding_dim=768, + num_heads=12, + num_layers=12, + hidden_dim=3072, + dropout_rate=0.1, + attn_dropout_rate=0.0, + use_representation=False, + conv_patch_representation=False, + positional_encoding_type="learned", + ) + + +def ViT_L16(dataset='imagenet'): + if dataset == 'imagenet': + img_dim = 224 + out_dim = 1000 + patch_dim = 16 + elif 'cifar' in dataset: + img_dim = 32 + out_dim = 10 + patch_dim = 4 + + return VisionTransformer( + img_dim=img_dim, + patch_dim=patch_dim, + out_dim=out_dim, + num_channels=3, + embedding_dim=1024, + num_heads=16, + num_layers=24, + hidden_dim=4096, + dropout_rate=0.1, + attn_dropout_rate=0.0, + use_representation=False, + conv_patch_representation=False, + positional_encoding_type="learned", + ) + + +def ViT_L32(dataset='imagenet'): + if dataset == 'imagenet': + img_dim = 224 + out_dim = 1000 + patch_dim = 32 + elif 'cifar' in dataset: + img_dim = 32 + out_dim = 10 + patch_dim = 4 + + return VisionTransformer( + img_dim=img_dim, + patch_dim=patch_dim, + out_dim=out_dim, + num_channels=3, + embedding_dim=1024, + num_heads=16, + num_layers=24, + hidden_dim=4096, + dropout_rate=0.1, + attn_dropout_rate=0.0, + use_representation=False, + conv_patch_representation=False, + positional_encoding_type="learned", + ) + + +def ViT_H14(dataset='imagenet'): + if dataset == 'imagenet': + img_dim = 224 + out_dim = 1000 + patch_dim = 14 + elif 'cifar' in dataset: + img_dim = 32 + out_dim = 10 + patch_dim = 4 + + return VisionTransformer( + img_dim=img_dim, + patch_dim=patch_dim, + out_dim=out_dim, + num_channels=3, + embedding_dim=1280, + num_heads=16, + num_layers=32, + hidden_dim=5120, + dropout_rate=0.1, + attn_dropout_rate=0.0, + use_representation=False, + conv_patch_representation=False, + positional_encoding_type="learned", + ) \ No newline at end of file diff --git a/third_party/EgoSpeak/MiniROAD/model/transformer_models/PositionalEncoding.py b/third_party/EgoSpeak/MiniROAD/model/transformer_models/PositionalEncoding.py new file mode 100755 index 0000000..e3f5019 --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/model/transformer_models/PositionalEncoding.py @@ -0,0 +1,44 @@ +import torch +import torch.nn as nn + + +class FixedPositionalEncoding(nn.Module): + def __init__(self, embedding_dim, max_length=5000): + super(FixedPositionalEncoding, self).__init__() + + pe = torch.zeros(max_length, embedding_dim) + position = torch.arange(0, max_length, dtype=torch.float).unsqueeze(1) + div_term = torch.exp( + torch.arange(0, embedding_dim, 2).float() + * (-torch.log(torch.tensor(10000.0)) / embedding_dim) + ) + pe[:, 0::2] = torch.sin(position * div_term) + pe[:, 1::2] = torch.cos(position * div_term) + pe = pe.unsqueeze(0).transpose(0, 1) + self.register_buffer('pe', pe) + + def forward(self, x): + x = x + self.pe[: x.size(0), :] + return x + + +class LearnedPositionalEncoding(nn.Module): + def __init__(self, max_position_embeddings, embedding_dim, seq_length): + super(LearnedPositionalEncoding, self).__init__() + self.pe = nn.Embedding(max_position_embeddings, embedding_dim) + self.seq_length = seq_length + + self.register_buffer( + "position_ids", + torch.arange(max_position_embeddings).expand((1, -1)), + ) + + def forward(self, x, position_ids=None): + if position_ids is None: + position_ids = self.position_ids[:, : self.seq_length] + + try: + position_embeddings = self.pe(position_ids) + return x + position_embeddings + except: + import pdb; pdb.set_trace() diff --git a/third_party/EgoSpeak/MiniROAD/model/transformer_models/Transformer.py b/third_party/EgoSpeak/MiniROAD/model/transformer_models/Transformer.py new file mode 100755 index 0000000..c1c9cb6 --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/model/transformer_models/Transformer.py @@ -0,0 +1,82 @@ +from torch import nn +from .Attention import SelfAttention + + +class Residual(nn.Module): + def __init__(self, fn): + super().__init__() + self.fn = fn + + def forward(self, x): + return self.fn(x) + x + + +class PreNorm(nn.Module): + def __init__(self, dim, fn): + super().__init__() + self.norm = nn.LayerNorm(dim) + self.fn = fn + + def forward(self, x): + return self.fn(self.norm(x)) + + +class PreNormDrop(nn.Module): + def __init__(self, dim, dropout_rate, fn): + super().__init__() + self.norm = nn.LayerNorm(dim) + self.dropout = nn.Dropout(p=dropout_rate) + self.fn = fn + + def forward(self, x): + return self.dropout(self.fn(self.norm(x))) + + +class FeedForward(nn.Module): + def __init__(self, dim, hidden_dim, dropout_rate): + super().__init__() + self.net = nn.Sequential( + nn.Linear(dim, hidden_dim), + nn.GELU(), + nn.Dropout(p=dropout_rate), + nn.Linear(hidden_dim, dim), + nn.Dropout(p=dropout_rate), + ) + + def forward(self, x): + return self.net(x) + + +class TransformerModel(nn.Module): + def __init__( + self, + dim, + depth, + heads, + mlp_dim, + dropout_rate=0.1, + attn_dropout_rate=0.1, + ): + super().__init__() + layers = [] + for _ in range(depth): + layers.extend( + [ + Residual( + PreNormDrop( + dim, + dropout_rate, + SelfAttention( + dim, heads=heads, dropout_rate=attn_dropout_rate + ), + ) + ), + Residual( + PreNorm(dim, FeedForward(dim, mlp_dim, dropout_rate)) + ), + ] + ) + self.net = nn.Sequential(*layers) + + def forward(self, x): + return self.net(x) diff --git a/third_party/EgoSpeak/MiniROAD/model/transformer_models/ViT.py b/third_party/EgoSpeak/MiniROAD/model/transformer_models/ViT.py new file mode 100755 index 0000000..16e1bbc --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/model/transformer_models/ViT.py @@ -0,0 +1,350 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from configs.feature_sizes import FEATURE_SIZES +from model.model_builder import META_ARCHITECTURES as registry + +from .attn import AttentionLayer, FullAttention, ProbAttention +from .decoder import Decoder, DecoderLayer +from .PositionalEncoding import (FixedPositionalEncoding, + LearnedPositionalEncoding) +from .Transformer import TransformerModel + + +# @registry.register('TRANSFORMER') +@registry.register('OADTR') +class OADTR(nn.Module): + def __init__(self, cfg, use_representation=True, + conv_patch_representation=False, + positional_encoding_type="learned" + ): + super(OADTR, self).__init__() + + self.img_dim = cfg["window_size"] # img_dim + self.out_dim = cfg["num_classes"] + + self.patch_dim = cfg["patch_dim"] + self.query_num = cfg["query_num"] + + self.embedding_dim = cfg["embedding_dim"] + self.num_heads = cfg["num_heads"] + self.num_layers = cfg["num_layers"] + self.hidden_dim = cfg["hidden_dim"] + + self.dropout_rate = cfg["dropout"] + self.attn_dropout_rate = cfg["attn_dropout_rate"] + + self.conv_patch_representation = conv_patch_representation + + self.decoder_embedding_dim = cfg["decoder_embedding_dim"] + self.decoder_num_heads = cfg["decoder_num_heads"] + self.decoder_layers = cfg["decoder_layers"] + self.decoder_embedding_dim_out = cfg["decoder_embedding_dim_out"] + self.decoder_attn_dropout_rate = cfg["decoder_attn_dropout_rate"] + + self.use_flow = not cfg['no_flow'] + self.use_rgb = not cfg['no_rgb'] + + self.num_channels = 0 + if self.use_rgb: + self.num_channels += FEATURE_SIZES[cfg['rgb_type']] + if self.use_flow: + self.num_channels += FEATURE_SIZES[cfg['flow_type']] + + assert self.embedding_dim % self.num_heads == 0 + assert self.img_dim % self.patch_dim == 0 + + self.num_patches = int(self.img_dim // self.patch_dim) + self.seq_length = self.num_patches + 1 + # self.seq_length = self.num_patches + self.img_dim + self.flatten_dim = self.patch_dim * self.patch_dim * self.num_channels + self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embedding_dim)) + # self.cls_token = nn.Parameter(torch.zeros(1, self.img_dim, self.embedding_dim)) + + self.linear_encoding = nn.Linear(self.flatten_dim, self.embedding_dim) + if positional_encoding_type == "learned": + self.position_encoding = LearnedPositionalEncoding( + self.seq_length, self.embedding_dim, self.seq_length + ) + elif positional_encoding_type == "fixed": + self.position_encoding = FixedPositionalEncoding( + self.embedding_dim, + ) + print('position encoding :', positional_encoding_type) + + self.pe_dropout = nn.Dropout(p=self.dropout_rate) + + self.encoder = TransformerModel( + self.embedding_dim, + self.num_layers, + self.num_heads, + self.hidden_dim, + self.dropout_rate, + self.attn_dropout_rate, + ) + self.pre_head_ln = nn.LayerNorm(self.embedding_dim) + + d_model = self.decoder_embedding_dim + use_representation = False # False + if use_representation: + self.mlp_head = nn.Sequential( + nn.Linear(self.embedding_dim , self.hidden_dim//2), + # nn.Tanh(), + nn.ReLU(), + nn.Linear(self.hidden_dim//2, self.out_dim), + ) + else: + self.mlp_head = nn.Linear(self.embedding_dim + d_model, self.out_dim) + + if self.conv_patch_representation: + # self.conv_x = nn.Conv2d( + # self.num_channels, + # self.embedding_dim, + # kernel_size=(self.patch_dim, self.patch_dim), + # stride=(self.patch_dim, self.patch_dim), + # padding=self._get_padding( + # 'VALID', (self.patch_dim, self.patch_dim), + # ), + # ) + self.conv_x = nn.Conv1d( + self.num_channels, + self.embedding_dim, + kernel_size=self.patch_dim, + stride=self.patch_dim, + padding=self._get_padding( + 'VALID', (self.patch_dim), + ), + ) + else: + self.conv_x = None + + self.to_cls_token = nn.Identity() + + # Decoder + factor = 1 # 5 + dropout = self.decoder_attn_dropout_rate + # d_model = self.decoder_embedding_dim + n_heads = self.decoder_num_heads + d_layers = self.decoder_layers + d_ff = self.decoder_embedding_dim_out # args.decoder_embedding_dim_out or 4*args.decoder_embedding_dim None + activation = 'gelu' # 'gelu' + self.decoder = Decoder( + [ + DecoderLayer( + AttentionLayer(FullAttention(True, factor, attention_dropout=dropout), # True + d_model, n_heads), # ProbAttention FullAttention + AttentionLayer(FullAttention(False, factor, attention_dropout=dropout), # False + d_model, n_heads), + d_model, + d_ff, + dropout=dropout, + activation=activation, + ) + for l in range(d_layers) + ], + norm_layer=torch.nn.LayerNorm(d_model) + ) + self.decoder_cls_token = nn.Parameter(torch.zeros(1, self.query_num, d_model)) + if positional_encoding_type == "learned": + self.decoder_position_encoding = LearnedPositionalEncoding( + self.query_num, self.embedding_dim, self.query_num + ) + elif positional_encoding_type == "fixed": + self.decoder_position_encoding = FixedPositionalEncoding( + self.embedding_dim, + ) + print('position decoding :', positional_encoding_type) + self.classifier = nn.Linear(d_model, self.out_dim) + self.after_dropout = nn.Dropout(p=self.dropout_rate) + # self.merge_fc = nn.Linear(d_model, 1) + # self.merge_sigmoid = nn.Sigmoid() + + + def forward(self, sequence_input_rgb, sequence_input_flow): + if self.use_rgb and self.use_flow: + x = torch.cat((sequence_input_rgb, sequence_input_flow), 2) + elif self.use_rgb: + x = sequence_input_rgb + elif self.use_flow: + x = sequence_input_flow + + x = self.linear_encoding(x) # B, 128, 1204 + cls_tokens = self.cls_token.expand(x.shape[0], -1, -1) # B, 1, 1024 + # x = torch.cat((cls_tokens, x), dim=1) + x = torch.cat((x, cls_tokens), dim=1) # B, seq+1, 1024 + try: + x = self.position_encoding(x) # B, seq+1, 1024 + except: + import pdb; pdb.set_trace() + x = self.pe_dropout(x) # not delete + + # apply transformer + x = self.encoder(x) + x = self.pre_head_ln(x) # B, seq+1, 1024 + + # decoder + decoder_cls_token = self.decoder_cls_token.expand(x.shape[0], -1, -1) # [B, 8, 1024] + # decoder_cls_token = self.after_dropout(decoder_cls_token) # add + # decoder_cls_token = self.decoder_position_encoding(decoder_cls_token) # [128, 8, 1024] + dec = self.decoder(decoder_cls_token, x) # [128, 8, 1024] + dec = self.after_dropout(dec) # add + + # merge_atte = self.merge_sigmoid(self.merge_fc(dec)) # [128, 8, 1] + # dec_for_token = (merge_atte*dec).sum(dim=1) # [128, 1024] + # dec_for_token = (merge_atte*dec).sum(dim=1)/(merge_atte.sum(dim=-2) + 0.0001) + # *** RuntimeError: mat1 and mat2 shapes cannot be multiplied (64x2048 and 1024x2) + dec_for_token = dec.mean(dim=1) # [64, 1024] + + # dec_for_token = dec.max(dim=1)[0] + dec_cls_out = self.classifier(dec) # [B, 8, NUM_CLASS?] + # set_trace() + # x = self.to_cls_token(x[:, 0]) + + x = torch.cat((self.to_cls_token(x[:, -1]), dec_for_token), dim=1) # [B, 2048] + x = self.mlp_head(x) # [B, NUM_CLASSES] + + # original OADTR + # return x, dec_cls_out + + # x = F.log_softmax(x, dim=-1) + out_dict = {} + out_dict['logits'] = x.unsqueeze(1) # [B, 1, NUM_CLASSES] ex) [64, 1, 2] + out_dict['dec_logits'] = dec_cls_out # [B, 8, NUM_CLASSES] ex) [64, 8, 2] + + # out_dict['logits'] = x #x.unsqueeze(1) + return out_dict + + def _get_padding(self, padding_type, kernel_size): + assert padding_type in ['SAME', 'VALID'] + if padding_type == 'SAME': + _list = [(k - 1) // 2 for k in kernel_size] + return tuple(_list) + return tuple(0 for _ in kernel_size) + +@registry.register('Transformer') +class ViTEnc(nn.Module): + def __init__(self, cfg, use_representation=True, + conv_patch_representation=False, + positional_encoding_type="learned" + ): + super(ViTEnc, self).__init__() + self.img_dim = cfg["window_size"] + self.out_dim = cfg["num_classes"] + self.embedding_dim = cfg["embedding_dim"] + self.patch_dim = cfg["patch_dim"] + self.num_heads = cfg["num_heads"] + self.num_layers = cfg["num_layers"] + self.hidden_dim = cfg["hidden_dim"] + self.dropout_rate = cfg["dropout"] + self.use_flow = not cfg['no_flow'] + self.use_rgb = not cfg['no_rgb'] + self.num_channels= 0 + if self.use_rgb: + self.num_channels += FEATURE_SIZES[cfg['rgb_type']] + if self.use_flow: + self.num_channels += FEATURE_SIZES[cfg['flow_type']] + self.attn_dropout_rate = cfg["attn_dropout_rate"] + assert self.embedding_dim % self.num_heads == 0 + assert self.img_dim % self.patch_dim == 0 + self.conv_patch_representation = conv_patch_representation + self.num_patches = int(self.img_dim // self.patch_dim) + self.seq_length = self.num_patches + 1 + # self.seq_length = self.num_patches + self.img_dim + self.flatten_dim = self.patch_dim * self.patch_dim * self.num_channels + self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embedding_dim)) + # self.cls_token = nn.Parameter(torch.zeros(1, self.img_dim, self.embedding_dim)) + + self.linear_encoding = nn.Linear(self.flatten_dim, self.embedding_dim) + if positional_encoding_type == "learned": + self.position_encoding = LearnedPositionalEncoding( + self.seq_length, self.embedding_dim, self.seq_length + ) + elif positional_encoding_type == "fixed": + self.position_encoding = FixedPositionalEncoding( + self.embedding_dim, + ) + print('position encoding :', positional_encoding_type) + + self.pe_dropout = nn.Dropout(p=self.dropout_rate) + + self.encoder = TransformerModel( + self.embedding_dim, + self.num_layers, + self.num_heads, + self.hidden_dim, + self.dropout_rate, + self.attn_dropout_rate, + ) + self.pre_head_ln = nn.LayerNorm(self.embedding_dim) + + use_representation = False # False + if use_representation: + self.mlp_head = nn.Sequential( + nn.Linear(self.embedding_dim , self.hidden_dim//2), + # nn.Tanh(), + nn.ReLU(), + nn.Linear(self.hidden_dim//2, self.out_dim), + ) + else: + self.mlp_head = nn.Linear(self.embedding_dim, self.out_dim) + + if self.conv_patch_representation: + # self.conv_x = nn.Conv2d( + # self.num_channels, + # self.embedding_dim, + # kernel_size=(self.patch_dim, self.patch_dim), + # stride=(self.patch_dim, self.patch_dim), + # padding=self._get_padding( + # 'VALID', (self.patch_dim, self.patch_dim), + # ), + # ) + self.conv_x = nn.Conv1d( + self.num_channels, + self.embedding_dim, + kernel_size=self.patch_dim, + stride=self.patch_dim, + padding=self._get_padding( + 'VALID', (self.patch_dim), + ), + ) + else: + self.conv_x = None + + self.to_cls_token = nn.Identity() + + + def forward(self, sequence_input_rgb, sequence_input_flow): + if self.use_rgb and self.use_flow: + x = torch.cat((sequence_input_rgb, sequence_input_flow), 2) + elif self.use_rgb: + x = sequence_input_rgb + elif self.use_flow: + x = sequence_input_flow + + x = self.linear_encoding(x) + cls_tokens = self.cls_token.expand(x.shape[0], -1, -1) # B, 1, 1024 + # x = torch.cat((cls_tokens, x), dim=1) + x = torch.cat((x, cls_tokens), dim=1) # B, seq+1, 1024 + + x = self.position_encoding(x) # B, seq+1, 1024 + x = self.pe_dropout(x) # not delete + + # apply transformer + x = self.encoder(x) + x = self.pre_head_ln(x) # B, seq+1, 1024 + + x = self.to_cls_token(x[:, 0]) # B, 1024 + # x = self.to_cls_token(x[:,0:self.img_dim]) # B, 1024 + x = self.mlp_head(x) + # x = F.log_softmax(x, dim=-1) + out_dict = {} + out_dict['logits'] = x.unsqueeze(1) + # out_dict['logits'] = x #x.unsqueeze(1) + return out_dict + + def _get_padding(self, padding_type, kernel_size): + assert padding_type in ['SAME', 'VALID'] + if padding_type == 'SAME': + _list = [(k - 1) // 2 for k in kernel_size] + return tuple(_list) + return tuple(0 for _ in kernel_size) diff --git a/third_party/EgoSpeak/MiniROAD/model/transformer_models/__init__.py b/third_party/EgoSpeak/MiniROAD/model/transformer_models/__init__.py new file mode 100755 index 0000000..62085d8 --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/model/transformer_models/__init__.py @@ -0,0 +1,4 @@ +from .ViT import ViTEnc +from .HybridViT import ResNetHybridViT, AxialNetHybridViT + +__all__ = ['ResNetHybridViT', 'AxialNetHybridViT', 'VisionTransformer_v3'] diff --git a/third_party/EgoSpeak/MiniROAD/model/transformer_models/attn.py b/third_party/EgoSpeak/MiniROAD/model/transformer_models/attn.py new file mode 100755 index 0000000..ef911b9 --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/model/transformer_models/attn.py @@ -0,0 +1,170 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +import numpy as np + +from math import sqrt + + +class TriangularCausalMask(): + def __init__(self, B, L, device="cpu"): + mask_shape = [B, 1, L, L] + with torch.no_grad(): + self._mask = torch.triu(torch.ones(mask_shape, dtype=torch.bool), diagonal=1).to(device) + + @property + def mask(self): + return self._mask + + +class ProbMask(): + def __init__(self, B, H, L, index, scores, device="cpu"): + _mask = torch.ones(L, scores.shape[-1], dytpe=torch.bool).to(device).triu(1) + _mask_ex = _mask[None, None, :].expand(B, H, L, scores.shape[-1]) + indicator = _mask_ex[torch.arange(B)[:, None, None], + torch.arange(H)[None, :, None], + index, :].to(device) + self._mask = indicator.view(scores.shape).to(device) + + @property + def mask(self): + return self._mask + + +class FullAttention(nn.Module): + def __init__(self, mask_flag=True, factor=5, scale=None, attention_dropout=0.1): + super(FullAttention, self).__init__() + self.scale = scale + self.mask_flag = mask_flag + self.dropout = nn.Dropout(attention_dropout) + + def forward(self, queries, keys, values, attn_mask): + B, L, H, E = queries.shape + _, S, _, D = values.shape + scale = self.scale or 1. / sqrt(E) + + scores = torch.einsum("blhe,bshe->bhls", queries, keys) + if self.mask_flag: + if attn_mask is None: + attn_mask = TriangularCausalMask(B, L, device=queries.device) + + scores.masked_fill_(attn_mask.mask, -np.inf) + + A = self.dropout(torch.softmax(scale * scores, dim=-1)) + V = torch.einsum("bhls,bshd->blhd", A, values) + + return V.contiguous() + + +class ProbAttention(nn.Module): + def __init__(self, mask_flag=True, factor=5, scale=None, attention_dropout=0.1): + super(ProbAttention, self).__init__() + self.factor = factor + self.scale = scale + self.mask_flag = mask_flag + self.dropout = nn.Dropout(attention_dropout) + + def _prob_QK(self, Q, K, sample_k, n_top): + # Q [B, H, L, D] + B, H, L, E = K.shape + _, _, S, _ = Q.shape + + # calculate the sampled Q_K + K_expand = K.unsqueeze(-3).expand(B, H, S, L, E) + indx_sample = torch.randint(L, (S, sample_k)) + K_sample = K_expand[:, :, torch.arange(S).unsqueeze(1), indx_sample, :] + Q_K_sample = torch.matmul(Q.unsqueeze(-2), K_sample.transpose(-2, -1)).squeeze() + + # find the Top_k query with sparisty measurement + M = Q_K_sample.max(-1)[0] - torch.div(Q_K_sample.sum(-1), L) + M_top = M.topk(n_top, sorted=False)[1] + + # use the reduced Q to calculate Q_K + Q_reduce = Q[torch.arange(B)[:, None, None], + torch.arange(H)[None, :, None], + M_top, :] + Q_K = torch.matmul(Q_reduce, K.transpose(-2, -1)) + + return Q_K, M_top + + def _get_initial_context(self, V, L_Q): + B, H, L_V, D = V.shape + if not self.mask_flag: + V_sum = V.sum(dim=-2) + contex = V_sum.unsqueeze(-2).expand(B, H, L_Q, V_sum.shape[-1]).clone() + else: # use mask + assert (L_Q == L_V) # requires that L_Q == L_V, i.e. for self-attention only + contex = V.cumsum(dim=-1) + return contex + + def _update_context(self, context_in, V, scores, index, L_Q, attn_mask): + B, H, L_V, D = V.shape + + if self.mask_flag: + attn_mask = ProbMask(B, H, L_Q, index, scores, device=V.device) + scores.masked_fill_(attn_mask.mask, -np.inf) + + attn = torch.softmax(scores, dim=-1) # nn.Softmax(dim=-1)(scores) + + context_in[torch.arange(B)[:, None, None], + torch.arange(H)[None, :, None], + index, :] = torch.matmul(attn, V) + return context_in + + def forward(self, queries, keys, values, attn_mask): + B, L, H, D = queries.shape + _, S, _, _ = keys.shape + + queries = queries.view(B, H, L, -1) + keys = keys.view(B, H, S, -1) + values = values.view(B, H, S, -1) + + U = self.factor * np.ceil(np.log(S)).astype('int').item() + u = self.factor * np.ceil(np.log(L)).astype('int').item() + + scores_top, index = self._prob_QK(queries, keys, u, U) + # add scale factor + scale = self.scale or 1. / sqrt(D) + if scale is not None: + scores_top = scores_top * scale + # get the context + context = self._get_initial_context(values, L) + # update the context with selected top_k queries + context = self._update_context(context, values, scores_top, index, L, attn_mask) + + return context.contiguous() + + +class AttentionLayer(nn.Module): + def __init__(self, attention, d_model, n_heads, d_keys=None, + d_values=None): + super(AttentionLayer, self).__init__() + + d_keys = d_keys or (d_model // n_heads) + d_values = d_values or (d_model // n_heads) + + self.inner_attention = attention + self.query_projection = nn.Linear(d_model, d_keys * n_heads) + self.key_projection = nn.Linear(d_model, d_keys * n_heads) + self.value_projection = nn.Linear(d_model, d_values * n_heads) + self.out_projection = nn.Linear(d_values * n_heads, d_model) + self.n_heads = n_heads + + def forward(self, queries, keys, values, attn_mask): + B, L, _ = queries.shape + _, S, _ = keys.shape + H = self.n_heads + + queries = self.query_projection(queries).view(B, L, H, -1) + keys = self.key_projection(keys).view(B, S, H, -1) + values = self.value_projection(values).view(B, S, H, -1) + + out = self.inner_attention( + queries, + keys, + values, + attn_mask + ).view(B, L, -1) + + return self.out_projection(out) \ No newline at end of file diff --git a/third_party/EgoSpeak/MiniROAD/model/transformer_models/decoder.py b/third_party/EgoSpeak/MiniROAD/model/transformer_models/decoder.py new file mode 100755 index 0000000..001e89e --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/model/transformer_models/decoder.py @@ -0,0 +1,53 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class DecoderLayer(nn.Module): + def __init__(self, self_attention, cross_attention, d_model, d_ff=None, + dropout=0.1, activation="relu"): + super(DecoderLayer, self).__init__() + d_ff = d_ff or 4*d_model + self.self_attention = self_attention + self.cross_attention = cross_attention + self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1) + self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1) + self.norm1 = nn.LayerNorm(d_model) + self.norm2 = nn.LayerNorm(d_model) + self.norm3 = nn.LayerNorm(d_model) + self.dropout = nn.Dropout(dropout) + self.activation = F.relu if activation == "relu" else F.gelu + + def forward(self, x, cross, x_mask=None, cross_mask=None): + x = x + self.dropout(self.self_attention( + x, x, x, + attn_mask=x_mask + )) + x = self.norm1(x) + + x = x + self.dropout(self.cross_attention( + x, cross, cross, + attn_mask=cross_mask + )) + + y = x = self.norm2(x) + y = self.dropout(self.activation(self.conv1(y.transpose(-1,1)))) + y = self.dropout(self.conv2(y).transpose(-1,1)) + + return self.norm3(x+y) + + +class Decoder(nn.Module): + def __init__(self, layers, norm_layer=None): + super(Decoder, self).__init__() + self.layers = nn.ModuleList(layers) + self.norm = norm_layer + + def forward(self, x, cross, x_mask=None, cross_mask=None): + for layer in self.layers: + x = layer(x, cross, x_mask=x_mask, cross_mask=cross_mask) + + if self.norm is not None: + x = self.norm(x) + + return x \ No newline at end of file diff --git a/third_party/EgoSpeak/MiniROAD/model/weights_init.py b/third_party/EgoSpeak/MiniROAD/model/weights_init.py new file mode 100755 index 0000000..016be34 --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/model/weights_init.py @@ -0,0 +1,35 @@ +# FROM LSTR +import math +import torch.nn as nn + +def weights_init(m): + if isinstance(m, nn.Linear): + nn.init.kaiming_uniform_(m.weight.data, a=math.sqrt(5)) + elif isinstance(m, nn.Conv1d): + nn.init.normal_(m.weight.data) + if m.bias is not None: + nn.init.normal_(m.bias.data) + elif isinstance(m, nn.Conv2d): + nn.init.xavier_normal_(m.weight.data) + if m.bias is not None: + nn.init.normal_(m.bias.data) + elif isinstance(m, nn.ConvTranspose1d): + nn.init.normal_(m.weight.data) + if m.bias is not None: + nn.init.normal_(m.bias.data) + elif isinstance(m, nn.ConvTranspose2d): + nn.init.xavier_normal_(m.weight.data) + if m.bias is not None: + nn.init.normal_(m.bias.data) + elif isinstance(m, nn.BatchNorm1d): + nn.init.normal_(m.weight.data, mean=1, std=0.02) + nn.init.constant_(m.bias.data, 0) + elif isinstance(m, nn.BatchNorm2d): + nn.init.normal_(m.weight.data, mean=1, std=0.02) + nn.init.constant_(m.bias.data, 0) + elif isinstance(m, nn.GRUCell): + for param in m.parameters(): + if len(param.shape) >= 2: + nn.init.orthogonal_(param.data) + else: + nn.init.normal_(param.data) \ No newline at end of file diff --git a/third_party/EgoSpeak/MiniROAD/run_mamba_inference.sh b/third_party/EgoSpeak/MiniROAD/run_mamba_inference.sh new file mode 100644 index 0000000..4246de5 --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/run_mamba_inference.sh @@ -0,0 +1,26 @@ +PATH_TO_CONFIG_FILE=configs/EasyCom/mamba_aa_easycom.yaml +PATH_TO_CHECKPOINT_FILE=checkpoints/mamba_easycom_AV.pth +shift # Remove the first argument so $@ contains all other arguments + +echo "Running MiniROAD inference" +echo "Config file: $PATH_TO_CONFIG_FILE" +echo "Additional options: $@" + +# get the basename of the config file, without the extension +CONFIG_FILENAME=$(basename -- "$PATH_TO_CONFIG_FILE") +CONFIG_FILENAME="${CONFIG_FILENAME%.*}" + +# --no_flow (for no audio) --no_rgb --wandb +# Pass all remaining arguments directly to the Python script +python -u main.py --config "$PATH_TO_CONFIG_FILE" \ + --save-as-cfg \ + --eval \ + --ckpt "$PATH_TO_CHECKPOINT_FILE" \ + $@ + +echo "Finished MiniROAD inference" +echo "python -u main.py --config '$PATH_TO_CONFIG_FILE' \ + --save-as-cfg \ + --ckpt '$PATH_TO_CHECKPOINT_FILE' \ + --eval \ + $@" \ No newline at end of file diff --git a/third_party/EgoSpeak/MiniROAD/run_miniroad_inference.sh b/third_party/EgoSpeak/MiniROAD/run_miniroad_inference.sh new file mode 100644 index 0000000..1de91d0 --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/run_miniroad_inference.sh @@ -0,0 +1,26 @@ +PATH_TO_CONFIG_FILE=configs/EasyCom/miniroad_aa_easycom.yaml +PATH_TO_CHECKPOINT_FILE=checkpoints/miniroad_easycom_AV.pth +shift # Remove the first argument so $@ contains all other arguments + +echo "Running MiniROAD inference" +echo "Config file: $PATH_TO_CONFIG_FILE" +echo "Additional options: $@" + +# get the basename of the config file, without the extension +CONFIG_FILENAME=$(basename -- "$PATH_TO_CONFIG_FILE") +CONFIG_FILENAME="${CONFIG_FILENAME%.*}" + +# --no_flow (for no audio) --no_rgb --wandb +# Pass all remaining arguments directly to the Python script +python -u main.py --config "$PATH_TO_CONFIG_FILE" \ + --save-as-cfg \ + --eval \ + --ckpt "$PATH_TO_CHECKPOINT_FILE" \ + $@ + +echo "Finished MiniROAD inference" +echo "python -u main.py --config '$PATH_TO_CONFIG_FILE' \ + --save-as-cfg \ + --ckpt '$PATH_TO_CHECKPOINT_FILE' \ + --eval \ + $@" \ No newline at end of file diff --git a/third_party/EgoSpeak/MiniROAD/run_miniroad_train.sh b/third_party/EgoSpeak/MiniROAD/run_miniroad_train.sh new file mode 100644 index 0000000..5ee9561 --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/run_miniroad_train.sh @@ -0,0 +1,13 @@ +#!/bin/bash +PATH_TO_CONFIG_FILE=configs/EasyCom/miniroad_aa_easycom.yaml + +echo "Running MiniROAD training" +echo "Config file: $PATH_TO_CONFIG_FILE" +echo "Additional options: $@" + +# Additional options: --no_flow (for no audio) --no_rgb --wandb +# Pass all remaining arguments directly to the Python script +python -u main.py --config "$PATH_TO_CONFIG_FILE" --save-as-cfg $@ + +echo "Finished MiniROAD training" +echo "Run python -u main.py --config $PATH_TO_CONFIG_FILE --save-as-cfg $@" \ No newline at end of file diff --git a/third_party/EgoSpeak/MiniROAD/trainer/__init__.py b/third_party/EgoSpeak/MiniROAD/trainer/__init__.py new file mode 100755 index 0000000..d9c43f5 --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/trainer/__init__.py @@ -0,0 +1,4 @@ +from .train_builder import build_trainer +from .train import train_one_epoch +from .eval_builder import build_eval +from .eval import Evaluate \ No newline at end of file diff --git a/third_party/EgoSpeak/MiniROAD/trainer/eval.py b/third_party/EgoSpeak/MiniROAD/trainer/eval.py new file mode 100755 index 0000000..280a292 --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/trainer/eval.py @@ -0,0 +1,288 @@ +import json +import time + +import numpy as np +import torch +import torch.nn as nn +import wandb +from tqdm import tqdm +from trainer.eval_builder import EVAL +from utils import * +from utils import (event_accuracy, perframe_average_precision, + thumos_postprocessing) + + +@EVAL.register("OAD") +class Evaluate(nn.Module): + + def __init__(self, cfg): + super(Evaluate, self).__init__() + self.data_processing = thumos_postprocessing if 'THUMOS' in cfg['data_name'] else None + self.metric = cfg['metric'] + self.eval_method = perframe_average_precision + self.event_eval_method = event_accuracy + self.all_class_names = json.load(open(cfg['video_list_path']))[cfg["data_name"].split('_')[0]]['class_index'] + self.cfg = cfg + self.ignore_index = cfg['ignore_index'] if 'ignore_index' in cfg else -100 + + def eval(self, model, dataloader, logger, criterion, use_wandb, epoch, writer, result_path): + device = "cuda:0" + model.eval() + eval_preds = {} + gt_targets_dict = {} + + with torch.no_grad(): + pred_scores, gt_targets = [], [] + start = time.time() + eval_loss = 0.0 + for it, (vid_name, rgb_input, flow_input, target) in enumerate(tqdm(dataloader, desc='Evaluation:', leave=False)): + rgb_input, flow_input, target = rgb_input.to(device), flow_input.to(device), target.to(device) + out_dict = model(rgb_input, flow_input) + + if criterion != None: + loss = criterion(out_dict, target) + eval_loss += loss.item() + + pred_logit = out_dict['logits'] + prob_val = pred_logit.squeeze().cpu().numpy() + target_batch = target.squeeze().cpu().numpy() + pred_scores += list(prob_val) + gt_targets += list(target_batch) + + eval_preds[vid_name] = prob_val + + end = time.time() + num_frames = len(gt_targets) + + # save prediction scores + # if result_path != None: + time_taken = end - start + logger.info(f'Processed {num_frames} frames in {time_taken:.1f} seconds ({num_frames / time_taken :.1f} FPS)') + + result_AP = self.eval_method(pred_scores, gt_targets, self.all_class_names, self.ignore_index, self.data_processing, 'AP') + result_cAP = self.eval_method(pred_scores, gt_targets, self.all_class_names, self.ignore_index, self.data_processing, 'cAP') + logger.info(f'per_class_ap: {result_AP["per_class_AP"]}') + logger.info(f'all_class_names: {self.all_class_names}') + logger.info(f'len(result_AP["per_class_AP"]: {len(result_AP["per_class_AP"])}') + + per_class_AP = result_AP["per_class_AP"] + if use_wandb == True: + wandb.log({"Validation Loss": eval_loss / len(dataloader)}) + wandb.log({"Validation per_class_AP": per_class_AP}) + + if writer != None: + writer.add_scalar("Validation Loss", eval_loss / len(dataloader), epoch) + + sum_class_AP = 0.0 + + for class_name in self.all_class_names: + # TVSeries + if class_name == 'background': + continue + # THUMOS + if class_name == 'Background' or class_name == 'Ambiguous': + continue + + # logger.info(f'{class_name}') + class_AP = result_AP["per_class_AP"][class_name] + sum_class_AP += class_AP + + class_cAP = result_cAP["per_class_AP"][class_name] + logger.info(f'AP @ {class_name}: {class_AP:.4f}') + logger.info(f'cAP @ {class_name}: {class_cAP:.4f}') + + logger.info(f'sum_AP: {sum_class_AP}, mean_AP: {sum_class_AP / (len(self.all_class_names) - 1):.4f}') + + logger.info(f'pred_scores: {pred_scores[:2]}') + logger.info(f'len(pred_scores): {len(pred_scores)}') + logger.info(f'len(pred_scores[0]): {len(pred_scores[0])}') + + return result_AP['mean_AP'], result_cAP['mean_AP'], eval_preds + + def forward(self, model, dataloader, logger, criterion=None, use_wandb=False, epoch=None, writer=None, result_path = None): + return self.eval(model, dataloader, logger, criterion, use_wandb, epoch, writer, result_path) + +@EVAL.register("OADTR") +class EvaluateOADTR(nn.Module): + + def __init__(self, cfg): + super(EvaluateOADTR, self).__init__() + self.data_processing = thumos_postprocessing if 'THUMOS' in cfg['data_name'] else None + self.metric = cfg['metric'] + self.eval_method = perframe_average_precision + self.all_class_names = json.load(open(cfg['video_list_path']))[cfg["data_name"].split('_')[0]]['class_index'] + self.cfg = cfg + + def eval(self, model, dataloader, logger, criterion, use_wandb, epoch, writer): + device = "cuda:0" + model.eval() + eval_preds = {} + gt_targets_dict = {} + + with torch.no_grad(): + pred_scores, gt_targets = [], [] + start = time.time() + eval_loss = 0.0 + for it, (vid_name, rgb_input, flow_input, enc_target, distance_target, class_h_target, dec_target) in enumerate(tqdm(dataloader, desc='Evaluation:', leave=False)): + rgb_input, flow_input, target = rgb_input.to(device), flow_input.to(device), target.to(device) + out_dict = model(rgb_input, flow_input) + + if criterion != None: + loss = criterion(out_dict, target) + eval_loss += loss.item() + + pred_logit = out_dict['logits'] + prob_val = pred_logit.squeeze().cpu().numpy() + target_batch = target.squeeze().cpu().numpy() + pred_scores += list(prob_val) + gt_targets += list(target_batch) + + # eval_preds[vid_name] = pred_scores + # gt_targets_dict[vid_name] = pred_scores + + if use_wandb == True: + wandb.log({"Validation Loss": eval_loss / len(dataloader)}) + + if writer != None: + writer.add_scalar("Validation Loss", eval_loss / len(dataloader), epoch) + + end = time.time() + num_frames = len(gt_targets) + result_AP = self.eval_method(pred_scores, gt_targets, self.all_class_names, self.data_processing, 'AP') + result_cAP = self.eval_method(pred_scores, gt_targets, self.all_class_names, self.data_processing, 'cAP') + time_taken = end - start + logger.info(f'Processed {num_frames} frames in {time_taken:.1f} seconds ({num_frames / time_taken :.1f} FPS)') + # logger.info(f'per_class_ap: {result["per_class_AP"]}') + logger.info(f'all_class_names: {self.all_class_names}') + logger.info(f'len(result_AP["per_class_AP"]: {len(result_AP["per_class_AP"])}') + + sum_class_AP = 0.0 + + for class_name in self.all_class_names: + # TVSeries + if class_name == 'background': + continue + # THUMOS + if class_name == 'Background' or class_name == 'Ambiguous': + continue + + # logger.info(f'{class_name}') + class_AP = result_AP["per_class_AP"][class_name] + sum_class_AP += class_AP + + class_cAP = result_cAP["per_class_AP"][class_name] + logger.info(f'AP @ {class_name}: {class_AP:.4f}') + logger.info(f'cAP @ {class_name}: {class_cAP:.4f}') + + logger.info(f'sum_AP: {sum_class_AP}, mean_AP: {sum_class_AP / (len(self.all_class_names) - 1):.4f}') + + logger.info(f'pred_scores: {pred_scores[:2]}') + logger.info(f'len(pred_scores): {len(pred_scores)}') + logger.info(f'len(pred_scores[0]): {len(pred_scores[0])}') + + return result_AP['mean_AP'], result_cAP['mean_AP'], eval_preds + + def forward(self, model, dataloader, logger, criterion=None, use_wandb=False, epoch=None, writer=None): + return self.eval(model, dataloader, logger, criterion, use_wandb, epoch, writer) + + +@EVAL.register("ANTICIPATION") +class ANT_Evaluate(nn.Module): + + def __init__(self, cfg): + super(ANT_Evaluate, self).__init__() + data_name = cfg["data_name"].split('_')[0] + self.data_processing = thumos_postprocessing if data_name == 'THUMOS' else None + self.metric = cfg['metric'] + self.eval_method = perframe_average_precision + self.event_eval_method = event_accuracy + self.all_class_names = json.load(open(cfg['video_list_path']))[data_name]['class_index'] + self.cfg = cfg + + if 'ignore_index' in cfg: + self.ignore_index = cfg['ignore_index'] + else: + raise NotImplementedError("Please set ignore_index in config") + + def eval(self, model, dataloader, logger, criterion, use_wandb, epoch, writer, result_path): + device = "cuda:0" + model.eval() + ant_eval_pred_scores_dict = {} + ant_gt_targets_dict = {} + + with torch.no_grad(): + pred_scores, gt_targets, ant_pred_scores, ant_gt_targets = [], [], [], [] + start = time.time() + anticipation_mAPs = [] + for vid_name, rgb_input, flow_input, target, ant_target in tqdm(dataloader, desc='Evaluation:', leave=False): + rgb_input, flow_input, target, ant_target = rgb_input.to(device), flow_input.to(device), target.to(device), ant_target.to(device) + + out_dict = model(rgb_input, flow_input) + pred_logit = out_dict['logits'] + ant_pred_logit = out_dict['anticipation_logits'] + + prob_val = pred_logit.squeeze().cpu().numpy() # [8390, 3] + target_batch = target.squeeze().cpu().numpy() + ant_prob_val = ant_pred_logit.squeeze().cpu().numpy() # [8390, ANTICIPATION_LENGTH, 3] + ant_target_batch = ant_target.squeeze().cpu().numpy() + + pred_scores += list(prob_val) + gt_targets += list(target_batch) + ant_pred_scores += list(ant_prob_val) + ant_gt_targets += list(ant_target_batch) + if isinstance(vid_name, tuple): + vid_name = vid_name[0] + + ant_eval_pred_scores_dict[vid_name] = ant_prob_val + ant_gt_targets_dict[vid_name] = ant_target_batch + + end = time.time() + num_frames = len(gt_targets) + + result = self.eval_method(pred_scores, gt_targets, self.all_class_names, self.ignore_index, self.data_processing, self.metric) + per_class_AP = result['per_class_AP'] + + ant_pred_scores = np.array(ant_pred_scores) + ant_gt_targets = np.array(ant_gt_targets) + logger.info(f'OAD mAP: {result["mean_AP"]*100:.2f}') + # Prepare the formatted string + formatted_output = "per_class_AP:\n" + "\n".join(f" {class_name:<25} : {ap_value:.4f}" for class_name, ap_value in per_class_AP.items()) + + # Log the formatted string + logger.info(formatted_output) + + if use_wandb == True: + wandb.log({"OAD mAP": result["mean_AP"]}) + + anticipation_mAPs_per_class = {key: [] for key in per_class_AP.keys()} + + for step in range(ant_gt_targets.shape[1]): + result[f'anticipation_{step+1}'] = self.eval_method(ant_pred_scores[:,step,:], ant_gt_targets[:,step,:], self.all_class_names, self.ignore_index, self.data_processing, self.metric) + anticipation_mAPs.append(result[f'anticipation_{step+1}']['mean_AP']) + formatted_output = "per_class_AP:\n" + "\n".join(f" {class_name:<25} : {ap_value:.4f}" for class_name, ap_value in result[f'anticipation_{step+1}']['per_class_AP'].items()) + logger.info(f"Anticipation at step {step+1}: {result[f'anticipation_{step+1}']['mean_AP']*100:.2f}") + logger.info(formatted_output) + + for class_name in self.all_class_names: + anticipation_mAPs_per_class[class_name].append(result[f'anticipation_{step+1}']['per_class_AP'][class_name]) + + if use_wandb == True: + wandb.log({f"Anticipation at step {step+1}": result[f'anticipation_{step+1}']['mean_AP']}) + + logger.info(f'Mean Anticipation mAP: {np.mean(anticipation_mAPs)*100:.2f}') + formatted_output = "Mean Anticipation per_class_AP:\n" + "\n".join(f" {class_name:<25} : {np.mean(anticipation_mAPs_per_class[class_name])*100:.4f}" for class_name in self.all_class_names) + # logger.info(f'Mean Anticipation mAP per class: {np.mean(list(anticipation_mAPs_per_class.values()), axis=1)}') + logger.info(formatted_output) + + if use_wandb == True: + wandb.log({"Mean Anticipation mAP": np.mean(anticipation_mAPs)}) + + time_taken = end - start + logger.info(f'Processed {num_frames} frames in {time_taken:.1f} seconds ({num_frames / time_taken :.1f} FPS)') + + return np.mean(anticipation_mAPs), np.mean(anticipation_mAPs), ant_eval_pred_scores_dict + + # def forward(self, model, dataloader, logger): + # return self.eval(model, dataloader, logger) + def forward(self, model, dataloader, logger, criterion=None, use_wandb=False, epoch=None, writer=None, result_path=None): + return self.eval(model, dataloader, logger, criterion, use_wandb, epoch, writer, result_path) diff --git a/third_party/EgoSpeak/MiniROAD/trainer/eval_builder.py b/third_party/EgoSpeak/MiniROAD/trainer/eval_builder.py new file mode 100755 index 0000000..5cd022a --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/trainer/eval_builder.py @@ -0,0 +1,11 @@ +__all__ = [ + 'build_eval' +] + +from utils import Registry + +EVAL = Registry() + +def build_eval(cfg): + eval = EVAL[cfg["task"]](cfg) + return eval diff --git a/third_party/EgoSpeak/MiniROAD/trainer/train.py b/third_party/EgoSpeak/MiniROAD/trainer/train.py new file mode 100755 index 0000000..1ec04ff --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/trainer/train.py @@ -0,0 +1,96 @@ +import torch +import wandb +from tqdm import tqdm +from trainer.train_builder import TRAINER + + +@TRAINER.register("OAD") +def train_one_epoch(trainloader, model, criterion, optimizer, scaler, epoch, use_wandb = False, writer=None, scheduler=None): + epoch_loss = 0 + for it, (vid_name, rgb_input, flow_input, target) in enumerate(tqdm(trainloader, desc=f'Epoch:{epoch} Training', postfix=f'lr: {optimizer.param_groups[0]["lr"]:.7f}')): + rgb_input, flow_input, target = rgb_input.cuda(), flow_input.cuda(), target.cuda() + model.train() + if scaler != None: + with torch.cuda.amp.autocast(): + out_dict = model(rgb_input, flow_input) + loss = criterion(out_dict, target) + optimizer.zero_grad(set_to_none=True) + scaler.scale(loss).backward() + scaler.step(optimizer) + scaler.update() + # scaler is None as argparse by default + else: + out_dict = model(rgb_input, flow_input) + loss = criterion(out_dict, target) + optimizer.zero_grad(set_to_none=True) + loss.backward() + optimizer.step() + + epoch_loss += loss.item() + if use_wandb == True: + wandb.log({"Train Loss": loss.item()}) + + if writer != None: + writer.add_scalar("Train Loss", loss.item(), it+epoch*len(trainloader)) + return epoch_loss + +@TRAINER.register("OADTR") +def train_one_epoch_oadtr(trainloader, model, criterion, optimizer, scaler, epoch, use_wandb = False, writer=None, scheduler=None): + epoch_loss = 0 + for it, (vid_name, rgb_input, flow_input, enc_target, distance_target, class_h_target, dec_target) in enumerate(tqdm(trainloader, desc=f'Epoch:{epoch} Training', postfix=f'lr: {optimizer.param_groups[0]["lr"]:.7f}')): + rgb_input, flow_input, enc_target, distance_target, class_h_target = rgb_input.cuda(), flow_input.cuda(), enc_target.cuda(), distance_target.cuda(), class_h_target.cuda() + model.train() + if scaler != None: + with torch.cuda.amp.autocast(): + out_dict = model(rgb_input, flow_input) + loss = criterion(out_dict, enc_target, dec_target) + optimizer.zero_grad(set_to_none=True) + scaler.scale(loss).backward() + scaler.step(optimizer) + scaler.update() + else: + out_dict = model(rgb_input, flow_input) + loss = criterion(out_dict, enc_target, dec_target) + optimizer.zero_grad(set_to_none=True) + loss.backward() + optimizer.step() + + epoch_loss += loss.item() + if use_wandb == True: + wandb.log({"Train Loss": loss.item()}) + + if writer != None: + writer.add_scalar("Train Loss", loss.item(), it+epoch*len(trainloader)) + + return epoch_loss + + +@TRAINER.register("ANTICIPATION") +def ant_train_one_epoch(trainloader, model, criterion, optimizer, scaler, epoch, use_wandb = False, writer=None, scheduler=None): + epoch_loss = 0 + for it, (vid_name, rgb_input, flow_input, target, ant_target) in enumerate(tqdm(trainloader, desc=f'Epoch:{epoch} Training', postfix=f'lr: {optimizer.param_groups[0]["lr"]:.7f}')): + rgb_input, flow_input, target, ant_target = rgb_input.cuda(), flow_input.cuda(), target.cuda(), ant_target.cuda() + model.train() + if scaler != None: + with torch.cuda.amp.autocast(): + out_dict = model(rgb_input, flow_input) + loss = criterion(out_dict, target, ant_target) + optimizer.zero_grad(set_to_none=True) + scaler.scale(loss).backward() + scaler.step(optimizer) + scaler.update() + else: + out_dict = model(rgb_input, flow_input) + loss = criterion(out_dict, target, ant_target) + optimizer.zero_grad(set_to_none=True) + loss.backward() + optimizer.step() + epoch_loss += loss.item() + + if use_wandb == True: + wandb.log({"Train Loss": loss.item()}) + + if writer != None: + writer.add_scalar("Train Loss", loss.item(), it+epoch*len(trainloader)) + + return epoch_loss \ No newline at end of file diff --git a/third_party/EgoSpeak/MiniROAD/trainer/train_builder.py b/third_party/EgoSpeak/MiniROAD/trainer/train_builder.py new file mode 100755 index 0000000..a7ba9e2 --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/trainer/train_builder.py @@ -0,0 +1,11 @@ +__all__ = [ + 'build_trainer' +] + +from utils import Registry + +TRAINER = Registry() + +def build_trainer(cfg): + trainer = TRAINER[cfg["task"]] + return trainer diff --git a/third_party/EgoSpeak/MiniROAD/utils/__init__.py b/third_party/EgoSpeak/MiniROAD/utils/__init__.py new file mode 100755 index 0000000..7fff467 --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/utils/__init__.py @@ -0,0 +1,7 @@ +from .metrics import perframe_average_precision, perstage_average_precision, event_accuracy +from .postprocessing import thumos_postprocessing +from .util import * +from .group_transforms import * +from .lr_scheduler import build_lr_scheduler +from .logger import get_logger +from .registry import Registry \ No newline at end of file diff --git a/third_party/EgoSpeak/MiniROAD/utils/group_transforms.py b/third_party/EgoSpeak/MiniROAD/utils/group_transforms.py new file mode 100755 index 0000000..aea7917 --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/utils/group_transforms.py @@ -0,0 +1,497 @@ +""" +This file is inherited from tsn-pytorch +""" + +import torchvision +import random +from PIL import Image, ImageOps +import numpy as np +import numbers +import math +import torch +from torchvision.transforms import functional as F + +# lstr +class Compose(object): + + def __init__(self, transforms): + self.transforms = transforms + + def __call__(self, data): + for t in self.transforms: + data = t(data) + return data + +#LSTR +class GroupCrop(object): + """perfrom group crop'""" + + def __init__(self, crop_size, crop_pos='random'): + super(GroupCrop, self).__init__() + self.crop_size = crop_size + self.crop_pos = crop_pos + assert crop_pos in ['random', 'top_left', 'center', 'bottom_right' + ], f'crop_pos: {self.crop_pos} not implemented' + + def __call__(self, img_group): + """ + Args: + img_group (List(PIL.Image) or List(torch.Tensor)): a list of image to crop. + Returns: The same type as input. cropped img_group. + + """ + crop_w, crop_h = self.crop_size + if isinstance(img_group, torch.Tensor): + assert img_group.dim( + ) == 4, f'input img_group should be dim-4(TCHW). Got shape: {img_group.shape}' + img_h, img_w = img_group.shape[-2:] + offset_h, offset_w = self.get_crop_offset(img_w, img_h, crop_w, + crop_h) + return F.crop(img_group, offset_h, offset_w, crop_h, crop_w) + elif isinstance(img_group[0], Image): + img_w, img_h = img_group[0].size + offset_h, offset_w = self.get_crop_offset(img_w, img_h, crop_w, + crop_h) + return [ + F.crop(img, offset_h, offset_w, crop_h, crop_w) + for img in img_group + ] + else: + raise NotImplementedError( + f'img_group of type: {type(img_group)} is not supported') + + def get_crop_offset(self, img_w, img_h, crop_w, crop_h): + """get the offset to crop. + Returns: (offset_h, offset_w). + """ + if self.crop_pos == 'random': + offset_h = int(torch.randint(img_h - crop_h, [])) + offset_w = int(torch.randint(img_w - crop_w, [])) + elif self.crop_pos == 'top_left': + offset_h = offset_w = 0 + elif self.crop_pos == 'center': + offset_h = int((img_h - crop_h) // 2) + offset_w = int((img_w - crop_w) // 2) + elif self.crop_pos == 'bottom_right': + offset_h = img_h - crop_h + offset_w = img_w - crop_w + else: + raise NotImplementedError( + f'crop_pos: {self.crop_pos} not implemented') + return (offset_h, offset_w) + +class GroupRandomCrop(object): + def __init__(self, size): + if isinstance(size, numbers.Number): + self.size = (int(size), int(size)) + else: + self.size = size + + def __call__(self, img_group): + + w, h = img_group[0].size + th, tw = self.size + + out_images = list() + + x1 = random.randint(0, w - tw) + y1 = random.randint(0, h - th) + + for img in img_group: + assert(img.size[0] == w and img.size[1] == h) + if w == tw and h == th: + out_images.append(img) + else: + out_images.append(img.crop((x1, y1, x1 + tw, y1 + th))) + + return out_images + + +class GroupCenterCrop(object): + def __init__(self, size): + self.worker = torchvision.transforms.CenterCrop(size) + + def __call__(self, img_group): + return [self.worker(img) for img in img_group] + + +class GroupRandomHorizontalFlip(object): + """Randomly horizontally flips the given PIL.Image with a probability of 0.5 + """ + def __init__(self, is_flow=False): + self.is_flow = is_flow + + def __call__(self, img_group, is_flow=False): + v = random.random() + if v < 0.5: + ret = [img.transpose(Image.FLIP_LEFT_RIGHT) for img in img_group] + if self.is_flow: + for i in range(0, len(ret), 2): + ret[i] = ImageOps.invert(ret[i]) # invert flow pixel values when flipping + return ret + else: + return img_group + + +# class GroupNormalize(object): +# def __init__(self, mean, std): +# self.mean = mean +# self.std = std + +# def __call__(self, tensor): +# rep_mean = self.mean * (tensor.size()[0]//len(self.mean)) +# rep_std = self.std * (tensor.size()[0]//len(self.std)) + +# # TODO: make efficient +# for t, m, s in zip(tensor, rep_mean, rep_std): +# t.sub_(m).div_(s) + +# return tensor + +# lstr +class GroupNormalize(object): + + def __init__(self, mean, std): + + if len(mean) == 1: # flow + assert len( + std + ) == 1, f'flow input must have std of length 1. got {len(std)}' + self.mean = mean[0] + self.std = std[0] + else: + mean = np.array(mean, dtype=np.float32) + self.mean = mean[np.newaxis, :, np.newaxis, + np.newaxis] # shape: (C,) ->(T,C,H,W) + std = np.array(std, dtype=np.float32) + self.std = std[np.newaxis, :, np.newaxis, + np.newaxis] # shape: (C,) ->(T,C,H,W) + + def __call__(self, tensor): + """ + Args: + tensor (torch.Tensor) : The stacked images. shape: (T, C, H, W). + Returns: (torch.Tensor). sub mean and div std along the channel dim. + + """ + return (tensor - self.mean) / self.std + +class GroupScale(object): + """ Rescales the input PIL.Image to the given 'size'. + 'size' will be the size of the smaller edge. + For example, if height > width, then image will be + rescaled to (size * height / width, size) + size: size of the smaller edge + interpolation: Default: PIL.Image.BILINEAR + """ + + def __init__(self, size, interpolation=Image.BILINEAR): + self.worker = torchvision.transforms.Resize(size, interpolation) + + def __call__(self, img_group): + return [self.worker(img) for img in img_group] + +# lstr +class GroupResize(object): + """perfrom group random resize""" + + def __init__(self, scale_range, keep_aspect_ratio=True): + """ + Args: + scale_range (tuple/list): The shorter side of the image will be resized to + 1) the range if `scale_range` is a tuple + 2) the specific size if `scale_range` contains only 1 element. + """ + super(GroupResize, self).__init__() + self.scale_range = scale_range + self.keep_aspect_ratio = keep_aspect_ratio + + def __call__(self, img_group): + """ + Args: + img_group (List(PIL.Image) or Tensor([T,C,H,W])): a list of image to resize. + Returns: The same type as input. resized img_group. + + """ + # resize + size = self.scale_range[0] if len( + self.scale_range) == 1 else torch.randint(self.scale_range[0], + self.scale_range[1], []) + size = int(size) + size = size if self.keep_aspect_ratio else (size, size) + + if isinstance(img_group, torch.Tensor): + assert img_group.dim( + ) == 4, f'input img_group should be dim-4(TCHW). Got shape: {img_group.shape}' + return F.resize(img_group, size) + elif isinstance(img_group[0], Image): + return [F.resize(img, size) for img in img_group] + else: + raise NotImplementedError( + f'img_group of type: {type(img_group)} is not supported') + +class GroupOverSample(object): + def __init__(self, crop_size, scale_size=None): + self.crop_size = crop_size if not isinstance(crop_size, int) else (crop_size, crop_size) + + if scale_size is not None: + self.scale_worker = GroupScale(scale_size) + else: + self.scale_worker = None + + def __call__(self, img_group): + + if self.scale_worker is not None: + img_group = self.scale_worker(img_group) + image_w, image_h = img_group[0].size + crop_w, crop_h = self.crop_size + + offsets = GroupMultiScaleCrop.fill_fix_offset(False, image_w, image_h, crop_w, crop_h) + oversample_group = list() + for o_w, o_h in offsets: + normal_group = list() + flip_group = list() + for i, img in enumerate(img_group): + crop = img.crop((o_w, o_h, o_w + crop_w, o_h + crop_h)) + normal_group.append(crop) + flip_crop = crop.copy().transpose(Image.FLIP_LEFT_RIGHT) + + if img.mode == 'L' and i % 2 == 0: + flip_group.append(ImageOps.invert(flip_crop)) + else: + flip_group.append(flip_crop) + + oversample_group.extend(normal_group) + oversample_group.extend(flip_group) + return oversample_group + + +class GroupMultiScaleCrop(object): + + def __init__(self, input_size, scales=None, max_distort=1, fix_crop=True, more_fix_crop=True): + self.scales = scales if scales is not None else [1, 875, .75, .66] + self.max_distort = max_distort + self.fix_crop = fix_crop + self.more_fix_crop = more_fix_crop + self.input_size = input_size if not isinstance(input_size, int) else [input_size, input_size] + self.interpolation = Image.BILINEAR + + def __call__(self, img_group): + + im_size = img_group[0].size + + crop_w, crop_h, offset_w, offset_h = self._sample_crop_size(im_size) + crop_img_group = [img.crop((offset_w, offset_h, offset_w + crop_w, offset_h + crop_h)) for img in img_group] + ret_img_group = [img.resize((self.input_size[0], self.input_size[1]), self.interpolation) + for img in crop_img_group] + return ret_img_group + + def _sample_crop_size(self, im_size): + image_w, image_h = im_size[0], im_size[1] + + # find a crop size + base_size = min(image_w, image_h) + crop_sizes = [int(base_size * x) for x in self.scales] + crop_h = [self.input_size[1] if abs(x - self.input_size[1]) < 3 else x for x in crop_sizes] + crop_w = [self.input_size[0] if abs(x - self.input_size[0]) < 3 else x for x in crop_sizes] + + pairs = [] + for i, h in enumerate(crop_h): + for j, w in enumerate(crop_w): + if abs(i - j) <= self.max_distort: + pairs.append((w, h)) + + crop_pair = random.choice(pairs) + if not self.fix_crop: + w_offset = random.randint(0, image_w - crop_pair[0]) + h_offset = random.randint(0, image_h - crop_pair[1]) + else: + w_offset, h_offset = self._sample_fix_offset(image_w, image_h, crop_pair[0], crop_pair[1]) + + return crop_pair[0], crop_pair[1], w_offset, h_offset + + def _sample_fix_offset(self, image_w, image_h, crop_w, crop_h): + offsets = self.fill_fix_offset(self.more_fix_crop, image_w, image_h, crop_w, crop_h) + return random.choice(offsets) + + @staticmethod + def fill_fix_offset(more_fix_crop, image_w, image_h, crop_w, crop_h): + w_step = (image_w - crop_w) // 4 + h_step = (image_h - crop_h) // 4 + + ret = list() + ret.append((0, 0)) # upper left + ret.append((4 * w_step, 0)) # upper right + ret.append((0, 4 * h_step)) # lower left + ret.append((4 * w_step, 4 * h_step)) # lower right + ret.append((2 * w_step, 2 * h_step)) # center + + if more_fix_crop: + ret.append((0, 2 * h_step)) # center left + ret.append((4 * w_step, 2 * h_step)) # center right + ret.append((2 * w_step, 4 * h_step)) # lower center + ret.append((2 * w_step, 0 * h_step)) # upper center + + ret.append((1 * w_step, 1 * h_step)) # upper left quarter + ret.append((3 * w_step, 1 * h_step)) # upper right quarter + ret.append((1 * w_step, 3 * h_step)) # lower left quarter + ret.append((3 * w_step, 3 * h_step)) # lower righ quarter + + return ret + + +class GroupRandomSizedCrop(object): + """Random crop the given PIL.Image to a random size of (0.08 to 1.0) of the original size + and and a random aspect ratio of 3/4 to 4/3 of the original aspect ratio + This is popularly used to train the Inception networks + size: size of the smaller edge + interpolation: Default: PIL.Image.BILINEAR + """ + def __init__(self, size, interpolation=Image.BILINEAR): + self.size = size + self.interpolation = interpolation + + def __call__(self, img_group): + for attempt in range(10): + area = img_group[0].size(0) * img_group[0].size(1) + target_area = random.uniform(0.08, 1.0) * area + aspect_ratio = random.uniform(3. / 4, 4. / 3) + + w = int(round(math.sqrt(target_area * aspect_ratio))) + h = int(round(math.sqrt(target_area / aspect_ratio))) + + if random.random() < 0.5: + w, h = h, w + + if w <= img_group[0].size(0) and h <= img_group[0].size(1): + x1 = random.randint(0, img_group[0].size(0) - w) + y1 = random.randint(0, img_group[0].size(1) - h) + found = True + break + else: + found = False + x1 = 0 + y1 = 0 + + if found: + out_group = list() + for img in img_group: + img = img.crop((x1, y1, x1 + w, y1 + h)) + assert(img.size == (w, h)) + out_group.append(img.resize((self.size, self.size), self.interpolation)) + return out_group + else: + # Fallback + scale = GroupScale(self.size, interpolation=self.interpolation) + crop = GroupRandomCrop(self.size) + return crop(scale(img_group)) + + +class Stack(object): + + def __init__(self, roll=False): + self.roll = roll + + def __call__(self, img_group): + if img_group[0].mode == 'L': + return np.concatenate([np.expand_dims(x, 2) for x in img_group], axis=2) + elif img_group[0].mode == 'RGB': + if self.roll: + return np.concatenate([np.array(x)[:, :, ::-1] for x in img_group], axis=2) + else: + return np.concatenate(img_group, axis=2) + +# LSTR +class GroupStack(object): + """Stack to numpy array and convert THWC to TCHW""" + + def __call__(self, img_group): + """ + Args: + img_group (List(PIL.Image)): a list of image. image shape: (H,W,C) for rgb, (H,W) for grayscale + Returns: np.ndarray. stack the images along the first axis. + shape: (T, C, H, W). + """ + if isinstance(img_group, (list, tuple)): + imgs = np.stack(img_group, axis=0) + else: + imgs = img_group + if imgs.ndim == 3: # THW + imgs = imgs[:, np.newaxis, :, :] # TCHW + else: # THWC + imgs = np.transpose(imgs, [0, 3, 1, 2]) # THWC -> TCHW + return imgs + +class ToTorchFormatTensor(object): + """ Converts a PIL.Image (RGB) or numpy.ndarray (H x W x C) in the range [0, 255] + to a torch.FloatTensor of shape (C x H x W) in the range [0.0, 1.0] """ + def __init__(self, div=True): + self.div = div + + def __call__(self, pic): + if isinstance(pic, np.ndarray): + # handle numpy array + img = torch.from_numpy(pic).permute(2, 0, 1).contiguous() + else: + # handle PIL Image + img = torch.ByteTensor(torch.ByteStorage.from_buffer(pic.tobytes())) + img = img.view(pic.size[1], pic.size[0], len(pic.mode)) + # put it from HWC to CHW format + # yikes, this transpose takes 80% of the loading time/CPU + img = img.transpose(0, 1).transpose(0, 2).contiguous() + return img.float().div(255) if self.div else img.float() + +# LSTR +class GroupToTensor(object): + + def __init__(self, div=True): + self.div = div + + def __call__(self, imgs): + """ + Args: + imgs (np.ndarray) : The stacked images. shape: (T, C, H, W). dtype: np.uint8 + Returns: (torch.Tensor). convert dtype to torch.float32. + If `div` is True, the images are divided by 255 to be normalized to [0,1] + + """ + imgs = torch.from_numpy(imgs).float() + if self.div: + imgs = imgs.div(255.) + return imgs + +class IdentityTransform(object): + + def __call__(self, data): + return data + + +if __name__ == "__main__": + trans = torchvision.transforms.Compose([ + GroupScale(256), + GroupRandomCrop(224), + Stack(), + ToTorchFormatTensor(), + GroupNormalize( + mean=[.485, .456, .406], + std=[.229, .224, .225] + )] + ) + + im = Image.open('../tensorflow-model-zoo.torch/lena_299.png') + + color_group = [im] * 3 + rst = trans(color_group) + + gray_group = [im.convert('L')] * 9 + gray_rst = trans(gray_group) + + trans2 = torchvision.transforms.Compose([ + GroupRandomSizedCrop(256), + Stack(), + ToTorchFormatTensor(), + GroupNormalize( + mean=[.485, .456, .406], + std=[.229, .224, .225]) + ]) + print(trans2(color_group)) \ No newline at end of file diff --git a/third_party/EgoSpeak/MiniROAD/utils/logger.py b/third_party/EgoSpeak/MiniROAD/utils/logger.py new file mode 100755 index 0000000..b50bc43 --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/utils/logger.py @@ -0,0 +1,17 @@ +import logging +import os + +def get_logger(output_path): + logger = logging.getLogger(__name__) + logger.setLevel(logging.DEBUG) + # formatter = logging.Formatter('%(asctime)s:%(message)s ', '%Y-%m-%d %H:%M:%S') + console_handler = logging.StreamHandler() + console_handler.setLevel(logging.INFO) + # console_handler.setFormatter(formatter) + logger.addHandler(console_handler) + log_file = os.path.join(output_path, 'log.txt') + fh = logging.FileHandler(log_file) + fh.setLevel(logging.INFO) + logger.addHandler(fh) + + return logger \ No newline at end of file diff --git a/third_party/EgoSpeak/MiniROAD/utils/lr_scheduler.py b/third_party/EgoSpeak/MiniROAD/utils/lr_scheduler.py new file mode 100755 index 0000000..8831ce1 --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/utils/lr_scheduler.py @@ -0,0 +1,90 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# code from detectron2 + +import math +from typing import List +import torch + +try: + from torch.optim.lr_scheduler import LRScheduler +except ImportError: + from torch.optim.lr_scheduler import _LRScheduler as LRScheduler + +class WarmupCosineLR(LRScheduler): + def __init__( + self, + optimizer: torch.optim.Optimizer, + max_iters: int, + warmup_factor: float = 0.001, + warmup_iters: int = 1000, + warmup_method: str = "linear", + last_epoch: int = -1, + ): + self.max_iters = max_iters + self.warmup_factor = warmup_factor + self.warmup_iters = warmup_iters + self.warmup_method = warmup_method + super().__init__(optimizer, last_epoch) + + def get_lr(self) -> List[float]: + warmup_factor = _get_warmup_factor_at_iter( + self.warmup_method, self.last_epoch, self.warmup_iters, self.warmup_factor + ) + # Different definitions of half-cosine with warmup are possible. For + # simplicity we multiply the standard half-cosine schedule by the warmup + # factor. An alternative is to start the period of the cosine at warmup_iters + # instead of at 0. In the case that warmup_iters << max_iters the two are + # very close to each other. + return [ + base_lr + * warmup_factor + * 0.5 + * (1.0 + math.cos(math.pi * self.last_epoch / self.max_iters)) + for base_lr in self.base_lrs + ] + + def _compute_values(self) -> List[float]: + # The new interface + return self.get_lr() + + +def _get_warmup_factor_at_iter( + method: str, iter: int, warmup_iters: int, warmup_factor: float +) -> float: + """ + Return the learning rate warmup factor at a specific iteration. + See :paper:`ImageNet in 1h` for more details. + + Args: + method (str): warmup method; either "constant" or "linear". + iter (int): iteration at which to calculate the warmup factor. + warmup_iters (int): the number of warmup iterations. + warmup_factor (float): the base warmup factor (the meaning changes according + to the method used). + + Returns: + float: the effective warmup factor at the given iteration. + """ + if iter >= warmup_iters: + return 1.0 + + if method == "constant": + return warmup_factor + elif method == "linear": + alpha = iter / warmup_iters + return warmup_factor * (1 - alpha) + alpha + else: + raise ValueError("Unknown warmup method: {}".format(method)) + +def build_lr_scheduler(cfg, optimizer, iters_per_epoch): + max_iters = cfg['num_epoch'] * iters_per_epoch + warmup_iters = int(cfg['warmup_epoch'] * iters_per_epoch) + # warmup_iters = int(max_iters * 1/6) + scheduler = WarmupCosineLR( + optimizer, + max_iters=max_iters, + warmup_factor=cfg['warmup_factor'], + warmup_iters=warmup_iters, + warmup_method=cfg['warmup_method'], + last_epoch=-1) + return scheduler \ No newline at end of file diff --git a/third_party/EgoSpeak/MiniROAD/utils/metrics.py b/third_party/EgoSpeak/MiniROAD/utils/metrics.py new file mode 100755 index 0000000..8fa5acf --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/utils/metrics.py @@ -0,0 +1,226 @@ +# Evaluation code from LSTR + +from collections import OrderedDict +from multiprocessing import Pool + +import numpy as np +from sklearn.metrics import average_precision_score + + +def calibrated_average_precision_score(y_true, y_score): + """Compute calibrated average precision (cAP), which is particularly + proposed for the TVSeries dataset. + """ + y_true_sorted = y_true[np.argsort(-y_score)] + tp = y_true_sorted.astype(float) + fp = np.abs(y_true_sorted.astype(float) - 1) + tps = np.cumsum(tp) + fps = np.cumsum(fp) + ratio = np.sum(tp == 0) / np.sum(tp) + cprec = tps / (tps + fps / (ratio + np.finfo(float).eps) + np.finfo(float).eps) + cap = np.sum(cprec[tp == 1]) / np.sum(tp) + return cap + +def perframe_average_precision(prediction, ground_truth, class_names, ignore_index, + postprocessing=None, metrics='AP'): + """Compute (frame-level) average precision between ground truth and + predictions data frames. + """ + result = OrderedDict() + ground_truth = np.array(ground_truth) + prediction = np.array(prediction) + + # Postprocessing + if postprocessing is not None: + ground_truth, prediction = postprocessing(ground_truth, prediction) + + # Build metrics + if metrics == 'AP': + compute_score = average_precision_score + elif metrics == 'cAP': + # print('cAP') + compute_score = calibrated_average_precision_score + else: + raise RuntimeError('Unknown metrics: {}'.format(metrics)) + + # Ignore backgroud class + ignore_set = set([ignore_index]) + + # print('ignore_set = set([{ignore_index}])') + + # Enable background according to ignore_index + # ignore_set = set([ignore_index]) + + # Compute average precision + result['per_class_AP'] = OrderedDict() + result['num'] = OrderedDict() + print(f"NUM FRAMES: {np.sum(ground_truth[:, 1:])}") + for idx, class_name in enumerate(class_names): + if idx not in ignore_set: + # ground_truth.shape [2901376, 2] + # prediction.shape [45334 ?] + if np.any(ground_truth[:, idx]): + ap_score = compute_score(ground_truth[:, idx], prediction[:, idx]) + result['per_class_AP'][class_name] = ap_score + result['num'][class_name] = f'[true: {int(np.sum(ground_truth[:, idx]))}, pred:{int(np.sum(prediction[:,idx]))}, AP:{ap_score*100:.1f}]' + result['mean_AP'] = np.mean(list(result['per_class_AP'].values())) + print(f"Per_class_AP: {result['per_class_AP']}") + return result + +def event_accuracy(pred_scores_dict, gt_targets_dict): + th_session_accuracies = {} + + threshold = 0.5 + # ERROR: something is wrong + # PLEASE EDIT: + # pred_scores_dict here has wrong dimension + # [64, 128, 10, 3] <- this is because I am using trainloader (do not evaluate on trainloader) + for threshold in np.arange(0.1, 1.0, 0.2): # 0.1, 0.3, 0.5, 0.7, 0.9 + num_events = 0 + th_total_speaking_start_correct = 0 + th_total_speaking_end_correct = 0 + + for session in pred_scores_dict.keys(): + th_session_speaking_start_correct = 0 + th_session_speaking_start_predictions = 0 + + th_session_speaking_end_correct = 0 + th_session_speaking_end_predictions = 0 + + argmax_session_speaking_start_correct = 0 + argmax_session_speaking_end_correct = 0 + + pred_scores = pred_scores_dict[session] + # pred_scores: [num_feature_frames, anticipation_steps, num_classes] + # classes: background, target speaker speaking, other speaker speaking + + # 12/06: To-do + # I forgot whether miniroad AA step 0 means anticipate 1 step future or current + # I did change something it would likely to be + ant_length = pred_scores.shape[1] + + gt_targets = gt_targets_dict[session] + event_labels = np.load(f'/home/junhyeok/projects/turn-taking/datasets/EasyCom/target_perframe/events/{session}.npy') + + num_speaking_start = 0 + num_speaking_end = 0 + print(f'{session} loaded, {len(event_labels)} events') + # 12/06: + # IndexError: index 8504 is out of bounds for axis 0 with size 8390 + # why the event_idx has much higher value than the length of pred_scores? + # check how the event label is made + + # event_label: [event_idx, event_class] + # ['18', 'target speaker speaking start'] + for event_label in event_labels: + event_idx, event_class = event_label + event_idx = int(event_idx) + + if event_class == 'target speaker speaking start': + num_speaking_start += 1 + num_speaking_end += 1 + # right before the event start, want to see how the model predict? + ant_prob = pred_scores[event_idx-1, 0, :] + if np.argmax(ant_prob) == 1: + argmax_session_speaking_start_correct += 1 + + ant_target_prob = pred_scores[event_idx-1, 0, 1] + if ant_target_prob >= threshold: + th_session_speaking_start_correct += 1 + + elif event_class == 'target speaker speaking end': + # right after the event end, want to see how the model predict? + ant_prob = pred_scores[event_idx, 0, :] + if np.argmax(ant_prob) == 0: + argmax_session_speaking_end_correct += 1 + + # target speaker will not speaker anymore + ant_target_prob = pred_scores[event_idx, 0, 1] + if ant_target_prob < threshold: + th_session_speaking_end_correct += 1 + + print(f'Threshold {threshold:.2f}) {session} speaking start accuracy: {(th_session_speaking_start_correct / (len(event_labels) / 2)):.4f}') + print(f'Threshold {threshold:.2f}) {session} speaking end accuracy: {(th_session_speaking_end_correct / (len(event_labels)/ 2 )):.4f}') + print(f'argmax {session} speaking start accuracy: {argmax_session_speaking_start_correct / (len(event_labels) / 2)}') + print(f'argmax {session} speaking end accuracy: {argmax_session_speaking_end_correct / (len(event_labels)/ 2 )}') + th_session_accuracies[session] = [th_session_speaking_start_correct / len(event_labels), th_session_speaking_end_correct / len(event_labels)] + th_total_speaking_start_correct += th_session_speaking_start_correct + th_total_speaking_end_correct += th_session_speaking_end_correct + num_events += len(event_labels) + + th_total_speaking_start_accuracy = th_total_speaking_start_correct / len(event_labels) + th_total_speaking_end_accuracy = th_total_speaking_end_correct / len(event_labels) + + th_total_accuracy = (th_total_speaking_start_correct + th_total_speaking_end_correct) / (num_events) + print(f'total speaking start accuracy: {th_total_accuracy}') + + return th_total_accuracy + +def get_stage_pred_scores(gt_targets, pred_scores, perc_s, perc_e): + starts = [] + ends = [] + stage_gt_targets = [] + stage_pred_scores = [] + for i in range(len(gt_targets)): + if gt_targets[i] == 0: + stage_gt_targets.append(gt_targets[i]) + stage_pred_scores.append(pred_scores[i]) + else: + if i == 0 or gt_targets[i - 1] == 0: + starts.append(i) + if i == len(gt_targets) - 1 or gt_targets[i + 1] == 0: + ends.append(i) + if len(starts) != len(ends): + raise ValueError('starts and ends cannot pair!') + + action_lens = [ends[i] - starts[i] for i in range(len(starts))] + stage_starts = [starts[i] + int(action_lens[i] * perc_s) for i in range(len(starts))] + stage_ends = [max(stage_starts[i] + 1, starts[i] + int(action_lens[i] * perc_e)) for i in range(len(starts))] + for i in range(len(starts)): + stage_gt_targets.extend(gt_targets[stage_starts[i]: stage_ends[i]]) + stage_pred_scores.extend(pred_scores[stage_starts[i]: stage_ends[i]]) + return np.array(stage_gt_targets), np.array(stage_pred_scores) + + +def perstage_average_precision(prediction, ground_truth, + class_names, postprocessing, + metrics='cAP'): + result = OrderedDict() + ground_truth = np.array(ground_truth) + prediction = np.array(prediction) + + # Postprocessing + if postprocessing is not None: + ground_truth, prediction = postprocessing(ground_truth, prediction) + + # Build metrics + if metrics == 'AP': + compute_score = average_precision_score + elif metrics == 'cAP': + compute_score = calibrated_average_precision_score + else: + raise RuntimeError('Unknown metrics: {}'.format(metrics)) + + raise NotImplementedError('This function is not implemented yet.') + # Ignore backgroud class + # ignore_index = set([0]) + + # Compute average precision + for perc_s in range(10): + perc_e = perc_s + 1 + stage_name = '{:2}%_{:3}%'.format(perc_s * 10, perc_e * 10) + result[stage_name] = OrderedDict({'per_class_AP': OrderedDict()}) + for idx, class_name in enumerate(class_names): + if idx not in ignore_index: + stage_gt_targets, stage_pred_scores = get_stage_pred_scores( + (ground_truth[:, idx] == 1).astype(int), + prediction[:, idx], + perc_s / 10, + perc_e / 10, + ) + result[stage_name]['per_class_AP'][class_name] = \ + compute_score(stage_gt_targets, stage_pred_scores) + result[stage_name]['mean_AP'] = \ + np.mean(list(result[stage_name]['per_class_AP'].values())) + + return result diff --git a/third_party/EgoSpeak/MiniROAD/utils/postprocessing.py b/third_party/EgoSpeak/MiniROAD/utils/postprocessing.py new file mode 100755 index 0000000..bace5ef --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/utils/postprocessing.py @@ -0,0 +1,29 @@ + +import numpy as np + +def thumos_postprocessing(ground_truth, prediction, smooth=False, switch=False): + """ + We follow (Shou et al., 2017) and adopt their perframe postprocessing method on THUMOS'14 datset. + Source: https://bitbucket.org/columbiadvmm/cdc/src/master/THUMOS14/eval/PreFrameLabeling/compute_framelevel_mAP.m + """ + + # Simple temporal smoothing via NMS of 5-frames window + if smooth: + prob = np.copy(prediction) + prob1 = prob.reshape(1, prob.shape[0], prob.shape[1]) + prob2 = np.append(prob[0, :].reshape(1, -1), prob[0: -1, :], axis=0).reshape(1, prob.shape[0], prob.shape[1]) + prob3 = np.append(prob[1:, :], prob[-1, :].reshape(1, -1), axis=0).reshape(1, prob.shape[0], prob.shape[1]) + prob4 = np.append(prob[0: 2, :], prob[0: -2, :], axis=0).reshape(1, prob.shape[0], prob.shape[1]) + prob5 = np.append(prob[2:, :], prob[-2:, :], axis=0).reshape(1, prob.shape[0], prob.shape[1]) + probsmooth = np.squeeze(np.max(np.concatenate((prob1, prob2, prob3, prob4, prob5), axis=0), axis=0)) + prediction = np.copy(probsmooth) + + # Assign cliff diving (5) as diving (8) + if switch: + switch_index = np.where(prediction[:, 5] > prediction[:, 8])[0] + prediction[switch_index, 8] = prediction[switch_index, 5] + + # Remove ambiguous (21) + valid_index = np.where(ground_truth[:, 21] != 1)[0] + + return ground_truth[valid_index], prediction[valid_index] \ No newline at end of file diff --git a/third_party/EgoSpeak/MiniROAD/utils/registry.py b/third_party/EgoSpeak/MiniROAD/utils/registry.py new file mode 100755 index 0000000..31ed83c --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/utils/registry.py @@ -0,0 +1,20 @@ +def _register_generic(module_dict, module_name, module): + assert module_name not in module_dict + module_dict[module_name] = module + + +class Registry(dict): + + def __init__(self, *args, **kwargs): + super(Registry, self).__init__(*args, **kwargs) + + def register(self, module_name, module=None): + if module is not None: + _register_generic(self, module_name, module) + return + + def register_fn(fn): + _register_generic(self, module_name, fn) + return fn + + return register_fn \ No newline at end of file diff --git a/third_party/EgoSpeak/MiniROAD/utils/util.py b/third_party/EgoSpeak/MiniROAD/utils/util.py new file mode 100755 index 0000000..3a6f9ad --- /dev/null +++ b/third_party/EgoSpeak/MiniROAD/utils/util.py @@ -0,0 +1,74 @@ +import pickle +import os +import os.path as osp +import random +import numpy as np +import torch + +def dump_pickle(lst, file_path, file_name): + with open(osp.join(file_path, file_name + '.pkl'), 'wb') as f: + pickle.dump(lst, f) + +def create_dir(dir_path): + if not osp.exists(dir_path): + os.makedirs(dir_path) + +def create_outdir(result_path, overwrite=False): + i = 1 + new_result_path = result_path + + if overwrite == False: + while osp.exists(new_result_path): + new_result_path = f'{result_path}_{i}' + i += 1 + + create_dir(osp.join(new_result_path, 'ckpts')) + create_dir(osp.join(new_result_path, 'runs')) + return new_result_path + +def set_seed(seed): + # os.environ['PYTHONHASHSEED'] = str(seed) + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + torch.backends.cudnn.benchmark = False + torch.backends.cudnn.deterministic = True + +def parse_rgb_type(rgb_type): + if "rgb_kinetics_resnet50" in rgb_type: + rgb_type = "rgb_kinetics_resnet50" + elif "rgb_farl" in rgb_type: + rgb_type = "rgb_farl" + + return rgb_type + +def parse_flow_type(flow_type): + if "audio_wav2vec2" in flow_type: + flow_type = "audio_wav2vec2" + + return flow_type + +def get_identifier(cfg): + mode = "eval" if cfg["eval"] == True else "train" + + modality = '' + if not cfg["no_flow"]: + modality += 'A' + if not cfg["no_rgb"]: + modality += 'V' + + identifier = f'{cfg["model"]}_{cfg["data_name"]}' + if 'A' in modality: + identifier += f"_{cfg['flow_type']}" + if 'V' in modality: + identifier += f"_{cfg['rgb_type']}" + + identifier += f'_{modality}_epoch{cfg["num_epoch"]}_{mode}' + if cfg["ckpt"] != None: + identifier += f'_ckpt_{".".join(osp.basename(cfg["ckpt"]).split(".")[:-1])}' + + return identifier + \ No newline at end of file diff --git a/third_party/EgoSpeak/README.md b/third_party/EgoSpeak/README.md new file mode 100644 index 0000000..36abd7e --- /dev/null +++ b/third_party/EgoSpeak/README.md @@ -0,0 +1,150 @@ +# EgoSpeak: Learning When to Speak for Egocentric Conversational Agents in the Wild + +This is the official codebase for our paper [EgoSpeak: Learning When to Speak for Egocentric Conversational Agents in the Wild](https://arxiv.org/abs/2502.14892), accepted to NAACL 2025 Findings. + +## Overview + +EgoSpeak is a framework for predicting speech initiation in egocentric perspective. +This repository contains two main components: + +1. **LSTR (Long-Short Term Transformer)**: A transformer-based model +2. **MiniROAD (Minimal RNN Framework for Online Action Detection)**: A lightweight RNN-based alternative + +## Environment Setup + +```bash +# Create a conda/micromamba environment +micromamba create -n egospeak python=3.10 -y +pip install -r requirements.txt + +# Install LSTR package +cd LSTR +pip install -e . +``` + +## Datasets + +This work uses multiple datasets for turn-taking prediction in conversations: + +```bash +# Create dataset directory +mkdir dataset +cd dataset + +# Download datasets +wget https://huggingface.co/datasets/kjunh/EgoSpeak/resolve/main/EasyCom.tar.gz +wget https://huggingface.co/datasets/kjunh/EgoSpeak/resolve/main/Ego4D.tar.gz + +# Set up symlinks +cd ../LSTR/data +ln -s ../../dataset/EasyCom EasyCom +ln -s ../../dataset/Ego4D Ego4D +``` + +Alternatively, download datasets via Hugging Face: + +```bash +pip install -U "huggingface_hub[cli]" +huggingface-cli download kjunh/EgoSpeak EasyCom.tar.gz Ego4D.tar.gz --repo-type dataset +``` + +If you want to download the YTConv dataset: + +```bash +huggingface-cli download kjunh/EgoSpeak YTConv.tar.gz --repo-type dataset +``` + +## Dataset Structure + +Each dataset follows a similar structure: +``` +EasyCom/ +├── audio_wav2vec2/ # Audio features extracted with wav2vec2 +├── rgb_kinetics_resnet50/ # Visual features extracted with ResNet50 +└── target_perframe/ # Frame-level annotations +``` + + +## Pre-trained Models + +We provide pre-trained checkpoints for both models: + +### LSTR Checkpoints + +```bash +cd LSTR +mkdir checkpoints +cd checkpoints +wget https://huggingface.co/kjunh/EgoSpeak/resolve/main/lstr_easycom_AV.pth +``` + +### MiniROAD & Mamba Checkpoints + +```bash +cd MiniROAD +mkdir checkpoints +cd checkpoints +wget https://huggingface.co/kjunh/EgoSpeak/resolve/main/miniroad_ego4dshuffle_AV.pth +wget https://huggingface.co/kjunh/EgoSpeak/resolve/main/mamba_easycom_AV.pth +``` + +## Usage + +### 1. LSTR + +Training: +```bash +bash run_LSTR_train.sh +``` + +Inference: +```bash +bash run_LSTR_inference.sh +``` + +To customize modalities: +```bash +bash run_LSTR_train.sh configs/EasyCom/LSTR/lstr_aa_easycom.yaml INPUT.MODALITY "visual+motion+object" +``` + +### 2. MiniROAD + +Training: +```bash +bash run_miniroad_train.sh +``` + +Inference: +```bash +bash run_miniroad_inference_cfg.sh +``` + +To customize modalities or enable wandb logging: +```bash +bash run_miniroad_train.sh configs/miniroad_aa_easycom.yaml --no_flow --no_rgb --wandb +``` + +## Citation + +If you find this code useful for your research, please cite our paper: + +```bibtex +@article{kim2025egospeak, + title={EgoSpeak: Learning When to Speak for Egocentric Conversational Agents in the Wild}, + author={Kim, Junhyeok and Kim, Min Soo and Chung, Jiwan and Cho, Jungbin and Kim, Jisoo and Kim, Sungwoong and Sim, Gyeongbo and Yu, Youngjae}, + journal={arXiv preprint arXiv:2502.14892}, + year={2025} +} +``` + +## Acknowledgements + +This codebase builds upon: +- [LSTR](https://github.com/amazon-research/long-short-term-transformer) +- [MiniROAD](https://github.com/joungbin/MiniROAD) + +We thank the authors of these repositories for making their code available. + +## License + +This project is licensed under the Apache-2.0 License. \ No newline at end of file diff --git a/third_party/EgoSpeak/requirements.txt b/third_party/EgoSpeak/requirements.txt new file mode 100644 index 0000000..bcfb222 --- /dev/null +++ b/third_party/EgoSpeak/requirements.txt @@ -0,0 +1,81 @@ +absl-py==2.1.0 +annotated-types==0.7.0 +appdirs==1.4.4 +argcomplete==3.0.8 +certifi==2025.1.31 +charset-normalizer==3.4.1 +click==8.1.8 +docker-pycreds==0.4.0 +einops==0.8.1 +filelock==3.17.0 +fsspec==2025.2.0 +gitdb==4.0.12 +GitPython==3.1.44 +grpcio==1.70.0 +huggingface-hub==0.28.1 +idna==3.10 +Jinja2==3.1.5 +joblib==1.4.2 +mamba-ssm==2.2.4 +Markdown==3.7 +MarkupSafe==3.0.2 +-e git+ssh://git@github.com/jun297/EgoSpeak-CAM.git@dcbfd115abd8ce6ecfe37cf4a40e71d0089c98a1#egg=MiniROAD&subdirectory=MiniROAD +mpmath==1.3.0 +networkx==3.4.2 +ninja==1.11.1.3 +numpy==2.2.2 +nvidia-cublas-cu12==12.4.2.65 +nvidia-cuda-cupti-cu12==12.4.99 +nvidia-cuda-nvrtc-cu12==12.4.99 +nvidia-cuda-runtime-cu12==12.4.99 +nvidia-cudnn-cu12==9.1.0.70 +nvidia-cufft-cu12==11.2.0.44 +nvidia-curand-cu12==10.3.5.119 +nvidia-cusolver-cu12==11.6.0.99 +nvidia-cusparse-cu12==12.3.0.142 +nvidia-cusparselt-cu12==0.6.2 +nvidia-nccl-cu12==2.20.5 +nvidia-nvjitlink-cu12==12.4.99 +nvidia-nvtx-cu12==12.4.99 +opencv-python==4.11.0.86 +packaging==24.2 +pandas==2.2.3 +pathtools==0.1.2 +pillow==11.0.0 +pipx==1.5.0 +platformdirs==4.3.6 +protobuf==5.29.3 +psutil==6.1.1 +pydantic==2.10.6 +pydantic_core==2.27.2 +python-dateutil==2.9.0.post0 +pytz==2025.1 +PyYAML==6.0.2 +regex==2024.11.6 +requests==2.32.3 +safetensors==0.5.2 +scikit-learn==1.6.1 +scipy==1.15.1 +sentry-sdk==2.21.0 +setproctitle==1.3.4 +six==1.17.0 +smmap==5.0.2 +sympy==1.13.1 +tensorboard==2.19.0 +tensorboard-data-server==0.7.2 +-e git+ssh://git@github.com/jun297/EgoSpeak-CAM.git@dcbfd115abd8ce6ecfe37cf4a40e71d0089c98a1#egg=TeSTra&subdirectory=LSTR +threadpoolctl==3.5.0 +tokenizers==0.21.0 +torch==2.4.1+cu124 +torchaudio==2.4.1+cu124 +torchvision==0.19.1+cu124 +tqdm==4.67.1 +transformers==4.48.3 +triton==3.0.0 +typing_extensions==4.12.2 +tzdata==2025.1 +urllib3==2.3.0 +userpath==1.8.0 +wandb==0.19.6 +Werkzeug==3.1.3 +yacs==0.1.8 diff --git a/third_party/MuVAP/.gitignore b/third_party/MuVAP/.gitignore new file mode 100644 index 0000000..cd8a12a --- /dev/null +++ b/third_party/MuVAP/.gitignore @@ -0,0 +1,13 @@ +__pycache__/ +*.py[cod] +*.egg-info/ +.pytest_cache/ +.ruff_cache/ +.venv/ +.vscode/ +artifacts/ +lightning_logs/ +wandb/ +build/ +dist/ +*.ckpt diff --git a/third_party/MuVAP/.upstream-revision b/third_party/MuVAP/.upstream-revision new file mode 100644 index 0000000..944aa35 --- /dev/null +++ b/third_party/MuVAP/.upstream-revision @@ -0,0 +1 @@ +a1ef9c8bde873196da2c37be5e74933c497b6a14 diff --git a/third_party/MuVAP/LICENSE b/third_party/MuVAP/LICENSE new file mode 100644 index 0000000..7d6b13b --- /dev/null +++ b/third_party/MuVAP/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 Haotian-Qi + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/third_party/MuVAP/README.md b/third_party/MuVAP/README.md new file mode 100644 index 0000000..947784f --- /dev/null +++ b/third_party/MuVAP/README.md @@ -0,0 +1,233 @@ +# MuVAP + +**MuVAP: Multimodal Multiparty Voice Activity Projection for Turn-taking +Prediction in the Wild** — Interspeech 2026. + +Demo: https://interspeech2026anonymous-eng.github.io/demo/ + +MuVAP is built from modules that share one codebase and are released one at a +time. Each is self-contained: it trains, evaluates, and publishes weights on its +own, without the modules that have not landed yet. + +| Module | Status | Entry point | +| --- | --- | --- | +| VAP — voice activity projection from audio | released | `train_vap.py` | +| ASD — audio-visual active speaker detection | released | `train_asd.py` | +| MuVAP — multiparty fusion of the two | to follow | — | + +## Install + +```bash +python -m pip install -e . # add ".[mimi]" for the Mimi frontend +``` + +Install the PyTorch build matching your CUDA version first. + +The frozen audio frontends download their pretrained weights the first time a +model is built. Fetch them ahead of time — worth doing before a batch job, on a +machine whose compute nodes have no outbound network: + +```bash +python tools/fetch_frontends.py # or: fetch_frontends.py cpc +``` + + +## VAP + +Voice Activity Projection: a causal model that predicts who will hold the floor +next, from audio alone. Every configuration shares one training program and +differs only in the config file, so a result difference between them comes from +the setup rather than from a different pipeline. + +| Configuration | Config | Input | Codebook | +| --- | --- | --- | --- | +| Role-based VAP | `config/yaml/vap_role.yaml` | 1 mixed channel | 136 role-relative classes | +| Speaker-based VAP | `config/yaml/vap_speaker.yaml` | 2 channels, one per speaker | 256 classes, one channel per row | +| Speaker-based, mono | `config/yaml/vap_speaker_mono_mimi.yaml` | 1 mixed channel + VAD | 256 classes, one channel per row | + +The mono setup is the original VAP arrangement: one channel carries the audio +and the causal voice activity behind each frame supplies the speaker identity a +second channel would have carried. `config/yaml/` also holds the frontend and +frame-rate variants of each - `*_mimi_12hz`, `*_cpc_50hz`. + +```bash +python train_vap.py --config config/yaml/vap_role.yaml --name role +``` + +### Positional encoding + +`vap.temporal.pos_encoding` selects `alibi`, causal `rope`, `sinusoidal`, or +`none`. The stack stays causal regardless — only how position reaches attention +changes. + +See [`docs/vap.md`](docs/vap.md) for what each setup means, what it logs, how to +prepare Fisher, and the exact causality and alignment contracts. + + +## ASD + +Audio-visual active speaker detection: given a face track and the audio around +it, decide frame by frame whether that face is speaking. The reported number is +the official AVA-ActiveSpeaker mAP. + +| Configuration | Config | Frontend | +| --- | --- | --- | +| ASD | `config/yaml/asd.yaml` | CPC | +| ASD + Mimi | `config/yaml/asd_mimi.yaml` | Mimi | + +The frontends are the same two the VAP module uses, and both emit features at +the 25 Hz rate the visual frontend runs at, so they substitute for each other +here too. + +```bash +python train_asd.py --config config/yaml/asd.yaml --name asd +``` + +Faces reach the model as 112x112 grayscale, the geometry the packs are written +in - see [preprocessing](docs/preprocessing.md) for where that is fixed and why. +On top of it, `asd.augment` applies one geometric transform per chunk, drawn +uniformly from none, horizontal flip, random crop rescaled back to 112, and +rotation; `crop_scale` is the side of that crop as a fraction of the frame. +Audio is augmented separately: with probability `audio_mix_prob` a donor clip +is mixed in at an SNR drawn from `audio_snr_db` while the label stays unchanged, +so the model has to reject speech it cannot see on the face. + +```yaml +augment: + audio_mix_prob: 0.5 # chance of mixing in off-screen speech + audio_snr_db: [-5.0, 5.0] # how loud that donor is against the real audio + visual: true # geometric transforms on the face track + crop_scale: [0.7, 1.0] # random-crop side, as a fraction of 112 + rotate_degrees: 15.0 # maximum rotation + substitute_audio_prob: 0.0 # ablation: replace the audio, relabel silent +``` + +Every source corpus is packed into one memory-mapped format first. A pack is a +decoded mirror of the corpus, not a model-ready tensor set: grayscale face +crops and one label per frame at the *native* video rate, plus 16 kHz audio at +its original gain. Sampling onto the 25 Hz model timeline and loudness +normalization happen in the dataloader, so one pack can serve models with +different conventions. + +```bash +python -m preprocess.asd.prepare ava --split train \ + --loader /data/AVA/csv/train_loader.csv \ + --visual-root /data/AVA/clips_videos --audio-root /data/AVA/clips_audios +python -m preprocess.asd.validate /data/AVA/packed/train +``` + +`python -m preprocess.asd.prepare --help` lists the accepted source schemas. +See [`docs/preprocessing.md`](docs/preprocessing.md) and +[`docs/data_alignment.md`](docs/data_alignment.md) for the storage and +alignment contracts. + +**Audio convention.** The model expects audio normalized to -25 dBFS RMS, +capped at 20 dB of gain and a 0.99 peak, with the level measured over the whole +recording rather than the evaluated slice. This is not the raw-waveform +convention other ASD codebases use, and a checkpoint evaluated on unnormalized +audio will not reproduce the mAP it was trained to. Packs store audio +unnormalized and the loader applies this; the constants live in +`data/media.py`. + + +## Running either module + +`--set KEY=VALUE` overrides any config value (dotted path, YAML-parsed) and can +be repeated; it fails on a key the config does not already define, so a typo +cannot silently do nothing. `--wandb` enables tracking, `--name` chooses the +run directory. + +`checkpoint_monitor` picks what `best.ckpt` selects on — `val/loss` for VAP, +`val/mAP_official` for ASD. Clear it where there is no held-out split to select +on, and the final epoch is the only artifact the run keeps: + +```bash +python train_vap.py --config config/yaml/vap_role.yaml --set vap.checkpoint_monitor= +``` + +`--test` then falls back to `last` instead of `best`. See +[Weight releases](#weight-releases) for what each checkpoint is. + +## Audio frontends + +`vap.audio_encoder` selects the frozen frontend. Both emit features at the rate +the projection window labels, so they are drop-in alternatives at a shared rate. + +| | `cpc` (default) | `mimi` | +| --- | --- | --- | +| Weights | CPC, LibriLight 60k | `kyutai/mimi`, from Moshi | +| Frozen params | 1.8M | 37.8M | +| Feature dim | 256 | 512 | +| Frame rates | 100 / 50 / 25 / 12.5 | 25 / 12.5 | + +CPC halves a 100 Hz convolutional stack down to the requested rate; Mimi taps +either its 25 Hz encoder output or its 12.5 Hz latent, so 50 Hz is a CPC-only +option. + +Both are causal, which is what makes the model usable as a streaming predictor. + + +## Weight releases + +Trained weights live on the Hub at +[Haotian-Qi/MuVAP](https://huggingface.co/Haotian-Qi/MuVAP). + +Two ASD releases, `asd-cpc` and `asd-mimi`, scoring 90.50 and 92.04 +`mAP_official` on AVA-ActiveSpeaker. The VAP releases span three codebooks, two +frontends and three frame rates - [docs/vap.md](docs/vap.md#released-models) +lists them with their scores. + +Fetch one, then evaluate it: + +```bash +huggingface-cli download Haotian-Qi/MuVAP --include 'vap-role-mimi/*' --local-dir weights + +python train_vap.py --config config/yaml/vap_role.yaml \ + --test --weights weights/vap-role-mimi \ + --set fisher_path=/path/to/fisher +``` + +That reproduces the score in the table exactly. All of them at once: + +```bash +huggingface-cli download Haotian-Qi/MuVAP --local-dir weights +``` + +A run produces two checkpoints. `last.ckpt` is the final epoch — the artifact +the run committed to. `best.ckpt` is the epoch a validation metric picked, and +exists only when `checkpoint_monitor` is set; where there is no held-out split +to select on, clear it and the final epoch is the only artifact. + +Publish either one with: + +```bash +python tools/export_checkpoint.py /last.ckpt --module vap +``` + +A release carries the trained weights alone: no optimizer state, and no copy of +the frozen audio frontend, which is fetched from its own pretrained source when +the model is built. It ships `config.yaml` (the architecture) and +`provenance.json`, which records the checkpoint it came from and what that +checkpoint scored. Point `--weights` at a release directory to evaluate it: + +```bash +python train_vap.py --config config/yaml/vap_role.yaml --test --weights +``` + +The architecture comes from the release's own `config.yaml`; dataset roots and +output directories come from the `--config` you pass, so a release runs anywhere +without reproducing the paths it was trained with. A release that does not fit +the config it is given fails immediately, naming the offending weights. + +Or load one directly: + +```python +import yaml +from models.release import load_weights, resolve +from models.vap import build_vap + +weights, config = resolve("path/to/release") +cfg = yaml.safe_load(open(config))["vap"] +model = load_weights(build_vap(cfg), weights).eval() +logits = model(waveform) # [batch, frames, classes] at the release's rate +``` diff --git a/third_party/MuVAP/config/__init__.py b/third_party/MuVAP/config/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/third_party/MuVAP/config/setup.py b/third_party/MuVAP/config/setup.py new file mode 100644 index 0000000..eee6d0e --- /dev/null +++ b/third_party/MuVAP/config/setup.py @@ -0,0 +1,17 @@ +import yaml + + +def load_config(path, module): + with open(path) as file: + config = yaml.safe_load(file) + + required = {"checkpoint_dir", module} + missing = required.difference(config) + if missing: + raise ValueError(f"Missing config keys: {sorted(missing)}") + return config + + +def init_yaml_config(path): + with open(path) as file: + return yaml.safe_load(file) diff --git a/third_party/MuVAP/config/yaml/asd.yaml b/third_party/MuVAP/config/yaml/asd.yaml new file mode 100644 index 0000000..22c0a66 --- /dev/null +++ b/third_party/MuVAP/config/yaml/asd.yaml @@ -0,0 +1,109 @@ +checkpoint_dir: runs/asd + +evaluation: + # The official val_orig.csv, shipped inside the pack it scores. + ava_groundtruth: data/AVA/packed/val/groundtruth.csv + ava_predictions: runs/asd/ava_predictions.csv + +# Each corpus keeps its pack beside its own media, so a dataset can be copied +# to fast local storage on its own. AVA val is disjoint from everything and is +# the only evaluation set. Measured overlap at the source-video level: +# +# AVA train / val 0 of 33 shared - a real held-out split +# WASD train / val 161 of 161 shared - never held out, free to train on +# MSDWild few.train / few.val 0 of 489 shared - a real held-out split +# MSDWild few.train / many.val 0 of 172 shared - a real held-out split +# +# WASD val and both MSDWild val sets are trained on deliberately. WASD val is +# not a held-out set to begin with. The MSDWild pair are: training on them +# spends two genuine evaluation sets (23.4 h, and 16.9 h of the harder +# many-speaker condition) to buy 40 h of training data. AVA val is untouched +# and is what the reported mAP is computed on. Move any of them back under +# `val` to recover them as evaluation sets. +packed_asd: + train: + - data/AVA/packed/train + - data/WASD/packed/train + - data/WASD/packed/val + - data/MSDWild/packed/few.train + - data/MSDWild/packed/few.val + - data/MSDWild/packed/many.val + val: + - data/AVA/packed/val + +asd: + num_workers: 16 + max_epochs: 5 + seed: 0 + + # Frozen audio frontend: cpc (256-dim) or mimi (512-dim, Kyutai). Both emit + # 25 Hz features, which is the rate the visual frontend runs at. + audio_encoder: cpc + + precision: bf16-mixed + gradient_clip_val: 1.0 + accumulate_grad_batches: 1 + matmul_precision: high + # Metric that selects best.ckpt. Leave empty when there is no held-out + # split to select on: last.ckpt is written either way. + checkpoint_monitor: val/mAP_official + checkpoint_mode: max + + # Batches are built to a frame budget rather than a fixed batch size, and + # tracks longer than train_window are split into near-equal chunks so one + # epoch still covers every frame. Each chunk after the first also reads the + # projection window's history as unscored context, so a causal model never + # has to predict from history it was not shown. See docs/preprocessing.md. + train_window: 250 + frame_budget: 2048 + min_batch: 4 + max_batch: 16 + + # Validation batches whole tracks, grouped by length and right-padded. It is + # forward-only, so it affords a much larger budget than training. + val_frame_budget: 8192 + val_max_batch: 64 + + # Train-time augmentation. Audio mixing is the one that attacks false + # positives: foreign speech is added while the label stays unchanged, so the + # model has to reject speech it cannot see on the face. Donors are drawn from + # a different clip, never from the batch. + augment: + audio_mix_prob: 0.5 + audio_snr_db: [-5.0, 5.0] + visual: true + crop_scale: [0.7, 1.0] + rotate_degrees: 15.0 + # Ablation: replace the audio and relabel the sample silent. Off by default. + substitute_audio_prob: 0.0 + + lr: 0.001 + weight_decay: 0.01 + warmup_ratio: 0.1 + min_lr: 0.00001 + + projection_window: + bin_sec: [0.8, 0.6, 0.4, 0.2, 0.2, 0.4] + num_hist_bins: 4 + frame_hz: 25 + threshold: 0.5 + mode: independent + cut_center_frame: false + + cross: + d_model: 256 + m_ff: 4 + n_heads: 4 + n_layers: 1 + dropout: 0.1 + activation: GELU + pos_encoding: alibi # alibi | rope | sinusoidal | none + + temporal: + d_model: 512 + m_ff: 4 + n_heads: 8 + n_layers: 1 + dropout: 0.1 + activation: GELU + pos_encoding: alibi # alibi | rope | sinusoidal | none diff --git a/third_party/MuVAP/config/yaml/asd_mimi.yaml b/third_party/MuVAP/config/yaml/asd_mimi.yaml new file mode 100644 index 0000000..f02ce9b --- /dev/null +++ b/third_party/MuVAP/config/yaml/asd_mimi.yaml @@ -0,0 +1,113 @@ +# ASD on a frozen Mimi frontend. Everything but the audio encoder matches +# asd.yaml, so a difference against that run is attributable to the encoder. +checkpoint_dir: runs/asd_mimi + +evaluation: + # The official val_orig.csv, shipped inside the pack it scores. + ava_groundtruth: data/AVA/packed/val/groundtruth.csv + ava_predictions: runs/asd_mimi/ava_predictions.csv + +# Each corpus keeps its pack beside its own media, so a dataset can be copied +# to fast local storage on its own. AVA val is disjoint from everything and is +# the only evaluation set. Measured overlap at the source-video level: +# +# AVA train / val 0 of 33 shared - a real held-out split +# WASD train / val 161 of 161 shared - never held out, free to train on +# MSDWild few.train / few.val 0 of 489 shared - a real held-out split +# MSDWild few.train / many.val 0 of 172 shared - a real held-out split +# +# WASD val and both MSDWild val sets are trained on deliberately. WASD val is +# not a held-out set to begin with. The MSDWild pair are: training on them +# spends two genuine evaluation sets (23.4 h, and 16.9 h of the harder +# many-speaker condition) to buy 40 h of training data. AVA val is untouched +# and is what the reported mAP is computed on. Move any of them back under +# `val` to recover them as evaluation sets. +packed_asd: + train: + - data/AVA/packed/train + - data/WASD/packed/train + - data/WASD/packed/val + - data/MSDWild/packed/few.train + - data/MSDWild/packed/few.val + - data/MSDWild/packed/many.val + val: + - data/AVA/packed/val + +asd: + num_workers: 16 + max_epochs: 5 + seed: 0 + + # Frozen audio frontend: cpc (256-dim) or mimi (512-dim, Kyutai). Both emit + # 25 Hz features, which is the rate the visual frontend runs at. + audio_encoder: + name: mimi + tap: encoder # 25 Hz, matching the visual timeline + + precision: bf16-mixed + gradient_clip_val: 1.0 + accumulate_grad_batches: 1 + matmul_precision: high + # Metric that selects best.ckpt. Leave empty when there is no held-out + # split to select on: last.ckpt is written either way. + checkpoint_monitor: val/mAP_official + checkpoint_mode: max + + # Batches are built to a frame budget rather than a fixed batch size, and + # tracks longer than train_window are split into near-equal chunks so one + # epoch still covers every frame. Each chunk after the first also reads the + # projection window's history as unscored context, so a causal model never + # has to predict from history it was not shown. See docs/preprocessing.md. + train_window: 250 + frame_budget: 2048 + min_batch: 4 + max_batch: 16 + + # Validation batches whole tracks, grouped by length and right-padded. It is + # forward-only, so it affords a much larger budget than training. + val_frame_budget: 8192 + val_max_batch: 64 + + # Train-time augmentation. Audio mixing is the one that attacks false + # positives: foreign speech is added while the label stays unchanged, so the + # model has to reject speech it cannot see on the face. Donors are drawn from + # a different clip, never from the batch. + augment: + audio_mix_prob: 0.5 + audio_snr_db: [-5.0, 5.0] + visual: true + crop_scale: [0.7, 1.0] + rotate_degrees: 15.0 + # Ablation: replace the audio and relabel the sample silent. Off by default. + substitute_audio_prob: 0.0 + + lr: 0.001 + weight_decay: 0.01 + warmup_ratio: 0.1 + min_lr: 0.00001 + + projection_window: + bin_sec: [0.8, 0.6, 0.4, 0.2, 0.2, 0.4] + num_hist_bins: 4 + frame_hz: 25 + threshold: 0.5 + mode: independent + cut_center_frame: false + + cross: + d_model: 256 + m_ff: 4 + n_heads: 4 + n_layers: 1 + dropout: 0.1 + activation: GELU + pos_encoding: alibi # alibi | rope | sinusoidal | none + + temporal: + d_model: 512 + m_ff: 4 + n_heads: 8 + n_layers: 1 + dropout: 0.1 + activation: GELU + pos_encoding: alibi # alibi | rope | sinusoidal | none diff --git a/third_party/MuVAP/config/yaml/vap_role.yaml b/third_party/MuVAP/config/yaml/vap_role.yaml new file mode 100644 index 0000000..843ac33 --- /dev/null +++ b/third_party/MuVAP/config/yaml/vap_role.yaml @@ -0,0 +1,49 @@ +# Role-based VAP: one mixed channel, 136 role-relative classes. +fisher_path: data/fisher +checkpoint_dir: runs/vap_role + +vap: + arch: mono # mono | stereo + batch_size: 64 + num_workers: 16 + max_epochs: 5 + lr: 0.001 + weight_decay: 0.01 + warmup_ratio: 0.1 + min_lr: 0.0001 + + precision: bf16-mixed + gradient_clip_val: 1.0 + accumulate_grad_batches: 1 + matmul_precision: high + # Metric that selects best.ckpt. Leave empty when there is no held-out + # split to select on: last.ckpt is written either way. + checkpoint_monitor: val/loss + checkpoint_mode: min + # npy reads the segments written by preprocessing; raw seeks the same window + # in the source recording and resamples it. + source: npy + swap_channels: false # stereo-only augmentation + + projection_window: + bin_sec: [1.4, 0.6, 0.6, 1.4] + num_hist_bins: 2 + frame_hz: 25 + threshold: 0.5 + mode: role_relative + cut_center_frame: false + + temporal: + d_model: 256 + m_ff: 3 + n_heads: 4 + n_layers: 4 + dropout: 0.1 + activation: GELU + causal: true + pos_encoding: alibi # alibi | rope | sinusoidal | none + ffn: mlp # mlp | swiglu + norm: layer # layer | rms + # Size of the sinusoidal position table. Rotary and ALiBi extend to any + # length, so this bounds the sinusoidal encoding only. + max_frames: 4096 diff --git a/third_party/MuVAP/config/yaml/vap_role_cpc_50hz.yaml b/third_party/MuVAP/config/yaml/vap_role_cpc_50hz.yaml new file mode 100644 index 0000000..1677809 --- /dev/null +++ b/third_party/MuVAP/config/yaml/vap_role_cpc_50hz.yaml @@ -0,0 +1,58 @@ +# Role-based VAP at 50 Hz on CPC - the original VAP's own configuration. +# +# CPC's conv stack runs at 100 Hz and `downsample` halves it once per block, so +# 50 Hz is one block instead of two. This is the rate and frontend of Ekstedt & +# Skantze (2022); Mimi cannot reach it, its encoder tops out at 25 Hz. +# Labels come from seg/vad50, written by 03_prep_vad_rate.py. +fisher_path: data/fisher +checkpoint_dir: runs/role_cpc_50hz + +vap: + audio_encoder: + name: cpc + frame_hz: 50 # one downsample block instead of two + + arch: mono # mono | stereo + batch_size: 64 + num_workers: 16 + max_epochs: 5 + lr: 0.001 + weight_decay: 0.01 + warmup_ratio: 0.1 + min_lr: 0.0001 + + precision: bf16-mixed + gradient_clip_val: 1.0 + accumulate_grad_batches: 1 + matmul_precision: high + # Metric that selects best.ckpt. Leave empty when there is no held-out + # split to select on: last.ckpt is written either way. + checkpoint_monitor: val/loss + checkpoint_mode: min + # npy reads the segments written by preprocessing; raw seeks the same window + # in the source recording and resamples it. + source: npy + swap_channels: false # stereo-only augmentation + + projection_window: + bin_sec: [1.4, 0.6, 0.6, 1.4] + num_hist_bins: 2 + frame_hz: 50 + threshold: 0.5 + mode: role_relative + cut_center_frame: false + + temporal: + d_model: 256 + m_ff: 3 + n_heads: 4 + n_layers: 4 + dropout: 0.1 + activation: GELU + causal: true + pos_encoding: alibi # alibi | rope | sinusoidal | none + ffn: mlp # mlp | swiglu + norm: layer # layer | rms + # Size of the sinusoidal position table. Rotary and ALiBi extend to any + # length, so this bounds the sinusoidal encoding only. + max_frames: 4096 diff --git a/third_party/MuVAP/config/yaml/vap_role_mimi_12hz.yaml b/third_party/MuVAP/config/yaml/vap_role_mimi_12hz.yaml new file mode 100644 index 0000000..807e826 --- /dev/null +++ b/third_party/MuVAP/config/yaml/vap_role_mimi_12hz.yaml @@ -0,0 +1,63 @@ +# Role-based VAP at 12.5 Hz, the rate Mimi natively codes at. +# +# `tap: latent` is Mimi's own frame rate, so this matches a full-duplex system +# built on Mimi codes frame for frame: same frontend, same 80 ms decision grid. +# Labels come from seg/vad12_5, written by 03_prep_vad_rate.py; the 25 Hz tree +# is untouched. bin_sec stays at the paper's times, so the rate is the only +# thing that changes against the 25 Hz runs - at 12.5 Hz they round to +# 0.16 / 0.4 / 0.64 / 0.8, still a 2.0 s horizon. +fisher_path: data/fisher +checkpoint_dir: runs/role_mimi_12hz + +vap: + audio_encoder: + name: mimi + tap: latent # 12.5 Hz, Mimi's own frame rate + + arch: mono # mono | stereo + batch_size: 64 + num_workers: 16 + max_epochs: 5 + lr: 0.001 + weight_decay: 0.01 + warmup_ratio: 0.1 + min_lr: 0.0001 + + precision: bf16-mixed + gradient_clip_val: 1.0 + accumulate_grad_batches: 1 + matmul_precision: high + # Metric that selects best.ckpt. Leave empty when there is no held-out + # split to select on: last.ckpt is written either way. + checkpoint_monitor: val/loss + checkpoint_mode: min + # npy reads the segments written by preprocessing; raw seeks the same window + # in the source recording and resamples it. + source: npy + swap_channels: false # stereo-only augmentation + + projection_window: + # 1.4 s is 17.5 frames at 12.5 Hz, so the paper's times cannot be + # represented here. These keep the 4.0 s window exact - [18, 7, 7, 18] - + # where rounding [1.4, 0.6] gives 26 frames of history for a 25-frame window. + bin_sec: [1.44, 0.56, 0.56, 1.44] + num_hist_bins: 2 + frame_hz: 12.5 + threshold: 0.5 + mode: role_relative + cut_center_frame: false + + temporal: + d_model: 256 + m_ff: 3 + n_heads: 4 + n_layers: 4 + dropout: 0.1 + activation: GELU + causal: true + pos_encoding: alibi # alibi | rope | sinusoidal | none + ffn: mlp # mlp | swiglu + norm: layer # layer | rms + # Size of the sinusoidal position table. Rotary and ALiBi extend to any + # length, so this bounds the sinusoidal encoding only. + max_frames: 4096 diff --git a/third_party/MuVAP/config/yaml/vap_speaker.yaml b/third_party/MuVAP/config/yaml/vap_speaker.yaml new file mode 100644 index 0000000..f36c0b3 --- /dev/null +++ b/third_party/MuVAP/config/yaml/vap_speaker.yaml @@ -0,0 +1,65 @@ +# Speaker-based VAP (Ekstedt & Skantze, 2022): two audio channels, one speaker per +# channel, and a 256-class codebook that projects one channel per row. +fisher_path: data/fisher +checkpoint_dir: runs/vap_speaker + +vap: + arch: stereo + batch_size: 32 + num_workers: 16 + max_epochs: 5 + lr: 0.001 + weight_decay: 0.01 + warmup_ratio: 0.1 + min_lr: 0.0001 + + precision: bf16-mixed + gradient_clip_val: 1.0 + accumulate_grad_batches: 2 + matmul_precision: high + # Metric that selects best.ckpt. Leave empty when there is no held-out + # split to select on: last.ckpt is written either way. + checkpoint_monitor: val/loss + checkpoint_mode: min + # npy reads the segments written by preprocessing; raw seeks the same window + # in the source recording and resamples it. + source: raw + swap_channels: true # channel order is arbitrary; swap both audio and VAD + + projection_window: + bin_sec: [0.2, 0.4, 0.6, 0.8] + num_hist_bins: 0 # every bin is a future bin, as in the original repo + frame_hz: 25 + threshold: 0.5 + mode: speaker_based + cut_center_frame: false + + temporal: + d_model: 256 + m_ff: 3 + n_heads: 4 + n_layers: 1 + dropout: 0.1 + activation: GELU + causal: true + pos_encoding: alibi # alibi | rope | sinusoidal | none + ffn: mlp + norm: layer + # Size of the sinusoidal position table. Rotary and ALiBi extend to any + # length, so this bounds the sinusoidal encoding only. + max_frames: 4096 + + cross: + d_model: 256 + m_ff: 3 + n_heads: 4 + n_layers: 3 + dropout: 0.1 + activation: GELU + causal: true + pos_encoding: alibi # alibi | rope | sinusoidal | none + ffn: mlp + norm: layer + # Size of the sinusoidal position table. Rotary and ALiBi extend to any + # length, so this bounds the sinusoidal encoding only. + max_frames: 4096 diff --git a/third_party/MuVAP/config/yaml/vap_speaker_cpc_50hz.yaml b/third_party/MuVAP/config/yaml/vap_speaker_cpc_50hz.yaml new file mode 100644 index 0000000..e694027 --- /dev/null +++ b/third_party/MuVAP/config/yaml/vap_speaker_cpc_50hz.yaml @@ -0,0 +1,73 @@ +# Speaker-based VAP at 50 Hz on CPC - the original VAP's own configuration. +# +# CPC's conv stack runs at 100 Hz and `downsample` halves it once per block, so +# 50 Hz is one block instead of two. This is the rate and frontend of Ekstedt & +# Skantze (2022); Mimi cannot reach it, its encoder tops out at 25 Hz. +# Labels come from seg/vad50, written by 03_prep_vad_rate.py. +fisher_path: data/fisher +checkpoint_dir: runs/speaker_cpc_50hz + +vap: + audio_encoder: + name: cpc + frame_hz: 50 # one downsample block instead of two + + arch: stereo + batch_size: 32 + num_workers: 16 + max_epochs: 5 + lr: 0.001 + weight_decay: 0.01 + warmup_ratio: 0.1 + min_lr: 0.0001 + + precision: bf16-mixed + gradient_clip_val: 1.0 + accumulate_grad_batches: 2 + matmul_precision: high + # Metric that selects best.ckpt. Leave empty when there is no held-out + # split to select on: last.ckpt is written either way. + checkpoint_monitor: val/loss + checkpoint_mode: min + # npy reads the segments written by preprocessing; raw seeks the same window + # in the source recording and resamples it. + source: raw + swap_channels: true # channel order is arbitrary; swap both audio and VAD + + projection_window: + bin_sec: [0.2, 0.4, 0.6, 0.8] + num_hist_bins: 0 # every bin is a future bin, as in the original repo + frame_hz: 50 + threshold: 0.5 + mode: speaker_based + cut_center_frame: false + + temporal: + d_model: 256 + m_ff: 3 + n_heads: 4 + n_layers: 1 + dropout: 0.1 + activation: GELU + causal: true + pos_encoding: alibi # alibi | rope | sinusoidal | none + ffn: mlp + norm: layer + # Size of the sinusoidal position table. Rotary and ALiBi extend to any + # length, so this bounds the sinusoidal encoding only. + max_frames: 4096 + + cross: + d_model: 256 + m_ff: 3 + n_heads: 4 + n_layers: 3 + dropout: 0.1 + activation: GELU + causal: true + pos_encoding: alibi # alibi | rope | sinusoidal | none + ffn: mlp + norm: layer + # Size of the sinusoidal position table. Rotary and ALiBi extend to any + # length, so this bounds the sinusoidal encoding only. + max_frames: 4096 diff --git a/third_party/MuVAP/config/yaml/vap_speaker_mimi_12hz.yaml b/third_party/MuVAP/config/yaml/vap_speaker_mimi_12hz.yaml new file mode 100644 index 0000000..7d6b570 --- /dev/null +++ b/third_party/MuVAP/config/yaml/vap_speaker_mimi_12hz.yaml @@ -0,0 +1,75 @@ +# Speaker-based VAP at 12.5 Hz, the rate Mimi natively codes at. +# +# `tap: latent` is Mimi's own frame rate, so this matches a full-duplex system +# built on Mimi codes frame for frame: same frontend, same 80 ms decision grid. +# Labels come from seg/vad12_5, written by 03_prep_vad_rate.py; the 25 Hz tree +# is untouched. bin_sec stays at the paper's times, so the rate is the only +# thing that changes against the 25 Hz runs - at 12.5 Hz they round to +# 0.16 / 0.4 / 0.64 / 0.8, still a 2.0 s horizon. +fisher_path: data/fisher +checkpoint_dir: runs/speaker_mimi_12hz + +vap: + audio_encoder: + name: mimi + tap: latent # 12.5 Hz, Mimi's own frame rate + + arch: stereo + batch_size: 32 + num_workers: 16 + max_epochs: 5 + lr: 0.001 + weight_decay: 0.01 + warmup_ratio: 0.1 + min_lr: 0.0001 + + precision: bf16-mixed + gradient_clip_val: 1.0 + accumulate_grad_batches: 2 + matmul_precision: high + # Metric that selects best.ckpt. Leave empty when there is no held-out + # split to select on: last.ckpt is written either way. + checkpoint_monitor: val/loss + checkpoint_mode: min + # npy reads the segments written by preprocessing; raw seeks the same window + # in the source recording and resamples it. + source: raw + swap_channels: true # channel order is arbitrary; swap both audio and VAD + + projection_window: + bin_sec: [0.2, 0.4, 0.6, 0.8] + num_hist_bins: 0 # every bin is a future bin, as in the original repo + frame_hz: 12.5 + threshold: 0.5 + mode: speaker_based + cut_center_frame: false + + temporal: + d_model: 256 + m_ff: 3 + n_heads: 4 + n_layers: 1 + dropout: 0.1 + activation: GELU + causal: true + pos_encoding: alibi # alibi | rope | sinusoidal | none + ffn: mlp + norm: layer + # Size of the sinusoidal position table. Rotary and ALiBi extend to any + # length, so this bounds the sinusoidal encoding only. + max_frames: 4096 + + cross: + d_model: 256 + m_ff: 3 + n_heads: 4 + n_layers: 3 + dropout: 0.1 + activation: GELU + causal: true + pos_encoding: alibi # alibi | rope | sinusoidal | none + ffn: mlp + norm: layer + # Size of the sinusoidal position table. Rotary and ALiBi extend to any + # length, so this bounds the sinusoidal encoding only. + max_frames: 4096 diff --git a/third_party/MuVAP/config/yaml/vap_speaker_mono_cpc_50hz.yaml b/third_party/MuVAP/config/yaml/vap_speaker_mono_cpc_50hz.yaml new file mode 100644 index 0000000..7180d8f --- /dev/null +++ b/third_party/MuVAP/config/yaml/vap_speaker_mono_cpc_50hz.yaml @@ -0,0 +1,73 @@ +# Anchored mono VAP at 50 Hz on CPC - the original VAP's own configuration. +# +# CPC's conv stack runs at 100 Hz and `downsample` halves it once per block, so +# 50 Hz is one block instead of two. This is the rate and frontend of Ekstedt & +# Skantze (2022); Mimi cannot reach it, its encoder tops out at 25 Hz. +# Labels come from seg/vad50, written by 03_prep_vad_rate.py. +fisher_path: data/fisher +checkpoint_dir: runs/mono_cpc_50hz + +vap: + audio_encoder: + name: cpc + frame_hz: 50 # one downsample block instead of two + + arch: anchored # one mixed channel, anchored by past VAD + batch_size: 32 + num_workers: 16 + max_epochs: 5 + lr: 0.001 + weight_decay: 0.01 + warmup_ratio: 0.1 + min_lr: 0.0001 + + precision: bf16-mixed + gradient_clip_val: 1.0 + accumulate_grad_batches: 2 + matmul_precision: high + # Metric that selects best.ckpt. Leave empty when there is no held-out + # split to select on: last.ckpt is written either way. + checkpoint_monitor: val/loss + checkpoint_mode: min + # npy reads the segments written by preprocessing; raw seeks the same window + # in the source recording and resamples it. + source: raw + swap_channels: true # channel order is arbitrary; swap both audio and VAD + + projection_window: + bin_sec: [0.2, 0.4, 0.6, 0.8] + num_hist_bins: 0 # every bin is a future bin, as in the original repo + frame_hz: 50 + threshold: 0.5 + mode: speaker_based + cut_center_frame: false + + temporal: + d_model: 256 + m_ff: 3 + n_heads: 4 + n_layers: 1 + dropout: 0.1 + activation: GELU + causal: true + pos_encoding: alibi # alibi | rope | sinusoidal | none + ffn: mlp + norm: layer + # Size of the sinusoidal position table. Rotary and ALiBi extend to any + # length, so this bounds the sinusoidal encoding only. + max_frames: 4096 + + cross: + d_model: 256 + m_ff: 3 + n_heads: 4 + n_layers: 3 + dropout: 0.1 + activation: GELU + causal: true + pos_encoding: alibi # alibi | rope | sinusoidal | none + ffn: mlp + norm: layer + # Size of the sinusoidal position table. Rotary and ALiBi extend to any + # length, so this bounds the sinusoidal encoding only. + max_frames: 4096 diff --git a/third_party/MuVAP/config/yaml/vap_speaker_mono_mimi.yaml b/third_party/MuVAP/config/yaml/vap_speaker_mono_mimi.yaml new file mode 100644 index 0000000..0e1aaa6 --- /dev/null +++ b/third_party/MuVAP/config/yaml/vap_speaker_mono_mimi.yaml @@ -0,0 +1,79 @@ +# Speaker-based VAP from ONE channel, anchored by voice activity. +# +# The 256-class codebook names channels, which mono audio cannot identify. Here +# the identity comes from the VAD instead: the model hears the mixture and is +# told who held the floor strictly before the frame it is labelling, so the two +# rows stay attached to real speakers without a second channel. The anchor is +# delayed by one frame because speaker_based bins the future from `t` inclusive. +# +# Frozen Mimi frontend, so this is comparable to vap-speaker-mimi on everything +# but the number of channels. +fisher_path: data/fisher +checkpoint_dir: runs/vap_speaker_mono_mimi + +vap: + # Frozen 37.8M Mimi encoder. Its 512-dim features are bridged onto d_model by + # the model's encoder projection. + audio_encoder: + name: mimi + tap: encoder # 25 Hz, matching the label grid + + arch: anchored # one mixed channel, anchored by past VAD + batch_size: 32 + num_workers: 16 + max_epochs: 5 + lr: 0.001 + weight_decay: 0.01 + warmup_ratio: 0.1 + min_lr: 0.0001 + + precision: bf16-mixed + gradient_clip_val: 1.0 + accumulate_grad_batches: 2 + matmul_precision: high + # Metric that selects best.ckpt. Leave empty when there is no held-out + # split to select on: last.ckpt is written either way. + checkpoint_monitor: val/loss + checkpoint_mode: min + # npy reads the segments written by preprocessing; raw seeks the same window + # in the source recording and resamples it. + source: raw + swap_channels: true # channel order is arbitrary; swap both audio and VAD + + projection_window: + bin_sec: [0.2, 0.4, 0.6, 0.8] + num_hist_bins: 0 # every bin is a future bin, as in the original repo + frame_hz: 25 + threshold: 0.5 + mode: speaker_based + cut_center_frame: false + + temporal: + d_model: 256 + m_ff: 3 + n_heads: 4 + n_layers: 1 + dropout: 0.1 + activation: GELU + causal: true + pos_encoding: alibi # alibi | rope | sinusoidal | none + ffn: mlp + norm: layer + # Size of the sinusoidal position table. Rotary and ALiBi extend to any + # length, so this bounds the sinusoidal encoding only. + max_frames: 4096 + + cross: + d_model: 256 + m_ff: 3 + n_heads: 4 + n_layers: 3 + dropout: 0.1 + activation: GELU + causal: true + pos_encoding: alibi # alibi | rope | sinusoidal | none + ffn: mlp + norm: layer + # Size of the sinusoidal position table. Rotary and ALiBi extend to any + # length, so this bounds the sinusoidal encoding only. + max_frames: 4096 diff --git a/third_party/MuVAP/config/yaml/vap_speaker_mono_mimi_12hz.yaml b/third_party/MuVAP/config/yaml/vap_speaker_mono_mimi_12hz.yaml new file mode 100644 index 0000000..e991708 --- /dev/null +++ b/third_party/MuVAP/config/yaml/vap_speaker_mono_mimi_12hz.yaml @@ -0,0 +1,77 @@ +# Anchored mono VAP at 12.5 Hz, the rate Mimi natively codes at. +# +# `tap: latent` is Mimi's own frame rate, so this matches a full-duplex system +# built on Mimi codes frame for frame: same frontend, same 80 ms decision grid. +# Labels come from seg/vad12_5, written by 03_prep_vad_rate.py; the 25 Hz tree +# is untouched. bin_sec stays at the paper's times, so the rate is the only +# thing that changes against the 25 Hz runs - at 12.5 Hz they round to +# 0.16 / 0.4 / 0.64 / 0.8, still a 2.0 s horizon. +fisher_path: data/fisher +checkpoint_dir: runs/mono_mimi_12hz + +vap: + # Frozen 37.8M Mimi encoder. Its 512-dim features are bridged onto d_model by + # the model's encoder projection. + audio_encoder: + name: mimi + tap: latent # 12.5 Hz, Mimi's own frame rate + + arch: anchored # one mixed channel, anchored by past VAD + batch_size: 32 + num_workers: 16 + max_epochs: 5 + lr: 0.001 + weight_decay: 0.01 + warmup_ratio: 0.1 + min_lr: 0.0001 + + precision: bf16-mixed + gradient_clip_val: 1.0 + accumulate_grad_batches: 2 + matmul_precision: high + # Metric that selects best.ckpt. Leave empty when there is no held-out + # split to select on: last.ckpt is written either way. + checkpoint_monitor: val/loss + checkpoint_mode: min + # npy reads the segments written by preprocessing; raw seeks the same window + # in the source recording and resamples it. + source: raw + swap_channels: true # channel order is arbitrary; swap both audio and VAD + + projection_window: + bin_sec: [0.2, 0.4, 0.6, 0.8] + num_hist_bins: 0 # every bin is a future bin, as in the original repo + frame_hz: 12.5 + threshold: 0.5 + mode: speaker_based + cut_center_frame: false + + temporal: + d_model: 256 + m_ff: 3 + n_heads: 4 + n_layers: 1 + dropout: 0.1 + activation: GELU + causal: true + pos_encoding: alibi # alibi | rope | sinusoidal | none + ffn: mlp + norm: layer + # Size of the sinusoidal position table. Rotary and ALiBi extend to any + # length, so this bounds the sinusoidal encoding only. + max_frames: 4096 + + cross: + d_model: 256 + m_ff: 3 + n_heads: 4 + n_layers: 3 + dropout: 0.1 + activation: GELU + causal: true + pos_encoding: alibi # alibi | rope | sinusoidal | none + ffn: mlp + norm: layer + # Size of the sinusoidal position table. Rotary and ALiBi extend to any + # length, so this bounds the sinusoidal encoding only. + max_frames: 4096 diff --git a/third_party/MuVAP/data/__init__.py b/third_party/MuVAP/data/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/third_party/MuVAP/data/alignment.py b/third_party/MuVAP/data/alignment.py new file mode 100644 index 0000000..3e42736 --- /dev/null +++ b/third_party/MuVAP/data/alignment.py @@ -0,0 +1,114 @@ +"""Time-based alignment utilities for audio, video frames, and frame labels. + +The model runs on a fixed 25 Hz timeline. Source videos remain at their native +frame rate; frames and labels are sampled onto the model timeline at load time. +""" + +from dataclasses import dataclass +from pathlib import Path +from typing import Sequence + +import numpy as np +import torch +import torch.nn.functional as F + + +@dataclass(frozen=True) +class AlignmentConfig: + model_fps: float = 25.0 + audio_sample_rate: int = 16_000 + + def __post_init__(self) -> None: + if self.model_fps <= 0: + raise ValueError("model_fps must be positive") + if self.audio_sample_rate <= 0: + raise ValueError("audio_sample_rate must be positive") + + +@dataclass(frozen=True) +class AVSample: + sample_id: str + audio_path: Path + visual_path: Path + source_fps: float + labels: Sequence[int] + start_sec: float = 0.0 + + def __post_init__(self) -> None: + if self.source_fps <= 0: + raise ValueError(f"{self.sample_id}: source_fps must be positive") + if self.start_sec < 0: + raise ValueError(f"{self.sample_id}: start_sec cannot be negative") + + +def model_frame_count(source_frames: int, source_fps: float, model_fps: float) -> int: + """Return the number of complete model intervals in a source-frame span.""" + if source_frames < 0 or source_fps <= 0 or model_fps <= 0: + raise ValueError("frame counts and frame rates must be valid") + return int(np.floor((source_frames / source_fps) * model_fps + 1e-9)) + + +def source_indices( + source_frames: int, source_fps: float, model_fps: float +) -> np.ndarray: + """Map model-frame centres to the source intervals that contain them.""" + count = model_frame_count(source_frames, source_fps, model_fps) + model_centres = (np.arange(count, dtype=np.float64) + 0.5) / model_fps + indices = np.floor(model_centres * source_fps).astype(np.int64) + return np.clip(indices, 0, max(source_frames - 1, 0)) + + +def model_indices( + source_frames: int, source_fps: float, model_fps: float +) -> np.ndarray: + """Map source-frame centres to the model intervals that contain them. + + The inverse direction of `source_indices`, used to lift model-rate + predictions back onto the native annotation timeline for evaluation. + Always returns exactly `source_frames` indices, so a caller can score one + prediction per native frame regardless of the source frame rate. + """ + if source_frames < 0: + raise ValueError("source_frames cannot be negative") + count = model_frame_count(source_frames, source_fps, model_fps) + source_centres = (np.arange(source_frames, dtype=np.float64) + 0.5) / source_fps + indices = np.floor(source_centres * model_fps).astype(np.int64) + return np.clip(indices, 0, max(count - 1, 0)) + + +def align_visual_and_labels( + visual: np.ndarray, + labels: Sequence[int], + source_fps: float, + model_fps: float, + strict: bool = True, +) -> tuple[np.ndarray, np.ndarray]: + """Sample faces and their labels with one shared timestamp mapping. + + A length disagreement means the face track and the annotation disagree + about how long the clip is, which silently shifts every later label. + `strict` reports that instead of trimming to the shorter of the two. + """ + labels_array = np.asarray(labels, dtype=np.float32) + if strict and len(visual) != len(labels_array): + raise ValueError( + f"visual has {len(visual)} frames but labels have {len(labels_array)}" + ) + usable_frames = min(len(visual), len(labels_array)) + if usable_frames == 0: + raise ValueError("visual and labels must contain at least one frame") + + indices = source_indices(usable_frames, source_fps, model_fps) + return visual[:usable_frames][indices], labels_array[:usable_frames][indices] + + +def fit_audio_to_frames( + waveform: torch.Tensor, + frame_count: int, + sample_rate: int, + model_fps: float, +) -> torch.Tensor: + """Pad/crop audio to the exact duration represented by model frames.""" + target_samples = round(frame_count * sample_rate / model_fps) + waveform = waveform[..., :target_samples] + return F.pad(waveform, (0, max(0, target_samples - waveform.shape[-1]))) diff --git a/third_party/MuVAP/data/augment.py b/third_party/MuVAP/data/augment.py new file mode 100644 index 0000000..7e1f2a9 --- /dev/null +++ b/third_party/MuVAP/data/augment.py @@ -0,0 +1,105 @@ +"""Train-time augmentation for packed ASD samples. + +Every transform here is label-preserving except `substitute_audio`, which +rewrites the activity labels to silence on purpose. Visual parameters are drawn +once per chunk and applied to every frame, because the model is temporal and a +per-frame jitter would destroy the motion it reads. +""" + +import math +import random +from dataclasses import dataclass, field + +import cv2 +import numpy as np + + +VISUAL_MODES = ("orig", "flip", "crop", "rotate") + + +@dataclass(frozen=True) +class AugmentConfig: + """How aggressively to augment; all probabilities are per chunk.""" + + audio_mix_prob: float = 0.5 + audio_snr_db: tuple[float, float] = (-5.0, 5.0) + visual: bool = True + crop_scale: tuple[float, float] = (0.7, 1.0) + rotate_degrees: float = 15.0 + # Ablation only: replace the audio outright and call the sample silent. + # Stronger than mixing for teaching rejection of off-screen speech, but it + # also lets the model key on audio plausibility instead of AV synchrony. + substitute_audio_prob: float = 0.0 + + @classmethod + def from_cfg(cls, cfg: dict | None) -> "AugmentConfig | None": + if not cfg: + return None + snr = cfg.get("audio_snr_db", (-5.0, 5.0)) + return cls( + audio_mix_prob=float(cfg.get("audio_mix_prob", 0.5)), + audio_snr_db=(float(snr[0]), float(snr[1])), + visual=bool(cfg.get("visual", True)), + crop_scale=tuple(cfg.get("crop_scale", (0.7, 1.0))), + rotate_degrees=float(cfg.get("rotate_degrees", 15.0)), + substitute_audio_prob=float(cfg.get("substitute_audio_prob", 0.0)), + ) + + +def _power_db(samples: np.ndarray) -> float: + return 10.0 * math.log10(float(np.mean(samples.astype(np.float64) ** 2)) + 1e-4) + + +def fit_length(donor: np.ndarray, samples: int) -> np.ndarray: + """Wrap or truncate donor audio to the target length.""" + if len(donor) == 0: + return np.zeros(samples, dtype=np.float32) + if len(donor) < samples: + donor = np.tile(donor, math.ceil(samples / len(donor))) + return donor[:samples] + + +def mix_audio( + clean: np.ndarray, donor: np.ndarray, snr_db: float +) -> np.ndarray: + """Add donor audio at a target signal-to-noise ratio. + + The label is deliberately unchanged. A negative face with foreign speech + added must stay negative, which is what teaches the model to reject speech + it cannot see; a positive face stays positive when a competing voice + arrives. + """ + donor = fit_length(donor, len(clean)) + scale = 10.0 ** ((_power_db(clean) - _power_db(donor) - snr_db) / 20.0) + return np.clip(clean + donor * scale, -1.0, 1.0) + + +def augment_visual( + faces: np.ndarray, config: AugmentConfig, rng: random.Random +) -> np.ndarray: + """Apply one geometric transform to every frame of a chunk.""" + mode = rng.choice(VISUAL_MODES) + if mode == "orig": + return faces + size = faces.shape[-1] + if mode == "flip": + return np.ascontiguousarray(faces[:, :, ::-1]) + if mode == "crop": + low, high = config.crop_scale + side = max(8, int(size * rng.uniform(low, high))) + if side >= size: + return faces + top = rng.randint(0, size - side) + left = rng.randint(0, size - side) + window = faces[:, top : top + side, left : left + side] + return np.stack( + [ + cv2.resize(frame, (size, size), interpolation=cv2.INTER_LINEAR) + for frame in window + ] + ) + angle = rng.uniform(-config.rotate_degrees, config.rotate_degrees) + matrix = cv2.getRotationMatrix2D((size / 2, size / 2), angle, 1.0) + return np.stack( + [cv2.warpAffine(frame, matrix, (size, size)) for frame in faces] + ) diff --git a/third_party/MuVAP/data/dataloaders/__init__.py b/third_party/MuVAP/data/dataloaders/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/third_party/MuVAP/data/dataloaders/fisher.py b/third_party/MuVAP/data/dataloaders/fisher.py new file mode 100644 index 0000000..f9c74c2 --- /dev/null +++ b/third_party/MuVAP/data/dataloaders/fisher.py @@ -0,0 +1,374 @@ +"""Fisher datasets for VAP training and turn-taking evaluation. + +Segments are stored as stereo audio beside the raw two-speaker VAD, not as +pre-encoded class indices. Labels are projected in the dataloader, so changing +`projection_window.mode` between the role-based and speaker-based setups is a +config edit rather than a re-run of preprocessing. + +Audio comes from one of two sources, selected by `vap.source`: + +* `npy` reads the pre-cut segments under `seg/` and `tune/`. +* `raw` seeks the same span in the 8 kHz mu-law source recording and resamples + it. The source is 97 GB where the pre-cut segments are 1.12 TB, and segment + length, stride, and sample rate stop being baked into the files, so they can + change without re-running preprocessing. + +Both sources are required to yield the same window, and a test pins them to +each other. +""" + +from functools import lru_cache +from glob import glob +from pathlib import Path + +import numpy as np +import soundfile as sf +import torch +import torchaudio.functional as AF +from torch.utils.data import Dataset + +from data.media import AudioSpec, load_waveform +from preprocess.vap.utils import EVENT_SEG_LEN_SEC, event_crop + +SOURCES = ("npy", "raw") +VAD_HZ = 25 + + +def vad_dirname(frame_hz): + """Where the VAD for this rate lives: `vad` at 25 Hz, `vad12_5` at 12.5, ... + + `00_prep_fisher.py` writes the 25 Hz labels to `seg/vad`; every other rate is + written beside it by `03_prep_vad_rate.py`, so the original tree is never + overwritten and two rates can coexist. + """ + if float(frame_hz) == float(VAD_HZ): + return "vad" + return "vad" + f"{float(frame_hz):g}".replace(".", "_") +TARGET_SAMPLE_RATE = 16_000 +# Resampling a slice is not the same as slicing a resampled recording: the +# filter needs signal beyond both edges. Read this much extra and trim it off. +RESAMPLE_PAD_SEC = 0.1 + + +@lru_cache(maxsize=1024) +def _source_rate(path): + return sf.info(path).samplerate + + +@lru_cache(maxsize=256) +def _speech_regions(root, part, group, conversation): + """Word-level speech intervals per speaker, as `((start, end), ...)` seconds. + + The event clips ship as audio only, but an anchored model needs the voice + activity behind them. It comes from the same word annotations preprocessing + rasterized to build the training VAD, so the two cannot disagree. + """ + pattern = Path(root) / part / "regions_vap" / group / f"{conversation}_*_words.txt" + speakers = [] + for path in sorted(glob(str(pattern))): + with open(path) as handle: + speakers.append( + tuple( + (float(line.split()[0]), float(line.split()[1])) + for line in handle + if line.strip() + ) + ) + if len(speakers) != 2: + raise FileNotFoundError( + f"expected two *_words.txt for {conversation} under {pattern.parent}, " + f"found {len(speakers)}" + ) + return tuple(speakers) + + +def window_vad(root, part, group, conversation, start_sec, frames, frame_hz=VAD_HZ): + """`[2, frames]` voice activity for the window starting at `start_sec`.""" + vad = torch.zeros(2, frames) + for speaker, regions in enumerate(_speech_regions(root, part, group, conversation)): + for begin, end in regions: + first = int(round((begin - start_sec) * frame_hz)) + last = int(round((end - start_sec) * frame_hz)) + if last <= 0 or first >= frames: + continue + vad[speaker, max(0, first) : min(frames, last)] = 1.0 + return vad + + +def read_source_window(path, start_sec, duration_sec, channels): + """Read `duration_sec` from a source recording, resampled to 16 kHz. + + Fisher ships as 8 kHz mu-law NIST Sphere, which `soundfile` reads and + `torchaudio` (2.11) no longer does. + """ + path = str(path) + rate = _source_rate(path) + pad = int(round(RESAMPLE_PAD_SEC * rate)) + start = int(round(start_sec * rate)) + frames = int(round(duration_sec * rate)) + + lead = min(pad, max(0, start)) + with sf.SoundFile(path) as handle: + handle.seek(start - lead) + block = handle.read(frames + lead + pad, dtype="float32", always_2d=True).T + + waveform = torch.from_numpy(np.ascontiguousarray(block)) + if channels == 1: + waveform = waveform.mean(dim=0, keepdim=True) + elif waveform.shape[0] != channels: + raise ValueError( + f"{path} holds {waveform.shape[0]} channel(s); {channels} requested" + ) + + waveform = AF.resample(waveform, rate, TARGET_SAMPLE_RATE) + scale = TARGET_SAMPLE_RATE / rate + offset = int(round(lead * scale)) + target = int(round(duration_sec * TARGET_SAMPLE_RATE)) + waveform = waveform[:, offset : offset + target] + if waveform.shape[1] < target: # the recording ended inside the window + waveform = torch.nn.functional.pad(waveform, (0, target - waveform.shape[1])) + # Clip after resampling, exactly where the `npy` path clips, so a checkpoint + # trained on one source evaluates identically on the other. The resampling + # filter overshoots at transients; mu-law input never does. + return waveform.clamp(-1.0, 1.0).contiguous() + + +def _check_source(source): + if source not in SOURCES: + raise ValueError(f"source must be one of {SOURCES}, got {source!r}") + return source + + +def _audio_spec(channels): + """Fisher segment spec. + + Segments written by earlier revisions of `00_prep_fisher.py` stored float + PCM without clipping, and about one in eight overshoots [-1, 1] by up to + 2.5 dB. Clipping them is what any encoder would do and keeps those shards + usable; freshly prepared segments are int16 and never reach this path. + """ + if channels not in (1, 2): + raise ValueError(f"channels must be 1 or 2, got {channels}") + return AudioSpec(sample_rate=16_000, mono=channels == 1, clamp_to_peak=True) + + +class Fisher(Dataset): + """Fixed-length conversation segments with labels projected on the fly.""" + + def __init__( + self, + fisher_path, + split_path, + projection, + channels=1, + frame_hz=25, + sample_rate=16_000, + swap_channels=False, + source="npy", + with_vad=False, + ): + self.root = Path(fisher_path) + self.projection = projection + self.channels = channels + self.source = _check_source(source) + self.spec = _audio_spec(channels) + # `frame_hz` can be fractional (12.5 Hz on Mimi's own grid), so this is + # a rounded int rather than floor division - a float here reaches a + # slice index and fails deep inside a dataloader worker. + self.samples_per_frame = int(round(sample_rate / frame_hz)) + self.with_vad = with_vad + self.vad_dir = vad_dirname(frame_hz) + # Swapping is label-preserving wherever speaker order is arbitrary: with + # two channels, and with one channel whose speakers are named by the VAD. + self.swap_channels = swap_channels and (channels == 2 or with_vad) + if projection.frame_hz != frame_hz: + raise ValueError( + f"projection is defined at {projection.frame_hz} Hz but the dataset runs at {frame_hz} Hz" + ) + self.samples = [] + self._build(split_path) + + def _paths(self, line): + # Paths are resolved without touching the filesystem: a split can list + # hundreds of thousands of segments on a network mount, and stat-ing + # each one at construction costs minutes before the first batch. + part, group, conversation, segment, start, end, *_ = line.split() + base = self.root / part / "seg" + audio = ( + base / "audio" / group / conversation / f"{segment}.npy" + if self.source == "npy" + else self.root / part / "audio" / group / f"{conversation}.wav" + ) + return ( + audio, + base / self.vad_dir / group / conversation / f"{segment}.npy", + float(start), + float(end) - float(start), + ) + + def _build(self, split_path): + with open(split_path) as file: + self.samples = [self._paths(line) for line in file if line.strip()] + if not self.samples: + raise ValueError(f"{split_path} lists no segments") + audio_path, vad_path, _, _ = self.samples[0] + for path in (audio_path, vad_path): + if not path.exists(): + raise FileNotFoundError( + f"missing preprocessed Fisher sample {path}; " + "run preprocess/vap/00_prep_fisher.py first" + ) + + def __len__(self): + return len(self.samples) + + def _labels(self, vad, num_frames): + """Project stored VAD, then drop the context kept only for the bins.""" + context = (vad.shape[-1] - num_frames) // 2 + if context < 0: + raise ValueError( + f"VAD has {vad.shape[-1]} frames, fewer than the {num_frames} audio frames" + ) + needed = max(self.projection.hist_frames, self.projection.fut_frames) + if context < needed: + raise ValueError( + f"projection needs {needed} context frames on each side but the segment stores {context}; " + "re-run preprocess/vap/00_prep_fisher.py with a wider --context-sec" + ) + labels = self.projection.get_labels(vad.unsqueeze(0))[0] + return labels[context : context + num_frames] + + def __getitem__(self, index): + audio_path, vad_path, start_sec, duration_sec = self.samples[index] + if self.source == "raw": + audio = read_source_window( + audio_path, start_sec, duration_sec, self.channels + ) + else: + audio, _ = load_waveform(audio_path, self.spec, npy_sample_rate=16_000) + vad = torch.from_numpy(np.load(vad_path, allow_pickle=False)).float() + + if self.channels == 2 and audio.shape[0] != 2: + raise ValueError( + f"{audio_path} holds {audio.shape[0]} channel(s); stereo VAP needs 2" + ) + if self.swap_channels and torch.rand(()) < 0.5: + # Which speaker sits on which channel is arbitrary, so swapping both + # the audio and its VAD together is a label-preserving augmentation. + audio = audio.flip(0) + vad = vad.flip(0) + + num_frames = int(audio.shape[-1] // self.samples_per_frame) + labels = self._labels(vad, num_frames) + if not self.with_vad: + return audio, labels + # The same frames the labels cover, so the anchor lines up with them. + context = (vad.shape[-1] - num_frames) // 2 + return audio, labels, vad[:, context : context + num_frames] + + +class FisherEvent(Dataset): + """Ten-second clips ending at an annotated hold or shift decision point.""" + + FIELDS = ("part", "group", "conversation", "start", "end", "prev_spk", "next_spk", + "timing", "gap_type", "label") + + def __init__(self, fisher_path, split_paths, channels=1, source="npy", + with_vad=False, frame_hz=VAD_HZ): + root = self.root = Path(fisher_path) + self.channels = channels + self.source = _check_source(source) + self.spec = _audio_spec(channels) + self.audio_dir = "audio" if channels == 1 else "audio_stereo" + self.with_vad = with_vad + self.frame_hz = frame_hz + self.samples = [] + if isinstance(split_paths, (str, Path)): + split_paths = [split_paths] + + for split_path in split_paths: + with open(split_path) as file: + for line in file: + if not line.strip(): + continue + fields = dict(zip(self.FIELDS, line.split())) + # The window is defined the same way for both sources; `raw` + # seeks it, `npy` reads a clip already cut to it. An anchored + # model needs its absolute start either way, to line the + # voice activity up with the audio. + crop = event_crop( + fields["timing"], + float(fields["start"]), + float(fields["end"]), + ) + if crop is None and (self.source == "raw" or self.with_vad): + continue + if self.source == "raw": + audio_path = ( + root + / fields["part"] + / "audio" + / fields["group"] + / f"{fields['conversation']}.wav" + ) + crop_start = crop[0] + else: + audio_path = ( + root + / fields["part"] + / "tune" + / self.audio_dir + / fields["group"] + / fields["conversation"] + / f"{fields['start']}.npy" + ) + crop_start = 0.0 + self.samples.append( + ( + audio_path, + 1.0 if fields["label"] == "SHIFT" else 0.0, + fields["conversation"], + fields["timing"], + fields["gap_type"], + int(fields["prev_spk"]), + crop_start, + (fields["part"], fields["group"], crop[0] if crop else None), + ) + ) + + def __len__(self): + return len(self.samples) + + def __getitem__(self, index): + ( + audio_path, + label, + conversation, + timing, + gap_type, + prev_spk, + crop_start, + window, + ) = self.samples[index] + + if self.source == "raw": + audio = read_source_window( + audio_path, crop_start, EVENT_SEG_LEN_SEC, self.channels + ) + else: + audio, _ = load_waveform(audio_path, self.spec, npy_sample_rate=16_000) + if self.channels == 2 and audio.shape[0] != 2: + raise ValueError( + f"{audio_path} holds {audio.shape[0]} channel(s); stereo evaluation needs 2. " + "Re-run preprocess/vap/02_prep_event.py with --channels 2, " + "or set vap.source: raw to read the source recording instead" + ) + if not self.with_vad: + return audio, label, conversation, timing, gap_type, prev_spk + part, group, window_start = window + vad = window_vad( + str(self.root), part, group, conversation, window_start, + int(round(audio.shape[-1] / TARGET_SAMPLE_RATE * self.frame_hz)), + frame_hz=self.frame_hz, + ) + return audio, label, conversation, timing, gap_type, prev_spk, vad diff --git a/third_party/MuVAP/data/dataloaders/packed.py b/third_party/MuVAP/data/dataloaders/packed.py new file mode 100644 index 0000000..b4a18a7 --- /dev/null +++ b/third_party/MuVAP/data/dataloaders/packed.py @@ -0,0 +1,410 @@ +"""Memory-mapped reader that turns native-rate packs into model-ready batches. + +Shards hold decoded source media: faces and labels at the native frame rate, +unnormalized 16 kHz audio. Everything MuVAP+ specific happens here - sampling +onto the 25 Hz model timeline, loudness normalization, chunking, augmentation - +so the same shards can serve a model with different conventions. +""" + +import json +import math +import random +from dataclasses import dataclass +from pathlib import Path +from typing import Sequence + +import numpy as np +import torch +import torch.nn.functional as F +from torch.utils.data import Dataset, Sampler + +from data.alignment import model_frame_count, source_indices +from data.augment import AugmentConfig, augment_visual, mix_audio +from data.media import WaveformStats, apply_normalization +from preprocess.asd.schema import FORMAT_NAME, MODEL_FPS, SAMPLE_RATE, SUPPORTED_VERSIONS + + +def _clip_id(sample_id: str) -> str: + """Strip the entity suffix so same-recording samples share a key. + + AVA and WASD both name tracks `__:`, where every + entity of a clip may share one audio recording. + """ + return sample_id.rsplit(":", 1)[0] + + +@dataclass(frozen=True) +class _Entry: + root: Path + shard: str + metadata: dict + + +class PackedASDDataset(Dataset): + """Random-access dataset backed by OS-page-cached NumPy memmaps.""" + + def __init__( + self, + roots: str | Path | Sequence[str | Path], + projection_window=None, + window: int | None = None, + context: int = 0, + jitter: bool = False, + augment: AugmentConfig | None = None, + model_fps: float = MODEL_FPS, + normalize: bool = True, + ): + """Read packed shards, optionally as fixed-length training chunks. + + `window` splits every track longer than it into near-equal chunks, so + one epoch still covers every frame while batches stay uniform enough to + pack tightly. + + `context` frames are read *before* each chunk and marked invalid in the + returned mask. The model is causal, so a frame near the start of a chunk + would otherwise be asked to predict from history the encoder was never + shown. Feeding that history as unscored context costs a little compute + and no storage, where overlapping the chunks would duplicate the data + and still train on the truncated copies. A track's real first frames get + no prefix: there the missing history is genuine, and the model has to + cope with it at inference too. + + `jitter` slides the chunk grid within its track on every access, which + varies the boundaries across epochs at no extra cost. All spans are + counted in model frames, not source frames. + """ + if isinstance(roots, (str, Path)): + roots = [roots] + self.projection_window = projection_window + self.window = window + self.context = max(0, context) + self.jitter = jitter + self.augment = augment + self.model_fps = model_fps + self.normalize = normalize + self.samples_per_frame = round(SAMPLE_RATE / model_fps) + self.entries: list[_Entry] = [] + self._arrays: dict[ + tuple[Path, str], tuple[np.ndarray, np.ndarray, np.ndarray] + ] = {} + + for raw_root in roots: + root = Path(raw_root) + with (root / "dataset.json").open() as handle: + manifest = json.load(handle) + if ( + manifest.get("format") != FORMAT_NAME + or manifest.get("version") not in SUPPORTED_VERSIONS + ): + raise ValueError(f"unsupported packed dataset format: {root}") + if manifest.get("timeline") != "native": + raise ValueError(f"pack is not on the native timeline: {root}") + if manifest.get("audio_sample_rate") != SAMPLE_RATE: + raise ValueError(f"unexpected audio sample rate: {root}") + if manifest.get("audio_normalization") is not None: + raise ValueError(f"pack has baked-in audio normalization: {root}") + + for shard_dir in sorted(root.glob("shard-*")): + with (shard_dir / "index.json").open() as handle: + for metadata in json.load(handle): + self.entries.append(_Entry(root, shard_dir.name, metadata)) + + self.model_frames = np.array( + [ + model_frame_count( + entry.metadata["source_frames"], + entry.metadata["source_fps"], + model_fps, + ) + for entry in self.entries + ], + dtype=np.int64, + ) + self.chunks = self._build_chunks() + self._clips = [_clip_id(entry.metadata["sample_id"]) for entry in self.entries] + + @property + def predict_window(self) -> int | None: + """Frames scored per chunk; the context prefix is read on top of this.""" + if self.window is None: + return None + return max(1, self.window - self.context) + + def _build_chunks(self) -> list[tuple[int, int, int]]: + """Split tracks into near-equal *scored* spans. + + A chunk is identified by the span it predicts. Whatever context the + loader prepends is read on top and never scored, so the spans still tile + the track exactly and every frame is predicted once per epoch. + """ + chunks: list[tuple[int, int, int]] = [] + span = self.predict_window + for index in range(len(self.entries)): + frames = int(self.model_frames[index]) + if frames == 0: + continue + if span is None or frames <= span: + chunks.append((index, 0, frames)) + continue + edges = np.linspace( + 0, frames, math.ceil(frames / span) + 1 + ).round().astype(int) + for start, end in zip(edges[:-1], edges[1:]): + chunks.append((index, int(start), int(end - start))) + return chunks + + def _fed_length(self, start: int, length: int) -> int: + return length + min(self.context, start) + + def chunk_lengths(self) -> np.ndarray: + """Frames actually fed per chunk - what bounds memory and batch size.""" + return np.fromiter( + (self._fed_length(start, length) for _, start, length in self.chunks), + dtype=np.int64, + count=len(self.chunks), + ) + + def __len__(self) -> int: + return len(self.chunks) + + def _open(self, entry: _Entry): + key = (entry.root, entry.shard) + arrays = self._arrays.get(key) + if arrays is None: + shard_dir = entry.root / entry.shard + faces = np.load(shard_dir / "faces.npy", mmap_mode="r", allow_pickle=False) + audio = np.load(shard_dir / "audio.npy", mmap_mode="r", allow_pickle=False) + labels = np.load(shard_dir / "labels.npy", mmap_mode="r", allow_pickle=False) + arrays = faces, audio, labels + self._arrays[key] = arrays + return arrays + + def _stats(self, entry: _Entry) -> WaveformStats: + return WaveformStats( + mean=entry.metadata["audio_mean"], + rms=entry.metadata["audio_rms"], + peak=entry.metadata["audio_peak"], + ) + + def _read_audio(self, entry: _Entry, start: int, length: int) -> torch.Tensor: + """Read one model-frame span and restore the recording's own gain.""" + _, audio, _ = self._open(entry) + offset = entry.metadata["audio_offset"] + start * self.samples_per_frame + available = entry.metadata["audio_samples"] - start * self.samples_per_frame + wanted = length * self.samples_per_frame + block = np.array(audio[offset : offset + min(wanted, max(available, 0))], copy=True) + waveform = torch.from_numpy(block).float().div_(32768.0).unsqueeze(0) + if waveform.shape[-1] < wanted: + waveform = F.pad(waveform, (0, wanted - waveform.shape[-1])) + if self.normalize: + waveform = apply_normalization(waveform, self._stats(entry)) + return waveform + + def _donor_audio(self, clip: str, frames: int, rng) -> np.ndarray: + """Read audio from a different clip than `clip`. + + Drawing from the whole dataset rather than the mini-batch matters here: + every face track in a WASD clip shares one clip-level recording, and the + batch sampler groups by length, which puts those tracks together. A + batch-local donor would often be the very same waveform. + """ + for _ in range(8): + index = rng.randrange(len(self.entries)) + if self._clips[index] == clip: + continue + entry = self.entries[index] + span = int(self.model_frames[index]) + start = rng.randrange(max(span - frames, 0) + 1) if span > frames else 0 + return self._read_audio(entry, start, min(frames, max(span, 1)))[0].numpy() + return np.zeros(frames * self.samples_per_frame, dtype=np.float32) + + def __getitem__(self, index: int): + entry_index, predict_start, predict_length = self.chunks[index] + entry = self.entries[entry_index] + faces, _, labels = self._open(entry) + frames = int(self.model_frames[entry_index]) + if self.jitter and predict_length < frames: + predict_start = random.randint(0, frames - predict_length) + + # Read a history prefix where one exists; at a track's real start there + # is none, and that truncation is genuine rather than an artefact. + prefix = min(self.context, predict_start) + start = predict_start - prefix + length = prefix + predict_length + scored = torch.zeros(length, dtype=torch.bool) + scored[prefix:] = True + + # One shared mapping puts faces and labels on the model timeline, which + # is what keeps a label attached to the face it describes. + rows = source_indices( + entry.metadata["source_frames"], entry.metadata["source_fps"], self.model_fps + ) + span = rows[start : start + length] + face_block = np.asarray(faces[entry.metadata["face_offset"] + span]) + # Labels are one byte a frame, so the whole track is cheap to read. The + # projection window needs 50 frames of history and 15 of future; taking + # them from the track rather than the chunk keeps a chunk boundary from + # fabricating silence in the targets. + track_labels = np.asarray( + labels[entry.metadata["label_offset"] + rows], dtype=np.uint8 + ).astype(np.float32) + waveform = self._read_audio(entry, start, length) + + if self.augment is not None: + # The module RNG is what Lightning seeds per worker, so augmentation + # is reproducible from the run seed. + rng = random + clip = self._clips[entry_index] + audio_block = waveform[0].numpy() + if rng.random() < self.augment.substitute_audio_prob: + audio_block = self._donor_audio(clip, length, rng) + track_labels = np.zeros_like(track_labels) + elif rng.random() < self.augment.audio_mix_prob: + audio_block = mix_audio( + audio_block, + self._donor_audio(clip, length, rng), + rng.uniform(*self.augment.audio_snr_db), + ) + waveform = torch.from_numpy(np.ascontiguousarray(audio_block)).unsqueeze(0) + if self.augment.visual: + face_block = augment_visual(face_block, self.augment, rng) + + visual = torch.from_numpy(np.ascontiguousarray(face_block)).float() + track_vad = torch.from_numpy(track_labels) + vad = track_vad[start : start + length] + if self.projection_window is None: + return waveform, visual, vad, scored, entry.metadata + with torch.no_grad(): + vap = self.projection_window.get_labels(track_vad.view(1, 1, frames)) + return ( + waveform, + visual, + vap[0, start : start + length], + vad, + scored, + entry.metadata, + ) + + +class ASDCollator: + """Bring a batch to one length, by cropping or by padding. + + Training crops to the shortest member. That keeps the model free of a + padding mask, and it is affordable because chunking bounds the spread of + lengths and the batch sampler groups similar lengths together; measured on + AVA plus WASD the pair loses 0.01% of frames per epoch, against 64% for + random batches. + + Evaluation cannot crop: the official export needs every frame of every + track, which is why it would otherwise run one track at a time. `pad=True` + right-pads to the longest member instead and marks the filler unscored. + Every stage of the model is causal - the CPC and Mimi frontends, the visual + TCN, and both transformer stacks - so appended frames cannot change any + earlier output, and the visual frontend's batch norm is on running + statistics in eval mode. A test pins both properties. + """ + + def __init__( + self, samples_per_frame: int = SAMPLE_RATE // int(MODEL_FPS), pad: bool = False + ): + self.samples_per_frame = samples_per_frame + self.pad = pad + + @staticmethod + def _to_length(tensor: torch.Tensor, length: int, axis: int) -> torch.Tensor: + """Crop or right-pad one stream along `axis` (counted from the end).""" + current = tensor.shape[-axis] + if current == length: + return tensor + if current > length: + index = [slice(None)] * tensor.ndim + index[-axis] = slice(0, length) + return tensor[tuple(index)] + padding = [0, 0] * (axis - 1) + [0, length - current] + return F.pad(tensor, padding) + + def __call__(self, batch): + lengths = [item[1].shape[0] for item in batch] + frames = max(lengths) if self.pad else min(lengths) + samples = frames * self.samples_per_frame + + audio = torch.stack([self._to_length(item[0], samples, 1) for item in batch]) + visual = torch.stack([self._to_length(item[1], frames, 3) for item in batch]) + vap = torch.stack([self._to_length(item[2], frames, 2) for item in batch]) + vad = torch.stack([self._to_length(item[3], frames, 1) for item in batch]) + scored = torch.stack([self._to_length(item[4], frames, 1) for item in batch]) + return audio, visual, vap, vad, scored, [item[5] for item in batch] + + +class FrameBudgetBatchSampler(Sampler): + """Group similar-length chunks into batches of roughly constant frame count. + + Holding `batch * frames` near a budget keeps three things stable that a + fixed batch size does not: activation memory, the `B*T` population the + visual frontend's batch norm sees, and the per-frame gradient noise of the + mean-reduced losses. The batch size is clamped because the contrastive term + is a `B`-way classification whose difficulty would otherwise swing with the + chunk length. + """ + + def __init__( + self, + lengths, + frame_budget: int, + min_batch: int = 4, + max_batch: int = 16, + shuffle: bool = True, + drop_last: bool = False, + seed: int = 0, + num_replicas: int = 1, + rank: int = 0, + ): + if min_batch < 1 or max_batch < min_batch: + raise ValueError("min_batch must be positive and not exceed max_batch") + self.lengths = np.asarray(lengths, dtype=np.int64) + self.frame_budget = frame_budget + self.min_batch = min_batch + self.max_batch = max_batch + self.shuffle = shuffle + self.drop_last = drop_last + self.seed = seed + self.num_replicas = num_replicas + self.rank = rank + self.epoch = 0 + self._batches = self._build() + + def set_epoch(self, epoch: int) -> None: + self.epoch = epoch + self._batches = self._build() + + def _build(self) -> list[list[int]]: + order = np.argsort(-self.lengths, kind="stable") + batches = [] + cursor = 0 + while cursor < len(order): + longest = int(self.lengths[order[cursor]]) + size = min( + max(self.frame_budget // max(longest, 1), self.min_batch), self.max_batch + ) + batch = order[cursor : cursor + size] + cursor += size + if self.drop_last and len(batch) < self.min_batch: + continue + batches.append([int(index) for index in batch]) + + if self.shuffle: + rng = random.Random(self.seed + self.epoch) + rng.shuffle(batches) + if self.num_replicas > 1: + # Every rank must run the same number of steps or DDP deadlocks. + remainder = len(batches) % self.num_replicas + if remainder: + batches += batches[: self.num_replicas - remainder] + batches = batches[self.rank :: self.num_replicas] + return batches + + def __iter__(self): + return iter(self._batches) + + def __len__(self) -> int: + return len(self._batches) diff --git a/third_party/MuVAP/data/media.py b/third_party/MuVAP/data/media.py new file mode 100644 index 0000000..b208fa5 --- /dev/null +++ b/third_party/MuVAP/data/media.py @@ -0,0 +1,211 @@ +"""Canonical media conversion shared by every MuVAP+ dataset. + +Dataset adapters resolve paths and metadata only. Signal conversion belongs +here so AVA, WASD, MSDWild, Fisher, and AVCC cannot silently diverge. +""" + +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import torch +import torch.nn.functional as F +import torchaudio +import torchaudio.functional as AF + + +@dataclass(frozen=True) +class AudioSpec: + sample_rate: int = 16_000 + mono: bool = True + peak_limit: float = 1.0 + 1e-4 + # Out-of-range float PCM is a scaling bug in most corpora, so it raises by + # default. Set this only for a source known to store real audio that + # overshoots, where clipping is the same thing any encoder would do. + clamp_to_peak: bool = False + + +@dataclass(frozen=True) +class VisualSpec: + height: int = 112 + width: int = 112 + + +CANONICAL_AUDIO = AudioSpec() +CANONICAL_VISUAL = VisualSpec() +ASD_TARGET_RMS_DBFS = -25.0 +ASD_MAX_GAIN_DB = 20.0 +ASD_PEAK_LIMIT = 0.99 + + +def _integer_audio_to_float(array: np.ndarray) -> np.ndarray: + info = np.iinfo(array.dtype) + scale = float(max(abs(info.min), abs(info.max))) + return array.astype(np.float32) / scale + + +def _channels_first(array: np.ndarray) -> np.ndarray: + """Canonicalize common `[N]`, `[C,N]`, and legacy `[N,C]` layouts.""" + if array.ndim == 1: + return array[None, :] + if array.ndim != 2: + raise ValueError(f"NumPy audio must have 1 or 2 dimensions, got {array.shape}") + if array.shape[0] <= 8: + return array + if array.shape[1] <= 8: + return array.T + raise ValueError( + f"cannot infer audio channel axis for shape {array.shape}; expected [C,N] or [N,C]" + ) + + +def canonicalize_waveform( + waveform: torch.Tensor, spec: AudioSpec = CANONICAL_AUDIO +) -> torch.Tensor: + """Return `[channels, samples]` float32 PCM without changing its gain.""" + if waveform.ndim == 1: + waveform = waveform.unsqueeze(0) + if waveform.ndim != 2: + raise ValueError( + f"audio must have 1 or 2 dimensions, got {tuple(waveform.shape)}" + ) + + waveform = waveform.to(torch.float32) + if not torch.isfinite(waveform).all(): + raise ValueError("audio contains NaN or infinite values") + if waveform.numel() and waveform.abs().max().item() > spec.peak_limit: + if not spec.clamp_to_peak: + raise ValueError( + "floating-point audio exceeds [-1, 1]; convert integer PCM to normalized " + "float during preprocessing instead of applying per-clip normalization" + ) + waveform = waveform.clamp(-1.0, 1.0) + if spec.mono and waveform.shape[0] > 1: + waveform = waveform.mean(dim=0, keepdim=True) + return waveform.contiguous() + + +@dataclass(frozen=True) +class WaveformStats: + """Level of a whole recording, measured after DC removal. + + Packed shards hold unnormalized audio, so these travel in the index and let + a reader reproduce the exact per-recording gain from any slice of it. Gain + derived from a slice would drift with the content of that slice. + """ + + mean: float + rms: float + peak: float + + +def waveform_stats(waveform: torch.Tensor) -> WaveformStats: + mean = waveform.mean().item() if waveform.numel() else 0.0 + centred = waveform - mean + return WaveformStats( + mean=mean, + rms=centred.square().mean().sqrt().item() if centred.numel() else 0.0, + peak=centred.abs().max().item() if centred.numel() else 0.0, + ) + + +def normalization_gain( + stats: WaveformStats, + target_rms_dbfs: float = ASD_TARGET_RMS_DBFS, + max_gain_db: float = ASD_MAX_GAIN_DB, + peak_limit: float = ASD_PEAK_LIMIT, +) -> float: + if stats.rms <= 1e-8 or stats.peak <= 1e-8: + return 1.0 + return min( + 10 ** (target_rms_dbfs / 20) / stats.rms, + 10 ** (max_gain_db / 20), + peak_limit / stats.peak, + ) + + +def apply_normalization( + waveform: torch.Tensor, + stats: WaveformStats, + target_rms_dbfs: float = ASD_TARGET_RMS_DBFS, + max_gain_db: float = ASD_MAX_GAIN_DB, + peak_limit: float = ASD_PEAK_LIMIT, +) -> torch.Tensor: + """Centre and scale a slice using its whole recording's statistics.""" + gain = normalization_gain(stats, target_rms_dbfs, max_gain_db, peak_limit) + return ((waveform - stats.mean) * gain).contiguous() + + +def normalize_waveform( + waveform: torch.Tensor, + target_rms_dbfs: float = ASD_TARGET_RMS_DBFS, + max_gain_db: float = ASD_MAX_GAIN_DB, + peak_limit: float = ASD_PEAK_LIMIT, +) -> tuple[torch.Tensor, float]: + stats = waveform_stats(waveform) + gain = normalization_gain(stats, target_rms_dbfs, max_gain_db, peak_limit) + return ((waveform - stats.mean) * gain).contiguous(), gain + + +def load_waveform( + path: str | Path, + spec: AudioSpec = CANONICAL_AUDIO, + *, + npy_sample_rate: int | None = None, +) -> tuple[torch.Tensor, int]: + """Load file or NumPy PCM into the canonical amplitude representation.""" + path = Path(path) + if path.suffix.lower() == ".npy": + array = _channels_first(np.load(path, allow_pickle=False)) + if np.issubdtype(array.dtype, np.integer): + array = _integer_audio_to_float(array) + waveform = torch.from_numpy(np.asarray(array)) + source_rate = npy_sample_rate + if source_rate is None: + raise ValueError(f"sample rate is required for NumPy audio: {path}") + else: + # torchaudio returns normalized float PCM for integer WAV/FLAC sources. + waveform, source_rate = torchaudio.load(path, normalize=True) + # Lossy codecs such as AAC can decode slightly outside nominal PCM bounds. + waveform = waveform.clamp(-1.0, 1.0) + + waveform = canonicalize_waveform(waveform, spec) + if source_rate != spec.sample_rate: + waveform = AF.resample(waveform, source_rate, spec.sample_rate) + return waveform.contiguous(), spec.sample_rate + + +def slice_audio( + waveform: torch.Tensor, + sample_rate: int, + *, + start_time: float | None = None, + end_time: float | None = None, + target_samples: int | None = None, +) -> torch.Tensor: + """Keep timeline origin: crop and zero-pad only at the right boundary.""" + start = round((start_time or 0.0) * sample_rate) + end = round(end_time * sample_rate) if end_time is not None else waveform.shape[-1] + if start < 0 or end < start: + raise ValueError("audio time range is invalid") + waveform = waveform[..., start:end] + + if target_samples is not None: + if target_samples < 0: + raise ValueError("target_samples cannot be negative") + waveform = waveform[..., :target_samples] + waveform = F.pad(waveform, (0, max(0, target_samples - waveform.shape[-1]))) + return waveform.contiguous() + + +def validate_visual( + array: np.ndarray, spec: VisualSpec = CANONICAL_VISUAL +) -> np.ndarray: + """Validate the tensor contract expected by `VisualEncoder`.""" + if array.ndim != 3 or tuple(array.shape[1:]) != (spec.height, spec.width): + raise ValueError( + f"visual must be [T, {spec.height}, {spec.width}], got {array.shape}" + ) + if array.dtype != np.uint8: + raise ValueError(f"visual must be uint8 in [0, 255], got {array.dtype}") + return array diff --git a/third_party/MuVAP/data/packed_asd_dm.py b/third_party/MuVAP/data/packed_asd_dm.py new file mode 100644 index 0000000..89172b3 --- /dev/null +++ b/third_party/MuVAP/data/packed_asd_dm.py @@ -0,0 +1,110 @@ +"""Lightning data module for canonical packed ASD datasets.""" + +import torch.distributed as dist +from lightning.pytorch import LightningDataModule +from torch.utils.data import DataLoader + +from data.augment import AugmentConfig +from data.dataloaders.packed import ( + ASDCollator, + FrameBudgetBatchSampler, + PackedASDDataset, +) + + +class PackedASDDataModule(LightningDataModule): + def __init__(self, cfg: dict, projection_window): + super().__init__() + self.cfg = cfg + self.projection_window = projection_window + asd = cfg["asd"] + self.num_workers = asd["num_workers"] + self.train_window = asd.get("train_window", 250) + self.frame_budget = asd.get("frame_budget", 2048) + self.min_batch = asd.get("min_batch", 4) + self.max_batch = asd.get("max_batch", 16) + self.seed = asd.get("seed", 0) + # Validation is forward-only, so it affords a larger budget than training. + self.val_frame_budget = asd.get("val_frame_budget", 4 * self.frame_budget) + self.val_max_batch = asd.get("val_max_batch", 64) + self.augment = AugmentConfig.from_cfg(asd.get("augment")) + + def _replicas(self): + if dist.is_available() and dist.is_initialized(): + return dist.get_world_size(), dist.get_rank() + return 1, 0 + + def setup(self, stage: str | None = None): + packed = self.cfg["packed_asd"] + if stage in ("fit", None): + # The context prefix is exactly the projection window's history, so + # every scored frame can see the past its target depends on. + self.train_dataset = PackedASDDataset( + packed["train"], + projection_window=self.projection_window, + window=self.train_window, + context=self.projection_window.hist_frames, + jitter=True, + augment=self.augment, + ) + if stage in ("fit", "validate", "test", None): + # Validation runs whole tracks so every native frame gets a score. + self.val_dataset = PackedASDDataset( + packed["val"], projection_window=self.projection_window + ) + + def train_dataloader(self): + replicas, rank = self._replicas() + sampler = FrameBudgetBatchSampler( + self.train_dataset.chunk_lengths(), + frame_budget=self.frame_budget, + min_batch=self.min_batch, + max_batch=self.max_batch, + shuffle=True, + drop_last=True, + seed=self.seed, + num_replicas=replicas, + rank=rank, + ) + return DataLoader( + self.train_dataset, + batch_sampler=sampler, + num_workers=self.num_workers, + collate_fn=ASDCollator(), + pin_memory=True, + persistent_workers=self.num_workers > 0, + prefetch_factor=2 if self.num_workers > 0 else None, + ) + + def val_dataloader(self): + """Batch whole tracks by grouping similar lengths and padding the rest. + + Evaluation needs every frame of every track, so the batch cannot be + cropped to its shortest member the way training's is. Grouping by + length first keeps the padding that replaces it small, which is what + makes batching worth doing at all. + + Lightning's automatic sampler injection is off because the train loader + shards itself, so this shards too. The last rank is padded by repeating + batches; the task de-duplicates by entity before scoring. + """ + replicas, rank = self._replicas() + sampler = FrameBudgetBatchSampler( + self.val_dataset.chunk_lengths(), + frame_budget=self.val_frame_budget, + min_batch=1, + max_batch=self.val_max_batch, + shuffle=False, + drop_last=False, + num_replicas=replicas, + rank=rank, + ) + return DataLoader( + self.val_dataset, + batch_sampler=sampler, + num_workers=self.num_workers, + collate_fn=ASDCollator(pad=True), + pin_memory=True, + persistent_workers=self.num_workers > 0, + prefetch_factor=2 if self.num_workers > 0 else None, + ) diff --git a/third_party/MuVAP/data/utils.py b/third_party/MuVAP/data/utils.py new file mode 100644 index 0000000..8dc3604 --- /dev/null +++ b/third_party/MuVAP/data/utils.py @@ -0,0 +1,64 @@ +import os +from glob import glob + +import torch + +from data.media import AudioSpec, load_waveform, slice_audio + + +def load_audio( + audio_path, + target_sr=16000, + mono=True, + start_time=None, + end_time=None, + target_samples=None, +): + """Load canonical float PCM while preserving the clip's time origin. + + Short clips are padded on the right and long clips are cropped on the + right. No peak, RMS, or dataset-specific volume normalization is applied. + """ + spec = AudioSpec(sample_rate=target_sr, mono=mono) + waveform, sr = load_waveform(audio_path, spec) + waveform = slice_audio( + waveform, + sr, + start_time=start_time, + end_time=end_time, + target_samples=target_samples, + ) + + duration = waveform.shape[1] / sr + + return waveform, duration + + +def load_fisher_vad(vad_path, conv_id, nframes, frame_hz=25): + + pattern = os.path.join(vad_path, f"{conv_id}_*_words.txt") + vad_files = sorted(glob(pattern)) + + vad_labels = [] + for vad_file in vad_files: + labels = torch.zeros(nframes) + with open(vad_file, "r") as f: + for line in f: + parts = line.strip().split() + start = float(parts[0]) + end = float(parts[1]) + start_idx = int(round(start * frame_hz)) + end_idx = int(round(end * frame_hz)) + start_idx = max(0, min(nframes, start_idx)) + end_idx = max(0, min(nframes, end_idx)) + labels[start_idx:end_idx] = 1 + vad_labels.append(labels) + + return torch.stack(vad_labels, dim=0) + + +def load_lit_state_dict(model_path): + device = "cuda" if torch.cuda.is_available() else "cpu" + checkpoint = torch.load(model_path, map_location=device) + state_dict = checkpoint.get("state_dict", checkpoint) + return {k.replace("model.", ""): v for k, v in state_dict.items()} diff --git a/third_party/MuVAP/data/vap_dm.py b/third_party/MuVAP/data/vap_dm.py new file mode 100644 index 0000000..0b0c22a --- /dev/null +++ b/third_party/MuVAP/data/vap_dm.py @@ -0,0 +1,100 @@ +from pathlib import Path + +from lightning.pytorch import LightningDataModule +from torch.utils.data import DataLoader + +from data.dataloaders.fisher import Fisher, FisherEvent +from models.vap import audio_channels, needs_vad + + +class VAPDataModule(LightningDataModule): + """Fisher segments for fitting, annotated turn-taking events for testing.""" + + def __init__(self, cfg, projection): + super().__init__() + vap_cfg = cfg["vap"] + self.root = Path(cfg["fisher_path"]) + self.projection = projection + self.batch_size = vap_cfg["batch_size"] + self.num_workers = vap_cfg["num_workers"] + self.channels = audio_channels(vap_cfg) + self.with_vad = needs_vad(vap_cfg) + self.swap_channels = bool(vap_cfg.get("swap_channels", False)) + self.source = str(vap_cfg.get("source", "npy")).lower() + + def _split(self, name): + path = self.root / f"{name}.txt" + if not path.exists(): + raise FileNotFoundError( + f"missing Fisher split {path}; run preprocess/vap/00_prep_fisher.py first" + ) + return path + + def _event_splits(self, names): + paths = [self.root / f"{name}_events.txt" for name in names] + found = [path for path in paths if path.exists()] + if not found: + raise FileNotFoundError( + f"none of {[p.name for p in paths]} exist; run preprocess/vap/02_prep_event.py first" + ) + return found + + def setup(self, stage=None): + if stage in ("fit", None): + self.train_dataset = self._segments("train", swap=self.swap_channels) + self.val_dataset = self._segments("val") + + if stage in ("test", None): + self.tune_dataset = FisherEvent( + self.root, + self._event_splits(["train", "val"]), + channels=self.channels, + source=self.source, + with_vad=self.with_vad, + frame_hz=self.projection.frame_hz, + ) + self.test_dataset = FisherEvent( + self.root, + self._event_splits(["test"]), + channels=self.channels, + source=self.source, + with_vad=self.with_vad, + frame_hz=self.projection.frame_hz, + ) + + def _segments(self, name, swap=False): + return Fisher( + self.root, + self._split(name), + self.projection, + channels=self.channels, + frame_hz=self.projection.frame_hz, + swap_channels=swap, + source=self.source, + with_vad=self.with_vad, + ) + + def _loader(self, dataset, shuffle=False): + workers = self.num_workers + return DataLoader( + dataset, + batch_size=self.batch_size, + shuffle=shuffle, + num_workers=workers, + pin_memory=True, + persistent_workers=bool(workers), + prefetch_factor=2 if workers else None, + drop_last=shuffle, + ) + + def train_dataloader(self): + return self._loader(self.train_dataset, shuffle=True) + + def val_dataloader(self): + return self._loader(self.val_dataset) + + def tune_dataloader(self): + return self._loader(self.tune_dataset) + + def test_dataloader(self): + return self._loader(self.test_dataset) diff --git a/third_party/MuVAP/docs/data_alignment.md b/third_party/MuVAP/docs/data_alignment.md new file mode 100644 index 0000000..88f57aa --- /dev/null +++ b/third_party/MuVAP/docs/data_alignment.md @@ -0,0 +1,96 @@ +# Audio-visual data contract + +## Canonical timeline + +MuVAP+ uses a **25 Hz model timeline**, because the CPC stack emits one audio +embedding per 640 samples at 16 kHz. Original videos are not transcoded or +rewritten to 25 fps. + +For model frame `k`, the timestamp is the centre of its interval: + +```text +t(k) = (k + 0.5) / 25 seconds +``` + +The visual face and label are selected from the native source interval that +contains `t(k)`: + +```text +source_index(k) = floor(t(k) * source_fps) +``` + +The exact same index is used for the face and its label. This prevents label +drift. At 30 fps some source frames are skipped; below 25 fps some are repeated. + +**This happens in the loader, not in the pack.** Packed shards store the native +frame rate, so the sampling above is applied when a batch is built and can be +changed - or replaced with a different `model_fps` - without re-packing. + +## Evaluation on the native timeline + +Metrics defined on the source annotations — the official AVA mAP above all — +need one score per native frame, not per model frame. `model_indices` reverses +the mapping by asking which model interval contains each source-frame centre: + +```text +model_index(j) = floor(((j + 0.5) / source_fps) * 25) +``` + +The result always has exactly `source_frames` entries, so every groundtruth row +receives a score at any frame rate. Below 25 fps this is an exact inverse of +`source_indices`. Above it, several source frames share a model frame, which is +information the 25 Hz timeline genuinely does not carry. Only the final source +frame can fall past the last complete model interval; it is clamped to that +interval, and no interior frame is ever mis-scored. + +## Audio + +Audio is loaded directly from WAV/FLAC (or NumPy waveforms), mixed to +mono, and resampled to 16 kHz for CPC. Packed shards keep the full span the face +frames cover, at the recording's own gain. + +The loader takes it from there. A batch of `T` model frames always carries +exactly `T * 640` samples, right-padded with silence if the recording is short. +Loudness normalization to -25 dBFS RMS is applied at that point, using level +statistics measured over the **whole recording** and stored in the index - +computing gain from a chunk alone would make a sample's level depend on which +chunk it landed in. + +## Required manifest fields + +Each row must resolve to: + +| Field | Meaning | +|---|---| +| `sample_id` | Stable, unique face-track/clip identifier | +| `audio_path` | Recording or clip audio path | +| `visual_path` | NumPy face array shaped `[source_frames, H, W]` | +| `source_fps` | Native FPS used by faces and labels | +| `labels` | One active-speaker label per native face frame | +| `start_sec` | Clip start on the audio recording timeline; default `0` only when audio is already clip-local | + +The canonical file is a tab-separated text file with a header. `labels` is a +JSON array; paths may be absolute or relative to the manifest directory: + +```text +sample_id\taudio_path\tvisual_path\tsource_fps\tlabels\tstart_sec +clip_001\taudio/recording.wav\tfaces/clip_001.npy\t29.97\t[0,0,1,1]\t12.48 +``` + +`len(labels)` and the face-frame count must match. Preprocessing rejects a +sample where they disagree rather than trimming to the shorter one: a length +mismatch means the face track and the annotation disagree about the clip, which +would silently shift every later label. + +## Dataset policy + +- AVA, WASD, and MSDWild use the same alignment implementation. +- FPS is never ignored and must never be replaced with a global constant. +- Packs are decode-only. Any transform whose parameters belong to the model - + frame rate, loudness, crop duration - happens in the loader, so a pack stays + usable by code that does not share those conventions. +- Labels are aligned before projection-window targets are computed. +- Train-time batching may crop examples to a shared duration, but audio, faces, + VAD labels, and projected labels must all use that same duration. +- A WASD face track using full-recording audio must provide `start_sec`. It is + mathematically impossible to recover that offset from FPS and labels alone. diff --git a/third_party/MuVAP/docs/preprocessing.md b/third_party/MuVAP/docs/preprocessing.md new file mode 100644 index 0000000..3b38c0d --- /dev/null +++ b/third_party/MuVAP/docs/preprocessing.md @@ -0,0 +1,219 @@ +# Canonical ASD preprocessing + +All datasets are converted to the same versioned `muvap-asd-mmap` format. +Dataset adapters only resolve source metadata; `preprocess.asd.writer` performs +all media conversion and alignment. + +## Packed format + +A pack is a **decoded mirror of the source corpus, not a model-ready tensor +set**. Faces and labels keep the native frame rate, audio keeps its full +duration and its original gain. Everything MuVAP+ specific - sampling onto the +25 Hz model timeline, loudness normalization, chunking - happens in +`data.dataloaders.packed` when a batch is built. That split is what lets a pack +serve a model whose conventions differ from this one. + +Each corpus keeps its pack beside its own media, so `--output` defaults to +`/packed/` and a dataset stays self-contained: + +```text +AVA/packed/train/ + dataset.json + shard-000000/ + faces.npy # uint8 [native_frames, 112, 112], grayscale + audio.npy # int16 mono 16 kHz, unnormalized, full clip duration + labels.npy # uint8, one label per NATIVE frame + index.json # offsets, source_fps, source_frames, audio level statistics +``` + +Each shard is approximately 1 GB. The arrays are uncompressed and opened with +`mmap_mode="r"`, so workers and DDP ranks share the operating-system page cache +without HDF5 handles or a per-process copy. Copy complete dataset directories +to node-local storage before training if the shared filesystem is the +bottleneck. + +The per-sample invariant is: + +```text +source_frames faces == source_frames labels +audio_samples == round(source_frames / source_fps * 16000) +``` + +### What the reader adds + +`index.json` carries `source_fps`, `source_frames`, and the recording's +`audio_mean`, `audio_rms`, and `audio_peak` measured after DC removal. Those +statistics are what let the loader reproduce a whole-recording normalization +from any slice: gain computed from a chunk alone would drift with that chunk's +content. The loader removes DC, scales to -25 dBFS RMS, caps gain at 20 dB and +the peak at 0.99, and samples faces and labels onto the 25 Hz timeline with one +shared index so a label stays attached to the face it describes. + +The one convention still baked into the pack is the face crop: 112x112 +grayscale, recorded in `dataset.json` as `face_shape` and `face_color`. It is +the ASD community's standard geometry rather than a MuVAP+ choice, but a reader +should check it rather than assume it. + +The size is a constant in two places, not a config key: `FACE_SIZE` in +`preprocess/asd/schema.py`, which the writer resizes every JPEG face to, and +`VisualSpec` in `data/media.py`, which is what the loader and the model expect. +Changing the geometry means editing both and re-running preprocessing - a pack +written at one size cannot be read at another. + +## AVA + +```bash +python -m preprocess.asd.prepare ava \ + --loader /data/AVA/csv/train_loader.csv \ + --visual-root /data/AVA/clips_videos \ + --audio-root /data/AVA/clips_audios \ + --split train # writes /data/AVA/packed/train +``` + +The adapter reads nested JPEG face directories and clip-local WAV files. +Missing inputs fail preprocessing. `--skip-missing` exists only for deliberately +partial development copies. + +## WASD + +WASD ships the same layout and the same loader schema as AVA, so it uses the +same reader: + +```bash +python -m preprocess.asd.prepare wasd \ + --loader /data/WASD/csv/train_loader.csv \ + --visual-root /data/WASD/clips_videos \ + --audio-root /data/WASD/clips_audios \ + --split train # writes /data/WASD/packed/train +``` + +Two differences from AVA are worth knowing when training on both. WASD labels +are strictly binary - `SPEAKING_NOT_AUDIBLE` never occurs - so the rule that +inaudible speech counts as negative is something only AVA teaches. And every +face track in a WASD clip shares one clip-level WAV, so audio alone cannot tell +its speakers apart; any augmentation that mixes in a second voice must draw it +from a different clip. + +## MSDWild + +Create 30-second native-FPS clips with 10 seconds of overlap directly from the +paired raw MP4 and bbox CSV files. The first window starts at the first frame +containing any face. Tracks must be present in at least 90% of a complete +window, which removes transient detections before face cropping. +Incomplete final windows are discarded rather than padded, because the model +does not use a padding mask and would otherwise learn from artificial silence. +Speaking labels come from `all.rttm`, which ships separately from the +bounding-box release and must be downloaded before this script can run; the +final bbox CSV column is not an activity label. `prepare.py` has no MSDWild +path - there is no loader TSV to read. + +```bash +python preprocess/asd/prep_msdwild.py \ + --input /data/MSDWild/msdwild_boundingbox_labels \ + --output /data/MSDWild/preprocessed +``` + +The script uses the shared audio loader and face cropper and writes the same +canonical memory-mapped shards as AVA and WASD. No intermediate per-clip files +or loader CSV are created. + +## AVCC + +AVCC uses the canonical headered TSV because conversation corpora vary in +their crop-generation layout: + +```text +sample_id\tvisual_path\taudio_path\tsource_fps\tlabels\tstart_sec\tvisual_kind\taudio_sample_rate +conv1_spk1\tfaces/conv1_spk1.npy\taudio/conv1.wav\t25.0\t[0,1,1,0]\t12.0\tnpy\t16000 +``` + +Paths may be absolute or relative to `--root`: + +```bash +python -m preprocess.asd.prepare avcc \ + --loader /data/AVCC/asd_manifest.tsv \ + --root /data/AVCC \ + --output /data/packed/avcc_train +``` + +Each row represents one speaker face track. AVCC conversation grouping, +bounding boxes, and speaker identities remain separate MuVAP metadata; its ASD +embeddings must be extracted from these canonical samples. + +## Lightning + +```yaml +packed_asd: + train: + - /scratch/local/muvap/AVA/packed/train + - /scratch/local/muvap/WASD/packed/train + - /scratch/local/muvap/WASD/packed/val + val: + - /scratch/local/muvap/AVA/packed/val +``` + +Use `PackedASDDataModule`. Start with four workers per GPU, +`persistent_workers=True`, `prefetch_factor=2`, and pinned memory, then measure +0/2/4/8 workers on the actual allocation. + +### Training batches + +Track lengths are wildly uneven - AVA's median is 64 frames against WASD's 746 - +and the collator crops a batch to its shortest member. Batching at random +therefore threw away most of the data: measured on the AVA pack, a shuffled +`batch_size: 8` used **36%** of its frames per epoch. + +Two settings fix that, and they only work together: + +- `train_window` splits any track longer than it into near-equal chunks, so one + epoch still covers every frame and no chunk is much longer than another. +- `frame_budget` builds each batch from similar-length chunks up to a target + `batch * frames`, clamped by `min_batch` and `max_batch`. + +Together they reach **99.95%** frame utilisation on AVA and 99.98% on AVA plus +WASD, which is why the model still needs no padding mask. Holding +`batch * frames` roughly constant also stabilises activation memory, the `B*T` +population the visual frontend's batch norm sees, and the per-frame gradient +noise of the mean-reduced losses. The clamp exists because the contrastive term +is a `B`-way classification whose difficulty would otherwise swing with the +chunk length. + +### Context prefix + +A chunk boundary creates a second, subtler problem. The model is causal, so a +frame near the start of a chunk can only attend to the few frames before it - +yet its projection target depends on 2 s of history. Without help, 31.6% of +scored frames on AVA plus WASD are asked to predict from history the encoder was +never shown. + +The loader therefore reads `hist_frames` extra frames *before* each chunk and +marks them unscored. They provide history and never contribute to a loss, so: + +- every frame is still scored exactly once per epoch, +- the artificial starvation disappears entirely - the remaining 21.5% are all + real track starts, where the missing history is genuine and the model has to + cope with it at inference too, +- the cost is 1.15x compute and no extra storage. + +Overlapping the chunks instead would duplicate the data on disk *and* keep +training on the truncated copies, since it adds well-formed frames beside the +starved ones rather than removing them. + +Validation ignores all of this: it runs whole tracks at batch size one, so every +native frame receives a score. + +Because the train loader shards itself across ranks, `Trainer` is constructed +with `use_distributed_sampler=False` and validation gets an explicit +`DistributedSampler`. + +Validate every completed dataset and compare amplitude distributions before +training: + +```bash +python -m preprocess.asd.validate /data/AVA/packed/train /data/WASD/packed/train +``` + +The report shows the stored source levels beside the levels the loader produces +after normalization. Large source differences between corpora are expected; +large differences in the *normalized* RMS usually mean peak limiting, +near-silence, or malformed input that should be inspected. diff --git a/third_party/MuVAP/docs/vap.md b/third_party/MuVAP/docs/vap.md new file mode 100644 index 0000000..2039ee7 --- /dev/null +++ b/third_party/MuVAP/docs/vap.md @@ -0,0 +1,232 @@ +# The VAP module + +Voice Activity Projection setups share one training program. They differ in +config only - how many audio channels the model reads, which codebook +`ProjectionWindow` projects the VAD onto, and which frozen audio frontend runs +in front - so any result difference between them comes from the setup rather +than from a different pipeline. + +| Experiment | Config | Frontend | Channels | `mode` | Classes | +| --- | --- | --- | --- | --- | --- | +| Speaker-based VAP | `vap_speaker.yaml` | CPC | 2 (one per speaker) | `speaker_based` | 256 | +| Role-based VAP | `vap_role.yaml` | CPC | 1 (downmix) | `role_relative` | 136 | + +Hold and shift come from the logits in both setups: `role_relative` names the +roles in the class, and `speaker_based` names the channels, which needs +`prev_spk` to say which one held the floor. + +## What each setup is + +**Role-Based** is the MuVAP+ default. Both speakers are mixed into one channel, +so the model is never told who is who. The class names the *roles* - current +speaker and next speaker - through an order-invariant pair codebook over two +history and two future bins. Hold and shift fall straight out of the logits +because the class already says which role continues. + +**Speaker-based VAP** reproduces Ekstedt & Skantze (2022): one audio channel per +speaker, a shared encoder and causal self-attention stack applied to each +channel, cross-attention between the two channels, and one head over their +concatenation. The projection window works per channel - one channel per row of +the code - giving `2 speakers x 4 future bins = 8 bits = 256` classes. Because +the classes name *channels* rather than roles, reading hold and shift needs to +know which channel held the floor, which is why the event manifest's +`prev_spk` column is carried into the test step. + +## Audio frontend + +`vap.audio_encoder` selects the frozen frontend. Both land on the same 25 Hz +grid the projection window labels, so they are drop-in alternatives: + +| | `cpc` (default) | `mimi` | +| --- | --- | --- | +| Weights | CPC, LibriLight 60k | `kyutai/mimi`, from Moshi | +| Frozen params | 1.8M | 37.8M | +| Feature dim | 256 | 512 | +| Native rate | 16 kHz | 24 kHz (resampled causally inside the encoder) | + +```yaml +audio_encoder: + name: mimi + tap: encoder # 25 Hz. `latent` is Mimi's own 12.5 Hz rate +``` + +Mimi is causal by construction - its SEANet convolutions left-pad and its +encoder transformer uses a sliding causal mask. Getting there causally +needs care in one place: `torchaudio`'s resampler convolves with a symmetric +sinc, so a plain 16 -> 24 kHz resample leaks about 9 output samples of future. +`causal_resample_16k_to_24k` delays the signal by 1 ms instead, which removes +the leak at the cost of a 1 ms lag against a 40 ms frame. + +The `latent` tap is refused at build time under a 25 Hz codebook rather than +failing later on a shape mismatch. + +Mimi's 512-dim features are bridged onto `d_model` by the model's encoder +projection, so `d_model` need not change. + +## Positional encoding + +`vap.temporal.pos_encoding` selects how position reaches attention. It never +changes what a frame can see: every temporal stack is causal, frame `t` attends +to frames `<= t` only, so a trained model stays usable as a streaming predictor. + +- `alibi` - per-head linear distance penalty on the logits (the previous default). +- `rope` - **causal rotary position embedding**. Query and key are rotated by an + angle proportional to the absolute frame index, so the attention logit depends + only on the distance between two frames. Under the causal mask that distance is + never negative, and no future frame can leak in through the encoding. A query + shorter than its key stream is offset to the end of that stream, so streaming + inference with a cached history lines up with training. +- `sinusoidal` - an additive input-side embedding, for ablation. +- `none` - no positional information beyond the causal mask itself. + +`ffn` (`mlp` / `swiglu`) and `norm` (`layer` / `rms`) are independent knobs on +the same stack. + +Switch encoding without editing the file: + +```bash +python train_vap.py --config config/yaml/vap_role.yaml --set vap.temporal.pos_encoding=rope +``` + +## Data + +Two paths in each config are the only machine-specific values: + +```yaml +fisher_path: data/fisher # the corpus +checkpoint_dir: runs/vap # checkpoints and logs +``` + +`checkpoint_dir` is also the trainer's root, so Lightning's logs land beside +the checkpoints. + +The corpus tree looks like this: + +```text +fisher/ + p1/ p2/ + audio//.wav # source, stereo + regions_vap//_*_words.txt + seg/audio///.npy # 30 s stereo segment + seg/vad///.npy # 2 x (750 + 2*context) VAD frames + tune/audio/... # 10 s mono event clips + tune/audio_stereo/... # 10 s stereo event clips (original VAP only) + train.txt val.txt test.txt + turn_events.txt + train_events.txt val_events.txt test_events.txt +``` + +**Labels are not stored.** Segments hold the raw two-speaker VAD with +`--context-sec` of extra frames on each side - everything the widest projection +window needs - and the dataloader projects them at load time. Changing +`projection_window.mode` between the three setups is therefore a config edit, +not a re-run of preprocessing. + +Build the tree: + +```bash +python preprocess/vap/00_prep_fisher.py --config config/yaml/vap_role.yaml # segments + VAD +python preprocess/vap/01_create_event.py --config config/yaml/vap_role.yaml # turn events +python preprocess/vap/02_prep_event.py --config config/yaml/vap_role.yaml # mono event clips +``` + +The speaker-based setup evaluates on two-channel clips. `config/yaml/vap_speaker.yaml` +ships with `source: raw`, which reads them straight from the source recording, +so nothing more is needed. To evaluate it from pre-cut clips instead, write a +two-channel copy - it lands in `tune/audio_stereo`, leaving the mono clips +untouched - and set `source: npy`: + +```bash +python preprocess/vap/02_prep_event.py --config config/yaml/vap_role.yaml --channels 2 +``` + +`--splits` defaults to `val test`, which is what evaluation reads. Adding +`train` cuts clips for roughly 2.5 M more events and is only worth it to +enlarge the probe's tuning pool. + +If a projection window needs more context than a segment stores, the dataloader +says so by name and asks for a wider `--context-sec`; the 2 s default covers all +three setups shipped here. + +## Running + +```bash +# Role-Based, the default +python train_vap.py --config config/yaml/vap_role.yaml --name role_alibi + +# Role-Based with causal RoPE +python train_vap.py --config config/yaml/vap_role.yaml --name role_rope \ + --set vap.temporal.pos_encoding=rope + +# Speaker-based VAP, two channels +python train_vap.py --config config/yaml/vap_speaker.yaml --name speaker +``` + +Evaluate a checkpoint without training: + +```bash +python train_vap.py --config config/yaml/vap_role.yaml --test --checkpoint runs/vap_role/role_alibi/best.ckpt +``` + +Add `--wandb` for tracking, `--seed` to change the seed, and `--name` to choose +the checkpoint directory. `--set KEY=VALUE` overrides any existing config key +(dotted path, value parsed as YAML) and can be repeated; it fails loudly on a +key the config does not already define, so a typo cannot silently do nothing. + +## Released models + +Weights are on the Hub at [Haotian-Qi/MuVAP](https://huggingface.co/Haotian-Qi/MuVAP). +`f1_macro` is zero-shot hold/shift on the Fisher events, decoded from the logits +over `p_future`, at the shift prior each codebook defaults to. + +| Release | Codebook | Audio in | Rate | Frontend | f1_macro | bacc | probe f1 | +| --- | --- | --- | --- | --- | --- | --- | --- | +| `vap-speaker-mimi-12hz` | speaker, 256 | 2 ch | 12.5 | Mimi | **0.7776** | 0.7495 | 0.7714 | +| `vap-speaker-mimi` | speaker, 256 | 2 ch | 25 | Mimi | 0.7755 | 0.7464 | 0.7716 | +| `vap-role-mimi` | role, 136 | 1 ch | 25 | Mimi | 0.7589 | 0.7698 | 0.7623 | +| `vap-role-mimi-12hz` | role, 136 | 1 ch | 12.5 | Mimi | 0.7565 | 0.7663 | 0.7570 | +| `vap-speaker-mono-mimi` | speaker, 256 | 1 ch + VAD | 25 | Mimi | 0.7517 | 0.7216 | 0.7619 | +| `vap-speaker-mono-mimi-12hz` | speaker, 256 | 1 ch + VAD | 12.5 | Mimi | 0.7484 | 0.7183 | 0.7585 | +| `vap-speaker-cpc-50hz` | speaker, 256 | 2 ch | 50 | CPC | 0.7444 | 0.7129 | 0.7520 | +| `vap-role-cpc-50hz` | role, 136 | 1 ch | 50 | CPC | 0.7358 | 0.7378 | 0.7382 | +| `vap-speaker-cpc` | speaker, 256 | 2 ch | 25 | CPC | 0.7310 | 0.6983 | 0.7508 | +| `vap-role-cpc` | role, 136 | 1 ch | 25 | CPC | 0.7289 | 0.7303 | 0.7304 | +| `vap-speaker-mono-cpc-50hz` | speaker, 256 | 1 ch + VAD | 50 | CPC | 0.7148 | 0.6823 | 0.7398 | + +Read it with three cautions. + +**The prior is a threshold, and it differs by codebook.** Scaling the shift slot +by `s` and renormalising is exactly a decision threshold at `1/(1+s)`, so the +speaker rows sit at 0.5 and the role rows at 0.333. The `bacc` column is where +that shows. Comparisons within a codebook are sound; across codebooks they are +comparing operating points as much as models. `probe f1`, a logistic probe on +the last-frame embedding, is threshold-free and closer to like-for-like. + +**The frontend is the only large effect.** Mimi over CPC is +0.045 on speaker and ++0.030 on role at matched rate. Frame rate is 3-6x smaller: 25 -> 50 Hz on CPC +gains 0.013 (speaker) and 0.007 (role), and on Mimi 25 -> 12.5 Hz costs nothing +measurable - +0.002, -0.002, -0.003 across the three codebooks. Matching a +12.5 Hz downstream system is therefore close to free. + +**The `-mono-` releases need VAD at inference.** They take one mixed channel plus +the voice activity behind each frame, `model(audio, vad)`, which is what supplies +the speaker identity a second channel would carry. Quality depends on the VAD you +feed them. + +## Metrics + +`trainer.test` runs two dataloaders: the tune events fit the linear probe, the +test events are scored. Events are pooled three ways - `all`, `SILENT`, +`ACTIVE` - and three prefixes keep the panels apart: + +- `test//f1_macro`, `test//bacc` - the headline. Hold and shift are + decoded from the logits, zero shot, over `p_future` at the shift prior the + codebook defaults to: 2.0 for `role_relative`, 1.0 for `speaker_based`, whose + rows name channels and so also need the event's `prev_spk`. +- `probe//f1_macro`, `probe//bacc` - a logistic probe on the + last-frame embedding. It bypasses the codebook, so the gap against `test/` + is what the readout costs rather than anything about the model. Being fitted, + it is a diagnostic and not a score. +- `ablation//f1_macro_scale_` - the shift prior swept 1.0 ... 3.0, the + way the VAP papers report it. `speaker_based` gains about 0.045 from a prior + of 2.0; the role setups are already at their optimum and gain 0.001. diff --git a/third_party/MuVAP/models/__init__.py b/third_party/MuVAP/models/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/third_party/MuVAP/models/asd.py b/third_party/MuVAP/models/asd.py new file mode 100644 index 0000000..2876308 --- /dev/null +++ b/third_party/MuVAP/models/asd.py @@ -0,0 +1,60 @@ +import torch +import torch.nn as nn + +from .modules.attention import TransformerModel +from .modules.audio import build_audio_encoder, check_frame_rate, encoder_projection +from .modules.projection import LinearHead +from .modules.visual_encoder import VisualEncoder + +#: The visual frontend samples faces onto this timeline, so the audio frontend +#: has to land on it too or the two streams cannot be fused frame for frame. +MODEL_FRAME_HZ = 25.0 + + +class AudioVisualASD(nn.Module): + def __init__(self, cfg): + super().__init__() + cross_dim = cfg["cross"]["d_model"] + temporal_dim = cfg["temporal"]["d_model"] + if temporal_dim != cross_dim * 2: + raise ValueError("ASD temporal d_model must be twice cross d_model") + + self.audio_encoder = build_audio_encoder(cfg) + check_frame_rate(self.audio_encoder, MODEL_FRAME_HZ, "the ASD visual timeline") + self.encoder_proj = encoder_projection(self.audio_encoder.out_dim, cross_dim) + self.visual_encoder = VisualEncoder() + if self.visual_encoder.out_dim != cross_dim: + raise ValueError( + f"ASD cross d_model must be {self.visual_encoder.out_dim} to match the " + f"visual frontend, got {cross_dim}" + ) + self.norm_audio = nn.LayerNorm(cross_dim) + self.norm_visual = nn.LayerNorm(cross_dim) + self.cross_v2a = TransformerModel(cfg["cross"], cross_attention=True) + self.cross_a2v = TransformerModel(cfg["cross"], cross_attention=True) + self.transformer = TransformerModel(cfg["temporal"]) + self.asd_head = LinearHead(temporal_dim, 6) + self.audio_head = LinearHead(cross_dim, 1) + self.visual_head = LinearHead(cross_dim, 1) + + def forward(self, audio, visual, return_embeddings=False): + audio_embedding = self.norm_audio( + self.encoder_proj(self.audio_encoder(audio)) + ) + visual_embedding = self.norm_visual(self.visual_encoder(visual)) + if audio_embedding.shape[1] != visual_embedding.shape[1]: + raise ValueError( + f"audio and visual lengths differ: {audio_embedding.shape[1]} and " + f"{visual_embedding.shape[1]}" + ) + + audio_cross = self.cross_v2a(audio_embedding, src=visual_embedding) + visual_cross = self.cross_a2v(visual_embedding, src=audio_embedding) + fused = self.transformer(torch.cat([audio_cross, visual_cross], dim=-1)) + projection_logits = self.asd_head(fused) + audio_logits = self.audio_head(audio_cross) + visual_logits = self.visual_head(visual_cross) + + if return_embeddings: + return projection_logits, audio_logits, visual_logits, fused, audio_embedding, visual_embedding + return projection_logits, audio_logits, visual_logits diff --git a/third_party/MuVAP/models/assets/checkpoints/cpc/60k_epoch4-d0f474de.pt b/third_party/MuVAP/models/assets/checkpoints/cpc/60k_epoch4-d0f474de.pt new file mode 100644 index 0000000..104a50c Binary files /dev/null and b/third_party/MuVAP/models/assets/checkpoints/cpc/60k_epoch4-d0f474de.pt differ diff --git a/third_party/MuVAP/models/modules/__init__.py b/third_party/MuVAP/models/modules/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/third_party/MuVAP/models/modules/attention.py b/third_party/MuVAP/models/modules/attention.py new file mode 100644 index 0000000..59a38ff --- /dev/null +++ b/third_party/MuVAP/models/modules/attention.py @@ -0,0 +1,294 @@ +"""Causal transformer stack. + +Frame `t` attends to frames `<= t` only, which is what makes a trained model +usable as a streaming turn-taking predictor. How position reaches attention is +configurable through `pos_encoding`; the visibility rule is not. +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .positional import ( + POS_ENCODINGS, + RotaryEmbedding, + alibi_slopes, + as_additive_mask, + build_input_position_embedding, + causal_alibi_bias, + causal_float_mask, +) + + +class TransformerModel(nn.Module): + """Stack of pre-norm transformer layers over a `[batch, frames, dim]` sequence.""" + + def __init__(self, cfg, cross_attention=False): + super().__init__() + self.dim = cfg["d_model"] + self.dff = cfg["m_ff"] * self.dim + self.num_layers = cfg["n_layers"] + self.num_heads = cfg["n_heads"] + self.activation = cfg["activation"] + self.dropout = cfg["dropout"] + self.cross_attention = cross_attention + self.causal = bool(cfg.get("causal", True)) + self.pos_encoding = str(cfg.get("pos_encoding", "alibi")).lower() + self.ffn_type = str(cfg.get("ffn", "mlp")).lower() + self.norm_type = str(cfg.get("norm", "layer")).lower() + self.max_frames = int(cfg.get("max_frames", 4096)) + if self.pos_encoding not in POS_ENCODINGS: + raise ValueError( + f"pos_encoding must be one of {POS_ENCODINGS}, got {self.pos_encoding!r}" + ) + + self.rope = ( + RotaryEmbedding(self.dim // self.num_heads, max_frames=self.max_frames) + if self.pos_encoding == "rope" + else None + ) + self.input_pos = build_input_position_embedding( + self.pos_encoding, self.dim, self.max_frames + ) + + self._build_layers() + self.apply(self._init_weights) + + def _build_layers(self): + self.layers = nn.ModuleList( + [ + TransformerLayer( + dim=self.dim, + ffn_dim=self.dff, + num_heads=self.num_heads, + ffn_activation=self.activation, + dropout=self.dropout, + cross_attention=self.cross_attention, + pos_encoding=self.pos_encoding, + causal=self.causal, + rope=self.rope, + ffn_type=self.ffn_type, + norm_type=self.norm_type, + ) + for _ in range(self.num_layers) + ] + ) + + def _init_weights(self, module): + if isinstance(module, (nn.Linear, nn.Embedding)): + torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) + if isinstance(module, nn.Linear) and module.bias is not None: + torch.nn.init.zeros_(module.bias) + elif isinstance(module, nn.LayerNorm): + torch.nn.init.zeros_(module.bias) + torch.nn.init.ones_(module.weight) + + def forward(self, x, src=None, mask=None): + if self.input_pos is not None: + x = self.input_pos(x) + for layer in self.layers: + x = layer(x, src, mask=mask) + return x + + +class TransformerLayer(nn.Module): + def __init__( + self, + dim=256, + ffn_dim=768, + num_heads=4, + ffn_activation="GELU", + dropout=0.1, + cross_attention=False, + pos_encoding="alibi", + causal=True, + rope=None, + ffn_type="mlp", + norm_type="layer", + ): + super().__init__() + self.cross_attention = cross_attention + self.dropout = nn.Dropout(dropout) + + self.ln_self_attn = build_norm(norm_type, dim) + self.mha = MultiHeadAttention( + dim=dim, + num_heads=num_heads, + dropout=dropout, + causal=causal, + pos_encoding=pos_encoding, + rope=rope, + ) + + if cross_attention: + self.ln_cross_query = build_norm(norm_type, dim) + self.ln_cross_source = build_norm(norm_type, dim) + self.mha_cross = MultiHeadAttention( + dim=dim, + num_heads=num_heads, + dropout=dropout, + causal=causal, + pos_encoding=pos_encoding, + rope=rope, + ) + + self.ln_ffnetwork = build_norm(norm_type, dim) + self.ffnetwork = ffn_block( + dim, ffn_dim, activation=ffn_activation, dropout=dropout, ffn_type=ffn_type + ) + + def forward(self, x, src=None, mask=None): + z = self.ln_self_attn(x) + x = x + self.dropout(self.mha(Q=z, K=z, V=z, mask=mask)) + + if self.cross_attention and src is not None: + query = self.ln_cross_query(x) + source = self.ln_cross_source(src) + x = x + self.dropout(self.mha_cross(Q=query, K=source, V=source, mask=mask)) + + x = x + self.dropout(self.ffnetwork(self.ln_ffnetwork(x))) + return x + + +class MultiHeadAttention(nn.Module): + """Scaled dot-product attention with a selectable positional encoding. + + `alibi` adds a per-head linear distance penalty to the logits, `rope` + rotates the query and key instead, and every other setting leaves position + to the input embedding. Causality is applied the same way in all cases. + """ + + def __init__( + self, + dim, + num_heads, + dropout, + bias=False, + causal=False, + pos_encoding="none", + rope=None, + ): + super().__init__() + assert dim % num_heads == 0 + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.causal = causal + self.pos_encoding = pos_encoding + self.dropout_p = dropout + self.rope = rope if pos_encoding == "rope" else None + + self.q_proj = nn.Linear(dim, dim, bias=bias) + self.k_proj = nn.Linear(dim, dim, bias=bias) + self.v_proj = nn.Linear(dim, dim, bias=bias) + self.proj = nn.Linear(dim, dim, bias=bias) + + self.resid_drop = nn.Dropout(dropout) + + if pos_encoding == "alibi": + self.register_buffer( + "slopes", torch.tensor(alibi_slopes(num_heads)), persistent=False + ) + self.register_buffer("alibi_bias", torch.empty(0), persistent=False) + + def _alibi_mask(self, num_frames, device, dtype): + stale = ( + self.alibi_bias.numel() == 0 + or self.alibi_bias.shape[-1] < num_frames + or self.alibi_bias.device != device + or self.alibi_bias.dtype != dtype + ) + if stale: + self.alibi_bias = causal_alibi_bias(num_frames, self.slopes, device, dtype) + return self.alibi_bias[..., :num_frames, :num_frames] + + def forward(self, Q, K, V, mask=None): + B, T, C = Q.size() + Tk = K.size(1) + + q = self.q_proj(Q).view(B, T, self.num_heads, self.head_dim).transpose(1, 2) + k = self.k_proj(K).view(B, Tk, self.num_heads, self.head_dim).transpose(1, 2) + v = self.v_proj(V).view(B, Tk, self.num_heads, self.head_dim).transpose(1, 2) + + if self.rope is not None: + # Query t sits at absolute frame Tk - T + t, so a query shorter than + # its key stream (streaming inference) still lines up with history. + q, k = self.rope(q, k, offset=Tk - T) + + # The causal mask is built on EVERY call when this layer is causal, and a + # caller's mask is *added* to it -- never substituted for it. Building it + # only when `mask is None` (as this did) meant that supplying any mask + # dropped both the causal bias and `is_causal`, leaving a silently + # bidirectional layer that still trained and still reported numbers. No + # caller passes a mask into a temporal stack today, so the defect was + # latent; it goes live the first time anyone adds padded batches, which + # is exactly how the sibling ASD codebase found it. + if self.causal and self.pos_encoding == "alibi" and T == Tk: + causal_bias = self._alibi_mask(T, q.device, q.dtype) + elif self.causal and mask is not None: + # ALiBi's helper only covers T == Tk; fall back to plain causality. + causal_bias = causal_float_mask(T, Tk, q.device, q.dtype) + else: + causal_bias = None + + if causal_bias is not None: + mask = ( + causal_bias + if mask is None + else causal_bias + as_additive_mask(mask, q.dtype) + ) + + y = F.scaled_dot_product_attention( + q, + k, + v, + attn_mask=mask, + dropout_p=self.dropout_p if self.training else 0.0, + is_causal=self.causal and mask is None, + ) + + y = y.transpose(1, 2).contiguous().view(B, T, C) + return self.resid_drop(self.proj(y)) + + +class RMSNorm(nn.Module): + def __init__(self, dim, eps=1e-6): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + normed = x * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps) + return normed.to(x.dtype) * self.weight + + +def build_norm(norm_type, dim): + if norm_type == "rms": + return RMSNorm(dim) + if norm_type == "layer": + return nn.LayerNorm(dim) + raise ValueError(f"norm must be 'layer' or 'rms', got {norm_type!r}") + + +class SwiGLU(nn.Module): + def __init__(self, din, dff, dropout=0.0, bias=False): + super().__init__() + self.gate = nn.Linear(din, dff, bias=bias) + self.up = nn.Linear(din, dff, bias=bias) + self.down = nn.Linear(dff, din, bias=bias) + self.drop = nn.Dropout(dropout) + + def forward(self, x): + return self.down(self.drop(F.silu(self.gate(x)) * self.up(x))) + + +def ffn_block(din, dff, activation="GELU", dropout=0.0, bias=False, ffn_type="mlp"): + if ffn_type == "swiglu": + return SwiGLU(din, dff, dropout=dropout, bias=bias) + if ffn_type != "mlp": + raise ValueError(f"ffn must be 'mlp' or 'swiglu', got {ffn_type!r}") + return nn.Sequential( + nn.Linear(din, dff, bias=bias), + getattr(nn, activation)(), + nn.Dropout(dropout), + nn.Linear(dff, din, bias=bias), + ) diff --git a/third_party/MuVAP/models/modules/audio.py b/third_party/MuVAP/models/modules/audio.py new file mode 100644 index 0000000..674202a --- /dev/null +++ b/third_party/MuVAP/models/modules/audio.py @@ -0,0 +1,49 @@ +"""Selection and dimension-matching of the frozen audio frontend. + +VAP and ASD both take a pretrained frozen encoder and run their own stack on +top, so which encoder that is belongs in one place rather than in each model. +""" + +import torch.nn as nn + +from .cpc_encoder import CPC +from .mimi_encoder import MimiEncoder + +AUDIO_ENCODERS = {"cpc": CPC, "mimi": MimiEncoder} + + +def build_audio_encoder(cfg): + """Frozen audio frontend named by `cfg['audio_encoder']` (default CPC). + + Accepts either a bare name or a block carrying the encoder's own options: + + audio_encoder: mimi + audio_encoder: {name: mimi, tap: encoder} + """ + spec = cfg.get("audio_encoder", "cpc") + if isinstance(spec, str): + spec = {"name": spec} + options = dict(spec) + name = str(options.pop("name", "cpc")).lower() + if name not in AUDIO_ENCODERS: + raise ValueError( + f"audio_encoder must be one of {sorted(AUDIO_ENCODERS)}, got {name!r}" + ) + return AUDIO_ENCODERS[name](**options) + + +def encoder_projection(encoder_dim, model_dim): + """Bridge a frozen encoder's output onto the stack width, free when equal.""" + if encoder_dim == model_dim: + return nn.Identity() + return nn.Sequential(nn.LayerNorm(encoder_dim), nn.Linear(encoder_dim, model_dim)) + + +def check_frame_rate(encoder, frame_hz, what): + """Refuse a frontend whose feature rate does not match the target timeline.""" + encoder_hz = getattr(encoder, "frame_hz", frame_hz) + if float(encoder_hz) != float(frame_hz): + raise ValueError( + f"audio encoder emits {encoder_hz} Hz features but {what} runs at " + f"{frame_hz} Hz; they must agree frame for frame" + ) diff --git a/third_party/MuVAP/models/modules/cpc_encoder.py b/third_party/MuVAP/models/modules/cpc_encoder.py new file mode 100644 index 0000000..8994e0d --- /dev/null +++ b/third_party/MuVAP/models/modules/cpc_encoder.py @@ -0,0 +1,566 @@ +import argparse +from os import makedirs +from os.path import dirname, exists, join +from typing import List + +import torch +import math + +import torch.nn as nn +import torch.nn.functional as F +from einops.layers.torch import Rearrange + +""" +Encoder should take a wave file, then return an embedding +""" + + +# 16 kHz / 160 in CPC's conv stack. `downsample` halves it once per block, so +# the rate the model sees is a choice: 50 Hz is the original VAP's, 25 Hz halves +# it again to meet Mimi on a shared grid. +CPC_CONV_HZ = 100.0 +CPC_RATES = (100.0, 50.0, 25.0, 12.5) + + +class CPC(nn.Module): + def __init__(self, frame_hz=25.0): + super().__init__() + frame_hz = float(frame_hz) + if frame_hz not in CPC_RATES: + raise ValueError( + f"cpc frame_hz must be one of {CPC_RATES} (the conv stack runs at " + f"{CPC_CONV_HZ:g} Hz and `downsample` halves it), got {frame_hz:g}" + ) + self.sample_rate = 16000 + self.frame_hz = frame_hz + self.encoder = load_CPC(True) + self.in_dim = self.encoder.gEncoder.conv4.out_channels + self.out_dim = 256 + + halvings = round(math.log2(CPC_CONV_HZ / frame_hz)) + self.downsample = nn.Sequential( + *[ + get_cnn_layer( + in_dim=self.in_dim if block == 0 else self.out_dim, + out_dim=self.out_dim, + kernel=[3], + stride=[2], + dilation=[1], + activation="GELU", + ) + for block in range(halvings) + ] + ) + if not halvings: # 100 Hz: still bridge the encoder width onto out_dim + self.downsample = nn.Sequential(nn.Linear(self.in_dim, self.out_dim)) + + self.freeze() + + def forward(self, waveform): + if waveform.ndim < 3: + waveform = waveform.unsqueeze(1) + z = self.encoder.gEncoder(waveform) + z = z.permute(0, 2, 1) + z = self.encoder.gAR(z) + z = self.downsample(z) + return z + + def freeze(self): + for p in self.encoder.parameters(): + p.requires_grad_(False) + + +""" + +############################################################# +############################################################# +WARNING - ATTENTION - HEAR YEE HEAAAR YEEEE +############################################################# +############################################################# + +Most of the code in this file are scaled down (and heavily copied) versions of + +https://github.com/facebookresearch/CPC_audio + +Please checkout their codebase if you are interested in CPC networks +---------------------------------------------------------------------------- + +torch.hub downloads to: `$HOME/.cache/torch/hub/checkpoints/` +Explicit checkpoint path saved manually in "assets/" see CHECKPOINTS below. +""" + + +def repo_root(): + """ + Returns the absolute path to the git repository + """ + root = dirname(__file__) + root = dirname(root) + return root + + +CHECKPOINTS = { + "cpc": join(repo_root(), "assets/checkpoints/cpc/60k_epoch4-d0f474de.pt") +} +NAMES = list(CHECKPOINTS.keys()) + + +class ChannelNorm(nn.Module): + """ + Most of the code in this file are scaled down (and heavily copied) versions of + https://github.com/facebookresearch/CPC_audio + """ + + def __init__(self, numFeatures, epsilon=1e-05, affine=True): + + super(ChannelNorm, self).__init__() + if affine: + self.weight = nn.parameter.Parameter(torch.Tensor(1, numFeatures, 1)) + self.bias = nn.parameter.Parameter(torch.Tensor(1, numFeatures, 1)) + else: + self.weight = None + self.bias = None + self.epsilon = epsilon + self.p = 0 + self.affine = affine + self.reset_parameters() + + def reset_parameters(self): + if self.affine: + torch.nn.init.ones_(self.weight) + torch.nn.init.zeros_(self.bias) + + def forward(self, x): + + cumMean = x.mean(dim=1, keepdim=True) + cumVar = x.var(dim=1, keepdim=True) + x = (x - cumMean) * torch.rsqrt(cumVar + self.epsilon) + + if self.weight is not None: + x = x * self.weight + self.bias + return x + + +class CPCEncoder(nn.Module): + """ + Most of the code in this file are scaled down (and heavily copied) versions of + https://github.com/facebookresearch/CPC_audio + """ + + def __init__(self, sizeHidden=512, normMode="layerNorm"): + super(CPCEncoder, self).__init__() + normLayer = ChannelNorm + self.dimEncoded = sizeHidden + self.conv0 = nn.Conv1d(1, sizeHidden, 10, stride=5, padding=3) + self.batchNorm0 = normLayer(sizeHidden) + self.conv1 = nn.Conv1d(sizeHidden, sizeHidden, 8, stride=4, padding=2) + self.batchNorm1 = normLayer(sizeHidden) + self.conv2 = nn.Conv1d(sizeHidden, sizeHidden, 4, stride=2, padding=1) + self.batchNorm2 = normLayer(sizeHidden) + self.conv3 = nn.Conv1d(sizeHidden, sizeHidden, 4, stride=2, padding=1) + self.batchNorm3 = normLayer(sizeHidden) + self.conv4 = nn.Conv1d(sizeHidden, sizeHidden, 4, stride=2, padding=1) + self.batchNorm4 = normLayer(sizeHidden) + self.DOWNSAMPLING = 160 + + def getDimOutput(self): + return self.conv4.out_channels + + def forward(self, x): + x = F.relu(self.batchNorm0(self.conv0(x))) + x = F.relu(self.batchNorm1(self.conv1(x))) + x = F.relu(self.batchNorm2(self.conv2(x))) + x = F.relu(self.batchNorm3(self.conv3(x))) + x = F.relu(self.batchNorm4(self.conv4(x))) + return x + + +class CPCAR(nn.Module): + """ + Most of the code in this file are scaled down (and heavily copied) versions of + https://github.com/facebookresearch/CPC_audio + """ + + def __init__( + self, dimEncoded, dimOutput, keepHidden, nLevelsGRU, mode="GRU", reverse=False + ): + + super(CPCAR, self).__init__() + self.RESIDUAL_STD = 0.1 + + if mode == "LSTM": + self.baseNet = nn.LSTM( + dimEncoded, dimOutput, num_layers=nLevelsGRU, batch_first=True + ) + elif mode == "RNN": + self.baseNet = nn.RNN( + dimEncoded, dimOutput, num_layers=nLevelsGRU, batch_first=True + ) + else: + self.baseNet = nn.GRU( + dimEncoded, dimOutput, num_layers=nLevelsGRU, batch_first=True + ) + + self.hidden = None + self.keepHidden = keepHidden + self.reverse = reverse + + def getDimOutput(self): + return self.baseNet.hidden_size + + def forward(self, x): + + if self.reverse: + x = torch.flip(x, [1]) + try: + self.baseNet.flatten_parameters() + except RuntimeError: + pass + x, h = self.baseNet(x, self.hidden) + if self.keepHidden: + if isinstance(h, tuple): + self.hidden = tuple(x.detach() for x in h) + else: + self.hidden = h.detach() + + # For better modularity, a sequence's order should be preserved + # by each module + if self.reverse: + x = torch.flip(x, [1]) + return x + + +class CPCModel(nn.Module): + """ + Most of the code in this file are scaled down (and heavily copied) versions of + https://github.com/facebookresearch/CPC_audio + """ + + def __init__(self, encoder, AR): + super(CPCModel, self).__init__() + self.gEncoder = encoder + self.gAR = AR + + def forward(self, batchData, label): + encodedData = self.gEncoder(batchData).permute(0, 2, 1) + cFeature = self.gAR(encodedData) + return cFeature, encodedData, label + + +def load_CPC(load_state_dict=True): + """ + Contrast predictive learning model for audio data + pretrained: if True, load a model trained on libri-light 60k + (https://arxiv.org/abs/1912.07875) + **kwargs : see cpc/cpc_default_config to get the list of possible arguments + + Most of the code in this file are scaled down (and heavily copied) versions of + https://github.com/facebookresearch/CPC_audio + """ + + def loadArgs(args, locArgs, forbiddenAttr=None): + for k, v in vars(locArgs).items(): + if forbiddenAttr is not None: + if k not in forbiddenAttr: + setattr(args, k, v) + else: + setattr(args, k, v) + + def get_default_cpc_config(): + parser = argparse.ArgumentParser() + + # Run parameters + group = parser.add_argument_group( + "Architecture configuration", + description="The arguments defining the model's architecture.", + ) + group.add_argument( + "--hiddenEncoder", + type=int, + default=256, + help="Hidden dimension of the encoder network.", + ) + group.add_argument( + "--hiddenGar", + type=int, + default=256, + help="Hidden dimension of the auto-regressive network", + ) + group.add_argument( + "--nPredicts", type=int, default=12, help="Number of steps to predict." + ) + group.add_argument( + "--negativeSamplingExt", + type=int, + default=128, + help="Number of negative samples to take.", + ) + group.add_argument("--learningRate", type=float, default=2e-4) + group.add_argument( + "--schedulerStep", + type=int, + default=-1, + help="Step of the learning rate scheduler: at each " + "step the learning rate is divided by 2. Default: " + "no scheduler.", + ) + group.add_argument( + "--schedulerRamp", + type=int, + default=None, + help="Enable a warm up phase for the learning rate: " + "adds a linear ramp of the given size.", + ) + group.add_argument( + "--beta1", + type=float, + default=0.9, + help="Value of beta1 for the Adam optimizer", + ) + group.add_argument( + "--beta2", + type=float, + default=0.999, + help="Value of beta2 for the Adam optimizer", + ) + group.add_argument( + "--epsilon", + type=float, + default=1e-08, + help="Value of epsilon for the Adam optimizer", + ) + group.add_argument( + "--sizeWindow", + type=int, + default=20480, + help="Number of frames to consider at each batch.", + ) + group.add_argument( + "--nEpoch", type=int, default=200, help="Number of epoch to run" + ) + group.add_argument( + "--samplingType", + type=str, + default="samespeaker", + choices=["samespeaker", "uniform", "samesequence", "sequential"], + help="How to sample the negative examples in the CPC loss.", + ) + group.add_argument( + "--nLevelsPhone", + type=int, + default=1, + help="(Supervised mode only). Number of layers in " + "the phone classification network.", + ) + group.add_argument( + "--cpc_mode", + type=str, + default=None, + choices=["reverse", "none"], + help="Some variations on CPC.", + ) + group.add_argument( + "--encoder_type", + type=str, + choices=["cpc", "mfcc", "lfb"], + default="cpc", + help="Replace the encoder network by mfcc features or learned filter banks", + ) + group.add_argument( + "--normMode", + type=str, + default="layerNorm", + choices=["instanceNorm", "ID", "layerNorm", "batchNorm"], + help="Type of normalization to use in the encoder " + "network (default is layerNorm).", + ) + group.add_argument( + "--onEncoder", + action="store_true", + help="(Supervised mode only) Perform the " + "classification on the encoder's output.", + ) + group.add_argument( + "--random_seed", type=int, default=None, help="Set a specific random seed." + ) + group.add_argument( + "--speakerEmbedding", + type=int, + default=0, + help="(Depreciated) Feed the prediction network with " + "speaker embeddings along with the usual sequence.", + ) + group.add_argument( + "--arMode", + default="LSTM", + choices=["GRU", "LSTM", "RNN", "no_ar", "transformer"], + help="Architecture to use for the auto-regressive " + "network (default is lstm).", + ) + group.add_argument( + "--nLevelsGRU", + type=int, + default=1, + help="Number of layers in the autoregressive network.", + ) + group.add_argument( + "--rnnMode", + type=str, + default="transformer", + choices=[ + "transformer", + "RNN", + "LSTM", + "linear", + "ffd", + "conv4", + "conv8", + "conv12", + ], + help="Architecture to use for the prediction network", + ) + group.add_argument( + "--dropout", + action="store_true", + help="Add a dropout layer at the output of the prediction network.", + ) + group.add_argument( + "--abspos", + action="store_true", + help="If the prediction network is a transformer, " + "active to use absolute coordinates.", + ) + return parser.parse_args([]) + + # from cpc.model import CPCModel as cpcmodel + # from cpc.cpc_default_config import get_default_cpc_config + # from cpc.feature_loader import getEncoder, getAR, loadArgs + # from cpc.feature_loader import loadArgs + + locArgs = get_default_cpc_config() + if exists(CHECKPOINTS["cpc"]): + checkpoint = torch.load( + CHECKPOINTS["cpc"], map_location="cpu", weights_only=True + ) + else: + checkpoint_url = "https://dl.fbaipublicfiles.com/librilight/CPC_checkpoints/60k_epoch4-d0f474de.pt" + checkpoint = torch.hub.load_state_dict_from_url( + checkpoint_url, progress=False, map_location="cpu" + ) + makedirs(dirname(CHECKPOINTS["cpc"]), exist_ok=True) + torch.save(checkpoint, CHECKPOINTS["cpc"]) + loadArgs(locArgs, argparse.Namespace(**checkpoint["config"])) + # encoderNet = getEncoder(locArgs) + encoderNet = CPCEncoder(locArgs.hiddenEncoder, locArgs.normMode) + # arNet = getAR(locArgs) + arNet = CPCAR( + locArgs.hiddenEncoder, + locArgs.hiddenGar, + locArgs.samplingType == "sequential", + locArgs.nLevelsGRU, + mode=locArgs.arMode, + reverse=locArgs.cpc_mode == "reverse", + ) + # model = cpcmodel(encoderNet, arNet) + model = CPCModel(encoderNet, arNet) + + # always load pretrained + if load_state_dict: + model.load_state_dict(checkpoint["weights"], strict=False) + model.name = "cpc" + return model + + +class LayerNorm(nn.Module): + def __init__(self, dim: int, rearrange_outputs: bool = True) -> None: + super().__init__() + self.ln = nn.LayerNorm(dim) + self.in_rearrange = Rearrange("b d t -> b t d") + if rearrange_outputs: + self.out_rearrange = Rearrange("b t d -> b d t") + else: + self.out_rearrange = nn.Identity() + + def __repr__(self): + return str(self.ln) + + def forward(self, x): + return self.out_rearrange(self.ln(self.in_rearrange(x))) + + +class CConv1d(nn.Conv1d): + """source: https://github.com/pytorch/pytorch/issues/1333""" + + def __init__( + self, + in_channels, + out_channels, + kernel_size, + stride=1, + dilation=1, + groups=1, + padding_value=0, + bias=True, + **kwargs, + ): + super().__init__( + in_channels, + out_channels, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + groups=groups, + bias=bias, + **kwargs, + ) + + ks = kernel_size if isinstance(kernel_size, int) else kernel_size[0] + pad_dim1_pre = ks - 1 + pad_dim1_post = 0 + if dilation > 0: + pad_dim1_pre *= dilation + pad = (pad_dim1_pre, pad_dim1_post) + self.pad = nn.ConstantPad1d(padding=pad, value=padding_value) + + def debug_weights(self, type="sum"): + w = 1.0 + if type == "mean": + w = 1.0 / self.kernel_size[0] + + elif type == "range": + k = self.kernel_size[0] + w = torch.arange(1, k + 1).float().pow(2) + w = w.repeat(self.out_channels, self.in_channels, 1) + self.weight.data = w + if self.bias: + self.bias.data = self.bias.data.fill_(0.0) + return None + + self.weight.data = self.weight.data.fill_(w) + if self.bias: + self.bias.data = self.bias.data.fill_(0.0) + + def forward(self, input): + return super().forward(self.pad(input)) + + +def get_cnn_layer( + in_dim: int, + out_dim: int, + kernel: List[int] = [3], + stride: List[int] = [2], + dilation: List[int] = [1], + padding: List[int] = [0], + activation: str = "GELU", +): + layers = [] + layers.append(Rearrange("b t d -> b d t")) + for k, s, d, p in zip(kernel, stride, dilation, padding): + layers.append( + CConv1d(in_dim, out_dim, kernel_size=k, stride=s, dilation=d, padding=p) + ) + layers.append(LayerNorm(out_dim)) + if activation is not None: + layers.append(getattr(torch.nn, activation)()) + layers.append(Rearrange("b d t -> b t d")) + return nn.Sequential(*layers) diff --git a/third_party/MuVAP/models/modules/mimi_encoder.py b/third_party/MuVAP/models/modules/mimi_encoder.py new file mode 100644 index 0000000..06e7377 --- /dev/null +++ b/third_party/MuVAP/models/modules/mimi_encoder.py @@ -0,0 +1,125 @@ +"""Kyutai Mimi as a frozen causal audio frontend. + +Mimi is the neural codec from Kyutai's Moshi. It is causal by construction - +its SEANet convolutions left-pad, and its encoder transformer uses a sliding +causal mask - which is what makes it usable here at all: a VAP model that saw +even a few milliseconds of future audio would not be a turn-taking predictor. + +The encoder is tapped *before* Mimi's own 2x downsample, which puts its +features on a 25 Hz grid - exactly the rate the projection window labels. The +12.5 Hz `latent` tap is available for completeness but does not line up with a +25 Hz label grid, so `build_vap` refuses the mismatch rather than let the loss +fail on a shape error deep in a run. + +Weights are the pretrained `kyutai/mimi` checkpoint and stay frozen; only the +projection onto `d_model` is learned. +""" + +import torch.nn as nn +import torch.nn.functional as F +import torchaudio + +SOURCE_RATE = 16_000 +MIMI_RATE = 24_000 +#: Mimi consumes 1920 samples (80 ms) per latent at 24 kHz; the encoder emits +#: one feature per 960 samples, which is 25 Hz. +MIMI_ENCODER_STRIDE = 960 +#: Left delay that buys back the resampler's lookahead, with margin. 1 ms. +RESAMPLE_DELAY = 16 + + +def causal_resample_16k_to_24k(waveform): + """16 -> 24 kHz without reading the future. + + `torchaudio.functional.resample` convolves with a symmetric sinc, so output + sample m depends on input past m*2/3 - measured, 9 output samples of + lookahead at this rate pair. Left-padding alone does not help, because + dropping the same number of output samples undoes the shift exactly. + Instead the signal is *delayed* by `RESAMPLE_DELAY` input samples while the + nominal output length is kept, so every output sample depends only on input + strictly behind it. The price is a 1 ms delay against a 40 ms frame. + """ + length = round(waveform.shape[-1] * MIMI_RATE / SOURCE_RATE) + delayed = F.pad(waveform, (RESAMPLE_DELAY, 0)) + return torchaudio.functional.resample(delayed, SOURCE_RATE, MIMI_RATE)[..., :length] + + +class MimiEncoder(nn.Module): + """Frozen Mimi encoder producing `[batch, frames, 512]` at 25 Hz.""" + + TAPS = ("encoder", "latent") + + def __init__(self, repo_id="kyutai/mimi", tap="encoder"): + super().__init__() + if tap not in self.TAPS: + raise ValueError(f"tap must be one of {self.TAPS}, got {tap!r}") + try: + from transformers import MimiModel + except ImportError as error: # pragma: no cover - dependency guard + raise ImportError( + "vap.audio_encoder: mimi needs transformers: pip install 'transformers>=4.45'" + ) from error + + mimi = MimiModel.from_pretrained(repo_id) + if not getattr(mimi.config, "use_causal_conv", False): + raise ValueError(f"{repo_id} was built non-causal; it cannot front a VAP model") + + # Nothing downstream reconstructs audio or uses discrete codes, and the + # decoder is over half the checkpoint. + for unused in ("decoder", "decoder_transformer", "upsample", "quantizer"): + if hasattr(mimi, unused): + setattr(mimi, unused, None) + if tap == "encoder": + mimi.downsample = None + + self.mimi = mimi + self.tap = tap + self._out_dim = int(mimi.config.hidden_size) + self.freeze() + + @property + def out_dim(self): + return self._out_dim + + @property + def frame_hz(self): + return 25.0 if self.tap == "encoder" else float(self.mimi.config.frame_rate) + + def freeze(self): + for parameter in self.mimi.parameters(): + parameter.requires_grad_(False) + + def train(self, mode=True): + """Keep the frozen frontend in eval mode. + + Its transformer has dropout; leaving it active would inject noise into + features that never adapt to it. + """ + super().train(mode) + self.mimi.eval() + return self + + def forward(self, waveform): + if waveform.ndim == 2: + waveform = waveform.unsqueeze(1) + if waveform.shape[1] != 1: + raise ValueError( + f"MimiEncoder takes one channel at a time, got {waveform.shape[1]}" + ) + + audio = causal_resample_16k_to_24k(waveform) + # Right-padding to a whole frame only ever adds future silence. + remainder = (-audio.shape[-1]) % MIMI_ENCODER_STRIDE + if remainder: + audio = F.pad(audio, (0, remainder)) + + features = self.mimi.encoder(audio) + encoded = self.mimi.encoder_transformer(features.transpose(1, 2)) + if hasattr(encoded, "last_hidden_state"): + encoded = encoded.last_hidden_state + elif isinstance(encoded, (tuple, list)): + encoded = encoded[0] + + if self.tap == "latent": + encoded = self.mimi.downsample(encoded.transpose(1, 2)).transpose(1, 2) + return encoded diff --git a/third_party/MuVAP/models/modules/positional.py b/third_party/MuVAP/models/modules/positional.py new file mode 100644 index 0000000..b125b89 --- /dev/null +++ b/third_party/MuVAP/models/modules/positional.py @@ -0,0 +1,156 @@ +"""Positional encodings shared by every MuVAP+ transformer stack. + +VAP is a causal model: a frame may only attend to frames at or before it. +Every encoding here is defined for that regime, so switching between them +changes how position is represented and never what is visible. +""" + +import math + +import torch +import torch.nn as nn + +#: A learned absolute table is deliberately absent. It is defined only up to +#: `max_frames` and cannot extrapolate past it, and these models are trained on +#: fixed-length clips but deployed on unbounded streams -- the position it would +#: need at inference is one it never saw a gradient for. ALiBi's distance bias +#: and RoPE's rotation are both defined at any offset; sinusoidal is at least +#: closed-form rather than a lookup. +POS_ENCODINGS = ("alibi", "rope", "sinusoidal", "none") + + +def alibi_slopes(num_heads): + """Per-head decay rates from the ALiBi paper (Press et al., 2022).""" + + def power_of_two_slopes(n): + start = 2 ** (-(2 ** -(math.log2(n) - 3))) + return [start * start**i for i in range(n)] + + if math.log2(num_heads).is_integer(): + return power_of_two_slopes(num_heads) + closest = 2 ** math.floor(math.log2(num_heads)) + extra = alibi_slopes(2 * closest)[0::2][: num_heads - closest] + return power_of_two_slopes(closest) + extra + + +def causal_alibi_bias(num_frames, slopes, device, dtype): + """Additive `[1, heads, T, T]` mask holding both ALiBi decay and causality.""" + positions = torch.arange(num_frames, device=device, dtype=dtype) + distance = positions.view(1, -1) - positions.view(-1, 1) + bias = distance.view(1, 1, num_frames, num_frames) * slopes.to( + device=device, dtype=dtype + ).view(1, -1, 1, 1) + causal = torch.ones(num_frames, num_frames, device=device, dtype=torch.bool).tril() + return bias.masked_fill(~causal.view(1, 1, num_frames, num_frames), float("-inf")) + + +def causal_float_mask(num_queries, num_keys, device, dtype): + """Additive `[1, 1, Tq, Tk]` mask: 0 where attendable, `-inf` ahead of the query. + + Query `t` sits at absolute frame `Tk - Tq + t` -- a query stream shorter than + its key stream is streaming inference over a longer history -- so it may + attend keys `<= Tk - Tq + t`. Additive rather than boolean on purpose: + `F.scaled_dot_product_attention` reads a bool mask as "True = attend" and a + float one as a pre-softmax bias, and mixing the two is silent. + """ + offset = num_keys - num_queries + queries = torch.arange(num_queries, device=device).view(-1, 1) + offset + keys = torch.arange(num_keys, device=device).view(1, -1) + mask = torch.zeros(num_queries, num_keys, device=device, dtype=dtype) + return mask.masked_fill(keys > queries, float("-inf")).view( + 1, 1, num_queries, num_keys + ) + + +def as_additive_mask(mask, dtype): + """A caller's mask as an additive float bias, whichever convention it used.""" + if mask.dtype == torch.bool: + return torch.zeros_like(mask, dtype=dtype).masked_fill_(~mask, float("-inf")) + return mask.to(dtype) + + +class RotaryEmbedding(nn.Module): + """Causal rotary position embedding (Su et al., 2024). + + Rotates the query and key of every head by an angle proportional to the + absolute frame index, which makes the attention logit depend only on the + distance between two frames. Under a causal mask that distance is always + non-negative, so the model sees a strictly backward-looking clock and no + future frame can leak through the encoding. + """ + + def __init__(self, head_dim, base=10_000.0, max_frames=4096): + super().__init__() + if head_dim % 2 != 0: + raise ValueError(f"rotary head_dim must be even, got {head_dim}") + self.head_dim = head_dim + self.base = float(base) + inv_freq = 1.0 / ( + self.base ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim) + ) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.register_buffer("cos", torch.empty(0), persistent=False) + self.register_buffer("sin", torch.empty(0), persistent=False) + self._build_cache(max_frames, inv_freq.device, torch.float32) + + def _build_cache(self, num_frames, device, dtype): + positions = torch.arange(num_frames, device=device, dtype=torch.float32) + angles = torch.outer(positions, self.inv_freq.to(device=device)) + angles = torch.cat([angles, angles], dim=-1) + self.cos = angles.cos().to(dtype) + self.sin = angles.sin().to(dtype) + + def _cached(self, num_frames, device, dtype): + stale = ( + self.cos.numel() == 0 + or self.cos.shape[0] < num_frames + or self.cos.device != device + or self.cos.dtype != dtype + ) + if stale: + self._build_cache(max(num_frames, 1), device, dtype) + return self.cos[:num_frames], self.sin[:num_frames] + + @staticmethod + def _rotate_half(x): + first, second = x.chunk(2, dim=-1) + return torch.cat([-second, first], dim=-1) + + def forward(self, query, key, offset=0): + """Rotate `[batch, heads, frames, head_dim]` query and key in place of no-op.""" + frames = query.shape[-2] + offset + cos, sin = self._cached(frames, query.device, query.dtype) + q_cos, q_sin = cos[offset:], sin[offset:] + k_cos, k_sin = cos[: key.shape[-2]], sin[: key.shape[-2]] + query = query * q_cos + self._rotate_half(query) * q_sin + key = key * k_cos + self._rotate_half(key) * k_sin + return query, key + + +class SinusoidalPositionEmbedding(nn.Module): + """Fixed additive sinusoidal embedding, kept for ablation against RoPE.""" + + def __init__(self, dim, max_frames=4096): + super().__init__() + position = torch.arange(max_frames, dtype=torch.float32).unsqueeze(1) + div = torch.exp( + torch.arange(0, dim, 2, dtype=torch.float32) * (-math.log(10_000.0) / dim) + ) + table = torch.zeros(max_frames, dim) + table[:, 0::2] = torch.sin(position * div) + table[:, 1::2] = torch.cos(position * div) + self.register_buffer("table", table, persistent=False) + + def forward(self, x): + if x.shape[1] > self.table.shape[0]: + raise ValueError( + f"sequence of {x.shape[1]} frames exceeds the {self.table.shape[0]}-frame table" + ) + return x + self.table[: x.shape[1]].to(dtype=x.dtype) + + +def build_input_position_embedding(pos_encoding, dim, max_frames=4096): + """Additive input-side encoding, or None when position lives in attention.""" + if pos_encoding == "sinusoidal": + return SinusoidalPositionEmbedding(dim, max_frames) + return None diff --git a/third_party/MuVAP/models/modules/projection.py b/third_party/MuVAP/models/modules/projection.py new file mode 100644 index 0000000..e9df79d --- /dev/null +++ b/third_party/MuVAP/models/modules/projection.py @@ -0,0 +1,27 @@ +import torch.nn as nn + + +class LinearHead(nn.Module): + def __init__(self, input_dim, output_dim): + super().__init__() + self.norm = nn.LayerNorm(input_dim) + self.linear = nn.Linear(input_dim, output_dim) + nn.init.normal_(self.linear.weight, std=0.02) + nn.init.zeros_(self.linear.bias) + + def forward(self, x): + return self.linear(self.norm(x)) + + +class MLPBlock(nn.Module): + def __init__(self, input_dim=512, hidden_dim=256, dropout=0.1): + super().__init__() + self.net = nn.Sequential( + nn.LayerNorm(input_dim), + nn.Linear(input_dim, hidden_dim), + nn.GELU(), + nn.Dropout(dropout), + ) + + def forward(self, x): + return self.net(x) diff --git a/third_party/MuVAP/models/modules/visual_encoder.py b/third_party/MuVAP/models/modules/visual_encoder.py new file mode 100644 index 0000000..0f8592b --- /dev/null +++ b/third_party/MuVAP/models/modules/visual_encoder.py @@ -0,0 +1,204 @@ +## +# ResNet18 Pretrained network to extract lip embedding +# This code is modified based on https://github.com/lordmartian/deep_avsr +## + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .cpc_encoder import CConv1d + + +class VisualEncoder(nn.Module): + #: Width of the final `visualConv1D`. Fixed by the architecture, so the + #: cross-attention stack has to be this wide too. + OUT_DIM = 256 + + def __init__(self): + super().__init__() + self.out_dim = self.OUT_DIM + self.frontend = visualFrontend() + self.tcn = visualTCN() + self.conv1d = visualConv1D() + + # Register normalization constants as buffers + self.register_buffer("mean", torch.tensor(0.4161)) + self.register_buffer("std", torch.tensor(0.1688)) + + def forward(self, x): + B, T, W, H = x.shape + # Normalize and Reshape + x = (x.view(B * T, 1, W, H) / 255.0 - self.mean) / self.std + + # Forward Pass + x = self.frontend(x) + x = x.view(B, T, -1).transpose(1, 2) + x = self.tcn(x) + x = self.conv1d(x) + return x.transpose(1, 2) + + +class ResNetLayer(nn.Module): + """ + A ResNet layer used to build the ResNet network. + Architecture: + --> conv-bn-relu -> conv -> + -> bn-relu -> conv-bn-relu -> conv -> + -> bn-relu --> + | | | | + -----> downsample ------> -------------------------------------> + """ + + def __init__(self, inplanes, outplanes, stride): + super(ResNetLayer, self).__init__() + self.conv1a = nn.Conv2d( + inplanes, outplanes, kernel_size=3, stride=stride, padding=1, bias=False + ) + self.bn1a = nn.BatchNorm2d(outplanes, momentum=0.01, eps=0.001) + self.conv2a = nn.Conv2d( + outplanes, outplanes, kernel_size=3, stride=1, padding=1, bias=False + ) + self.stride = stride + self.downsample = nn.Conv2d( + inplanes, outplanes, kernel_size=(1, 1), stride=stride, bias=False + ) + self.outbna = nn.BatchNorm2d(outplanes, momentum=0.01, eps=0.001) + + self.conv1b = nn.Conv2d( + outplanes, outplanes, kernel_size=3, stride=1, padding=1, bias=False + ) + self.bn1b = nn.BatchNorm2d(outplanes, momentum=0.01, eps=0.001) + self.conv2b = nn.Conv2d( + outplanes, outplanes, kernel_size=3, stride=1, padding=1, bias=False + ) + self.outbnb = nn.BatchNorm2d(outplanes, momentum=0.01, eps=0.001) + return + + def forward(self, inputBatch): + batch = F.relu(self.bn1a(self.conv1a(inputBatch))) + batch = self.conv2a(batch) + if self.stride == 1: + residualBatch = inputBatch + else: + residualBatch = self.downsample(inputBatch) + batch = batch + residualBatch + intermediateBatch = batch + batch = F.relu(self.outbna(batch)) + + batch = F.relu(self.bn1b(self.conv1b(batch))) + batch = self.conv2b(batch) + residualBatch = intermediateBatch + batch = batch + residualBatch + outputBatch = F.relu(self.outbnb(batch)) + return outputBatch + + +class ResNet(nn.Module): + """ + An 18-layer ResNet architecture. + """ + + def __init__(self): + super(ResNet, self).__init__() + self.layer1 = ResNetLayer(64, 64, stride=1) + self.layer2 = ResNetLayer(64, 128, stride=2) + self.layer3 = ResNetLayer(128, 256, stride=2) + self.layer4 = ResNetLayer(256, 512, stride=2) + self.avgpool = nn.AvgPool2d(kernel_size=(4, 4), stride=(1, 1)) + + return + + def forward(self, inputBatch): + batch = self.layer1(inputBatch) + batch = self.layer2(batch) + batch = self.layer3(batch) + batch = self.layer4(batch) + outputBatch = self.avgpool(batch) + return outputBatch + + +class CausalLayerNorm(nn.Module): + def __init__(self, channel_size, eps=1e-8): + super().__init__() + self.gamma = nn.Parameter(torch.ones(1, channel_size, 1)) + self.beta = nn.Parameter(torch.zeros(1, channel_size, 1)) + self.eps = eps + + def forward(self, y): + mean = y.mean(dim=1, keepdim=True) + var = ((y - mean) ** 2).mean(dim=1, keepdim=True) + y_norm = (y - mean) / torch.sqrt(var + self.eps) + return self.gamma * y_norm + self.beta + + +class visualFrontend(nn.Module): + """ + A visual feature extraction module. Generates a 512-dim feature vector per video frame. + Architecture: A 2D convolution block followed by an 18-layer ResNet. + """ + + def __init__(self): + super(visualFrontend, self).__init__() + self.frontend2D = nn.Sequential( + nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False), + nn.BatchNorm2d(64, momentum=0.01, eps=0.001), + nn.ReLU(), + nn.MaxPool2d(kernel_size=3, stride=2, padding=1), + ) + self.resnet = ResNet() + return + + def forward(self, inputBatch): + batchsize = inputBatch.shape[0] + batch = self.frontend2D(inputBatch) + + outputBatch = self.resnet(batch) + outputBatch = outputBatch.reshape(batchsize, -1, 512) + outputBatch = outputBatch.transpose(1, 2) + outputBatch = outputBatch.transpose(1, 2).transpose(0, 1) + return outputBatch + + +# With Causal Convolution 1D +class DSConv1d(nn.Module): + def __init__(self, dilation): + super(DSConv1d, self).__init__() + self.net = nn.Sequential( + nn.ReLU(), + CausalLayerNorm(512), + CConv1d(512, 512, 3, stride=1, dilation=dilation, groups=512, bias=False), + nn.PReLU(), + CausalLayerNorm(512), + CConv1d(512, 512, 1, bias=False), + ) + + def forward(self, x): + out = self.net(x) + return out + x + + +class visualTCN(nn.Module): + def __init__(self): + super(visualTCN, self).__init__() + stacks = [] + for i in range(5): + stacks += [DSConv1d(dilation=2**i)] + self.net = nn.Sequential(*stacks) # Visual Temporal Network V-TCN + + def forward(self, x): + out = self.net(x) + return out + + +class visualConv1D(nn.Module): + def __init__(self): + super(visualConv1D, self).__init__() + self.net = nn.Sequential( + CConv1d(512, 256, 5, stride=1), + CausalLayerNorm(256), + nn.ReLU(), + CConv1d(256, 256, 1), + ) + + def forward(self, x): + out = self.net(x) + return out diff --git a/third_party/MuVAP/models/release.py b/third_party/MuVAP/models/release.py new file mode 100644 index 0000000..e152470 --- /dev/null +++ b/third_party/MuVAP/models/release.py @@ -0,0 +1,60 @@ +"""Reading a published weight release back into a model. + +A release holds the trained weights alone: no optimizer state, and no copy of +the frozen frontend, which is restored from its own pretrained source when the +model is built. It also ships the model half of the config it was trained with, +so a caller supplies only where the data lives and where results go. +""" + +from pathlib import Path + +import torch +import yaml + +WEIGHTS_NAME = "weights.pt" +CONFIG_NAME = "config.yaml" +PROVENANCE_NAME = "provenance.json" + + +def resolve(path): + """Accept either a release directory or the weights file inside one.""" + path = Path(path) + if path.is_dir(): + weights = path / WEIGHTS_NAME + if not weights.exists(): + raise SystemExit(f"{path} is not a release: no {WEIGHTS_NAME}") + return weights, path / CONFIG_NAME + return path, path.parent / CONFIG_NAME + + +def merge_config(cfg, release_config, module): + """Take the architecture from the release, keep local paths from `cfg`. + + A release must be rebuilt with the architecture it was trained with, but + dataset roots and output directories belong to whoever is running it. + """ + if not Path(release_config).exists(): + return cfg + with open(release_config) as handle: + published = yaml.safe_load(handle) or {} + if module in published: + cfg[module] = {**cfg.get(module, {}), **published[module]} + return cfg + + +def load_weights(model, weights_path): + """Load a release into `model`, allowing only the frozen frontend to be absent.""" + state = torch.load(weights_path, map_location="cpu", weights_only=True) + missing, unexpected = model.load_state_dict(state, strict=False) + if unexpected: + raise SystemExit( + f"{weights_path} holds weights this model does not want, so it was trained " + f"with a different architecture: {sorted(unexpected)[:5]}" + ) + frozen = {name for name, p in model.named_parameters() if not p.requires_grad} + unexplained = set(missing) - frozen + if unexplained: + raise SystemExit( + f"{weights_path} is missing trainable weights: {sorted(unexplained)[:5]}" + ) + return model diff --git a/third_party/MuVAP/models/vap.py b/third_party/MuVAP/models/vap.py new file mode 100644 index 0000000..f4c1c90 --- /dev/null +++ b/third_party/MuVAP/models/vap.py @@ -0,0 +1,193 @@ +"""Voice Activity Projection encoders. + +Three architectures share one head contract: a `[batch, frames, classes]` logit +stream over whatever codebook the configured `ProjectionWindow` defines. + +* `AudioVAP` reads one mixed-mono channel. Speaker identity is not given to + the model, so it only works with the `role_relative` codebook, which names + speakers by their role in the conversation. +* `StereoVAP` reproduces the original VAP model: one channel per speaker, + a shared encoder and self-attention stack, cross-attention between the two + channels, and one head over their concatenation. Channels are ordered, so it + pairs with the `speaker_based` codebook that projects one channel per row. +* `AnchoredVAP` also uses the `speaker_based` codebook, but from one mixed + channel: the identity a second channel would carry comes instead from the + voice activity strictly before the frame being labelled, conditioned in. +""" + +import torch +import torch.nn as nn + +from projection_window import ProjectionWindow + +from .modules.attention import TransformerModel +from .modules.audio import build_audio_encoder, check_frame_rate, encoder_projection +from .modules.projection import LinearHead + + +class AudioVAP(nn.Module): + """Single-channel VAP over a mixed recording.""" + + def __init__(self, cfg): + super().__init__() + model_dim = cfg["temporal"]["d_model"] + self.audio_encoder = build_audio_encoder(cfg) + self.encoder_proj = encoder_projection(self.audio_encoder.out_dim, model_dim) + self.transformer = TransformerModel(cfg["temporal"]) + output_dim = ProjectionWindow(**cfg["projection_window"]).n_classes + self.vap_head = LinearHead(model_dim, output_dim) + + def forward(self, audio, return_embeddings=False): + if audio.ndim == 3 and audio.shape[1] != 1: + raise ValueError( + f"AudioVAP expects one audio channel, got {audio.shape[1]}; " + "use arch: stereo for two-channel input" + ) + embedding = self.transformer(self.encoder_proj(self.audio_encoder(audio))) + logits = self.vap_head(embedding) + if return_embeddings: + return logits, embedding + return logits + + +class StereoVAP(nn.Module): + """Two-channel VAP, one channel per speaker, as in the original VAP paper.""" + + def __init__(self, cfg): + super().__init__() + model_dim = cfg["temporal"]["d_model"] + cross_cfg = cfg.get("cross", cfg["temporal"]) + if cross_cfg["d_model"] != model_dim: + raise ValueError("StereoVAP temporal and cross d_model must match") + + self.audio_encoder = build_audio_encoder(cfg) + self.encoder_proj = encoder_projection(self.audio_encoder.out_dim, model_dim) + # One set of weights sees both channels: which speaker sits on which + # channel is arbitrary, so the encoder must not learn a per-channel bias. + self.transformer = TransformerModel(cfg["temporal"]) + self.cross = TransformerModel(cross_cfg, cross_attention=True) + self.channel_embed = nn.Embedding(2, model_dim) + nn.init.zeros_(self.channel_embed.weight) + + output_dim = ProjectionWindow(**cfg["projection_window"]).n_classes + self.vap_head = LinearHead(model_dim * 2, output_dim) + + def encode_channel(self, waveform, channel): + index = torch.tensor(channel, device=waveform.device) + encoded = self.encoder_proj(self.audio_encoder(waveform)) + return encoded + self.channel_embed(index) + + def forward(self, audio, return_embeddings=False): + if audio.ndim != 3 or audio.shape[1] != 2: + raise ValueError( + f"StereoVAP expects [batch, 2, samples] audio, got {tuple(audio.shape)}" + ) + + first = self.transformer(self.encode_channel(audio[:, 0:1], 0)) + second = self.transformer(self.encode_channel(audio[:, 1:2], 1)) + + first_cross = self.cross(first, src=second) + second_cross = self.cross(second, src=first) + + embedding = torch.cat([first_cross, second_cross], dim=-1) + logits = self.vap_head(embedding) + if return_embeddings: + return logits, embedding + return logits + + +class AnchoredVAP(nn.Module): + """One mixed channel, anchored by the voice activity behind each frame. + + This is `VapGPTMono` from the original VAP repository, on this codebase's + parts. `speaker_based` classes name channels, which mono audio cannot + identify - which is why `AudioVAP` is barred from that codebook. The voice + activity supplies the identity instead: a `Linear(2, d_model)`, layer + normalised, added to the encoder output, then the same two causal stacks the + stereo model uses, run back to back on the single channel. + + One departure, forced by the label convention. The reference conditions on + the activity at frame `t` because its target starts at `t + 1` + (`ObjectiveVAP.projection` shifts with `va[..., 1:, :]`). `speaker_based` + here bins the future from `t` inclusive, so the anchor is delayed by one + frame to keep the same property: the model is never shown any part of the + window it is being asked to predict. + """ + + def __init__(self, cfg): + super().__init__() + model_dim = cfg["temporal"]["d_model"] + cross_cfg = cfg.get("cross", cfg["temporal"]) + if cross_cfg["d_model"] != model_dim: + raise ValueError("AnchoredVAP temporal and cross d_model must match") + + self.audio_encoder = build_audio_encoder(cfg) + self.encoder_proj = encoder_projection(self.audio_encoder.out_dim, model_dim) + # `init_va_conditioning` in the reference: orthogonal weights, and a + # layer norm on the conditioning before it is added to the audio. + self.anchor = nn.Linear(2, model_dim) + nn.init.orthogonal_(self.anchor.weight) + self.anchor_norm = nn.LayerNorm(model_dim) + # `ar_channel` then `ar`: both causal self-attention over one channel. + self.transformer = TransformerModel(cfg["temporal"]) + self.cross = TransformerModel(cross_cfg) + self.vap_head = LinearHead(model_dim, ProjectionWindow(**cfg["projection_window"]).n_classes) + + @staticmethod + def past_activity(vad): + """`[batch, 2, frames]` voice activity delayed by one frame.""" + return torch.cat([torch.zeros_like(vad[:, :, :1]), vad[:, :, :-1]], dim=2) + + def encode_activity(self, vad, dtype): + return self.anchor_norm(self.anchor(self.past_activity(vad).transpose(1, 2).to(dtype))) + + def forward(self, audio, vad, return_embeddings=False): + if audio.ndim == 3 and audio.shape[1] != 1: + raise ValueError( + f"AnchoredVAP expects one audio channel, got {audio.shape[1]}" + ) + features = self.encoder_proj(self.audio_encoder(audio)) + if vad.shape[-1] != features.shape[1]: + raise ValueError( + f"anchor has {vad.shape[-1]} frames, audio has {features.shape[1]}" + ) + embedding = self.cross(self.transformer(features + self.encode_activity(vad, features.dtype))) + logits = self.vap_head(embedding) + if return_embeddings: + return logits, embedding + return logits + + +ARCHITECTURES = {"mono": AudioVAP, "stereo": StereoVAP, "anchored": AnchoredVAP} + + +def build_vap(cfg): + """Instantiate the VAP architecture named by `cfg['arch']` (default mono).""" + arch = str(cfg.get("arch", "mono")).lower() + if arch not in ARCHITECTURES: + raise ValueError(f"arch must be one of {sorted(ARCHITECTURES)}, got {arch!r}") + + projection = ProjectionWindow(**cfg["projection_window"]) + if arch == "anchored" and projection.mode != "speaker_based": + raise ValueError( + "anchored conditioning names channels, so it pairs with speaker_based; " + "a role codebook needs no anchor" + ) + if arch == "mono" and projection.mode == "speaker_based": + raise ValueError( + "speaker_based labels name a channel, which mono audio cannot identify; " + "use arch: stereo or a role-based projection mode" + ) + model = ARCHITECTURES[arch](cfg) + check_frame_rate(model.audio_encoder, projection.frame_hz, "the projection window") + return model + + +def audio_channels(cfg): + """Channels the configured architecture reads from disk.""" + return 2 if str(cfg.get("arch", "mono")).lower() == "stereo" else 1 + + +def needs_vad(cfg): + """Whether the architecture is conditioned on voice activity as well as audio.""" + return str(cfg.get("arch", "mono")).lower() == "anchored" diff --git a/third_party/MuVAP/preprocess/__init__.py b/third_party/MuVAP/preprocess/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/third_party/MuVAP/preprocess/asd/__init__.py b/third_party/MuVAP/preprocess/asd/__init__.py new file mode 100644 index 0000000..72a4115 --- /dev/null +++ b/third_party/MuVAP/preprocess/asd/__init__.py @@ -0,0 +1 @@ +"""Canonical active-speaker preprocessing for MuVAP+.""" diff --git a/third_party/MuVAP/preprocess/asd/adapters.py b/third_party/MuVAP/preprocess/asd/adapters.py new file mode 100644 index 0000000..b3b6e1e --- /dev/null +++ b/third_party/MuVAP/preprocess/asd/adapters.py @@ -0,0 +1,112 @@ +"""Dataset-specific metadata adapters. + +Adapters do not normalize signals. They only convert source layouts into +`SourceRecord`; the shared writer owns every media transformation. +""" + +import ast +import csv +import json +from pathlib import Path +from typing import Iterable + +from preprocess.asd.schema import SAMPLE_RATE, SourceRecord + + +def _labels(text: str) -> list[int]: + values = ast.literal_eval(text) + if not isinstance(values, (list, tuple)): + raise ValueError("labels must be a list") + result = [int(value) for value in values] + if any(value not in (0, 1) for value in result): + raise ValueError("labels must contain only 0 and 1") + return result + + +def _existing(path: Path, sample_id: str) -> Path: + if not path.exists(): + raise FileNotFoundError(f"{sample_id}: missing {path}") + return path + + +def clip_directory_records( + loader_path: Path, + visual_root: Path, + audio_root: Path, + split: str, + dataset: str = "ava", + skip_missing: bool = False, +) -> Iterable[SourceRecord]: + """Read the AVA-style loader TSV and nested JPEG/WAV layout. + + AVA and WASD ship the identical layout - `clips_videos/// + /.jpg` beside `clips_audios///.wav` + - and the identical five-column loader TSV, so one reader serves both. + """ + face_by_id = { + path.name: path for path in (visual_root / split).glob("*/*") if path.is_dir() + } + audio_by_id = {path.stem: path for path in (audio_root / split).glob("*/*.wav")} + with loader_path.open(newline="") as handle: + for row_number, row in enumerate(csv.reader(handle, delimiter="\t"), start=1): + if len(row) < 4: + raise ValueError( + f"{loader_path}:{row_number}: expected at least 4 columns" + ) + sample_id, _, fps, label_text = row[:4] + face_path = face_by_id.get(sample_id) + audio_path = audio_by_id.get(sample_id) + if face_path is None or audio_path is None: + if skip_missing: + continue + raise FileNotFoundError( + f"{sample_id}: missing face directory or clip-local WAV" + ) + yield SourceRecord( + dataset=dataset, + sample_id=sample_id, + visual_path=face_path, + audio_path=audio_path, + source_fps=float(fps), + labels=_labels(label_text), + visual_kind="jpeg_directory", + ) + + +def manifest_records( + manifest_path: Path, + dataset: str, + root: Path | None = None, + default_npy_audio_sample_rate: int = SAMPLE_RATE, +) -> Iterable[SourceRecord]: + """Read the canonical TSV adapter used for AVCC and future datasets. + + Required header fields: sample_id, visual_path, audio_path, source_fps, + labels. Optional: start_sec, visual_kind, audio_sample_rate. + """ + root = root or manifest_path.parent + with manifest_path.open(newline="") as handle: + reader = csv.DictReader(handle, delimiter="\t") + required = {"sample_id", "visual_path", "audio_path", "source_fps", "labels"} + missing = required.difference(reader.fieldnames or []) + if missing: + raise ValueError(f"{manifest_path}: missing columns {sorted(missing)}") + for row_number, row in enumerate(reader, start=2): + try: + yield SourceRecord( + dataset=dataset, + sample_id=row["sample_id"], + visual_path=_existing(root / row["visual_path"], row["sample_id"]), + audio_path=_existing(root / row["audio_path"], row["sample_id"]), + source_fps=float(row["source_fps"]), + labels=json.loads(row["labels"]), + start_sec=float(row.get("start_sec") or 0.0), + npy_audio_sample_rate=int( + row.get("audio_sample_rate") or default_npy_audio_sample_rate + ), + visual_kind=row.get("visual_kind") or "npy", + ) + except (TypeError, ValueError, json.JSONDecodeError) as error: + raise ValueError( + f"{manifest_path}:{row_number}: invalid row" + ) from error diff --git a/third_party/MuVAP/preprocess/asd/prep_msdwild.py b/third_party/MuVAP/preprocess/asd/prep_msdwild.py new file mode 100644 index 0000000..366c1a9 --- /dev/null +++ b/third_party/MuVAP/preprocess/asd/prep_msdwild.py @@ -0,0 +1,270 @@ +"""Crop raw MSDWild videos into aligned 30-second speaker-track samples.""" + +import argparse +import csv +import shutil +import sys +from collections import defaultdict +from dataclasses import dataclass +from pathlib import Path + +import cv2 +import numpy as np + +PROJECT_ROOT = Path(__file__).resolve().parents[2] +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) + +from data.utils import load_audio # noqa: E402 +from preprocess.asd.video import crop_face, load_visual # noqa: E402 +from preprocess.asd.writer import ShardWriter # noqa: E402 + + +@dataclass(frozen=True) +class Detection: + bbox: tuple[float, float, float, float] + + +def read_annotations(path: Path) -> dict[int, dict[str, Detection]]: + annotations: dict[int, dict[str, Detection]] = defaultdict(dict) + with path.open(newline="") as handle: + for row_number, row in enumerate(csv.reader(handle), start=1): + if len(row) < 8: + raise ValueError(f"{path}:{row_number}: expected 8 columns") + try: + frame = int(row[0]) + speaker_id = row[2] + bbox = tuple(float(value) for value in row[3:7]) + int(row[7]) + except ValueError as error: + raise ValueError(f"{path}:{row_number}: invalid annotation") from error + annotations[frame][speaker_id] = Detection(bbox) + return dict(annotations) + + +def read_rttm(path: Path) -> dict[str, dict[str, list[tuple[float, float]]]]: + activity: dict[str, dict[str, list[tuple[float, float]]]] = defaultdict( + lambda: defaultdict(list) + ) + with path.open() as handle: + for line_number, line in enumerate(handle, start=1): + parts = line.split() + if not parts: + continue + if len(parts) < 8 or parts[0] != "SPEAKER": + raise ValueError(f"{path}:{line_number}: invalid RTTM row") + try: + start = float(parts[3]) + duration = float(parts[4]) + except ValueError as error: + raise ValueError(f"{path}:{line_number}: invalid time") from error + activity[parts[1]][parts[7]].append((start, start + duration)) + return { + video_id: {speaker: sorted(intervals) for speaker, intervals in speakers.items()} + for video_id, speakers in activity.items() + } + + +def retained_speakers( + annotations: dict[int, dict[str, Detection]], + start_frame: int, + frame_count: int, + min_face_ratio: float, +) -> list[str]: + counts: dict[str, int] = defaultdict(int) + for frame in range(start_frame, start_frame + frame_count): + for speaker_id in annotations.get(frame, {}): + counts[speaker_id] += 1 + minimum = min_face_ratio * frame_count + return sorted(speaker_id for speaker_id, count in counts.items() if count >= minimum) + + +def clip_starts(first_face_frame: int, last_face_frame: int, stride_frames: int): + return range(first_face_frame, last_face_frame + 1, stride_frames) + + +def video_metadata(path: Path) -> tuple[float, int]: + cap = cv2.VideoCapture(str(path)) + if not cap.isOpened(): + raise FileNotFoundError(f"could not open video: {path}") + fps = float(cap.get(cv2.CAP_PROP_FPS)) + frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + cap.release() + if fps <= 0 or frame_count <= 0: + raise ValueError(f"invalid video metadata: {path}") + return fps, frame_count + + +def process_video( + video_path: Path, + annotation_path: Path, + activity: dict[str, list[tuple[float, float]]], + writer: ShardWriter, + clip_seconds: float, + overlap_seconds: float, + min_face_ratio: float, + min_clip_seconds: float, +) -> int: + annotations = read_annotations(annotation_path) + if not annotations: + return 0 + + fps, video_frames = video_metadata(video_path) + clip_frames = round(clip_seconds * fps) + stride_frames = round((clip_seconds - overlap_seconds) * fps) + min_clip_frames = max(1, round(min_clip_seconds * fps)) + if clip_frames <= 0 or stride_frames <= 0: + raise ValueError("clip and stride must each contain at least one frame") + waveform, _ = load_audio(video_path) + first_face = min(annotations) + last_face = min(max(annotations), video_frames - 1) + written = 0 + + for segment_id, start_frame in enumerate( + clip_starts(first_face, last_face, stride_frames) + ): + # Take whatever is left at the tail. Windows used to be dropped unless + # they were complete, which discarded every video shorter than one clip + # - about a third of MSDWild. The loader pads and chunks, so a short + # sample costs nothing. + window_frames = min(clip_frames, video_frames - start_frame) + if window_frames < min_clip_frames: + break + speakers = retained_speakers( + annotations, start_frame, window_frames, min_face_ratio + ) + if not speakers: + continue + + frames, decoded_fps = load_visual(video_path, start_frame, window_frames) + if not np.isclose(decoded_fps, fps): + raise ValueError(f"frame-rate changed while decoding {video_path}") + if len(frames) < min_clip_frames: + continue + window_frames = len(frames) + + start_time = start_frame / fps + start_sample = round(start_time * 16_000) + target_samples = round(window_frames / fps * 16_000) + audio = waveform[:, start_sample : start_sample + target_samples] + + for speaker_id in speakers: + faces = [] + labels = np.zeros(window_frames, dtype=np.uint8) + for offset, frame in enumerate(frames): + detection = annotations.get(start_frame + offset, {}).get(speaker_id) + faces.append(crop_face(frame, detection.bbox if detection else None)) + intervals = activity.get(speaker_id, []) + for offset in range(window_frames): + timestamp = start_time + (offset + 0.5) / fps + labels[offset] = any(start <= timestamp < end for start, end in intervals) + faces = np.stack(faces).astype(np.uint8, copy=False) + sample_id = f"{video_path.stem}_{segment_id:05d}_{speaker_id}" + writer.add_arrays( + dataset="msdwild", + sample_id=sample_id, + faces=faces, + waveform=audio, + labels=labels, + source_fps=fps, + start_sec=start_time, + ) + written += 1 + return written + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--input", type=Path, required=True) + parser.add_argument("--rttm", type=Path) + parser.add_argument( + "--output", + type=Path, + help="defaults to /packed/, beside the source media", + ) + parser.add_argument( + "--split", + default="all", + help="official split to pack: all, few.train, few.val, many.val", + ) + parser.add_argument("--clip-seconds", type=float, default=30.0) + # Windows no longer overlap: the loader computes projection targets over a + # whole track before chunking, so duplicated frames buy nothing and would + # both inflate the pack and over-weight this corpus in an epoch. + parser.add_argument("--overlap-seconds", type=float, default=0.0) + parser.add_argument("--min-clip-seconds", type=float, default=2.0) + parser.add_argument("--min-face-ratio", type=float, default=0.9) + parser.add_argument("--max-shard-gb", type=float, default=1.0) + parser.add_argument("--overwrite", action="store_true") + return parser + + +def main() -> None: + args = build_parser().parse_args() + if args.clip_seconds <= 0: + raise ValueError("clip-seconds must be positive") + if not 0 <= args.overlap_seconds < args.clip_seconds: + raise ValueError("overlap-seconds must be in [0, clip-seconds)") + if not 0 < args.min_face_ratio <= 1: + raise ValueError("min-face-ratio must be in (0, 1]") + if args.max_shard_gb <= 0: + raise ValueError("max-shard-gb must be positive") + if not 0 < args.min_clip_seconds <= args.clip_seconds: + raise ValueError("min-clip-seconds must be in (0, clip-seconds]") + if args.output is None: + args.output = args.input.parent / "packed" / args.split + print(f"packing into {args.output}") + if args.output.exists(): + if not args.overwrite and any(args.output.iterdir()): + raise FileExistsError(f"output directory is not empty: {args.output}") + if args.overwrite: + shutil.rmtree(args.output) + + args.output.mkdir(parents=True) + videos = sorted(args.input.glob("*.mp4")) + if not videos: + raise FileNotFoundError(f"no MP4 files found in {args.input}") + + rttm_path = args.rttm or args.input.parent / "all.rttm" + if not rttm_path.exists(): + raise FileNotFoundError(f"missing RTTM labels: {rttm_path}") + activity = read_rttm(rttm_path) + + if args.split != "all": + # The official split lists are themselves RTTMs; their video ids are + # the split membership. + split_path = args.input.parent / f"{args.split}.rttm" + if not split_path.exists(): + raise FileNotFoundError(f"missing split list: {split_path}") + wanted = set(read_rttm(split_path)) + videos = [video for video in videos if video.stem in wanted] + if not videos: + raise ValueError(f"no videos matched split {args.split}") + print(f"{args.split}: {len(videos)} videos") + total = 0 + writer = ShardWriter( + args.output, max_shard_bytes=round(args.max_shard_gb * 1_000_000_000) + ) + try: + for index, video_path in enumerate(videos, start=1): + annotation_path = video_path.with_suffix(".csv") + if not annotation_path.exists(): + raise FileNotFoundError(f"missing annotations: {annotation_path}") + total += process_video( + video_path, + annotation_path, + activity.get(video_path.stem, {}), + writer, + args.clip_seconds, + args.overlap_seconds, + args.min_face_ratio, + args.min_clip_seconds, + ) + print(f"[{index}/{len(videos)}] {video_path.stem}: {total} samples") + finally: + writer.close() + print(f"Wrote {total} samples to {args.output}") + + +if __name__ == "__main__": + main() diff --git a/third_party/MuVAP/preprocess/asd/prepare.py b/third_party/MuVAP/preprocess/asd/prepare.py new file mode 100644 index 0000000..820b941 --- /dev/null +++ b/third_party/MuVAP/preprocess/asd/prepare.py @@ -0,0 +1,93 @@ +"""Build the canonical fast-training format for AVA, WASD, and AVCC.""" + +import argparse +from pathlib import Path + +from preprocess.asd.adapters import clip_directory_records, manifest_records +from preprocess.asd.writer import write_records + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("dataset", choices=["ava", "wasd", "msdwild", "avcc"]) + parser.add_argument( + "--loader", type=Path, required=True, help="source TSV/manifest" + ) + parser.add_argument( + "--output", + type=Path, + help="defaults to /packed/, beside the source media", + ) + parser.add_argument("--visual-root", type=Path) + parser.add_argument("--audio-root", type=Path) + parser.add_argument("--root", type=Path, help="base for AVCC manifest paths") + parser.add_argument("--split", default="train") + parser.add_argument("--audio-sample-rate", type=int, default=16_000) + parser.add_argument("--max-shard-gb", type=float, default=1.0) + parser.add_argument( + "--skip-missing", + action="store_true", + help="development subsets only; production preprocessing should fail on missing files", + ) + return parser + + +def require_roots(args: argparse.Namespace) -> None: + if args.visual_root is None or args.audio_root is None: + raise ValueError(f"{args.dataset} requires --visual-root and --audio-root") + + +def default_output(args: argparse.Namespace) -> Path: + """Keep each corpus's pack beside that corpus. + + `/clips_videos/` becomes `/packed/`, so a dataset + stays self-contained and can be copied to fast local storage on its own. + """ + if args.visual_root is not None: + return args.visual_root.parent / "packed" / args.split + if args.root is not None: + return args.root / "packed" / args.split + raise ValueError("--output is required when no dataset root is given") + + +def main() -> None: + args = build_parser().parse_args() + if args.output is None: + args.output = default_output(args) + print(f"packing into {args.output}") + if args.output.exists() and any(args.output.iterdir()): + raise FileExistsError(f"output directory is not empty: {args.output}") + + if args.dataset in ("ava", "wasd"): + require_roots(args) + records = clip_directory_records( + args.loader, + args.visual_root, + args.audio_root, + args.split, + dataset=args.dataset, + skip_missing=args.skip_missing, + ) + elif args.dataset == "msdwild": + raise SystemExit( + "MSDWild is packed from raw media by preprocess/asd/prep_msdwild.py, " + "which needs all.rttm for labels; there is no loader TSV to read here" + ) + else: + records = manifest_records( + args.loader, + dataset="avcc", + root=args.root, + default_npy_audio_sample_rate=args.audio_sample_rate, + ) + + count = write_records( + records, + args.output, + max_shard_bytes=round(args.max_shard_gb * 1_000_000_000), + ) + print(f"Wrote {count} {args.dataset} samples to {args.output}") + + +if __name__ == "__main__": + main() diff --git a/third_party/MuVAP/preprocess/asd/schema.py b/third_party/MuVAP/preprocess/asd/schema.py new file mode 100644 index 0000000..204be0d --- /dev/null +++ b/third_party/MuVAP/preprocess/asd/schema.py @@ -0,0 +1,44 @@ +"""Shared source record and packed-shard schema. + +The packed format is a *decoded mirror* of the source corpus, not a +model-ready tensor set. Faces, labels, and audio are stored on the native +timeline with no gain applied, so any codebase can consume a pack without +inheriting MuVAP+'s 25 Hz model timeline or its loudness convention. Every +model-specific transform lives in `data.dataloaders.packed` instead. +""" + +from dataclasses import dataclass +from pathlib import Path +from typing import Sequence + +FORMAT_NAME = "muvap-asd-mmap" +FORMAT_VERSION = 4 +SUPPORTED_VERSIONS = (4,) + +# The model timeline. It belongs to the loader, not to the stored data, and is +# repeated here only because preprocessing reports frame counts against it. +MODEL_FPS = 25.0 + +SAMPLE_RATE = 16_000 +FACE_SIZE = 112 + + +@dataclass(frozen=True) +class SourceRecord: + dataset: str + sample_id: str + visual_path: Path + audio_path: Path + source_fps: float + labels: Sequence[int] + start_sec: float = 0.0 + npy_audio_sample_rate: int | None = None + visual_kind: str = "npy" + + def __post_init__(self) -> None: + if self.source_fps <= 0: + raise ValueError(f"{self.sample_id}: source_fps must be positive") + if self.start_sec < 0: + raise ValueError(f"{self.sample_id}: start_sec cannot be negative") + if self.visual_kind not in {"npy", "jpeg_directory"}: + raise ValueError(f"unsupported visual_kind: {self.visual_kind}") diff --git a/third_party/MuVAP/preprocess/asd/validate.py b/third_party/MuVAP/preprocess/asd/validate.py new file mode 100644 index 0000000..7668d97 --- /dev/null +++ b/third_party/MuVAP/preprocess/asd/validate.py @@ -0,0 +1,93 @@ +"""Validate packed samples and report per-dataset signal statistics.""" + +import argparse +import json +import math +from collections import defaultdict +from pathlib import Path +from typing import Sequence + +import numpy as np + +from data.alignment import model_frame_count +from data.dataloaders.packed import PackedASDDataset +from data.media import WaveformStats, normalization_gain +from preprocess.asd.schema import MODEL_FPS +from preprocess.asd.writer import native_audio_samples + + +def percentile(values: list[float], points=(0, 1, 50, 95, 99, 100)) -> dict: + return {str(point): float(np.percentile(values, point)) for point in points} + + +def validate(roots: Path | Sequence[Path]) -> None: + dataset = PackedASDDataset(roots) + grouped = defaultdict(lambda: defaultdict(list)) + + for index in range(len(dataset)): + audio, visual, vad, _, metadata = dataset[index] + sample_id = metadata["sample_id"] + expected = model_frame_count( + metadata["source_frames"], metadata["source_fps"], MODEL_FPS + ) + frames = visual.shape[0] + if frames == 0: + raise ValueError(f"{sample_id}: empty sample") + if frames != expected: + raise ValueError( + f"{sample_id}: {metadata['source_frames']} source frames at " + f"{metadata['source_fps']} fps imply {expected} model frames, " + f"loader produced {frames}" + ) + if vad.shape[0] != frames or audio.shape[-1] != frames * 640: + raise ValueError(f"{sample_id}: stream-length mismatch") + if visual.shape[1:] != (112, 112) or visual.min() < 0 or visual.max() > 255: + raise ValueError(f"{sample_id}: invalid visual tensor") + if not np.isfinite(metadata["audio_peak"]): + raise ValueError(f"{sample_id}: invalid audio statistics") + if metadata["audio_samples"] != native_audio_samples( + metadata["source_frames"], metadata["source_fps"] + ): + raise ValueError(f"{sample_id}: stored audio does not span the faces") + + stats = grouped[metadata["dataset"]] + stats["source_audio_peak"].append(metadata["audio_peak"]) + stats["source_audio_rms"].append(metadata["audio_rms"]) + stats["source_audio_mean"].append(metadata["audio_mean"]) + stats["speaking_ratio"].append(metadata["speaking_ratio"]) + stats["normalized_rms"].append( + audio.square().mean().sqrt().item() + ) + stats["normalized_peak"].append(audio.abs().max().item()) + gain = normalization_gain( + WaveformStats( + mean=metadata["audio_mean"], + rms=metadata["audio_rms"], + peak=metadata["audio_peak"], + ) + ) + stats["gain_db"].append(20 * math.log10(max(gain, 1e-8))) + + report = {"samples": len(dataset), "datasets": {}} + for name, stats in grouped.items(): + report["datasets"][name] = { + "samples": len(stats["source_audio_peak"]), + "source_audio_peak_percentiles": percentile(stats["source_audio_peak"]), + "source_audio_rms_percentiles": percentile(stats["source_audio_rms"]), + "normalized_rms_percentiles": percentile(stats["normalized_rms"]), + "normalized_peak_percentiles": percentile(stats["normalized_peak"]), + "gain_db_percentiles": percentile(stats["gain_db"]), + "max_abs_source_mean": max(abs(v) for v in stats["source_audio_mean"]), + "speaking_ratio_percentiles": percentile(stats["speaking_ratio"]), + } + print(json.dumps(report, indent=2)) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("roots", type=Path, nargs="+") + validate(parser.parse_args().roots) + + +if __name__ == "__main__": + main() diff --git a/third_party/MuVAP/preprocess/asd/video.py b/third_party/MuVAP/preprocess/asd/video.py new file mode 100644 index 0000000..915cc74 --- /dev/null +++ b/third_party/MuVAP/preprocess/asd/video.py @@ -0,0 +1,84 @@ +"""Reading face crops off a source video. + +Shared by the dataset adapters that pack from video rather than from +pre-extracted frames. +""" + +import cv2 +import numpy as np + + +def crop_face(frame, bbox, gray=True, dim=(112, 112)): + if bbox is None: + return np.full(dim if gray else (dim[1], dim[0], 3), 128, dtype=np.uint8) + + x1, y1, x2, y2 = int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3]) + + if x1 >= x2 or y1 >= y2: + return np.full(dim if gray else (dim[1], dim[0], 3), 128, dtype=np.uint8) + + h, w = frame.shape[:2] + cx, cy = (x1 + x2) // 2, (y1 + y2) // 2 + size = max(x2 - x1, y2 - y1) + half = size // 2 + + x1s, y1s = cx - half, cy - half + x2s, y2s = x1s + size, y1s + size + + x1c, y1c = max(0, x1s), max(0, y1s) + x2c, y2c = min(w, x2s), min(h, y2s) + + if x1c >= x2c or y1c >= y2c: + return np.full(dim if gray else (dim[1], dim[0], 3), 128, dtype=np.uint8) + + cropped = frame[y1c:y2c, x1c:x2c] + + if gray: + cropped = cv2.cvtColor(cropped, cv2.COLOR_BGR2GRAY) + border_val = 128 + else: + cropped = cv2.cvtColor(cropped, cv2.COLOR_BGR2RGB) + border_val = (128, 128, 128) + + pad_top, pad_left = y1c - y1s, x1c - x1s + pad_bottom, pad_right = y2s - y2c, x2s - x2c + + if pad_top > 0 or pad_bottom > 0 or pad_left > 0 or pad_right > 0: + cropped = cv2.copyMakeBorder( + cropped, + pad_top, + pad_bottom, + pad_left, + pad_right, + cv2.BORDER_CONSTANT, + value=border_val, + ) + + return cv2.resize(cropped, dim) + + +def load_visual(video_path, start_frame=0, target_frames=None): + """Decode a native-FPS video interval without changing its timeline.""" + if start_frame < 0: + raise ValueError("start_frame cannot be negative") + if target_frames is not None and target_frames < 0: + raise ValueError("target_frames cannot be negative") + + cap = cv2.VideoCapture(str(video_path)) + if not cap.isOpened(): + raise FileNotFoundError(f"could not open video: {video_path}") + + fps = float(cap.get(cv2.CAP_PROP_FPS)) + if fps <= 0: + cap.release() + raise ValueError(f"video has an invalid frame rate: {video_path}") + + cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame) + frames = [] + while target_frames is None or len(frames) < target_frames: + ok, frame = cap.read() + if not ok: + break + frames.append(frame) + cap.release() + return frames, fps diff --git a/third_party/MuVAP/preprocess/asd/writer.py b/third_party/MuVAP/preprocess/asd/writer.py new file mode 100644 index 0000000..4c03cf0 --- /dev/null +++ b/third_party/MuVAP/preprocess/asd/writer.py @@ -0,0 +1,204 @@ +"""Pack decoded native-rate samples into memory-mapped NumPy shards. + +Nothing here is specific to the model that will read the shards. Faces and +labels keep the source frame rate, audio keeps its full duration and its +original gain, and the per-recording level statistics needed to reproduce a +normalization travel in the index. Resampling onto a model timeline and +applying loudness normalization are the reader's job. +""" + +import json +from pathlib import Path +from typing import Iterable + +import cv2 +import numpy as np +import torch + +from data.alignment import model_frame_count +from data.media import ( + AudioSpec, + load_waveform, + slice_audio, + validate_visual, + waveform_stats, +) +from preprocess.asd.schema import ( + FACE_SIZE, + FORMAT_NAME, + FORMAT_VERSION, + MODEL_FPS, + SAMPLE_RATE, + SourceRecord, +) + + +def _load_faces(record: SourceRecord, face_size: int = FACE_SIZE) -> np.ndarray: + if record.visual_kind == "npy": + return validate_visual(np.load(record.visual_path, allow_pickle=False)) + + paths = sorted(record.visual_path.glob("*.jpg"), key=lambda path: float(path.stem)) + if not paths: + raise FileNotFoundError(f"{record.sample_id}: no JPEG faces") + frames = [] + for path in paths: + frame = cv2.imread(str(path), cv2.IMREAD_GRAYSCALE) + if frame is None: + raise ValueError(f"{record.sample_id}: could not decode {path}") + frames.append( + cv2.resize(frame, (face_size, face_size), interpolation=cv2.INTER_LINEAR) + ) + return np.stack(frames).astype(np.uint8, copy=False) + + +def _pcm16(waveform: torch.Tensor) -> np.ndarray: + values = waveform.squeeze(0).clamp(-1.0, 1.0) + return (values * 32767.0).round().to(torch.int16).cpu().numpy() + + +def native_audio_samples(source_frames: int, source_fps: float) -> int: + """Audio length covering exactly the span of the stored face frames.""" + return round(source_frames / source_fps * SAMPLE_RATE) + + +class ShardWriter: + def __init__(self, output_dir: Path, max_shard_bytes: int = 1_000_000_000): + self.output_dir = output_dir + self.max_shard_bytes = max_shard_bytes + self.shard_index = 0 + self.samples_written = 0 + self._reset() + output_dir.mkdir(parents=True, exist_ok=True) + + def _reset(self) -> None: + self.faces: list[np.ndarray] = [] + self.audio: list[np.ndarray] = [] + self.labels: list[np.ndarray] = [] + self.entries: list[dict] = [] + self.byte_count = 0 + self.face_count = 0 + self.audio_count = 0 + self.label_count = 0 + + def add_arrays( + self, + *, + dataset: str, + sample_id: str, + faces: np.ndarray, + waveform: torch.Tensor, + labels, + source_fps: float, + start_sec: float = 0.0, + ) -> None: + faces = validate_visual(faces) + labels = np.asarray(labels, dtype=np.uint8) + if len(faces) != len(labels): + raise ValueError( + f"{sample_id}: {len(faces)} faces but {len(labels)} labels" + ) + source_frames = len(labels) + if source_frames == 0: + raise ValueError(f"{sample_id}: sample is empty") + + waveform = slice_audio( + waveform, + SAMPLE_RATE, + target_samples=native_audio_samples(source_frames, source_fps), + ) + stats = waveform_stats(waveform) + audio = _pcm16(waveform) + + sample_bytes = faces.nbytes + audio.nbytes + labels.nbytes + if self.entries and self.byte_count + sample_bytes > self.max_shard_bytes: + self.flush() + + self.entries.append( + { + "dataset": dataset, + "sample_id": sample_id, + "source_fps": source_fps, + "source_frames": source_frames, + "start_sec": start_sec, + # Reported for convenience; the loader recomputes it. + "model_frames": model_frame_count(source_frames, source_fps, MODEL_FPS), + "face_offset": self.face_count, + "audio_offset": self.audio_count, + "audio_samples": len(audio), + "label_offset": self.label_count, + "audio_mean": stats.mean, + "audio_rms": stats.rms, + "audio_peak": stats.peak, + "speaking_ratio": float(labels.mean()), + } + ) + self.faces.append(faces) + self.audio.append(audio) + self.labels.append(labels) + self.byte_count += sample_bytes + self.face_count += len(faces) + self.audio_count += len(audio) + self.label_count += len(labels) + + def add(self, record: SourceRecord) -> None: + faces = _load_faces(record) + waveform, sample_rate = load_waveform( + record.audio_path, + AudioSpec(sample_rate=SAMPLE_RATE, mono=True), + npy_sample_rate=record.npy_audio_sample_rate, + ) + waveform = slice_audio(waveform, sample_rate, start_time=record.start_sec) + self.add_arrays( + dataset=record.dataset, + sample_id=record.sample_id, + faces=faces, + waveform=waveform, + labels=record.labels, + source_fps=record.source_fps, + start_sec=record.start_sec, + ) + + def flush(self) -> None: + if not self.entries: + return + shard_dir = self.output_dir / f"shard-{self.shard_index:06d}" + shard_dir.mkdir(parents=False, exist_ok=False) + np.save(shard_dir / "faces.npy", np.concatenate(self.faces), allow_pickle=False) + np.save(shard_dir / "audio.npy", np.concatenate(self.audio), allow_pickle=False) + np.save( + shard_dir / "labels.npy", np.concatenate(self.labels), allow_pickle=False + ) + with (shard_dir / "index.json").open("w") as handle: + json.dump(self.entries, handle, separators=(",", ":")) + self.samples_written += len(self.entries) + self.shard_index += 1 + self._reset() + + def close(self) -> None: + self.flush() + manifest = { + "format": FORMAT_NAME, + "version": FORMAT_VERSION, + "timeline": "native", + "audio_sample_rate": SAMPLE_RATE, + "audio_normalization": None, + "face_shape": [FACE_SIZE, FACE_SIZE], + "face_color": "grayscale", + "face_dtype": "uint8", + "audio_dtype": "int16", + "label_dtype": "uint8", + "shards": self.shard_index, + "samples": self.samples_written, + } + with (self.output_dir / "dataset.json").open("w") as handle: + json.dump(manifest, handle, indent=2) + + +def write_records( + records: Iterable[SourceRecord], output_dir: Path, max_shard_bytes: int +) -> int: + writer = ShardWriter(output_dir, max_shard_bytes) + for record in records: + writer.add(record) + writer.close() + return writer.samples_written diff --git a/third_party/MuVAP/preprocess/vap/00_prep_fisher.py b/third_party/MuVAP/preprocess/vap/00_prep_fisher.py new file mode 100644 index 0000000..cb18a4f --- /dev/null +++ b/third_party/MuVAP/preprocess/vap/00_prep_fisher.py @@ -0,0 +1,196 @@ +"""Cut Fisher into fixed segments of stereo audio beside their raw VAD. + +Class labels are *not* baked in here. Each segment stores the two-speaker VAD +padded with `--context-sec` of extra frames on both sides, which is everything +`ProjectionWindow` needs to produce role-based or speaker-based labels in the +dataloader. Switching projection setups therefore never requires re-running +this script. +""" + +import argparse +import os +from dataclasses import dataclass +from glob import glob + +import numpy as np +from tqdm import tqdm + +from config.setup import init_yaml_config +from data.utils import load_audio, load_fisher_vad +from preprocess.vap.utils import ( + SAMPLE_RATE, + TEST_GROUPS, + VAL_GROUPS, + ProcessConfig, + get_dataset_split, +) + +SEG_SEC: int = 30 +OVERLAP: int = 10 +VAD_HZ: int = 25 +CONTEXT_SEC: float = 2.0 + +STRIDE_SEC: int = SEG_SEC - OVERLAP +SEG_SAMPLES: int = int(SEG_SEC * SAMPLE_RATE) +STRIDE_SAMPLES: int = int(STRIDE_SEC * SAMPLE_RATE) +SEG_VAD_FRAMES: int = int(SEG_SEC * VAD_HZ) + + +@dataclass() +class SegmentMeta: + part: str + group: str + cid: str + seg_id: int + start_t: float + end_t: float + nframes: int + + +def format_meta(s: SegmentMeta) -> str: + return f"{s.part} {s.group} {s.cid} {s.seg_id} {s.start_t:.2f} {s.end_t:.2f} {s.nframes}\n" + + +def extract_and_save_chunks( + audio: np.ndarray, + vad: np.ndarray, + part: str, + group: str, + cid: str, + cfg: ProcessConfig, + context_frames: int, +): + total_samples = audio.shape[-1] + total_vad_frames = vad.shape[-1] + + audio_out = os.path.join(cfg.fisher_path, part, "seg", "audio", group, cid) + vad_out = os.path.join(cfg.fisher_path, part, "seg", "vad", group, cid) + os.makedirs(audio_out, exist_ok=True) + os.makedirs(vad_out, exist_ok=True) + + segments = [] + start_sample = 0 + seg_idx = 0 + + while start_sample + SEG_SAMPLES <= total_samples: + end_sample = start_sample + SEG_SAMPLES + + chunk_audio = audio[:, start_sample:end_sample] + pcm16 = np.round(np.clip(chunk_audio, -1.0, 1.0) * 32767).astype(np.int16) + audio_file = os.path.join(audio_out, f"{seg_idx}.npy") + if not os.path.exists(audio_file): + np.save(audio_file, pcm16) + + curr_vad_start = int((start_sample / SAMPLE_RATE) * VAD_HZ) + curr_vad_end = curr_vad_start + SEG_VAD_FRAMES + + ext_start = max(0, curr_vad_start - context_frames) + ext_end = min(total_vad_frames, curr_vad_end + context_frames) + + chunk_vad = vad[:, ext_start:ext_end] + + pad_left = max(0, context_frames - curr_vad_start) + pad_right = max(0, (curr_vad_end + context_frames) - total_vad_frames) + + if pad_left > 0 or pad_right > 0: + chunk_vad = np.pad( + chunk_vad, ((0, 0), (pad_left, pad_right)), mode="constant" + ) + + vad_file = os.path.join(vad_out, f"{seg_idx}.npy") + if not os.path.exists(vad_file): + np.save(vad_file, chunk_vad.astype(np.float32)) + + segments.append( + SegmentMeta( + part=part, + group=group, + cid=cid, + seg_id=seg_idx, + start_t=start_sample / SAMPLE_RATE, + end_t=end_sample / SAMPLE_RATE, + nframes=total_vad_frames, + ) + ) + + start_sample += STRIDE_SAMPLES + seg_idx += 1 + + return segments + + +def process_entry( + part: str, + group: str, + cid: str, + cfg: ProcessConfig, + context_frames: int, +): + audio_path = os.path.join(cfg.fisher_path, part, "audio", group, f"{cid}.wav") + vad_dir = os.path.join(cfg.fisher_path, part, "regions_vap", group) + + if not os.path.exists(audio_path): + return [] + + # Both channels are kept: the speaker-based setup needs one per speaker and + # the mono setups downmix in the dataloader. + audio_tensor, duration = load_audio(audio_path, mono=False) + audio = audio_tensor.numpy() + + nframes = int(round(duration * VAD_HZ)) + vad_tensor = load_fisher_vad(vad_dir, cid, nframes) + vad = vad_tensor.numpy() + + return extract_and_save_chunks(audio, vad, part, group, cid, cfg, context_frames) + + +def run_pipeline(cfg: ProcessConfig, context_frames: int): + pattern = os.path.join(cfg.fisher_path, "*", "regions_vap", "*", "*_words.txt") + all_files = glob(pattern) + + entries = {"train": [], "val": [], "test": []} + seen = set() + + for vad_file in tqdm(all_files, desc="Building Dataset"): + parts = vad_file.split(os.sep) + part, group = parts[-4], parts[-2] + cid = "_".join(os.path.basename(vad_file).split("_")[:3]) + + if (part, group, cid) in seen: + continue + seen.add((part, group, cid)) + + split = get_dataset_split(group, cfg) + if split not in cfg.active_splits: + continue + + segments = process_entry(part, group, cid, cfg, context_frames) + entries[split].extend(format_meta(s) for s in segments) + + for split in cfg.active_splits: + lines = entries.get(split, []) + if lines: + out_file = os.path.join(cfg.fisher_path, f"{split}.txt") + with open(out_file, "w") as f: + f.writelines(lines) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--config", default="config/yaml/vap_role.yaml") + parser.add_argument( + "--context-sec", + type=float, + default=CONTEXT_SEC, + help="VAD kept on each side of a segment; must cover the widest projection window", + ) + args = parser.parse_args() + yaml_cfg = init_yaml_config(args.config) + + config = ProcessConfig( + fisher_path=yaml_cfg["fisher_path"], + val_groups=VAL_GROUPS, + test_groups=TEST_GROUPS, + active_splits={"train", "val", "test"}, + ) + run_pipeline(config, int(round(args.context_sec * VAD_HZ))) diff --git a/third_party/MuVAP/preprocess/vap/01_create_event.py b/third_party/MuVAP/preprocess/vap/01_create_event.py new file mode 100644 index 0000000..7692e73 --- /dev/null +++ b/third_party/MuVAP/preprocess/vap/01_create_event.py @@ -0,0 +1,268 @@ +import argparse +import os +from dataclasses import dataclass +from glob import glob +from typing import List, Set, Tuple + +from tqdm import tqdm + +from config.setup import init_yaml_config +from preprocess.vap.utils import print_dataset_summary + + +@dataclass +class LogicConfig: + pause_threshold: float = 0.1 + long_pause_threshold: float = 1.0 + max_pause_length: float = 3.0 + active_window: float = 0.5 + purity_win: float = 1.0 + min_speech_density: float = 1.0 + + +@dataclass +class SpeechSegment: + start_t: float + end_t: float + speaker: str + + +@dataclass +class TurnEvent: + part: str + group: str + conv_id: str + region_start: float + region_end: float + prev_spk: str + next_spk: str + label: str + timing: str + gap_type: str + + +def parse_vad_file(file_path: str, speaker_id: str): + segments = [] + with open(file_path, "r") as f: + for line in f: + parts = line.strip().split() + if len(parts) >= 2: + segments.append( + SpeechSegment(float(parts[0]), float(parts[1]), speaker_id) + ) + return segments + + +def merge_into_ipus(segments: List[SpeechSegment], cfg: LogicConfig): + if not segments: + return [] + + sorted_segs = sorted(segments, key=lambda s: s.start_t) + ipus = [] + + first = sorted_segs[0] + c_start, c_end, c_spk = first.start_t, first.end_t, first.speaker + + for seg in sorted_segs[1:]: + if seg.speaker == c_spk and (seg.start_t - c_end) <= cfg.pause_threshold: + c_end = max(c_end, seg.end_t) + else: + ipus.append(SpeechSegment(c_start, c_end, c_spk)) + c_start, c_end, c_spk = seg.start_t, seg.end_t, seg.speaker + + ipus.append(SpeechSegment(c_start, c_end, c_spk)) + return ipus + + +def get_pure_speaker( + ipus: List[SpeechSegment], win_start: float, win_end: float, min_density: float +): + speaker_time = {} + + for ipu in ipus: + if ipu.end_t <= win_start: + continue + if ipu.start_t >= win_end: + break + + o_start = max(ipu.start_t, win_start) + o_end = min(ipu.end_t, win_end) + + if o_end > o_start: + speaker_time[ipu.speaker] = speaker_time.get(ipu.speaker, 0.0) + ( + o_end - o_start + ) + + active_speakers = [spk for spk, duration in speaker_time.items() if duration > 0] + + if len(active_speakers) == 1: + spk = active_speakers[0] + if speaker_time[spk] >= min_density: + return spk + + return None + + +def get_mutual_silences(ipus: List[SpeechSegment]): + if not ipus: + return [] + + sorted_all = sorted(ipus, key=lambda x: x.start_t) + merged_speech = [] + + c_s, c_e = sorted_all[0].start_t, sorted_all[0].end_t + for ipu in sorted_all[1:]: + if ipu.start_t <= c_e: + c_e = max(c_e, ipu.end_t) + else: + merged_speech.append((c_s, c_e)) + c_s, c_e = ipu.start_t, ipu.end_t + merged_speech.append((c_s, c_e)) + + gaps = [] + for i in range(len(merged_speech) - 1): + gaps.append((merged_speech[i][1], merged_speech[i + 1][0])) + + return gaps + + +def get_overlaps(ipus: List[SpeechSegment]): + overlaps = [] + sorted_ipus = sorted(ipus, key=lambda x: x.start_t) + active = [] + + for ipu in sorted_ipus: + active = [a for a in active if a.end_t > ipu.start_t] + + for a in active: + if a.speaker != ipu.speaker: + o_start = ipu.start_t + o_end = min(a.end_t, ipu.end_t) + if o_end > o_start: + overlaps.append((o_start, o_end)) + active.append(ipu) + + return overlaps + + +def extract_events( + ipus: List[SpeechSegment], part: str, group: str, conv_id: str, cfg: LogicConfig +): + events = [] + sorted_ipus = sorted(ipus, key=lambda s: s.start_t) + + gaps = get_mutual_silences(sorted_ipus) + + for g_start, g_end in gaps: + g_dur = g_end - g_start + + prev_spk = get_pure_speaker( + sorted_ipus, g_start - cfg.purity_win, g_start, cfg.min_speech_density + ) + next_spk = get_pure_speaker( + sorted_ipus, g_end, g_end + cfg.purity_win, cfg.min_speech_density + ) + + if prev_spk is not None and next_spk is not None: + label = "HOLD" if prev_spk == next_spk else "SHIFT" + g_type = "LONG" if g_dur >= cfg.long_pause_threshold else "SHORT" + + if g_dur <= cfg.max_pause_length: + events.append( + TurnEvent( + part, + group, + conv_id, + g_start, + g_end, + prev_spk, + next_spk, + label, + "SILENT", + g_type, + ) + ) + + region_start = g_start - cfg.active_window + region_end = g_start + predicted_gap_len = g_dur + cfg.active_window + + if predicted_gap_len <= cfg.max_pause_length: + events.append( + TurnEvent( + part, + group, + conv_id, + region_start, + region_end, + prev_spk, + next_spk, + label, + "ACTIVE", + g_type, + ) + ) + + return events + + +def collect_vad(fisher_path: str, part: str): + seen: Set[Tuple[str, str]] = set() + conv = [] + pattern = os.path.join(fisher_path, part, "regions_vap", "*", "*_words.txt") + + for vad_file in glob(pattern): + group = vad_file.split(os.sep)[-2] + file_name = os.path.basename(vad_file) + if not file_name.endswith("_words.txt"): + continue + + conv_id = "_".join(file_name.split("_")[:3]) + key = (group, conv_id) + if key not in seen: + seen.add(key) + conv.append(key) + + return conv + + +def run_pipeline(fisher_path: str): + cfg = LogicConfig() + lines = [] + all_events = [] + + for partition in ["p1", "p2"]: + conversations = collect_vad(fisher_path, partition) + + for group, conv_id in tqdm(conversations, desc=f"Gen labels {partition}"): + vad_dir = os.path.join(fisher_path, partition, "regions_vap", group) + vad_files = sorted(glob(os.path.join(vad_dir, f"{conv_id}_*_words.txt"))) + if not vad_files: + continue + + all_ipus = [] + for i, file in enumerate(vad_files): + raw_segments = parse_vad_file(file, str(i)) + all_ipus.extend(merge_into_ipus(raw_segments, cfg)) + + events = extract_events(all_ipus, partition, group, conv_id, cfg) + all_events.extend(events) + + for e in events: + lines.append( + f"{e.part} {e.group} {e.conv_id} {e.region_start:.2f} {e.region_end:.2f} " + f"{e.prev_spk} {e.next_spk} {e.timing} {e.gap_type} {e.label}" + ) + + out_path = os.path.join(fisher_path, "turn_events.txt") + with open(out_path, "w") as f: + f.write("\n".join(lines)) + + print_dataset_summary(all_events) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Create Fisher turn events") + parser.add_argument("--config", default="config/yaml/vap_role.yaml") + args = parser.parse_args() + run_pipeline(init_yaml_config(args.config)["fisher_path"]) diff --git a/third_party/MuVAP/preprocess/vap/02_prep_event.py b/third_party/MuVAP/preprocess/vap/02_prep_event.py new file mode 100644 index 0000000..f6a949f --- /dev/null +++ b/third_party/MuVAP/preprocess/vap/02_prep_event.py @@ -0,0 +1,217 @@ +import argparse +import os +from collections import defaultdict +from dataclasses import dataclass +from typing import List + +import numpy as np +import torch +from tqdm import tqdm + +from config.setup import init_yaml_config +from data.utils import load_audio +from preprocess.vap.utils import ( + EVENT_PADDING_SEC, + EVENT_SEG_LEN_SEC, + SAMPLE_RATE, + TEST_GROUPS, + VAL_GROUPS, + ProcessConfig, + event_crop, + get_dataset_split, + print_dataset_summary, +) + +SEG_LEN_SEC: int = EVENT_SEG_LEN_SEC +PADDING_SEC: float = EVENT_PADDING_SEC + + +AUDIO_DIRS = {1: "audio", 2: "audio_stereo"} + + +@dataclass(slots=True) +class TurnEvent: + part: str + group: str + conv_id: str + region_start: float + region_end: float + prev_spk: str + next_spk: str + timing: str + gap_type: str + label: str + crop_start: float = 0.0 + crop_end: float = 0.0 + + +def parse_line(line: str) -> TurnEvent: + tokens = line.strip().split() + if len(tokens) != 10: + return None + + entry = TurnEvent( + part=tokens[0], + group=tokens[1], + conv_id=tokens[2], + region_start=float(tokens[3]), + region_end=float(tokens[4]), + prev_spk=tokens[5], + next_spk=tokens[6], + timing=tokens[7], + gap_type=tokens[8], + label=tokens[9], + ) + + crop = event_crop(entry.timing, entry.region_start, entry.region_end) + if crop is None: + return None + entry.crop_start, entry.crop_end = crop + + return entry + + +def group_by_conversation(entries: List[TurnEvent]): + grouped = defaultdict(list) + for entry in entries: + grouped[(entry.part, entry.group, entry.conv_id)].append(entry) + return dict(grouped) + + +def format_output_line(entry: TurnEvent): + location = f"{entry.part} {entry.group} {entry.conv_id} {entry.region_start:.2f} {entry.region_end:.2f}" + event = f"{entry.prev_spk} {entry.next_spk} {entry.timing} {entry.gap_type} {entry.label}" + return f"{location} {event}" + + +def load_entries_from_disk(cfg: ProcessConfig): + splits = {"train": [], "val": [], "test": []} + path = os.path.join(cfg.fisher_path, cfg.input_file) + + with open(path, "r") as f: + parsed_entries = (parse_line(line) for line in f) + + valid_entries = (e for e in parsed_entries if e is not None) + + for entry in valid_entries: + split_name = get_dataset_split(entry.group, cfg) + if split_name in cfg.active_splits: + splits[split_name].append(entry) + + return splits + + +def save_segment(entry: TurnEvent, audio: np.ndarray, cfg: ProcessConfig, channels: int): + seg_samples = int(SEG_LEN_SEC * SAMPLE_RATE) + start_sample = int(entry.crop_start * SAMPLE_RATE) + end_sample = start_sample + seg_samples + + if end_sample > audio.shape[1]: + return False + + seg = audio[:, start_sample:end_sample] + if seg.shape[1] != seg_samples: + return False + + out_dir = os.path.join( + cfg.fisher_path, + entry.part, + "tune", + AUDIO_DIRS[channels], + entry.group, + entry.conv_id, + ) + os.makedirs(out_dir, exist_ok=True) + + out_path = os.path.join(out_dir, f"{entry.region_start:.2f}.npy") + if not os.path.exists(out_path): + pcm16 = (seg.clamp(-1.0, 1.0) * 32767).round().to(torch.int16) + np.save(out_path, pcm16.numpy()) + + return True + + +def process_split( + split_name: str, entries: List[TurnEvent], cfg: ProcessConfig, channels: int +): + if not entries: + return [], [] + + out_lines: List[str] = [] + successful_events: List[TurnEvent] = [] + grouped_convos = group_by_conversation(entries) + + for (part, group, cid), convo_entries in tqdm( + grouped_convos.items(), desc=f"Processing {split_name}" + ): + audio_sub_dir = "p1" if part == "p1" else "p2" + audio_path = os.path.join( + cfg.fisher_path, audio_sub_dir, "audio", group, f"{cid}.wav" + ) + + if not os.path.exists(audio_path): + continue + + audio, _ = load_audio(audio_path, mono=channels == 1) + + for entry in convo_entries: + if save_segment(entry, audio, cfg, channels): + out_lines.append(format_output_line(entry)) + successful_events.append(entry) + + return out_lines, successful_events + + +def run_pipeline(cfg: ProcessConfig, channels: int = 1): + splits = load_entries_from_disk(cfg) + all_saved_events = [] + + for split_name in cfg.active_splits: + entries = splits.get(split_name, []) + if not entries: + continue + + out_lines, success_events = process_split(split_name, entries, cfg, channels) + all_saved_events.extend(success_events) + + # The manifest lists which events have audio, so it is written once per + # channel layout and stays identical between them. + if out_lines: + out_path = os.path.join(cfg.fisher_path, f"{split_name}_events.txt") + with open(out_path, "w") as f: + f.write("\n".join(out_lines) + "\n") + + if all_saved_events: + print_dataset_summary(all_saved_events) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--config", default="config/yaml/vap_role.yaml") + parser.add_argument( + "--splits", + nargs="+", + choices=("train", "val", "test"), + default=["val", "test"], + help="splits to cut clips for; training reads segments, not events, so " + "train clips are only needed to enlarge the probe's tuning pool", + ) + parser.add_argument( + "--channels", + type=int, + choices=(1, 2), + default=1, + help="1 writes tune/audio for the mono setups, 2 writes tune/audio_stereo for original VAP", + ) + args = parser.parse_args() + yaml_cfg = init_yaml_config(args.config) + + config = ProcessConfig( + fisher_path=yaml_cfg["fisher_path"], + input_file="turn_events.txt", + val_groups=VAL_GROUPS, + test_groups=TEST_GROUPS, + active_splits=set(args.splits), + ) + + run_pipeline(config, channels=args.channels) diff --git a/third_party/MuVAP/preprocess/vap/03_prep_vad_rate.py b/third_party/MuVAP/preprocess/vap/03_prep_vad_rate.py new file mode 100644 index 0000000..e82a7c4 --- /dev/null +++ b/third_party/MuVAP/preprocess/vap/03_prep_vad_rate.py @@ -0,0 +1,112 @@ +"""Write the segment VAD again at a different frame rate. + +`00_prep_fisher.py` rasterizes the word annotations at `VAD_HZ` and stores one +`.npy` per segment beside the audio. A projection window running at another rate +needs its labels on that grid: resampling the 25 Hz array would inherit its +quantisation, so this rasterizes the annotations again from their original +continuous times. + +Only the VAD is rewritten. The audio is rate-independent and stays where it is, +so this touches ~1 GB rather than the 1.12 TB of segments, and it writes to +`seg/vad{rate}/` rather than over `seg/vad/`. + + python preprocess/vap/03_prep_vad_rate.py --frame-hz 12.5 \ + --fisher-path /path/to/fisher --splits train val test +""" + +import argparse +from collections import defaultdict +from pathlib import Path + +import numpy as np + +from preprocess.vap.utils import SAMPLE_RATE + +SEG_SEC = 30 +CONTEXT_SEC = 2.0 + + +def rasterize(words_dir, conversation, frames, frame_hz): + """`[2, frames]` activity for a whole conversation, from the word times.""" + vad = np.zeros((2, frames), dtype=np.float32) + paths = sorted(words_dir.glob(f"{conversation}_*_words.txt")) + if len(paths) != 2: + raise FileNotFoundError( + f"expected two *_words.txt for {conversation} in {words_dir}, found {len(paths)}" + ) + for speaker, path in enumerate(paths): + with open(path) as handle: + for line in handle: + if not line.strip(): + continue + begin, end, *_ = line.split() + first = max(0, int(round(float(begin) * frame_hz))) + last = min(frames, int(round(float(end) * frame_hz))) + if last > first: + vad[speaker, first:last] = 1.0 + return vad + + +def segment_window(vad, start_sec, frame_hz, context_frames): + """The stored slice for one segment: the segment plus context, zero-padded. + + Mirrors `extract_and_save_chunks`, so a reader can treat these files exactly + as it treats the 25 Hz ones. + """ + seg_frames = int(round(SEG_SEC * frame_hz)) + start = int(round(start_sec * frame_hz)) + end = start + seg_frames + lo, hi = max(0, start - context_frames), min(vad.shape[-1], end + context_frames) + chunk = vad[:, lo:hi] + pad_left = max(0, context_frames - start) + pad_right = max(0, (end + context_frames) - vad.shape[-1]) + if pad_left or pad_right: + chunk = np.pad(chunk, ((0, 0), (pad_left, pad_right))) + return chunk + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--fisher-path", required=True) + parser.add_argument("--frame-hz", type=float, required=True) + parser.add_argument("--splits", nargs="+", default=["train", "val", "test"]) + parser.add_argument("--overwrite", action="store_true") + args = parser.parse_args() + + root = Path(args.fisher_path) + label = f"{args.frame_hz:g}".replace(".", "_") + context_frames = int(round(CONTEXT_SEC * args.frame_hz)) + + # One rasterization per conversation, then slice every segment out of it. + by_conversation = defaultdict(list) + for split in args.splits: + for line in open(root / f"{split}.txt"): + if not line.strip(): + continue + part, group, conversation, seg_id, start, *_ = line.split() + by_conversation[(part, group, conversation)].append((int(seg_id), float(start))) + + written = skipped = 0 + for index, ((part, group, conversation), segments) in enumerate(sorted(by_conversation.items()), 1): + out = root / part / "seg" / f"vad{label}" / group / conversation + out.mkdir(parents=True, exist_ok=True) + targets = [(s, p) for s, p in ((s, out / f"{s}.npy") for s, _ in segments) + if args.overwrite or not p.exists()] + if not targets: + skipped += len(segments) + continue + last_start = max(start for _, start in segments) + frames = int(round((last_start + SEG_SEC + CONTEXT_SEC) * args.frame_hz)) + 1 + vad = rasterize(root / part / "regions_vap" / group, conversation, frames, args.frame_hz) + starts = dict(segments) + for seg_id, path in targets: + np.save(path, segment_window(vad, starts[seg_id], args.frame_hz, context_frames)) + written += 1 + if index % 500 == 0: + print(f" {index}/{len(by_conversation)} conversations, {written} written", flush=True) + print(f"done: {written} written, {skipped} already present, " + f"{len(by_conversation)} conversations -> seg/vad{label}/") + + +if __name__ == "__main__": + main() diff --git a/third_party/MuVAP/preprocess/vap/__init__.py b/third_party/MuVAP/preprocess/vap/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/third_party/MuVAP/preprocess/vap/utils.py b/third_party/MuVAP/preprocess/vap/utils.py new file mode 100644 index 0000000..a5da210 --- /dev/null +++ b/third_party/MuVAP/preprocess/vap/utils.py @@ -0,0 +1,98 @@ +from collections import defaultdict +from dataclasses import dataclass +from typing import List + +SAMPLE_RATE: int = 16000 +VAL_GROUPS = {"000", "016", "032", "048", "096", "112"} +TEST_GROUPS = {"064", "080"} + + +@dataclass(slots=True) +class ProcessConfig: + fisher_path: str + val_groups: set[str] + test_groups: set[str] + active_splits: set[str] + input_file: str = "" + + +def get_dataset_split(group: str, cfg: ProcessConfig): + if group in cfg.val_groups: + return "val" + if group in cfg.test_groups: + return "test" + return "train" + + +def print_dataset_summary(events: List, is_turn_event=True): + stats = defaultdict(lambda: defaultdict(int)) + categories = set() + + for e in events: + split = get_dataset_split( + e.group, ProcessConfig("", VAL_GROUPS, TEST_GROUPS, set()) + ) + + cat = f"{e.timing} {e.label}" if is_turn_event else e.label + categories.add(cat) + stats[split][cat] += 1 + stats["TOTAL"][cat] += 1 + + sorted_cats = sorted(list(categories)) + + header = f"{'SPLIT':<10}" + for cat in sorted_cats: + header += f" | {cat:<15}" + header += f" | {'TOTAL'}" + + width = len(header) + print(f"\n{'=' * width}") + print(f"{'DATASET EVENT SUMMARY':^{width}}") + print(f"{'=' * width}") + print(header) + print(f"{'-' * width}") + + for split in ["train", "val", "test", "TOTAL"]: + row = f"{split.upper():<10}" + total = 0 + for cat in sorted_cats: + val = stats[split][cat] + total += val + row += f" | {val:<15}" + row += f" | {total}" + print(row) + print(f"{'=' * width}\n") + + +EVENT_SEG_LEN_SEC: int = 10 +EVENT_PADDING_SEC: float = 0.1 + + +def event_crop( + timing: str, + region_start: float, + region_end: float, + seg_len_sec: float = EVENT_SEG_LEN_SEC, + padding_sec: float = EVENT_PADDING_SEC, +): + """Audio window an evaluation clip covers, or None if the event has none. + + The window ends just before the decision point: a hair into the silence for + a SILENT event, a hair before the speech ends for an ACTIVE one. This is the + single definition of that window - `02_prep_event.py` writes clips with it + and the raw dataloader reads the same span straight from the source audio, + so the two cannot drift apart. + """ + if timing == "SILENT": + crop_end = region_start + padding_sec + if crop_end > region_end: + return None + elif timing == "ACTIVE": + crop_end = region_end - padding_sec + else: + return None + + crop_start = crop_end - seg_len_sec + if crop_start < 0: + return None + return crop_start, crop_end diff --git a/third_party/MuVAP/projection_window.py b/third_party/MuVAP/projection_window.py new file mode 100644 index 0000000..4bb4f82 --- /dev/null +++ b/third_party/MuVAP/projection_window.py @@ -0,0 +1,282 @@ +"""Voice activity projection windows: VAD frames in, one class index per frame. + +A class is an ordered pair of `row_bins`-bit rows. What a row is, per mode: + +* `speaker_based` - one row per channel, in channel order. The original VAP: + 2 speakers x 4 future bins = 256 classes. The rows name channels, not roles, + so hold/shift needs the event's floor holder. +* `role_relative` - N speakers ranked down to a pair on the 1.4, 0.6, 0.6, 1.4 + window and encoded `Scurr` first, then folded onto `16 * 17 / 2 = 136` + classes. Row 1 is the next speaker, so its `p_future` is the shift. +* `independent` - no codebook, the raw per-speaker bins. +""" + +import torch +import torch.nn.functional as F + +P_FUTURE_SEC = 1.4 # hold/shift is decided over the trailing 1.4s +MODES = ("role_relative", "speaker_based", "independent") + + +def bin_times_to_frames(bin_times, frame_hz): + frames = [round(seconds * frame_hz) for seconds in bin_times] + if any(frame <= 0 for frame in frames): + raise ValueError("Projection bins must contain at least one frame") + return frames + + +class ProjectionWindow: + def __init__( + self, + bin_sec=(1.4, 0.6, 0.6, 1.4), + num_hist_bins=2, + frame_hz=25, + threshold=0.5, + mode="role_relative", + cut_center_frame=False, + ): + mode = str(mode).lower() + if mode not in MODES: + raise ValueError(f"mode must be one of {MODES}, got {mode!r}") + if mode == "role_relative" and (len(bin_sec) != 4 or num_hist_bins != 2): + raise ValueError("role_relative needs two history and two future bins") + + self.mode = mode + self.bin_sec = list(bin_sec) + self.frame_hz = frame_hz + self.threshold = threshold + self.num_hist_bins = num_hist_bins + self.cut_center_frame = cut_center_frame + self.bin_frames = bin_times_to_frames(self.bin_sec, frame_hz) + self.n_bins = len(self.bin_frames) + self.hist_frames = sum(self.bin_frames[:num_hist_bins]) + self.fut_frames = sum(self.bin_frames[num_hist_bins:]) + + keeps_history = mode == "role_relative" + self.row_bins = self.n_bins if keeps_history else self.n_bins - num_hist_bins + self.row_hist_bins = num_hist_bins if keeps_history else 0 + self.row_sec = self.bin_sec[self.n_bins - self.row_bins :] + self._row_sec = torch.tensor(self.row_sec) + + if self.is_role_based: + self.codebook, self.fold = self._role_codebook() + elif mode == "speaker_based": + self.codebook, self.fold = self._codebook(2 * self.row_bins), None + else: + self.codebook, self.fold = None, None + self.n_classes = 0 if self.codebook is None else len(self.codebook) + + self._lead = 0 if cut_center_frame or self.hist_frames == 0 else 1 + self._skip = 1 if cut_center_frame else 0 + self._offsets = self._bin_offsets(self.bin_frames) + self._hist_sec = torch.tensor(self.bin_sec[:num_hist_bins]) + self._fut_sec = torch.tensor(self.bin_sec[num_hist_bins:]) + self._weights = 2 ** torch.arange(self.row_bins - 1, -1, -1) + self._pair_weights = 2 ** torch.arange(2 * self.row_bins - 1, -1, -1) + self._slots = {} + + def __repr__(self): + return ( + f"ProjectionWindow(mode={self.mode!r}, frame_hz={self.frame_hz}, " + f"bin_frames={self.bin_frames}, n_classes={self.n_classes})" + ) + + @property + def is_role_based(self): + return self.mode == "role_relative" + + # ---------------------------------------------------------------- codebook + + @staticmethod + def _codebook(bits): + values = torch.arange(2**bits) + return values.unsqueeze(1).bitwise_right_shift( + torch.arange(bits - 1, -1, -1) + ).bitwise_and(1).float() + + def _role_codebook(self): + """`(a, b)` and `(b, a)` share a class: 256 ordered codes -> 136. + + Indices follow the paper's enumeration, by pattern value; the pair each + index *stores* is the ranked one, so `decode` returns `(Scurr, Snext)` + with nothing to recompute. That works because the ranking is a function + of the rows - each role keeps its own history bins - so a pair whose top + row ranks below its bottom row never occurs: exactly 136 of the 256 + ordered states are reachable and the fold is a bijection. + """ + patterns = self._codebook(self.row_bins) + size = len(patterns) + split = self.row_hist_bins + rank = self._activity(patterns[:, :split], self._row_sec[:split]) * 10 + rank += self._activity(patterns[:, split:], self._row_sec[split:]) + + fold = torch.empty(size * size, dtype=torch.long) + codebook = torch.empty(size * (size + 1) // 2, 2 * self.row_bins) + for low in range(size): + for high in range(low, size): + index = low * size - low * (low - 1) // 2 + high - low + fold[low * size + high] = fold[high * size + low] = index + top, bottom = (high, low) if rank[high] >= rank[low] else (low, high) + codebook[index] = torch.cat([patterns[top], patterns[bottom]]) + return codebook, fold + + # ------------------------------------------------------------------ labels + + def _bin_offsets(self, bin_frames): + offsets, start = [], 0 + for width in bin_frames: + shift = self._lead + (self._skip if start >= self.hist_frames else 0) + offsets.append((start + shift, start + shift + width, width)) + start += width + return offsets + + def _totals(self, activity): + """Running frame count; the extra left pad makes the sum exclusive.""" + return F.pad( + activity, + (0, 0, self.hist_frames + 1, self.fut_frames + self._lead + self._skip), + ).cumsum(dim=1) + + def _bins(self, totals, offsets, frames, dtype): + """A bin is two slices of the running count subtracted, any width.""" + return torch.stack( + [ + (totals[:, end : end + frames] - totals[:, start : start + frames]) + >= self.threshold * width + for start, end, width in offsets + ], + dim=-1, + ).to(dtype) + + def extract_bins(self, activity): + """Thresholded `[B, T, S, n_bins]` from `[B, T, S]` activity.""" + return self._bins( + self._totals(activity), self._offsets, activity.shape[1], activity.dtype + ) + + def _activity(self, bins, seconds): + """Seconds of activity: the 2, 1 / 1, 2 weighting of the original code.""" + return (bins * seconds.to(bins.device)).sum(dim=-1) + + def _rank_roles(self, bins): + """History picks `Scurr`; the remaining speakers' future picks `Snext`. + + When the top two tie on history outright, stage 1's runner-up takes the + second slot rather than the speaker with the most future. + """ + history = self._activity(bins[..., : self.num_hist_bins], self._hist_sec) + future = self._activity(bins[..., self.num_hist_bins :], self._fut_sec) + by_history = history * 10 + future + + if bins.shape[2] == 2: # stage 2 has one candidate left + current = (by_history[..., 1:] > by_history[..., :1]).long() + return torch.cat([current, 1 - current], dim=2) + + by_future = future * 10 + history + pair = torch.topk(by_history, k=2, dim=2).indices + current, runner_up = pair[..., :1], pair[..., 1:2] + tied = history.gather(2, current) == history.gather(2, runner_up) + by_future.scatter_(2, current, -torch.inf) + following = torch.where(tied, runner_up, by_future.argmax(dim=2, keepdim=True)) + return torch.cat([current, following], dim=2) + + @torch.no_grad() + def get_labels(self, activity): + """`[B, S, frames]` voice activity to `[B, frames]` class indices.""" + activity = activity.transpose(1, 2) + frames = activity.shape[1] + totals = self._totals(activity) + bins = self._bins(totals, self._offsets, frames, activity.dtype) + + if self.mode == "independent": + if bins.shape[2] != 1: + raise ValueError("independent projection expects one speaker per item") + return bins.squeeze(2) + if bins.shape[2] < 2: + raise ValueError(f"{self.mode} projection requires at least two speakers") + + if self.mode == "speaker_based": # rows are channels, in order + pair = bins[..., self.num_hist_bins :].flatten(start_dim=2) + return (pair * self._pair_weights.to(pair.dtype)).sum(dim=-1).long() + + # Index every row first: gathering two integers beats gathering two rows. + patterns = (bins * self._weights.to(bins.dtype)).sum(dim=-1).long() + speakers = self._rank_roles(bins) + ordered = patterns.gather(2, speakers[..., :1]) * 2**self.row_bins + ordered += patterns.gather(2, speakers[..., 1:]) + return self.fold.to(ordered.device)[ordered.squeeze(-1)] + + # ----------------------------------------------------------------- readout + + def decode(self, indices): + """`(Scurr, Snext)` for the role modes, `(channel 1, channel 2)` otherwise.""" + values = self.codebook.to(indices.device)[indices] + return values.view(*indices.shape, 2, self.row_bins) + + def _default_future_window(self): + """The `p_future` bins: the trailing bins covering the last 1.4s.""" + span, start = 0.0, self.row_bins + while start > self.row_hist_bins and span < P_FUTURE_SEC - 1e-6: + start -= 1 + span += self.row_sec[start] + return start, self.row_bins - 1 + + def _default_now_window(self): + """The `p_now` bins: the rest of the future, ahead of `p_future`.""" + start = self._default_future_window()[0] + return (self.row_hist_bins, start - 1) if start > self.row_hist_bins else None + + def _rows_over(self, probs, window): + """Class probabilities collapsed onto the two rows over one bin window.""" + if window not in self._slots: + states = self.decode(torch.arange(self.n_classes)) + self._slots[window] = states[:, :, window[0] : window[1] + 1].sum(dim=-1) + rows = torch.einsum("...d,dc->...c", probs, self._slots[window].to(probs.device)) + return rows / rows.sum(dim=-1, keepdim=True).clamp_min(1e-8) + + def get_probs(self, logits): + """Row probabilities over `p_now` and `p_future`. + + The two entries are the codebook's two rows: `(Scurr, Snext)` for the + role modes, `(speaker 1, speaker 2)` for `speaker_based`. This is the + plain readout - turning it into hold and shift is `get_shift_hold`. + """ + probs = logits.softmax(dim=-1) + now = self._default_now_window() + return { + "probs": probs, + "p_now": None if now is None else self._rows_over(probs, now), + "p_future": self._rows_over(probs, self._default_future_window()), + } + + @staticmethod + def as_hold_shift(rows, prev_speaker): + """Channel probabilities to (hold, shift), given who held the floor.""" + prev = prev_speaker.to(device=rows.device).long().unsqueeze(-1) + return torch.cat([rows.gather(-1, prev), rows.gather(-1, 1 - prev)], dim=-1) + + def get_shift_hold(self, logits, shift_scale=None, prev_speaker=None): + """Shift probability over `p_future`, and its 0.5-threshold prediction. + + The role modes rank into `(Scurr, Snext)` at encode time, so their two + rows already are (hold, shift). `speaker_based` rows are the channels, + and which one is the hold is a property of the event, so the caller + supplies the floor holder. + """ + output = self.get_probs(logits) + rows = output["p_future"] + + if self.mode == "speaker_based": + if prev_speaker is None: + raise ValueError( + "speaker_based rows are speaker 1 and speaker 2, not hold and " + "shift; pass prev_speaker to say which one held the floor" + ) + rows = self.as_hold_shift(rows, prev_speaker) + + if shift_scale is None: + shift_scale = 2.0 if self.is_role_based else 1.0 + scaled = rows.clone() + scaled[..., 1] *= shift_scale + shift = scaled[..., 1] / scaled.sum(dim=-1).clamp_min(1e-8) + return {**output, "pred": (shift > 0.5).long(), "p_shift": shift} diff --git a/third_party/MuVAP/pyproject.toml b/third_party/MuVAP/pyproject.toml new file mode 100644 index 0000000..7ae01b7 --- /dev/null +++ b/third_party/MuVAP/pyproject.toml @@ -0,0 +1,43 @@ +[build-system] +requires = ["setuptools>=77"] +build-backend = "setuptools.build_meta" + +[project] +name = "muvap" +version = "0.1.0" +description = "MuVAP: Multimodal Multiparty Voice Activity Projection for Turn-taking Prediction in the Wild" +readme = "README.md" +requires-python = ">=3.10" +license = "MIT" +dependencies = [ + "einops>=0.8", + "lightning>=2.5", + "numpy>=1.26", + "opencv-python-headless>=4.10", + "pyyaml>=6.0", + "scikit-learn>=1.5", + "soundfile>=0.12", + "torch>=2.6", + "torchaudio>=2.6", + "torchmetrics>=1.6", + "tqdm>=4.66", +] + +[project.optional-dependencies] +# The Mimi frontend. CPC needs nothing beyond the base install. +mimi = ["transformers>=4.45"] +wandb = ["wandb>=0.19"] + +[tool.setuptools] +# One entry point per module. They share the packages found below. +py-modules = ["projection_window", "train_vap", "train_asd"] + +[tool.setuptools.packages.find] +include = ["config*", "data*", "models*", "preprocess*", "tasks*", "tools*"] + +[tool.ruff] +line-length = 120 +target-version = "py310" + +[tool.ruff.lint] +select = ["E", "F", "I"] diff --git a/third_party/MuVAP/tasks/__init__.py b/third_party/MuVAP/tasks/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/third_party/MuVAP/tasks/asd_metrics.py b/third_party/MuVAP/tasks/asd_metrics.py new file mode 100644 index 0000000..eb00bce --- /dev/null +++ b/third_party/MuVAP/tasks/asd_metrics.py @@ -0,0 +1,25 @@ +import torch + + +def ava_active_speaker_map(scores, labels): + scores = torch.as_tensor(scores, dtype=torch.float64).flatten() + labels = torch.as_tensor(labels, device=scores.device).bool().flatten() + if scores.shape != labels.shape: + raise ValueError("scores and labels must have the same shape") + if not torch.isfinite(scores).all(): + raise ValueError("scores contain NaN or infinite values") + + positives = labels.sum() + if not len(labels) or positives == 0: + return 0.0 + + order = torch.argsort(scores, descending=True, stable=True) + true_positives = labels[order].cumsum(0) + precision = true_positives / torch.arange(1, len(labels) + 1, dtype=torch.float64, device=scores.device) + recall = true_positives / positives + precision = torch.cat([precision.new_zeros(1), precision, precision.new_zeros(1)]) + recall = torch.cat([recall.new_zeros(1), recall, recall.new_ones(1)]) + precision = torch.flip(torch.cummax(torch.flip(precision, dims=[0]), dim=0).values, dims=[0]) + indices = torch.where(recall[1:] != recall[:-1])[0] + 1 + average_precision = ((recall[indices] - recall[indices - 1]) * precision[indices]).sum() + return 100 * average_precision.item() diff --git a/third_party/MuVAP/tasks/asd_task.py b/third_party/MuVAP/tasks/asd_task.py new file mode 100644 index 0000000..cf67e3c --- /dev/null +++ b/third_party/MuVAP/tasks/asd_task.py @@ -0,0 +1,217 @@ +import torch +import torch.distributed as dist +import torch.nn.functional as F +from lightning.pytorch import LightningModule +from torchmetrics.classification import ( + BinaryConfusionMatrix, + BinaryF1Score, + BinaryPrecision, + BinaryRecall, +) + +from models.asd import AudioVisualASD +from tasks.asd_metrics import ava_active_speaker_map +from tasks.ava_evaluation import export_ava_predictions +from tasks.setup import ( + get_adamw_optimizer, + get_cosine_warmup_scheduler, + get_optimizer_cfg, + get_param_groups, +) + + +class ASDTask(LightningModule): + def __init__(self, cfg): + super().__init__() + self.save_hyperparameters(cfg) + self.model = AudioVisualASD(cfg["asd"]) + self.precision = BinaryPrecision() + self.recall = BinaryRecall() + self.f1 = BinaryF1Score() + self.confusion = BinaryConfusionMatrix() + self.validation_entities = [] + evaluation = cfg.get("evaluation", {}) + self.ava_groundtruth = evaluation.get("ava_groundtruth") + self.ava_predictions = evaluation.get("ava_predictions") + + def forward(self, audio, visual): + return self.model(audio, visual, return_embeddings=True) + + @staticmethod + def _masked_bce(logits, targets, scored): + """Average binary cross-entropy over scored frames only.""" + loss = F.binary_cross_entropy_with_logits( + logits, targets.float(), reduction="none" + ) + weight = scored.float() + if loss.dim() == 3: + weight = weight.unsqueeze(-1) + weight = weight.expand_as(loss) + return (loss * weight).sum() / weight.sum().clamp_min(1.0) + + def _shared_step(self, batch): + audio, visual, projection_labels, activity_labels, scored = batch[:5] + projection, audio_logits, visual_logits, _, audio_emb, visual_emb = self( + audio, visual + ) + audio_logits = audio_logits.squeeze(-1) + visual_logits = visual_logits.squeeze(-1) + activity_logits = (audio_logits + visual_logits) / 2 + + projection_loss = self._masked_bce(projection, projection_labels, scored) + audio_loss = self._masked_bce(audio_logits, activity_labels, scored) + visual_loss = self._masked_bce(visual_logits, activity_labels, scored) + contrastive_loss = self._talknce( + audio_emb, visual_emb, activity_labels, scored + ) + loss = projection_loss + 0.4 * audio_loss + 0.4 * visual_loss + 0.3 * contrastive_loss + losses = { + "loss": loss, + "projection_loss": projection_loss, + "audio_loss": audio_loss, + "visual_loss": visual_loss, + "contrastive_loss": contrastive_loss, + } + return activity_logits, losses + + def training_step(self, batch, batch_idx): + _, losses = self._shared_step(batch) + for name, value in losses.items(): + self.log(f"train/{name}", value, prog_bar=name == "loss", sync_dist=True) + return losses["loss"] + + def validation_step(self, batch, batch_idx): + logits, losses = self._shared_step(batch) + labels = batch[3].int() + probabilities = logits.sigmoid() + if len(batch) == 6: + scores = probabilities.detach().cpu() + targets = labels.detach().cpu() + for index, metadata in enumerate(batch[5]): + frames = min(int(metadata["model_frames"]), scores.shape[-1]) + self.validation_entities.append( + ( + metadata["sample_id"], + scores[index, :frames], + targets[index, :frames], + float(metadata["source_fps"]), + ) + ) + scored = batch[4].bool() + flat_scores, flat_labels = probabilities[scored], labels[scored] + self.precision.update(flat_scores, flat_labels) + self.recall.update(flat_scores, flat_labels) + self.f1.update(flat_scores, flat_labels) + self.confusion.update(flat_scores, flat_labels) + for name, value in losses.items(): + self.log(f"val/{name}", value, prog_bar=name == "loss", sync_dist=True) + + def _validation_is_partial(self): + """True when this pass deliberately saw only part of the val set. + + The official export needs one score for every groundtruth entity, so it + cannot run during the sanity check or a truncated debug pass. An + unexpectedly incomplete set still raises inside the exporter. + """ + return bool( + self.trainer.sanity_checking + or self.trainer.fast_dev_run + or self.trainer.limit_val_batches != 1.0 + ) + + def on_validation_epoch_end(self): + if dist.is_available() and dist.is_initialized(): + gathered = [None] * dist.get_world_size() + dist.all_gather_object(gathered, self.validation_entities) + entities = [entity for item in gathered for entity in item] + else: + entities = self.validation_entities + + # The distributed val sampler pads the epoch by repeating samples, so + # the same entity can arrive from several ranks. Ranks are not required + # to agree bit for bit, so the first copy wins. + scores_by_entity = {} + labels_by_entity = {} + fps_by_entity = {} + for entity_id, entity_scores, entity_labels, source_fps in entities: + if entity_id in scores_by_entity: + continue + scores_by_entity[entity_id] = entity_scores + labels_by_entity[entity_id] = entity_labels + fps_by_entity[entity_id] = source_fps + if scores_by_entity: + scores = torch.cat(list(scores_by_entity.values())) + labels = torch.cat(list(labels_by_entity.values())) + else: + scores = torch.empty(0) + labels = torch.empty(0, dtype=torch.int) + map_25hz = ava_active_speaker_map(scores, labels) + + matrix = self.confusion.compute() + false_positive_rate = matrix[0, 1] / matrix[0].sum().clamp_min(1) + self.log("val/mAP_25hz", map_25hz, prog_bar=True) + if ( + self.ava_groundtruth + and self.ava_predictions + and scores_by_entity + and not self._validation_is_partial() + ): + # Only rank zero writes the prediction file, but every rank has to + # log the metric or a ModelCheckpoint monitoring it will hang. + official_map = None + if self.trainer.is_global_zero: + official_map = export_ava_predictions( + self.ava_groundtruth, + self.ava_predictions, + scores_by_entity, + fps_by_entity=fps_by_entity, + ) + if dist.is_available() and dist.is_initialized(): + payload = [official_map] + dist.broadcast_object_list(payload, src=0) + official_map = payload[0] + self.log("val/mAP_official", official_map, prog_bar=True) + self.log("val/precision", self.precision.compute(), sync_dist=True) + self.log("val/recall", self.recall.compute(), sync_dist=True) + self.log("val/f1", self.f1.compute(), prog_bar=True, sync_dist=True) + self.log("val/fpr", false_positive_rate, sync_dist=True) + for metric in ( + self.precision, + self.recall, + self.f1, + self.confusion, + ): + metric.reset() + self.validation_entities.clear() + + @staticmethod + def _talknce(audio, visual, labels, scored=None, temperature=0.07): + positive = labels.bool() + if scored is not None: + positive = positive & scored.bool() + if audio.shape[0] < 2 or not positive.any(): + return audio.sum() * 0 + + audio = F.normalize(audio, dim=-1) + visual = F.normalize(visual, dim=-1) + logits = torch.einsum("btd,ctd->tbc", audio, visual) / temperature + batch_size = audio.shape[0] + targets = torch.arange(batch_size, device=audio.device).repeat(audio.shape[1]) + active = positive.transpose(0, 1).reshape(-1) + return F.cross_entropy(logits.reshape(-1, batch_size)[active], targets[active]) + + def configure_optimizers(self): + cfg = self.hparams["asd"] + weight_decay = float(cfg["weight_decay"]) + optimizer = get_adamw_optimizer( + get_param_groups(self.named_parameters(), weight_decay), + float(cfg["lr"]), + weight_decay, + ) + scheduler = get_cosine_warmup_scheduler( + optimizer, + self.trainer.estimated_stepping_batches, + float(cfg["warmup_ratio"]), + float(cfg["min_lr"]) if cfg.get("min_lr") else None, + ) + return get_optimizer_cfg(optimizer, scheduler) diff --git a/third_party/MuVAP/tasks/ava_evaluation.py b/third_party/MuVAP/tasks/ava_evaluation.py new file mode 100644 index 0000000..a466f69 --- /dev/null +++ b/third_party/MuVAP/tasks/ava_evaluation.py @@ -0,0 +1,130 @@ +import csv +import logging +from collections import Counter, defaultdict +from pathlib import Path + +import torch + +from data.alignment import model_frame_count, model_indices +from tasks.asd_metrics import ava_active_speaker_map + +logger = logging.getLogger(__name__) + +MODEL_FPS = 25.0 + +PREDICTION_FIELDS = [ + "video_id", + "frame_timestamp", + "entity_box_x1", + "entity_box_y1", + "entity_box_x2", + "entity_box_y2", + "label", + "entity_id", + "score", +] + + +def _lift_to_native(entity_id, scores, native_frames, source_fps, model_fps): + """Expand model-rate scores to one score per native groundtruth row.""" + expected_frames = model_frame_count(native_frames, source_fps, model_fps) + if len(scores) != expected_frames: + raise ValueError( + f"{entity_id}: {native_frames} groundtruth rows at {source_fps} fps imply " + f"{expected_frames} model frames, got {len(scores)} scores" + ) + indices = model_indices(native_frames, source_fps, model_fps) + return scores[torch.as_tensor(indices, dtype=torch.long)] + + +def _count_drifting_entities(rows, fps_by_entity, model_fps): + """Count entities whose groundtruth timestamps are not uniformly spaced. + + The lift assumes native frame `j` sits at `t0 + j / source_fps`, matching + the positional mapping used when the sample was packed. A few AVA entities + have gaps in their annotation timeline; they are scored positionally + anyway, but the count belongs in the log so the number stays auditable. + """ + timestamps = defaultdict(list) + for row in rows: + timestamps[row["entity_id"]].append(float(row["frame_timestamp"])) + tolerance = 1.0 / model_fps + drifting = 0 + for entity_id, values in timestamps.items(): + source_fps = fps_by_entity.get(entity_id) + if not source_fps: + continue + origin = values[0] + if any( + abs((value - origin) - index / source_fps) > tolerance + for index, value in enumerate(values) + ): + drifting += 1 + return drifting + + +def export_ava_predictions( + groundtruth_path, + output_path, + scores_by_entity, + fps_by_entity=None, + model_fps=MODEL_FPS, +): + """Write official AVA predictions and return the mAP over every row. + + `scores_by_entity` holds one score per model frame. When `fps_by_entity` is + given, scores are lifted onto each entity's native annotation timeline + first, so the exported file has exactly one row per groundtruth row. + """ + groundtruth_path = Path(groundtruth_path) + output_path = Path(output_path) + rows = list(csv.DictReader(groundtruth_path.open())) + expected = Counter(row["entity_id"] for row in rows) + missing = expected.keys() - scores_by_entity.keys() + extra = scores_by_entity.keys() - expected.keys() + if missing or extra: + raise ValueError(f"AVA entity mismatch: {len(missing)} missing, {len(extra)} extra") + + if fps_by_entity is not None: + absent = expected.keys() - fps_by_entity.keys() + if absent: + raise ValueError(f"AVA source fps missing for {len(absent)} entities") + scores_by_entity = { + entity_id: _lift_to_native( + entity_id, + scores_by_entity[entity_id], + expected[entity_id], + fps_by_entity[entity_id], + model_fps, + ) + for entity_id in expected + } + drifting = _count_drifting_entities(rows, fps_by_entity, model_fps) + if drifting: + logger.info( + "%d of %d AVA entities have non-uniform groundtruth timestamps; " + "scored positionally", + drifting, + len(expected), + ) + + for entity_id, count in expected.items(): + if len(scores_by_entity[entity_id]) != count: + raise ValueError(f"{entity_id}: expected {count} scores, got {len(scores_by_entity[entity_id])}") + + positions = defaultdict(int) + labels = [] + scores = [] + output_path.parent.mkdir(parents=True, exist_ok=True) + with output_path.open("w", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=PREDICTION_FIELDS) + writer.writeheader() + for row in rows: + entity_id = row["entity_id"] + score = float(scores_by_entity[entity_id][positions[entity_id]]) + positions[entity_id] += 1 + prediction = {field: row[field] for field in PREDICTION_FIELDS[:-2]} + writer.writerow({**prediction, "label": "SPEAKING_AUDIBLE", "entity_id": entity_id, "score": score}) + labels.append(row["label"] == "SPEAKING_AUDIBLE") + scores.append(score) + return ava_active_speaker_map(torch.tensor(scores), torch.tensor(labels)) diff --git a/third_party/MuVAP/tasks/setup.py b/third_party/MuVAP/tasks/setup.py new file mode 100644 index 0000000..aa90e50 --- /dev/null +++ b/third_party/MuVAP/tasks/setup.py @@ -0,0 +1,82 @@ +import torch +from sklearn.linear_model import LogisticRegression +from sklearn.preprocessing import StandardScaler + + +class LogisticProber: + def __init__(self): + self.scaler = StandardScaler() + # `penalty="l2"` is the default and was deprecated in scikit-learn 1.8; + # naming it explicitly stops working in 1.10. + self.clf = LogisticRegression( + C=1.0, class_weight="balanced", max_iter=1000, random_state=42 + ) + + def fit_predict(self, X_train, y_train, X_test): + """Predictions on the test events, or None if either pool is too small. + + Returning the predictions rather than one score lets the caller report + the probe with the same metrics as the zero-shot readout. + """ + if len(y_train) < 10 or len(X_test) < 10: + return None + + self.clf.fit(self.scaler.fit_transform(X_train), y_train) + return self.clf.predict(self.scaler.transform(X_test)) + + +def get_param_groups(named_parameters, weight_decay): + """Decay matrices only. + + Norm gains, biases, and embedding tables are 1-D; shrinking them toward + zero costs capacity without the regularizing effect it has on weight + matrices. Frozen parameters are dropped so they never reach the optimizer. + """ + decay, no_decay = [], [] + for _, parameter in named_parameters: + if not parameter.requires_grad: + continue + (decay if parameter.ndim >= 2 else no_decay).append(parameter) + return [ + {"params": decay, "weight_decay": weight_decay}, + {"params": no_decay, "weight_decay": 0.0}, + ] + + +def get_adamw_optimizer(params, lr, weight_decay): + return torch.optim.AdamW(params, lr=lr, weight_decay=weight_decay) + + +def get_cosine_warmup_scheduler(optimizer, total_steps, warmup_ratio=0.1, min_lr=None): + total_steps = max(1, int(total_steps)) + warmup_steps = min(int(total_steps * warmup_ratio), total_steps - 1) + decay_steps = max(1, total_steps - warmup_steps) + + max_lr = optimizer.param_groups[0]["lr"] + min_lr = max_lr * 0.01 if min_lr is None else min_lr + + cosine = torch.optim.lr_scheduler.CosineAnnealingLR( + optimizer, T_max=decay_steps, eta_min=min_lr + ) + + if warmup_steps == 0: + return cosine + + warmup = torch.optim.lr_scheduler.LinearLR( + optimizer, start_factor=0.01, end_factor=1.0, total_iters=warmup_steps + ) + + return torch.optim.lr_scheduler.SequentialLR( + optimizer, schedulers=[warmup, cosine], milestones=[warmup_steps] + ) + + +def get_optimizer_cfg(optimizer, scheduler=None): + cfg = {"optimizer": optimizer} + if scheduler: + cfg["lr_scheduler"] = { + "scheduler": scheduler, + "interval": "step", + "frequency": 1, + } + return cfg diff --git a/third_party/MuVAP/tasks/vap_task.py b/third_party/MuVAP/tasks/vap_task.py new file mode 100644 index 0000000..a0dd9ce --- /dev/null +++ b/third_party/MuVAP/tasks/vap_task.py @@ -0,0 +1,151 @@ +"""Lightning task wrapping any VAP architecture and projection setup. + +Training is a single next-state classification over the configured codebook. +Testing is the papers' turn-taking protocol, zero shot: hold and shift are +decoded straight from the logits at each annotated event and scored on the +SILENT and ACTIVE pools. Two supporting views log beside it under their own +prefixes, so `test/` stays the headline: `probe/` fits a logistic probe on the +last-frame embedding, an upper bound on what the representation carries, and +`ablation/` sweeps the shift prior the VAP papers report. +""" + +from collections import defaultdict + +import numpy as np +import torch +import torch.nn.functional as F +from lightning.pytorch import LightningModule +from sklearn.metrics import balanced_accuracy_score, f1_score + +from models.vap import build_vap, needs_vad +from tasks.setup import ( + LogisticProber, + get_adamw_optimizer, + get_cosine_warmup_scheduler, + get_optimizer_cfg, + get_param_groups, +) + + +# The shift prior the VAP papers sweep over, reported as an ablation. +SHIFT_SCALES = tuple(round(0.1 * step, 1) for step in range(10, 31, 2)) + + +class VAPTask(LightningModule): + def __init__(self, cfg, proj_win): + super().__init__() + self.save_hyperparameters(cfg, ignore=["proj_win"]) + self.model = build_vap(cfg["vap"]) + # The anchored architecture is conditioned on voice activity as well as + # audio, so its batches carry one more tensor than the others'. + self.anchored = needs_vad(cfg["vap"]) + self.proj_win = proj_win + self.tune_outputs = [] + self.test_outputs = [] + + def forward(self, audio, vad=None): + return self.model(audio, vad) if self.anchored else self.model(audio) + + def _shared_step(self, batch): + audio, labels, *rest = batch + logits = self(audio, rest[0] if self.anchored else None) + if self.proj_win.mode == "independent": + return F.binary_cross_entropy_with_logits(logits, labels.float()) + return F.cross_entropy(logits.transpose(1, 2), labels.long()) + + def training_step(self, batch, batch_idx): + loss = self._shared_step(batch) + self.log("train/loss", loss, prog_bar=True, sync_dist=True) + return loss + + def validation_step(self, batch, batch_idx): + loss = self._shared_step(batch) + self.log("val/loss", loss, prog_bar=True, sync_dist=True) + + def test_step(self, batch, batch_idx, dataloader_idx=0): + audio, labels, _, timing, _, prev_spk, *rest = batch + logits, embedding = ( + self.model(audio, rest[0], return_embeddings=True) + if self.anchored + else self.model(audio, return_embeddings=True) + ) + output = { + "embedding": embedding[:, -1].float().detach().cpu(), + "logits": logits[:, -1].float().detach().cpu(), + "labels": labels.detach().cpu(), + "timing": timing, + "prev_spk": prev_spk.detach().cpu(), + } + # The probe is fitted on held-out events and scored on the test events, + # so the two never share a conversation group. + target = self.tune_outputs if dataloader_idx == 0 else self.test_outputs + target.append(output) + + def on_test_epoch_end(self): + tune = self._pools(self.tune_outputs) + for pool, values in self._pools(self.test_outputs).items(): + prev = values["prev_spk"] if self.proj_win.mode == "speaker_based" else None + labels = values["labels"] + + prediction = self.proj_win.get_shift_hold( + values["logits"], prev_speaker=prev + )["pred"].numpy() + self.log(f"test/{pool}/f1_macro", f1_score(labels, prediction, average="macro")) + self.log(f"test/{pool}/bacc", balanced_accuracy_score(labels, prediction)) + + for scale in SHIFT_SCALES: + swept = self.proj_win.get_shift_hold( + values["logits"], shift_scale=scale, prev_speaker=prev + )["pred"].numpy() + self.log( + f"ablation/{pool}/f1_macro_scale_{scale}", + f1_score(labels, swept, average="macro"), + ) + + if pool in tune: + probe = LogisticProber().fit_predict( + tune[pool]["embedding"], tune[pool]["labels"], values["embedding"] + ) + if probe is not None: + self.log(f"probe/{pool}/f1_macro", f1_score(labels, probe, average="macro")) + self.log(f"probe/{pool}/bacc", balanced_accuracy_score(labels, probe)) + self.tune_outputs.clear() + self.test_outputs.clear() + + @staticmethod + def _pools(outputs): + """Every event, then the SILENT and ACTIVE pools on their own.""" + grouped = defaultdict( + lambda: {"embedding": [], "labels": [], "logits": [], "prev_spk": []} + ) + for output in outputs: + for index, label in enumerate(output["labels"]): + for pool in ("all", output["timing"][index]): + values = grouped[pool] + values["embedding"].append(output["embedding"][index].numpy()) + values["labels"].append(float(label)) + values["logits"].append(output["logits"][index]) + values["prev_spk"].append(output["prev_spk"][index]) + + for values in grouped.values(): + values["embedding"] = np.asarray(values["embedding"]) + values["labels"] = np.asarray(values["labels"], dtype=np.int64) + values["logits"] = torch.stack(values["logits"]) + values["prev_spk"] = torch.stack(values["prev_spk"]).long() + return grouped + + def configure_optimizers(self): + cfg = self.hparams["vap"] + weight_decay = float(cfg["weight_decay"]) + optimizer = get_adamw_optimizer( + get_param_groups(self.named_parameters(), weight_decay), + float(cfg["lr"]), + weight_decay, + ) + scheduler = get_cosine_warmup_scheduler( + optimizer, + self.trainer.estimated_stepping_batches, + float(cfg["warmup_ratio"]), + float(cfg["min_lr"]), + ) + return get_optimizer_cfg(optimizer, scheduler) diff --git a/third_party/MuVAP/tools/export_checkpoint.py b/third_party/MuVAP/tools/export_checkpoint.py new file mode 100644 index 0000000..3cc6f09 --- /dev/null +++ b/third_party/MuVAP/tools/export_checkpoint.py @@ -0,0 +1,141 @@ +"""Turn a training checkpoint into something publishable. + +A Lightning checkpoint is a resume artifact: most of it is optimizer moments, +scheduler and loop state, and - because the frozen frontend is an ordinary +submodule - a copy of pretrained weights that carry their own licence. + +This writes the trained weights alone, keyed as the model expects them, beside +the config needed to rebuild it. The frozen frontend is dropped by default and +restored from its own pretrained source when the model is constructed. +""" + +import argparse +import json +import sys +from pathlib import Path + +import torch +import yaml + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from models.release import CONFIG_NAME, PROVENANCE_NAME, WEIGHTS_NAME # noqa: E402 + +MODULE_PREFIX = "model." + + +def strip_prefix(state_dict, prefix=MODULE_PREFIX): + """Drop the LightningModule attribute name from every key.""" + return { + key[len(prefix) :] if key.startswith(prefix) else key: value + for key, value in state_dict.items() + } + + +def build_model(cfg, module): + if module == "asd": + from models.asd import AudioVisualASD + + return AudioVisualASD(cfg["asd"]) + from models.vap import build_vap + + return build_vap(cfg["vap"]) + + +def frozen_keys(model): + return {name for name, parameter in model.named_parameters() if not parameter.requires_grad} + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("checkpoint") + parser.add_argument("output", help="destination directory") + parser.add_argument("--module", choices=("vap", "asd"), default="vap") + parser.add_argument( + "--keep-frozen", + action="store_true", + help="also ship the frozen frontend, making the file self-contained " + "but redistributing pretrained weights under their own licence", + ) + args = parser.parse_args() + + checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=False) + if "state_dict" not in checkpoint: + raise SystemExit(f"{args.checkpoint} is not a Lightning checkpoint") + + cfg = dict(checkpoint.get("hyper_parameters") or {}) + if args.module not in cfg: + raise SystemExit( + f"checkpoint carries no '{args.module}' config; was it trained by train_{args.module}.py?" + ) + + model = build_model(cfg, args.module) + # Ship exactly what the model asks for. A checkpoint also carries buffers + # the model derives at build time, and those are not weights. + wanted = set(model.state_dict()) + weights = {k: v for k, v in strip_prefix(checkpoint["state_dict"]).items() if k in wanted} + + dropped = set() + if not args.keep_frozen: + dropped = frozen_keys(model) & set(weights) + weights = {k: v for k, v in weights.items() if k not in dropped} + + # A checkpoint can also carry buffers this build does not register - ALiBi + # slopes, say, which are recomputed from the head count. They are not trained + # state, and shipping them makes `load_weights` reject the release, so drop + # them; the `missing` check below still proves nothing needed was lost. + recomputed = sorted(set(weights) - set(model.state_dict())) + weights = {k: v for k, v in weights.items() if k not in set(recomputed)} + + # Loading into a freshly built model proves the export is usable: the frozen + # frontend arrives from its own pretrained source, everything else from here. + missing, _ = model.load_state_dict(weights, strict=False) + if set(missing) - dropped: + raise SystemExit(f"weights missing from the export: {sorted(set(missing) - dropped)[:5]}") + if recomputed: + print( + f"dropped {len(recomputed)} buffer(s) the model recomputes: " + f"{', '.join(recomputed[:3])}{' ...' if len(recomputed) > 3 else ''}" + ) + + out = Path(args.output) + out.mkdir(parents=True, exist_ok=True) + torch.save(weights, out / WEIGHTS_NAME) + with (out / CONFIG_NAME).open("w") as handle: + yaml.safe_dump({args.module: cfg[args.module]}, handle, sort_keys=False) + selection = next( + (v for k, v in (checkpoint.get("callbacks") or {}).items() if "ModelCheckpoint" in k), + {}, + ) + score = selection.get("best_model_score") + epoch = checkpoint.get("epoch") + max_epochs = cfg[args.module].get("max_epochs") + provenance = { + "module": args.module, + "selected_on": selection.get("monitor"), + "selected_score": float(score) if score is not None else None, + "epoch": epoch, + "epochs_planned": max_epochs, + # Whether these are the weights training ended on. A checkpoint chosen + # by a validation metric need not be, and a release should say so. + "final_epoch": None if (epoch is None or not max_epochs) else epoch == max_epochs - 1, + "global_step": checkpoint.get("global_step"), + "lightning_version": checkpoint.get("pytorch-lightning_version"), + "parameters": sum(v.numel() for v in weights.values()), + "frozen_frontend_omitted": sorted({k.split(".")[0] for k in dropped}) or None, + } + with (out / PROVENANCE_NAME).open("w") as handle: + json.dump(provenance, handle, indent=2) + + source_mb = Path(args.checkpoint).stat().st_size / 1e6 + exported_mb = (out / WEIGHTS_NAME).stat().st_size / 1e6 + print(f"{args.checkpoint} {source_mb:8.1f} MB") + print(f"{out / WEIGHTS_NAME} {exported_mb:8.1f} MB " + f"({provenance['parameters'] / 1e6:.2f}M params, {exported_mb / source_mb:.0%} of the original)") + if dropped: + print(f"omitted frozen frontend: {provenance['frozen_frontend_omitted']} " + f"({len(dropped)} tensors), restored from its pretrained source on load") + + +if __name__ == "__main__": + main() diff --git a/third_party/MuVAP/tools/fetch_frontends.py b/third_party/MuVAP/tools/fetch_frontends.py new file mode 100644 index 0000000..9094720 --- /dev/null +++ b/third_party/MuVAP/tools/fetch_frontends.py @@ -0,0 +1,44 @@ +"""Download the frozen audio frontends ahead of a run. + +Both frontends fetch their pretrained weights the first time a model is built, +which is fine interactively but not on a compute node with no outbound network, +and not when a typo in a repository id should fail in seconds rather than after +the dataset has been indexed. This does that fetch on its own. +""" + +import argparse +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from models.modules.audio import AUDIO_ENCODERS, build_audio_encoder # noqa: E402 + + +def describe(encoder): + parameters = sum(p.numel() for p in encoder.parameters()) + return f"{parameters / 1e6:.1f}M parameters, {encoder.out_dim}-dim at {encoder.frame_hz} Hz" + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "encoders", + nargs="*", + default=[], + choices=sorted(AUDIO_ENCODERS), + help="which frontends to fetch (default: all)", + ) + args = parser.parse_args() + + for name in args.encoders or sorted(AUDIO_ENCODERS): + print(f"{name}: fetching") + try: + encoder = build_audio_encoder({"audio_encoder": name}) + except Exception as error: # network, missing extra, bad repository id + raise SystemExit(f"{name}: {error}") from error + print(f"{name}: ready - {describe(encoder)}") + + +if __name__ == "__main__": + main() diff --git a/third_party/MuVAP/train_asd.py b/third_party/MuVAP/train_asd.py new file mode 100644 index 0000000..a509409 --- /dev/null +++ b/third_party/MuVAP/train_asd.py @@ -0,0 +1,180 @@ +import argparse +from pathlib import Path + +import torch +from lightning.pytorch import Trainer, seed_everything +from lightning.pytorch.callbacks import ( + Callback, + LearningRateMonitor, + ModelCheckpoint, + RichProgressBar, +) + +from config.setup import load_config +from data.packed_asd_dm import PackedASDDataModule +from models.release import load_weights, merge_config, resolve +from projection_window import ProjectionWindow +from tasks.asd_task import ASDTask + + +class BatchSamplerEpoch(Callback): + """Advance the train batch sampler's epoch so its shuffle changes. + + Lightning only forwards `set_epoch` to `dataloader.sampler` and + `dataloader.batch_sampler.sampler`, so a bare batch sampler never hears + about the epoch and would replay one fixed batch order forever. + """ + + def on_train_epoch_start(self, trainer, pl_module): + sampler = getattr(trainer.train_dataloader, "batch_sampler", None) + set_epoch = getattr(sampler, "set_epoch", None) + if callable(set_epoch): + set_epoch(trainer.current_epoch) + + +def parse_args(): + parser = argparse.ArgumentParser(description="Train or evaluate ASD") + parser.add_argument( + "--config", default=Path(__file__).resolve().parent / "config/yaml/asd.yaml" + ) + parser.add_argument("--test", action="store_true") + parser.add_argument("--checkpoint", help="Lightning .ckpt to resume or evaluate") + parser.add_argument( + "--weights", + help="published release directory (or its weights.pt) to evaluate; " + "the model architecture comes from the release, data paths from --config", + ) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--wandb", action="store_true") + parser.add_argument( + "--wandb-project", + default="MuVAP-ASD", + help="W&B project the run is logged to (with --wandb)", + ) + parser.add_argument("--name", help="run name used for the checkpoint directory") + parser.add_argument( + "--set", + action="append", + default=[], + metavar="KEY=VALUE", + help="override one config value, e.g. --set asd.audio_encoder=mimi", + ) + return parser.parse_args() + + +def apply_overrides(cfg, overrides): + """Apply `--set a.b.c=value` onto a loaded config, parsed as YAML scalars.""" + import yaml + + for override in overrides: + key, _, raw = override.partition("=") + if not _: + raise SystemExit(f"--set expects KEY=VALUE, got {override!r}") + target = cfg + *path, leaf = key.split(".") + for step in path: + if step not in target: + raise SystemExit(f"unknown config key in --set {override!r}") + target = target[step] + if leaf not in target: + raise SystemExit(f"unknown config key in --set {override!r}") + target[leaf] = yaml.safe_load(raw) + return cfg + + +def build_checkpoints(asd_cfg, checkpoint_dir): + """Always keep `last.ckpt`; keep a selected `best.ckpt` only if asked to. + + Two callbacks, not one: a single `ModelCheckpoint` given both a monitor and + `save_last=True` writes a `last.ckpt` that tracks the selected checkpoint + rather than the final epoch, so the final epoch would never be kept. + + `val/mAP_official` is the reported metric; `val/loss` mixes the projection + and contrastive terms and does not track it closely. Setting + `checkpoint_monitor` to nothing says there is no held-out split to select + on, and the final epoch becomes the model the run produced. + """ + last = ModelCheckpoint(dirpath=checkpoint_dir, save_top_k=0, save_last=True) + monitor = asd_cfg.get("checkpoint_monitor", "val/mAP_official") + if not monitor: + return [last] + best = ModelCheckpoint( + monitor=monitor, + mode=asd_cfg.get("checkpoint_mode", "max"), + save_top_k=1, + dirpath=checkpoint_dir, + filename="best", + ) + return [best, last] + + +def main(): + args = parse_args() + if args.test and not (args.checkpoint or args.weights): + raise SystemExit("--test needs --checkpoint (a .ckpt) or --weights (a release)") + + cfg = apply_overrides(load_config(args.config, "asd"), args.set) + weights_path = None + if args.weights: + weights_path, release_config = resolve(args.weights) + cfg = merge_config(cfg, release_config, "asd") + + seed_everything(args.seed, workers=True) + asd_cfg = cfg["asd"] + # Ampere and newer trade a little float32 mantissa for a large matmul speedup. + torch.set_float32_matmul_precision(asd_cfg.get("matmul_precision", "high")) + + logger = None + run_id = args.name or "local" + if args.wandb: + from lightning.pytorch.loggers import WandbLogger + + logger = WandbLogger(project=args.wandb_project, name=args.name) + run_id = args.name or logger.experiment.id + + checkpoint_dir = Path(cfg["checkpoint_dir"]) / run_id + callbacks = [ + *build_checkpoints(asd_cfg, checkpoint_dir), + BatchSamplerEpoch(), + RichProgressBar(), + LearningRateMonitor("step"), + ] + projection = ProjectionWindow(**cfg["asd"]["projection_window"]) + encoder = asd_cfg.get("audio_encoder", "cpc") + print(f"ASD setup: audio_encoder={encoder} {projection}") + + data = PackedASDDataModule(cfg, projection) + model = ASDTask(cfg) + if weights_path is not None: + load_weights(model.model, weights_path) + print(f"loaded release weights from {weights_path}") + trainer = Trainer( + # Checkpoints and Lightning's fallback CSV logs both land here rather + # than as `lightning_logs/` inside the code checkout. + default_root_dir=checkpoint_dir, + max_epochs=asd_cfg["max_epochs"], + logger=logger, + callbacks=callbacks, + log_every_n_steps=100, + accelerator="auto", + devices="auto", + precision=asd_cfg.get("precision", "bf16-mixed"), + gradient_clip_val=asd_cfg.get("gradient_clip_val", 1.0), + accumulate_grad_batches=asd_cfg.get("accumulate_grad_batches", 1), + # The train loader owns a frame-budget batch sampler that already + # shards across ranks; Lightning must not wrap it in its own sampler. + use_distributed_sampler=False, + ) + if args.test: + # A release is already in the model; a ckpt_path would overwrite it. + trainer.validate( + model, + datamodule=data, + ckpt_path=None if weights_path is not None else args.checkpoint, + ) + else: + trainer.fit(model, datamodule=data) + + +if __name__ == "__main__": + main() diff --git a/third_party/MuVAP/train_vap.py b/third_party/MuVAP/train_vap.py new file mode 100644 index 0000000..d174a85 --- /dev/null +++ b/third_party/MuVAP/train_vap.py @@ -0,0 +1,164 @@ +import argparse +from pathlib import Path + +import torch +from lightning.pytorch import Trainer, seed_everything +from lightning.pytorch.callbacks import ( + LearningRateMonitor, + ModelCheckpoint, + RichProgressBar, +) + +from config.setup import load_config +from data.vap_dm import VAPDataModule +from models.release import load_weights, merge_config, resolve +from projection_window import ProjectionWindow +from tasks.vap_task import VAPTask + + +def parse_args(): + parser = argparse.ArgumentParser(description="Train or evaluate VAP") + parser.add_argument( + "--config", default=Path(__file__).resolve().parent / "config/yaml/vap_role.yaml" + ) + parser.add_argument("--test", action="store_true") + parser.add_argument("--checkpoint", help="Lightning .ckpt to resume or evaluate") + parser.add_argument( + "--weights", + help="published release directory (or its weights.pt) to evaluate; " + "the model architecture comes from the release, data paths from --config", + ) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--wandb", action="store_true") + parser.add_argument( + "--wandb-project", + default="VAP", + help="W&B project the run is logged to (with --wandb)", + ) + parser.add_argument("--name", help="run name used for the checkpoint directory") + parser.add_argument( + "--set", + action="append", + default=[], + metavar="KEY=VALUE", + help="override one config value, e.g. --set vap.temporal.pos_encoding=rope", + ) + return parser.parse_args() + + +def apply_overrides(cfg, overrides): + """Apply `--set a.b.c=value` onto a loaded config, parsed as YAML scalars.""" + import yaml + + for override in overrides: + key, _, raw = override.partition("=") + if not _: + raise SystemExit(f"--set expects KEY=VALUE, got {override!r}") + target = cfg + *path, leaf = key.split(".") + for step in path: + if step not in target: + raise SystemExit(f"unknown config key in --set {override!r}") + target = target[step] + if leaf not in target: + raise SystemExit(f"unknown config key in --set {override!r}") + target[leaf] = yaml.safe_load(raw) + return cfg + + +def build_checkpoints(vap_cfg, checkpoint_dir): + """Always keep `last.ckpt`; keep a selected `best.ckpt` only if asked to. + + Two callbacks, not one: a single `ModelCheckpoint` given both a monitor and + `save_last=True` writes a `last.ckpt` that tracks the selected checkpoint + rather than the final epoch, so the final epoch would never be kept. + + Selecting on a validation metric is only meaningful with a held-out split. + Setting `checkpoint_monitor` to nothing says there is none, and the final + epoch becomes the model the run produced. + """ + last = ModelCheckpoint(dirpath=checkpoint_dir, save_top_k=0, save_last=True) + monitor = vap_cfg.get("checkpoint_monitor", "val/loss") + if not monitor: + return [last] + best = ModelCheckpoint( + monitor=monitor, + mode=vap_cfg.get("checkpoint_mode", "min"), + save_top_k=1, + dirpath=checkpoint_dir, + filename="best", + ) + return [best, last] + + +def build_trainer(cfg, logger, checkpoint_dir): + """Trainer whose checkpoints and logs all land under `checkpoint_dir`.""" + vap_cfg = cfg["vap"] + return Trainer( + default_root_dir=checkpoint_dir, + max_epochs=vap_cfg["max_epochs"], + logger=logger, + callbacks=[ + *build_checkpoints(vap_cfg, checkpoint_dir), + RichProgressBar(), + LearningRateMonitor("step"), + ], + log_every_n_steps=100, + accelerator="auto", + devices="auto", + precision=vap_cfg.get("precision", "bf16-mixed"), + gradient_clip_val=vap_cfg.get("gradient_clip_val", 1.0), + accumulate_grad_batches=vap_cfg.get("accumulate_grad_batches", 1), + ) + + +def main(): + args = parse_args() + if args.test and not (args.checkpoint or args.weights): + raise SystemExit("--test needs --checkpoint (a .ckpt) or --weights (a release)") + + cfg = apply_overrides(load_config(args.config, "vap"), args.set) + weights_path = None + if args.weights: + weights_path, release_config = resolve(args.weights) + cfg = merge_config(cfg, release_config, "vap") + seed_everything(args.seed, workers=True) + torch.set_float32_matmul_precision(cfg["vap"].get("matmul_precision", "high")) + + logger = None + run_id = args.name or "local" + if args.wandb: + from lightning.pytorch.loggers import WandbLogger + + logger = WandbLogger(project=args.wandb_project, name=args.name) + run_id = args.name or logger.experiment.id + + projection = ProjectionWindow(**cfg["vap"]["projection_window"]) + print(f"VAP setup: arch={cfg['vap'].get('arch', 'mono')} {projection}") + + data = VAPDataModule(cfg, projection) + model = VAPTask(cfg, projection) + if weights_path is not None: + load_weights(model.model, weights_path) + print(f"loaded release weights from {weights_path}") + trainer = build_trainer(cfg, logger, Path(cfg["checkpoint_dir"]) / run_id) + + if not args.test: + trainer.fit(model, datamodule=data) + data.setup("test") + # A release is already loaded into the model; a ckpt_path would replace it. + if weights_path is not None: + ckpt_path = None + else: + # Without a selected best there is only the final epoch to fall back on. + fallback = "best" if cfg["vap"].get("checkpoint_monitor", "val/loss") else "last" + ckpt_path = args.checkpoint or fallback + trainer.test( + model, + dataloaders=[data.tune_dataloader(), data.test_dataloader()], + ckpt_path=ckpt_path, + ) + + +if __name__ == "__main__": + main() diff --git a/third_party/MuVAP/weights/asd-cpc/config.yaml b/third_party/MuVAP/weights/asd-cpc/config.yaml new file mode 100644 index 0000000..f38d806 --- /dev/null +++ b/third_party/MuVAP/weights/asd-cpc/config.yaml @@ -0,0 +1,62 @@ +asd: + num_workers: 16 + max_epochs: 5 + seed: 0 + audio_encoder: cpc + precision: bf16-mixed + gradient_clip_val: 1.0 + accumulate_grad_batches: 1 + matmul_precision: high + early_stopping_patience: 0 + limit_train_batches: 1.0 + limit_val_batches: 1.0 + train_window: 250 + frame_budget: 2048 + min_batch: 4 + max_batch: 16 + val_frame_budget: 8192 + val_max_batch: 64 + augment: + audio_mix_prob: 0.5 + audio_snr_db: + - -5.0 + - 5.0 + visual: true + crop_scale: + - 0.7 + - 1.0 + rotate_degrees: 15.0 + substitute_audio_prob: 0.0 + lr: 0.001 + weight_decay: 0.01 + warmup_ratio: 0.1 + min_lr: 1.0e-05 + projection_window: + bin_sec: + - 0.8 + - 0.6 + - 0.4 + - 0.2 + - 0.2 + - 0.4 + num_hist_bins: 4 + frame_hz: 25 + threshold: 0.5 + mode: independent + cut_center_frame: false + cross: + d_model: 256 + m_ff: 4 + n_heads: 4 + n_layers: 1 + dropout: 0.1 + activation: GELU + pos_encoding: alibi + temporal: + d_model: 512 + m_ff: 4 + n_heads: 8 + n_layers: 1 + dropout: 0.1 + activation: GELU + pos_encoding: alibi diff --git a/third_party/MuVAP/weights/asd-cpc/provenance.json b/third_party/MuVAP/weights/asd-cpc/provenance.json new file mode 100644 index 0000000..3cd84d6 --- /dev/null +++ b/third_party/MuVAP/weights/asd-cpc/provenance.json @@ -0,0 +1,14 @@ +{ + "module": "asd", + "selected_on": "val/mAP_official", + "selected_score": 90.498291015625, + "epoch": 4, + "epochs_planned": 5, + "final_epoch": true, + "global_step": 87200, + "lightning_version": "2.6.1", + "parameters": 18881376, + "frozen_frontend_omitted": [ + "audio_encoder" + ] +} \ No newline at end of file diff --git a/third_party/MuVAP/weights/asd-cpc/weights.pt b/third_party/MuVAP/weights/asd-cpc/weights.pt new file mode 100644 index 0000000..8bceace Binary files /dev/null and b/third_party/MuVAP/weights/asd-cpc/weights.pt differ diff --git a/third_party/MuVAP/weights/vap-role-cpc/config.yaml b/third_party/MuVAP/weights/vap-role-cpc/config.yaml new file mode 100644 index 0000000..1dfd82e --- /dev/null +++ b/third_party/MuVAP/weights/vap-role-cpc/config.yaml @@ -0,0 +1,42 @@ +vap: + arch: mono + batch_size: 64 + num_workers: 16 + max_epochs: 5 + lr: 0.001 + weight_decay: 0.01 + warmup_ratio: 0.1 + min_lr: 0.0001 + precision: bf16-mixed + gradient_clip_val: 1.0 + accumulate_grad_batches: 1 + matmul_precision: high + limit_train_batches: 1.0 + limit_val_batches: 1.0 + limit_test_batches: 1.0 + early_stopping_patience: 0 + source: npy + swap_channels: false + projection_window: + bin_sec: + - 1.4 + - 0.6 + - 0.6 + - 1.4 + num_hist_bins: 2 + frame_hz: 25 + threshold: 0.5 + mode: role_relative + cut_center_frame: false + temporal: + d_model: 256 + m_ff: 3 + n_heads: 4 + n_layers: 4 + dropout: 0.1 + activation: GELU + causal: true + pos_encoding: alibi + ffn: mlp + norm: layer + max_frames: 4096 diff --git a/third_party/MuVAP/weights/vap-role-cpc/provenance.json b/third_party/MuVAP/weights/vap-role-cpc/provenance.json new file mode 100644 index 0000000..44fee8a --- /dev/null +++ b/third_party/MuVAP/weights/vap-role-cpc/provenance.json @@ -0,0 +1,14 @@ +{ + "module": "vap", + "selected_on": "val/loss", + "selected_score": 1.56204092502594, + "epoch": 4, + "epochs_planned": 5, + "final_epoch": true, + "global_step": 23190, + "lightning_version": "2.6.1", + "parameters": 3055752, + "frozen_frontend_omitted": [ + "audio_encoder" + ] +} \ No newline at end of file diff --git a/third_party/MuVAP/weights/vap-role-cpc/weights.pt b/third_party/MuVAP/weights/vap-role-cpc/weights.pt new file mode 100644 index 0000000..caaf336 Binary files /dev/null and b/third_party/MuVAP/weights/vap-role-cpc/weights.pt differ diff --git a/third_party/README.md b/third_party/README.md index ced1e49..0b8b26c 100644 --- a/third_party/README.md +++ b/third_party/README.md @@ -1,7 +1,9 @@ # Third-party reference code -Reference repositories are inspected here but are not vendored into this -project. They are not runtime dependencies. +Audited reference repositories are vendored here so a normal clone of this +project contains the exact source used by the experiments. The nested +upstream `.git` directories are deliberately removed; each checkout records +its immutable source commit in `.upstream-revision`. The track-GRU design was checked against the official EgoSpeak repository at revision `b5bee35e796d8f2d6aa7da2bf1b07836270f8bfb`: @@ -11,5 +13,33 @@ git clone https://github.com/jun297/EgoSpeak.git third_party/EgoSpeak git -C third_party/EgoSpeak checkout b5bee35e796d8f2d6aa7da2bf1b07836270f8bfb ``` -`third_party/EgoSpeak/` is intentionally ignored. No EgoSpeak source file is -copied or imported by the project implementation. +The project implementation does not import EgoSpeak at runtime. Its source is +retained for implementation comparison and license/provenance auditing. + +The high-rate local AV extractor imports the vendored official MuVAP ASD +implementation at audited revision `a1ef9c8` and its published `asd-cpc` +weights. The MuVAP-only Stage-3 reconstruction additionally uses the published +25 Hz `vap-role-cpc` release. To refresh this vendored copy intentionally, +clone and audit the new revision outside the project first, then update the +source, weights and `.upstream-revision` together: + +```bash +git clone https://github.com/Haotian-Qi/MuVAP.git third_party/MuVAP +git -C third_party/MuVAP checkout a1ef9c8 +python third_party/MuVAP/tools/fetch_frontends.py cpc +huggingface-cli download Haotian-Qi/MuVAP \ + --include 'asd-cpc/*' 'vap-role-cpc/*' \ + --local-dir third_party/MuVAP/weights +``` + +The published `asd-cpc`/`vap-role-cpc` releases and CPC frontend are retained +in this directory. The extractors hash each consumed release into every +feature contract. +It imports this exact checkout directly, so an editable install is not +required. In an established Qwen container, avoid letting MuVAP's package +metadata replace the working PyTorch/Transformers stack. The additional import +requirements are `torchaudio`, `einops`, and `PyYAML`; verify/install only the +missing packages. + +See [`docs/easycom_highrate_person_av.md`](../docs/easycom_highrate_person_av.md) +for the causality contract and experiment commands. diff --git a/train_muvap_egospeak_stage3.py b/train_muvap_egospeak_stage3.py new file mode 100644 index 0000000..0240fd4 --- /dev/null +++ b/train_muvap_egospeak_stage3.py @@ -0,0 +1,423 @@ +"""Train only a MuVAP-style Stage-3 fusion module on EasyCom targets. + +The published VAP and ASD backbones have already been frozen into sidecars by +``extract_muvap_egospeak_features.py``. This program trains only projections, +frame person attention, one causal temporal Transformer and the EgoSpeak +10-horizon head. It is a paper reconstruction because upstream Stage-3 code +is not public. +""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import random +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from torch.optim.lr_scheduler import LambdaLR +from torch.utils.data import DataLoader, Subset +from tqdm.auto import tqdm + +from easycom_head import ( + CLASS_NAMES, + LOSS_NAMES, + NUM_HORIZONS, + egospeak_multihorizon_loss, + four_state_multihorizon_loss, + head_output_names_for_loss, + independent_bce_multihorizon_loss, + miniroad_multihorizon_loss, + normalize_sample_name, + output_classes_for_loss, + output_space_for_loss, +) +from muvap_egospeak_stage3 import ( + MODEL_KIND, + MuVAPEgoSpeakStage3, + MuVAPWindowDataset, + collate_stage3_windows, + evaluate_valid_predictions, + load_muvap_session, + model_config, + predict_loader, +) + + +OFFICIAL_TRAIN = tuple(f"Session_{index}" for index in range(4, 13)) +OFFICIAL_TEST = {f"Session_{index}" for index in range(1, 4)} + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--feature-root", type=Path, required=True) + parser.add_argument("--target-root", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--sessions", nargs="+", default=list(range(4, 13))) + parser.add_argument("--dev-sessions", nargs="+", default=[12]) + parser.add_argument("--epochs", type=int, default=5) + parser.add_argument("--batch-size", type=int, default=16) + parser.add_argument("--learning-rate", type=float, default=1e-3) + parser.add_argument("--min-learning-rate", type=float, default=1e-4) + parser.add_argument("--weight-decay", type=float, default=1e-2) + parser.add_argument("--warmup-ratio", type=float, default=0.1) + parser.add_argument("--gradient-clip", type=float, default=1.0) + parser.add_argument("--history-seconds", type=float, default=8.0) + parser.add_argument("--model-dim", type=int, default=256) + parser.add_argument("--num-heads", type=int, default=4) + parser.add_argument("--temporal-layers", type=int, default=1) + parser.add_argument("--dropout", type=float, default=0.1) + parser.add_argument("--loss", choices=LOSS_NAMES, default="lstr-mce") + parser.add_argument( + "--person-evidence", + choices=("aligned", "none"), + default="aligned", + help="aligned is Stage 3; none is its capacity-matched VAP-only control", + ) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--num-workers", type=int, default=0) + parser.add_argument( + "--max-train-examples", + type=int, + help="Deterministic prefix used only for a smoke test", + ) + parser.add_argument("--no-progress", action="store_true") + args = parser.parse_args(argv) + if args.epochs <= 0 or args.batch_size <= 0 or args.num_workers < 0: + parser.error("epochs/batch-size must be positive and num-workers nonnegative") + if args.learning_rate <= 0 or not 0 <= args.min_learning_rate <= args.learning_rate: + parser.error("Require 0 <= min-learning-rate <= learning-rate") + if args.weight_decay < 0 or not 0 <= args.warmup_ratio < 1: + parser.error("weight-decay/warmup-ratio are invalid") + if args.gradient_clip <= 0 or args.history_seconds <= 0: + parser.error("gradient-clip/history-seconds must be positive") + if args.max_train_examples is not None and args.max_train_examples <= 0: + parser.error("--max-train-examples must be positive") + return args + + +def seed_everything(seed: int) -> None: + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + + +def resolve_device(name: str) -> torch.device: + device = torch.device(name) + if device.type == "cuda" and not torch.cuda.is_available(): + raise RuntimeError(f"CUDA requested but unavailable: {name}") + return device + + +def make_scheduler( + optimizer: torch.optim.Optimizer, + *, + total_steps: int, + warmup_ratio: float, + min_learning_rate: float, + learning_rate: float, +) -> LambdaLR: + warmup_steps = round(total_steps * warmup_ratio) + minimum = min_learning_rate / learning_rate + + def multiplier(step: int) -> float: + if warmup_steps and step < warmup_steps: + return max(1.0 / warmup_steps, (step + 1) / warmup_steps) + progress = (step - warmup_steps) / max(1, total_steps - warmup_steps) + progress = min(1.0, max(0.0, progress)) + return minimum + (1.0 - minimum) * 0.5 * (1.0 + math.cos(math.pi * progress)) + + return LambdaLR(optimizer, multiplier) + + +def checkpoint_payload( + model: MuVAPEgoSpeakStage3, + *, + epoch: int, + best_epoch: int, + best_dev_map: float | None, + history: list[dict[str, Any]], + extraction_contract: dict[str, Any], + train_sessions: list[str], + dev_sessions: list[str], + args: argparse.Namespace, +) -> dict[str, Any]: + return { + "format_version": 1, + "model_kind": MODEL_KIND, + "implementation_scope": ( + "clean-room MuVAP Stage-3 paper reconstruction; frozen released " + "VAP/ASD; EgoSpeak activity head; not upstream Stage-3 code" + ), + "model_state": model.state_dict(), + "model_config": model_config(model), + "feature_contract": extraction_contract, + "epoch": epoch, + "best_epoch": best_epoch, + "best_dev_map": best_dev_map, + "train_sessions": train_sessions, + "dev_sessions": dev_sessions, + "classes": list(CLASS_NAMES), + "loss": args.loss, + "output_space": output_space_for_loss(args.loss), + "head_output_names": list(head_output_names_for_loss(args.loss)), + "alignment": { + "input_hz": 25.0, + "target_hz": 5.0, + "current_cutoff": "task row t ends at (t+1)/5 seconds", + "future_offsets": list(range(1, NUM_HORIZONS + 1)), + }, + "training_config": { + "epochs": args.epochs, + "batch_size": args.batch_size, + "learning_rate": args.learning_rate, + "min_learning_rate": args.min_learning_rate, + "weight_decay": args.weight_decay, + "warmup_ratio": args.warmup_ratio, + "gradient_clip": args.gradient_clip, + "history_seconds": args.history_seconds, + "seed": args.seed, + "max_train_examples": args.max_train_examples, + "person_evidence": args.person_evidence, + }, + "history": history, + } + + +def write_json_summary(path: Path, payload: dict[str, Any]) -> None: + summary = {key: value for key, value in payload.items() if key not in ("model_state", "feature_contract")} + summary["feature_contract"] = payload["feature_contract"] + path.write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n") + + +def main(argv: list[str] | None = None) -> None: + args = parse_args(argv) + seed_everything(args.seed) + device = resolve_device(args.device) + requested = [normalize_sample_name(value) for value in args.sessions] + dev_sessions = [normalize_sample_name(value) for value in args.dev_sessions] + if len(requested) != len(set(requested)) or len(dev_sessions) != len(set(dev_sessions)): + raise ValueError("Duplicate train/dev sessions") + unknown = (set(requested) | set(dev_sessions)) - set(OFFICIAL_TRAIN) + if unknown: + raise ValueError(f"Only EasyCom Sessions 4--12 may train/select: {sorted(unknown)}") + train_sessions = [name for name in requested if name not in set(dev_sessions)] + if not train_sessions: + raise ValueError("No training sessions remain after removing dev sessions") + if set(train_sessions) & set(dev_sessions) or (set(train_sessions) | set(dev_sessions)) & OFFICIAL_TEST: + raise ValueError("Train/dev/test split leakage") + if args.output.exists() or args.output.with_suffix(args.output.suffix + ".json").exists(): + raise FileExistsError(f"Refusing to overwrite {args.output} or its JSON summary") + + print(f"device: {device}") + print(f"model: {MODEL_KIND}") + print(f"train sessions: {', '.join(train_sessions)}") + print(f"dev sessions: {', '.join(dev_sessions) if dev_sessions else 'none'}") + contract = None + train_features = [] + for name in train_sessions: + session = load_muvap_session( + args.feature_root, + args.target_root, + name, + expected_contract=contract, + ) + contract = session.contract + train_features.append(session) + print( + f"{name}: task={len(session.targets)}, 25Hz={len(session.vap)}, " + f"max_people={session.people.shape[1]}" + ) + assert contract is not None + dev_features = [ + load_muvap_session( + args.feature_root, + args.target_root, + name, + expected_contract=contract, + ) + for name in dev_sessions + ] + dimensions = {(item.vap_dim, item.asd_dim) for item in train_features + dev_features} + if len(dimensions) != 1: + raise ValueError(f"Feature dimensions differ across sessions: {dimensions}") + vap_dim, asd_dim = dimensions.pop() + + train_dataset: Any = MuVAPWindowDataset( + train_features, + history_seconds=args.history_seconds, + require_full_future=True, + ) + if args.max_train_examples is not None: + keep = min(len(train_dataset), args.max_train_examples) + train_dataset = Subset(train_dataset, range(keep)) + print(f"SMOKE: restricting training to first {keep} examples") + generator = torch.Generator().manual_seed(args.seed) + train_loader = DataLoader( + train_dataset, + batch_size=args.batch_size, + shuffle=True, + num_workers=args.num_workers, + pin_memory=device.type == "cuda", + persistent_workers=args.num_workers > 0, + generator=generator, + collate_fn=collate_stage3_windows, + ) + dev_dataset = ( + MuVAPWindowDataset( + dev_features, + history_seconds=args.history_seconds, + require_full_future=True, + ) + if dev_features + else None + ) + dev_loader = ( + DataLoader( + dev_dataset, + batch_size=args.batch_size * 2, + shuffle=False, + num_workers=args.num_workers, + pin_memory=device.type == "cuda", + persistent_workers=args.num_workers > 0, + collate_fn=collate_stage3_windows, + ) + if dev_dataset is not None + else None + ) + history_frames = round(args.history_seconds * 25) + model = MuVAPEgoSpeakStage3( + vap_dim=vap_dim, + asd_dim=asd_dim, + model_dim=args.model_dim, + num_heads=args.num_heads, + dropout=args.dropout, + temporal_layers=args.temporal_layers, + history_frames=history_frames, + num_classes=output_classes_for_loss(args.loss), + person_evidence=args.person_evidence, + ).to(device) + trainable = sum(parameter.numel() for parameter in model.parameters() if parameter.requires_grad) + print( + f"training examples: {len(train_dataset)}; 25Hz history={history_frames} " + f"({args.history_seconds:g}s); trainable Stage-3 parameters={trainable:,}" + ) + print( + f"frozen inputs: Z_VAP={vap_dim}D, Z_ASD={asd_dim}D; " + f"person_evidence={args.person_evidence}; loss={args.loss}; " + f"output={output_space_for_loss(args.loss)}" + ) + + optimizer = torch.optim.AdamW( + model.parameters(), + lr=args.learning_rate, + weight_decay=args.weight_decay, + ) + total_steps = max(1, args.epochs * len(train_loader)) + scheduler = make_scheduler( + optimizer, + total_steps=total_steps, + warmup_ratio=args.warmup_ratio, + min_learning_rate=args.min_learning_rate, + learning_rate=args.learning_rate, + ) + loss_function = { + "lstr-mce": egospeak_multihorizon_loss, + "miniroad-normalized": miniroad_multihorizon_loss, + "independent-bce": independent_bce_multihorizon_loss, + "four-state-ce": four_state_multihorizon_loss, + }[args.loss] + args.output.parent.mkdir(parents=True, exist_ok=True) + history: list[dict[str, Any]] = [] + best_dev_map = -math.inf if dev_loader is not None else None + best_epoch = 0 + for epoch in range(1, args.epochs + 1): + model.train() + loss_sum = 0.0 + example_count = 0 + progress = tqdm( + train_loader, + desc=f"MuVAP Stage3 {epoch}/{args.epochs}", + unit="batch", + disable=args.no_progress, + ) + for batch in progress: + vap = batch["vap"].to(device=device, dtype=torch.float32, non_blocking=True) + people = batch["people"].to(device=device, dtype=torch.float32, non_blocking=True) + person_valid = batch["person_valid"].to(device=device, non_blocking=True) + history_valid = batch["history_valid"].to(device=device, non_blocking=True) + targets = batch["targets"].to(device=device, dtype=torch.float32, non_blocking=True) + optimizer.zero_grad(set_to_none=True) + logits = model(vap, people, person_valid, history_valid) + loss = loss_function(logits, targets) + loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), args.gradient_clip) + optimizer.step() + scheduler.step() + count = len(vap) + loss_sum += float(loss.detach()) * count + example_count += count + progress.set_postfix( + loss=f"{float(loss.detach()):.5f}", + lr=f"{scheduler.get_last_lr()[0]:.3g}", + ) + record: dict[str, Any] = { + "epoch": epoch, + "train_loss": loss_sum / example_count, + "learning_rate": scheduler.get_last_lr()[0], + } + if dev_loader is not None: + probabilities, dev_targets, _, _ = predict_loader( + model, + dev_loader, + device=device, + loss_name=args.loss, + ) + dev_metrics = evaluate_valid_predictions(probabilities, dev_targets) + record["dev_valid_mean_map"] = dev_metrics["mean_map"] + record["dev_valid_0.2s_map"] = dev_metrics["horizons"][0]["map"] + improved = dev_metrics["mean_map"] > best_dev_map + if improved: + best_dev_map = dev_metrics["mean_map"] + best_epoch = epoch + else: + improved = epoch == args.epochs + best_epoch = epoch + history.append(record) + if improved: + payload = checkpoint_payload( + model, + epoch=epoch, + best_epoch=best_epoch, + best_dev_map=best_dev_map, + history=history, + extraction_contract=contract, + train_sessions=train_sessions, + dev_sessions=dev_sessions, + args=args, + ) + temporary = args.output.with_suffix(args.output.suffix + ".tmp") + torch.save(payload, temporary) + os.replace(temporary, args.output) + write_json_summary(args.output.with_suffix(args.output.suffix + ".json"), payload) + dev_text = ( + "" + if dev_loader is None + else f", dev_mean_mAP={record['dev_valid_mean_map'] * 100:.5f}%" + ) + print( + f"epoch {epoch}: train_loss={record['train_loss']:.6f}{dev_text}; " + f"best_epoch={best_epoch}" + ) + print(f"Saved best Stage-3 checkpoint: {args.output} (epoch {best_epoch})") + + +if __name__ == "__main__": + main() diff --git a/train_turn_head.py b/train_turn_head.py index 0ccc9ae..542c473 100644 --- a/train_turn_head.py +++ b/train_turn_head.py @@ -8,6 +8,7 @@ decay, linear warmup for 10 epochs, cosine decay, dropout 0.2, and MCE loss. from __future__ import annotations import argparse +import hashlib import json import math import random @@ -23,6 +24,7 @@ from tqdm.auto import tqdm from easycom_head import ( CLASS_NAMES, DATASET_NAMES, + HIGHRATE_PERSON_AV_FUSION_MODES, LAYER_FUSION_MODES, LOCAL_FUSION_MODES, LOCAL_FUSION_SCOPES, @@ -30,6 +32,7 @@ from easycom_head import ( NUM_HORIZONS, OFFICIAL_TRAIN_SESSIONS, STEP_SECONDS, + TRACK_PERSON_FUSION_MODES, TRANSITION_OBJECTIVE_NAMES, TurnPredictionHead, average_precision_metrics, @@ -53,6 +56,12 @@ from easycom_head import ( from local_turn_features import FramePersonFeatures +INITIALIZED_PERSON_RESIDUAL_MODES = ( + *TRACK_PERSON_FUSION_MODES, + *HIGHRATE_PERSON_AV_FUSION_MODES, +) + + def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--feature-root", type=Path, required=True) @@ -69,7 +78,15 @@ def parse_args() -> argparse.Namespace: default="pooled", help=( "pooled keeps the existing baseline; frame-set-attention is F1; " - "track-gru-residual adds a zero-gated causal track correction; " + "track-gru-residual preserves the legacy zero-gated screen; " + "causal-person-adapter-v2 adds a dynamic-gated, zero-output " + "causal track correction to a matched pooled AV checkpoint; " + "causal-person-adapter-egospeak-gru keeps that adapter but uses " + "EgoSpeak MiniROADA's Linear-LN-ReLU-Dropout-GRU-ReLU temporal " + "stem on each anonymous person track; " + "highrate-person-av-attention lets global Omni plus one shared " + "traditional CPC audio stream query causal 25 Hz person visual/ASD " + "states sampled at the 5 Hz cutoffs; " "audio-person-frame-set and gated-person-state-adapter are " "current-sidecar F5/F6 screening variants, not full F5/F6." ), @@ -96,8 +113,8 @@ def parse_args() -> argparse.Namespace: "--initialize-pooled-checkpoint", type=Path, help=( - "Required by track-gru-residual: copy the exact pooled G+A+V base " - "before training the zero-gated track correction." + "Required by person residual modes: copy an exact pooled G+A(+V) " + "base before training the person correction." ), ) parser.add_argument( @@ -223,15 +240,15 @@ def _validate_args(args: argparse.Namespace) -> None: "Person-set --local-fusion modes require --local-feature-root and " "--local-modalities av" ) - if args.local_fusion == "track-gru-residual": + if args.local_fusion in INITIALIZED_PERSON_RESIDUAL_MODES: if args.initialize_pooled_checkpoint is None: raise ValueError( - "track-gru-residual requires --initialize-pooled-checkpoint" + f"{args.local_fusion} requires --initialize-pooled-checkpoint" ) elif args.initialize_pooled_checkpoint is not None or args.train_pooled_base: raise ValueError( "--initialize-pooled-checkpoint/--train-pooled-base require " - "--local-fusion track-gru-residual" + "a person-residual --local-fusion mode" ) if args.epochs <= 0: raise ValueError("--epochs must be positive") @@ -511,7 +528,7 @@ def _evaluate_valid_horizons( person_mask = None if people is not None: rows = torch.arange(start, min(start + batch_size, len(states))) - if head.local_fusion == "track-gru-residual": + if head.local_fusion in TRACK_PERSON_FUSION_MODES: person_batch, person_mask = people.padded_history( rows, max_observations=head.person_history_frames, @@ -635,7 +652,21 @@ def _checkpoint( if args.initialize_pooled_checkpoint is not None else None ), + "initialize_pooled_checkpoint_sha256": getattr( + args, + "initialize_pooled_checkpoint_sha256", + None, + ), + "initialize_pooled_checkpoint_epoch": getattr( + args, + "initialize_pooled_checkpoint_epoch", + None, + ), "train_pooled_base": args.train_pooled_base, + "pooled_base_frozen": bool( + args.local_fusion in INITIALIZED_PERSON_RESIDUAL_MODES + and not args.train_pooled_base + ), "learnable_layer_mixture": args.layer_fusion == "scalar-mixture", "supervision_policy": ( "transition-sparse-manifest" @@ -651,7 +682,7 @@ def _checkpoint( } -def _initialize_track_residual_from_pooled( +def _initialize_person_residual_from_pooled( head: TurnPredictionHead, checkpoint: Path, *, @@ -660,7 +691,7 @@ def _initialize_track_residual_from_pooled( dev_sessions: list[str], local_contract: dict[str, Any] | None, ) -> None: - """Load and optionally freeze the exact pooled base inside a track model.""" + """Load and optionally freeze the exact pooled base in a person model.""" path = checkpoint.expanduser().resolve() if not path.is_file(): @@ -671,18 +702,28 @@ def _initialize_track_residual_from_pooled( config = checkpoint_head_config(payload) if config["local_fusion"] != "pooled": raise ValueError( - "Track residual initialization requires a pooled checkpoint, got " + "Person residual initialization requires a pooled checkpoint, got " f"{config['local_fusion']!r}: {path}" ) if checkpoint_loss_name(payload) != args.loss: - raise ValueError("Pooled checkpoint loss differs from the track experiment") + raise ValueError("Pooled checkpoint loss differs from the person experiment") if payload.get("dataset", args.dataset) != args.dataset: - raise ValueError("Pooled checkpoint dataset differs from the track experiment") + raise ValueError("Pooled checkpoint dataset differs from the person experiment") if payload.get("train_sessions") != train_sessions: raise ValueError("Pooled checkpoint train sessions differ from this run") if payload.get("dev_sessions", []) != dev_sessions: raise ValueError("Pooled checkpoint dev sessions differ from this run") - if payload.get("local_contract") != local_contract: + if args.local_fusion in HIGHRATE_PERSON_AV_FUSION_MODES: + expected_base_contract = ( + {**local_contract, "modalities": "audio"} + if local_contract is not None + else None + ) + expected_local_visual_dim = 0 + else: + expected_base_contract = local_contract + expected_local_visual_dim = head.local_visual_dim + if payload.get("local_contract") != expected_base_contract: raise ValueError("Pooled checkpoint local feature contract differs from this run") expected = { @@ -691,9 +732,11 @@ def _initialize_track_residual_from_pooled( "num_classes": head.num_classes, "num_input_layers": head.num_input_layers, "layer_fusion": head.layer_fusion, - "local_visual_dim": head.local_visual_dim, + "local_visual_dim": expected_local_visual_dim, "local_audio_dim": head.local_audio_dim, "local_projection_dim": head.local_projection_dim, + "dropout": head.dropout, + "transition_auxiliary": head.transition_auxiliary, } mismatches = { key: (config.get(key), value) @@ -702,12 +745,17 @@ def _initialize_track_residual_from_pooled( } if mismatches: raise ValueError( - f"Pooled checkpoint architecture differs from track base: {mismatches}" + f"Pooled checkpoint architecture differs from person base: {mismatches}" ) def is_base_parameter(name: str) -> bool: return name == "layer_logits" or name.startswith( - ("layer_fusion_net.", "local_projections.", "net.") + ( + "layer_fusion_net.", + "local_projections.", + "net.", + "transition_head.", + ) ) source = payload["head"] @@ -728,6 +776,13 @@ def _initialize_track_residual_from_pooled( current[name] = source[name] head.load_state_dict(current) + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + args.initialize_pooled_checkpoint_sha256 = digest.hexdigest() + args.initialize_pooled_checkpoint_epoch = payload.get("epoch") + if not args.train_pooled_base: for name, parameter in head.named_parameters(): if is_base_parameter(name): @@ -932,9 +987,9 @@ def main() -> None: person_track_max_gap_seconds=args.person_track_max_gap_seconds, **local_head_dimensions(local_contract), ).to(device) - if args.local_fusion == "track-gru-residual": + if args.local_fusion in INITIALIZED_PERSON_RESIDUAL_MODES: assert args.initialize_pooled_checkpoint is not None - _initialize_track_residual_from_pooled( + _initialize_person_residual_from_pooled( head, args.initialize_pooled_checkpoint, args=args, @@ -1023,13 +1078,21 @@ def main() -> None: best_epoch = args.epochs if not dev_sessions else 0 for epoch in range(1, args.epochs + 1): head.train() - if args.local_fusion == "track-gru-residual" and not args.train_pooled_base: + if ( + args.local_fusion in INITIALIZED_PERSON_RESIDUAL_MODES + and not args.train_pooled_base + ): # Keep dropout/norm behavior in the frozen pooled checkpoint fixed - # while the new track branch learns a residual correction. + # while the new person branch learns a residual correction. This + # applies to both track-history adapters and the high-rate person + # set; otherwise head.train() would silently turn dropout back on + # inside the supposedly frozen matched base. if head.layer_fusion_net is not None: head.layer_fusion_net.eval() head.local_projections.eval() head.net.eval() + if head.transition_head is not None: + head.transition_head.eval() loss_sum = 0.0 activity_loss_sum = 0.0 transition_loss_sum = 0.0 @@ -1069,7 +1132,7 @@ def main() -> None: person_batch = None person_mask = None if person_features is not None: - if head.local_fusion == "track-gru-residual": + if head.local_fusion in TRACK_PERSON_FUSION_MODES: person_batch, person_mask = person_features.padded_history( example_index_batch, max_observations=head.person_history_frames, @@ -1155,6 +1218,9 @@ def main() -> None: interaction_gates = head.person_interaction_gates() if interaction_gates: record["person_interaction_gates"] = interaction_gates + adapter_diagnostics = head.person_adapter_diagnostics() + if adapter_diagnostics: + record["person_adapter_diagnostics"] = adapter_diagnostics weight_text = ( "" if layer_weights is None @@ -1181,6 +1247,12 @@ def main() -> None: else ", person_gates=" + json.dumps(interaction_gates, sort_keys=True) ) + adapter_text = ( + "" + if not adapter_diagnostics + else ", person_adapter=" + + json.dumps(adapter_diagnostics, sort_keys=True) + ) if dev_sessions: assert dev_states is not None and dev_targets is not None @@ -1199,7 +1271,7 @@ def main() -> None: print( f"epoch {epoch:02d}: train_loss={train_loss:.6f}, " f"dev_valid_mean_mAP={dev_map * 100:.3f}" - f"{transition_text}{weight_text}{gate_text}" + f"{transition_text}{weight_text}{gate_text}{adapter_text}" ) if dev_map > float(best_dev_map): best_dev_map = dev_map @@ -1225,7 +1297,7 @@ def main() -> None: else: print( f"epoch {epoch:02d}: train_loss={train_loss:.6f}" - f"{transition_text}{weight_text}{gate_text}" + f"{transition_text}{weight_text}{gate_text}{adapter_text}" ) history.append(record)