Initial commit: Omni-only turn-taking experiment scaffold.

Add EasyCom prep, Qwen2.5-Omni feature extraction, streaming turn model prototype, and docs for the ICASSP 2027 ego-speak study.

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Xin Wang
2026-08-16 00:02:58 +08:00
commit bde72487bb
9 changed files with 4349 additions and 0 deletions

752
README.md Normal file
View File

@@ -0,0 +1,752 @@
# Omni-only Turn-Taking Experiment
This project tests whether **Qwen2.5-Omni-3B can replace both the handcrafted
audio/visual feature extractors and the downstream RNN/Transformer used by
EgoSpeak** for egocentric turn-taking anticipation.
The main experiment is deliberately simple:
```text
causal audio + video up to time t
|
v
Qwen2.5-Omni-3B Thinker
|
fused state q_t [2048]
|
v
trainable prediction head
|
v
future logits [10 horizons, 3 classes]
```
There is **no RNN after Qwen**. The Omni Thinker is the temporal model. The
only supervised component in the first experiment is a small `10 x 3`
classification head.
## Current status
| File | Status | Purpose |
|---|---|---|
| `smoke_test.py` | Verified on one NVIDIA L20 | Text and 12-second EasyCom audio-video inference |
| `easycom_manifest.py` | Implemented, awaiting dataset-side run | Discover one-minute clips, preserve session timecodes, and check 5 Hz targets |
| `qwen_omni_features.py` | Implemented, not executed here | Load only the 3B Thinker and extract fused states from one AV window |
| `streaming_turn_model.py` | Prototype implemented, not executed here | Produce one causal Omni state every 200 ms, with reference and KV-cache modes |
| `TurnPredictionHead` | Defined, untrained | Map `[2048]` to `[10, 3]` logits |
| Head training and EgoSpeak evaluation script | Not yet implemented | Align saved states with `target_perframe`, train the head, and reproduce Table 2 mAP |
The official full model, Thinker input path, embedded audio, video decoding,
and FlashAttention-2 have passed a 12-second EasyCom smoke test on one NVIDIA
L20. Streaming extraction, cache equivalence, target alignment, training, and
paper metrics remain unverified.
## Why Omni-only instead of Omni + RNN
EgoSpeak separately extracts visual and audio features and then uses GRU,
LSTR, or Mamba for temporal modeling:
```text
ResNet50 video feature + wav2vec2 audio feature -> temporal model -> prediction
```
Qwen2.5-Omni already contains:
- audio and visual encoders;
- joint audio-video temporal positions;
- a causal decoder-only Thinker with KV cache;
- pretrained acoustic, visual, and semantic knowledge.
Adding a GRU after a continuously processed Omni sequence would model history
twice and would make it harder to attribute an improvement to Omni. The primary
comparison is therefore:
```text
EgoSpeak: ResNet50 + wav2vec2 + GRU
Ours: Qwen2.5-Omni Thinker + small prediction head
```
An Omni-feature-plus-GRU model can be kept as an ablation, but it is not the
main architecture.
## EgoSpeak contract that this project reproduces
The inspected reference repository is `/Volumes/760P/Code/EgoSpeak`.
It contains no Qwen or Omni implementation and no raw-media feature extraction
scripts. It distributes precomputed arrays instead:
```text
DATASET/
|- rgb_kinetics_resnet50/ # [T, 2048]
|- audio_wav2vec2/ # [T, 5120]
`- target_perframe/ # [T, 3]
```
The useful part for this project is the exact task interface.
### Time and classes
Both EasyCom and the EgoSpeak `EGO4DSHUFFLE` split use **5 FPS**, so one task
timestep is **0.2 seconds**. The three output columns are:
```text
0: background
1: camera wearer speaking
2: other people speaking
```
They are not strictly one-hot. Background is exclusive, but simultaneous
camera-wearer and other-person speech is stored as `[0, 1, 1]`. The inspected
EasyCom features contain 7,711 such overlap steps. Preserve the full
multi-positive target rather than applying `argmax`.
Relevant EgoSpeak files:
- `LSTR/data/data_info.json`: FPS, classes, and session splits;
- `MiniROAD/configs/EasyCom/miniroad_aa_easycom.yaml`: 128-step training
windows, stride 4, three classes, and anticipation length 10;
- `MiniROAD/datasets/dataset.py`: feature loading and future-label indexing;
- `MiniROAD/criterions/loss.py`: anticipation loss;
- `MiniROAD/trainer/eval.py`: per-horizon AP and mean anticipation mAP.
### Future target alignment
For current task index `i`, the prediction target is:
```python
target[i + 1 : i + 11]
```
Thus the ten output rows represent:
```text
t+0.2, t+0.4, t+0.6, ..., t+2.0 seconds
```
If `states[i]` is verified to correspond to EgoSpeak task index `i`, the
offline head-training tensors should be constructed as follows:
```python
target = target_perframe # [T, 3], multi-positive
usable = min(len(states), len(target)) - 10
x = states[:usable] # [N, 2048]
y = stack(
[target[h : h + usable] for h in range(1, 11)],
axis=1,
) # [N, 10, 3]
```
Do not silently assume that the first saved Omni state matches label index
zero. EgoSpeak does not publish the script that generated its original
ResNet/wav2vec2 features, so the exact zero-time convention is not documented.
Before training, inspect a short clip with known speech boundaries and test the
candidate mappings `states[i] -> target[i]` and `states[i] -> target[i+1]`.
Store the selected offset in every feature file/checkpoint.
### Splits
Use the published split, not a random frame split.
EasyCom:
- train, preserving EgoSpeak's list: `Session_7`, `Session_11`, `Session_10`,
`Session_6`, `Session_9`, `Session_4`, `Session_8`, `Session_5`, `Session_12`;
- test, preserving EgoSpeak's list: `Session_3`, `Session_2`, `Session_1`.
This is the **EgoSpeak comparison split**. The EasyCom dataset repository
itself recommends a different train/validation/test split; do not substitute
that split when reproducing EgoSpeak Table 2.
Ego4D Audio-Visual Diarization:
- use the `EGO4DSHUFFLE` key;
- 346 training clips;
- 87 test clips;
- each clip is approximately five minutes.
Never mix timesteps from the same session across training and validation.
## EasyCom raw-media layout and manifest
The extracted release used here has this relevant structure:
```text
extracted_easycom/
`- Main/
|- Video_Compressed/
| |- Session_1/
| | |- 00-00-000.mp4
| | |- 01-00-000.mp4
| | `- ...
| `- Session_2 ... Session_12/
|- Glasses_Microphone_Array_Audio/
|- Close_Microphone_Audio/
|- Speech_Transcriptions/
`- Metadata/
```
`Video_Compressed` is the AR-glasses wearer's egocentric stream with embedded
stereo AAC audio. Each session is divided into approximately one-minute clips;
the `mm-ss-SSS` filename is the clip's start time on the original recording
timeline. Some intervals are absent because EasyCom redacted one-minute
segments. These raw time gaps are expected.
EgoSpeak does not preserve those empty intervals. Its appendix states that it
crops every EasyCom clip to retain only the 1920-pixel-wide video region and
then merges the retained clips into one dense video per session. Therefore the
EgoSpeak-compatible label time for a clip is its cumulative position in the
concatenation, not the `mm-ss-SSS` source time. The manifest stores both.
Build a raw-video manifest and, when the EgoSpeak labels are available, check
their duration at 5 Hz:
```bash
python easycom_manifest.py \
--dataset-root /data1/yfzx/wangx/Datasets/easycom/extracted_easycom \
--target-root /path/to/EasyCom/target_perframe \
--output manifests/easycom.json
```
Omit `--target-root` for the first media-only inventory. The manifest reports
raw-source discontinuities as informational and computes each clip's dense
EgoSpeak concatenation offset and corresponding 5 Hz target index. For the
audio-only ablation, use the same embedded glasses audio as A+V (extract it to
mono WAV if needed); do not switch to close microphones, because that changes
the sensor and makes the modality ablation unfair.
Audit downloaded EgoSpeak arrays, and optionally compare every session length
against the raw manifest:
```bash
python audit_egospeak_features.py \
--feature-root /path/to/EgoSpeak/dataset/EasyCom \
--raw-manifest manifests/easycom.json \
--output reports/egospeak_easycom_audit.json
```
The inspected release contains 95,453 synchronized steps across all 12
sessions: 5.302944 hours at 5 Hz. Every session has identical RGB, flow, audio,
and target length. Shapes are `[T,2048]`, `[T,2048]`, `[T,5120]`, and `[T,3]`.
## Recommended comparison matrix
| Experiment | Visual sampling | Audio | Output rate |
|---|---:|---|---:|
| EgoSpeak official A+V | 5 Hz ResNet features | 5 Hz wav2vec2 features | 5 Hz |
| Omni audio-only | none | raw glasses-wearer waveform | 5 Hz |
| Omni A+V primary | 2 FPS | raw embedded waveform | 5 Hz |
| Omni A+V FPS ablation | 5 FPS | raw embedded waveform | 5 Hz |
Two FPS is the primary Omni setting because it matches Qwen's default video
sampling and permits longer context. Five FPS is the visual-rate ablation.
Input feature rates need not match EgoSpeak; the split, causal prediction
times, three classes, ten target horizons, and evaluator must match exactly.
## Prepare an EgoSpeak-compatible EasyCom session
Do not pass the raw one-minute files independently to Omni and do not preserve
the gaps implied by their filenames. EgoSpeak concatenates the retained clips
into a dense session timeline and crops the participant-ID banner on the left.
`prepare_easycom_session.py` reproduces this policy and creates videos whose
stored frames already lie on a fixed session-global 2 FPS or 5 FPS grid. This
prevents Qwen's per-window video loader from choosing a new adaptive set of
frames for every overlapping prefix.
Session 1 has been checked against the local release:
```text
24 retained clips
27,768 source frames at 20 FPS
/ 4 frames per task step
= 6,942 steps at 5 Hz
= target_perframe/Session_1.npy length exactly
```
Preview the commands and validate frames against the target without encoding:
```bash
python prepare_easycom_session.py \
--session-dir /Volumes/MX500B/Code/easycom/Main/Video_Compressed/Session_1 \
--output-dir prepared_easycom/Session_1 \
--target /Volumes/MX500B/Code/EgoSpeak/dataset/EasyCom/target_perframe/Session_1.npy \
--fps 2 5 \
--dry-run
```
Remove `--dry-run` to run FFmpeg. The outputs are:
```text
Session_1_2fps.mp4 primary Omni A+V input
Session_1_5fps.mp4 visual-rate ablation
Session_1.json clip order, exact T, crop and sampling provenance
Session_1.ffconcat reproducible dense concatenation order
```
The script fails rather than silently padding or truncating when
`sum(source_frames) / 4 != target.shape[0]`.
Each clip's decoded video and audio are trimmed to the exact duration implied
by its video frame count before concatenation. This prevents AAC priming/tail
padding from accumulating at the 24 clip boundaries and causing progressive
audio/video drift.
On the remote L20 machine, use its own dataset paths. For example:
```bash
python prepare_easycom_session.py \
--session-dir /data1/yfzx/wangx/Datasets/easycom/extracted_easycom/Main/Video_Compressed/Session_1 \
--output-dir /data1/yfzx/wangx/Datasets/easycom/omni_prepared/Session_1 \
--target /path/to/EgoSpeak/dataset/EasyCom/target_perframe/Session_1.npy \
--fps 2 5
```
This is the command that performs the potentially long FFmpeg processing; it
has not been run by the code authoring workflow.
If encoding completed but a later validation/metadata step failed, update the
script and reuse the existing MP4 without encoding it again:
```bash
python prepare_easycom_session.py ... --fps 2 --reuse-existing
```
## What Qwen's native 2-second chunk means
Qwen2.5-Omni is **not restricted to receiving audio once every two seconds**.
It can receive an arbitrary-duration AV input. The two-second value describes
internal processing:
1. the audio encoder divides acoustic features into approximately two-second
attention blocks;
2. audio and video tokens are interleaved in two-second temporal groups before
entering the Thinker.
The audio encoder attention inside a block is non-causal. Therefore this is
invalid for online evaluation:
```text
encode the complete [0, 2.0 s] block
-> read a purported state at 0.2 s
```
That early state can contain audio from 0.2--2.0 seconds. At prediction time
`t`, the model must receive only audio and video with timestamps `<= t`.
## Causal 5 Hz extraction modes
`streaming_turn_model.py` implements three modes.
### 1. `causal-prefix`: correctness and semantic baseline
At every 200 ms, independently process the causal context ending at the
current time:
```text
[max(0, t-20 s), t] AV -> Omni -> q_t
```
Properties:
- strict no-future input;
- one continuous media window per prediction;
- no cross-forward cache;
- very expensive because overlapping context is recomputed;
- preferred reference when checking semantics and causality.
Example:
```bash
python streaming_turn_model.py \
--video /path/to/easycom_session.mp4 \
--use-audio-in-video \
--mode causal-prefix \
--start 0 \
--end 20 \
--memory-seconds 20 \
--output features/easycom_20s_prefix.pt
```
### 2. `commit-replay`: closed-turn compatibility baseline
The extractor retains Thinker KV for completed two-second media blocks. During
the unfinished block it branches from the same committed cache and evaluates
only the available prefix:
```text
committed historical KV + current block [0.0, 0.2] -> q_t
committed historical KV + current block [0.0, 0.4] -> q_t
...
committed historical KV + current block [0.0, 2.0] -> q_t
-> commit media KV
```
The temporary readout/assistant query is never committed. Only the completed
media block enters long-term cache. This avoids recomputing the full history,
but still re-encodes at most the current two-second prefix.
Example with numerical checking:
```bash
python streaming_turn_model.py \
--video /path/to/easycom_session.mp4 \
--use-audio-in-video \
--mode commit-replay \
--start 0 \
--end 20 \
--memory-seconds 20 \
--verify-every 5 \
--verify-continuous-every 5 \
--output features/easycom_20s_stream.pt
```
`--verify-every 5` compares cached processing with a from-scratch replay of the
same sequence of two-second media turns. The saved cosine similarity should be
extremely close to 1. If it is not, the TMRoPE/cache path is invalid and must
not be used for a paper result.
`--verify-continuous-every 5` is a different check: it compares the cached
sequence of two-second media turns with a single continuous causal window.
This comparison need not be numerically identical, but it measures whether the
input organization materially changes the representation. Save and report the
two check types separately (`cache_replay` and `continuous_window_semantics`).
There is an important experimental distinction: `commit-replay` represents
history as a sequence of two-second media turns. Its numerical check proves
cache equivalence to that same turn sequence, not equivalence to one continuous
20-second video placeholder. Compare it against `causal-prefix` on identical
clips and report the representation if it changes mAP.
### 3. `open-stream-replay`: recommended efficient mode
This mode keeps one user/media turn open across the whole bounded-memory
segment. The first block contributes the chat and AV opening markers once;
later completed two-second blocks contribute only interleaved audio/video
payload tokens:
```text
open AV prefix
completed [0,2] payload -> commit KV
current [2,2.2] payload -> temporary branch -> close/read q_t
current [2,2.4] payload -> temporary branch -> close/read q_t
...
completed [2,4] payload -> commit KV
AV/chat closure -> never committed; readout branch only
```
The audio encoder accepts a shorter final chunk with padding and an
actual-length mask. For the two-frame visual patch requirement, strict-grid
mode duplicates only the most recent past frame when necessary.
Run a four-second validation before using it for extraction:
```bash
python streaming_turn_model.py \
--model-id /path/to/Qwen2.5-Omni-3B \
--video prepared_easycom/Session_1/Session_1_2fps.mp4 \
--use-audio-in-video \
--mode open-stream-replay \
--start 0 --end 4 \
--video-fps 2 \
--strict-fixed-video-grid \
--verify-every 5 \
--verify-continuous-every 5 \
--flash-attention-2 \
--output features/Session_1_open_first_4s.pt
```
`cache_replay` must be identical to replaying the same open stream from an
empty cache. `continuous_window_semantics` compares it with one normal
continuous Qwen AV input. Do not assume equivalence: the modality encoders are
still invoked per tail block in the efficient path, so this numerical check is
required even though artificial user-turn boundaries have been removed.
## Whole-session prediction and the 20-second limit
Twenty seconds is a practical memory/context setting for the 3090, not the
maximum video duration. The complete session is processed sequentially and a
prediction is emitted every 0.2 seconds.
The legacy `commit-replay` prototype resets Thinker KV at each non-overlapping
20-second boundary:
```text
session: 0 -------- 20 -------- 40 -------- 60 ...
cache: [segment 1] [segment 2] [segment 3] ...
output: continuous predictions across the entire session
```
This preserves complete output coverage but limits available history after a
reset. Report the memory length and additionally evaluate the first two seconds
after each reset versus all other timesteps. Rolling KV eviction is a possible
future optimization, but requires correct multimodal position rebasing.
`open-stream-replay` instead uses blockwise rolling memory. With a 20-second
budget it rebuilds the open cache from the latest nine completed blocks
(18 seconds) when a new block starts, then adds the current 0--2 second tail.
Its effective history therefore varies from 18 to 20 seconds and never drops
to an empty 0.2-second context at a 20-second boundary.
For the paper experiment, drive the output length from the preparation sidecar
rather than a floating MP4 duration:
```bash
python streaming_turn_model.py \
--session-metadata prepared_easycom/Session_1/Session_1.json \
--use-audio-in-video \
--mode open-stream-replay \
--memory-seconds 20 \
--video-fps 2 \
--strict-fixed-video-grid \
--verify-every 25 \
--verify-continuous-every 25 \
--output features/Session_1_omni_2fps.pt
```
The sidecar selects the matching fixed-FPS MP4 and sets `N` to its exact
`num_steps`. Alternatively pass `--video ... --target ...` explicitly. A raw
20 FPS MP4 with `--video-fps 2` is rejected when
`--strict-fixed-video-grid` is enabled.
In strict A+V mode the extractor does not delegate temporal frame selection to
the generic video utility. It reads the already-fixed MP4 grid explicitly,
uses only frames with timestamps `<= t`, and duplicates the most recent past
frame when the visual tower needs an even frame count. Audio is loaded once at
16 kHz and sliced exactly at each causal boundary. The processor still receives
`use_audio_in_video=True`, so Qwen's audio/video token interleaving is retained.
Before a full session, run only the first four seconds and inspect both check
types:
```bash
python streaming_turn_model.py \
--model-id /data1/yfzx/wangx/Models/Qwen2.5-Omni-3B \
--video /data1/yfzx/wangx/Datasets/easycom/omni_prepared/Session_1/Session_1_2fps.mp4 \
--use-audio-in-video \
--mode open-stream-replay \
--start 0 --end 4 \
--video-fps 2 \
--strict-fixed-video-grid \
--verify-every 5 \
--verify-continuous-every 5 \
--flash-attention-2 \
--output features/Session_1_first_4s_check.pt
```
Do not launch all sessions unless `cache_replay` cosine is nearly 1. The
`continuous_window_semantics` cosine answers a different question and may be
lower because media-turn boundaries change the representation.
Video can be sampled at 2 FPS while predictions remain at 5 Hz. Audio and the
causal boundary update every 200 ms; consecutive prediction times may use the
same most-recent visual frame. A frame with timestamp later than `t` must never
be included.
## Saved streaming tensors
Each `.pt` file contains:
- `times`: `[N]`, prediction times in seconds;
- `states`: `[N, 2048]`, causal fused Thinker states;
- `label_indices`: `[0, ..., N-1]`, direct `target_perframe` row alignment;
- `valid_10_horizon_steps`: `N-10`, because the last ten current rows do not
have all future labels;
- `mode`: `causal-prefix`, `commit-replay`, or `open-stream-replay`;
- `stream_representation`: explicit input-organization metadata;
- `config`: step, native block, memory, FPS, and pixel settings;
- `verification`: cached-versus-exact numerical checks;
- `input_video_fps` and `video_sampling_grid`: sampling provenance;
- `logits`: optional `[N, 10, 3]` output if a trained head is supplied;
- `probabilities`: optional softmax probabilities.
## Training the `10 x 3` head
The head must be trained. It is not part of pretrained Qwen.
The first experiment should freeze all Qwen parameters and train only:
```text
LayerNorm(2048)
-> Linear(2048, 1024)
-> GELU
-> Dropout
-> Linear(1024, 30)
-> reshape [10, 3]
```
`streaming_turn_model.py` exposes:
```python
from streaming_turn_model import (
TurnPredictionHead,
egospeak_multihorizon_loss,
)
```
Recommended initial procedure:
1. Freeze Qwen and extract all training-session states offline.
2. Verify state/label time offset and truncate only the final ten invalid
current positions.
3. Keep EgoSpeak targets as float multi-positive tensors `[N, 10, 3]`; never
convert them with `argmax`.
4. Train with `egospeak_multihorizon_loss`. It exactly mirrors MiniROAD's
anticipation objective: L2-normalize each three-class target, then compute
its weighted negative log-softmax.
5. Select hyperparameters using only training sessions. For EasyCom, reserve a
training session for development or use session-level cross-validation;
never tune on Sessions 1--3.
6. After choosing the epoch/hyperparameters, optionally retrain on all nine
EasyCom training sessions and evaluate the fixed model once on the test set.
Class-balanced/focal loss, a learned turn-query token, and LoRA are later
ablations. They should not replace the frozen-Omni linear/MLP-head baseline.
## EgoSpeak-compatible evaluation
For every valid current timestep, produce probabilities shaped:
```text
[number of timesteps, 10 future horizons, 3 classes]
```
Then reproduce EgoSpeak's Table 2 procedure:
1. for each horizon independently, concatenate test-session predictions and
ground truth;
2. compute per-class frame-level average precision;
3. average the three class AP values to obtain horizon mAP;
4. report all ten horizon mAP values;
5. average the ten horizon mAP values for the final score.
With EgoSpeak's `ignore_index=-100`, the released evaluator does not remove
class index zero, so background is included. Reuse the same implementation or
verify a reimplementation numerically.
Reference Table 2 targets from the EgoSpeak paper:
- EasyCom GRU A+V: **60.6 mean mAP**;
- EasyCom Transformer A+V: **58.7 mean mAP**;
- Ego4D Transformer audio: **69.2 mean mAP**;
- Ego4D Transformer A+V: **69.0 mean mAP**;
- Ego4D GRU audio: **69.2 mean mAP**;
- Ego4D GRU A+V: **68.2 mean mAP**.
Do not compare these numbers with EgoSpeak Table 3 target-speaker-only AP;
Table 3 uses a different metric.
## Experiments that demonstrate Omni's semantic advantage
The headline comparison alone cannot establish that improvements come from
semantics. Run the following controlled experiments:
1. **Omni audio vs Omni audio+video**: measure whether visual grounding helps.
2. **Per-horizon curve**: semantic completion should be especially visible at
longer horizons such as 1.2--2.0 seconds.
3. **Context length**: compare 4, 10, 15, and 20 seconds without adding an RNN.
4. **Layer ablation**: compare audio-encoder output and several Thinker layers,
because an intermediate layer may be more discriminative than the final
generation-oriented layer.
5. **Semantic-completion subset**: contrast completed utterances with
hesitation/incomplete pauses while matching pause duration and basic
prosody as closely as possible.
6. **Question/statement/backchannel subsets**: report target-speaker AP and
premature-interruption rate on linguistically meaningful events.
7. **Causal transcript baseline**: compare Omni audio with causal ASR text plus
a language model to isolate lexical/semantic contribution.
The frozen-Omni head is the cleanest representation test. LoRA can be reported
as a second-stage ceiling, not as the only result.
## Single-window feature inspection
`qwen_omni_features.py` is useful for inspecting a short window or layer, but
it is not the main whole-session experiment.
If audio is embedded in the video:
```bash
python qwen_omni_features.py \
--video /path/to/window.mp4 \
--use-audio-in-video \
--output features/window.pt
```
If audio and video are separate:
```bash
python qwen_omni_features.py \
--video /path/to/window.mp4 \
--audio /path/to/window.wav \
--output features/window.pt
```
Do not provide the same audio both as a separate file and via
`--use-audio-in-video`.
Saved single-window fields are:
- `token_features`: `[B, token_length, 2048]`;
- `last_token`: `[B, 2048]`;
- `tail_mean`: `[B, 2048]`;
- `attention_mask`: valid-token mask.
## Installation
Install `ffmpeg`, then:
```bash
pip install -r requirements.txt
```
The implementation is pinned to `transformers==4.52.3` because the
commit/replay path uses version-specific Thinker internals. On Linux,
`qwen-omni-utils[decord]` is recommended for faster video decoding. On a
compatible NVIDIA GPU, optionally install `flash-attn` and pass
`--flash-attention-2`.
## Ego4D Audio-Visual Diarization download
Do not download the complete Ego4D release for this experiment. Download only
the compressed Audio-Visual Diarization benchmark clips and annotations:
```bash
ego4d \
--output_directory=/path/to/ego4d \
--datasets clips annotations \
--benchmarks AV
```
The approximately 1 TB figure on the Ego4D download page refers to all
benchmark clips combined, not only the AV subset. The annotated AVD subset has
about 572 five-minute clips (roughly 50 hours). EgoSpeak uses a custom 346/87
train/test split, about 36.1 hours total. Reproducing its Table 2 requires the
exact IDs in EgoSpeak's `EGO4DSHUFFLE` split.
## Required validation before a full run
Do not launch full EasyCom/Ego4D extraction until all of these pass:
- no input audio sample or video frame has timestamp later than prediction
time `t`;
- `states` and `target_perframe` lengths/time origins are reconciled and the
chosen label offset is saved;
- commit/replay cosine equivalence is near 1 on short clips;
- commit/replay and causal-prefix predictions are compared on identical clips;
- a 20-second run fits the 3090 with the selected pixel budget and attention
implementation;
- saved outputs contain exactly ten future horizons in the order 0.2--2.0 s;
- the local metric reproduces EgoSpeak's evaluator on known predictions;
- test session IDs never enter training, validation, normalization statistics,
or hyperparameter selection.
Once these checks pass, the first publishable experiment is:
```text
frozen Qwen2.5-Omni-3B Thinker
+ trained 10x3 head
+ 20-second causal memory
+ full-session 5 Hz output
+ exact EgoSpeak split and Table 2 metric
```

169
audit_egospeak_features.py Normal file
View File

@@ -0,0 +1,169 @@
#!/usr/bin/env python3
"""Audit EgoSpeak EasyCom feature rates, shapes, labels, and raw alignment.
Example:
python audit_egospeak_features.py \
--feature-root /path/to/EgoSpeak/dataset/EasyCom \
--raw-manifest manifests/easycom.json \
--output reports/egospeak_easycom_audit.json
"""
from __future__ import annotations
import argparse
import json
from pathlib import Path
from typing import Any
import numpy as np
from easycom_manifest import ALL_SESSIONS, TASK_FPS, TEST_SESSIONS, TRAIN_SESSIONS
FEATURE_SPECS = {
"rgb_kinetics_resnet50": 2048,
"flow_kinetics_resnet50": 2048,
"audio_wav2vec2": 5120,
"target_perframe": 3,
}
def load_raw_steps(path: Path | None) -> dict[str, int]:
if path is None:
return {}
payload = json.loads(path.read_text())
return {
session["session"]: int(session["expected_5hz_steps"])
for session in payload["sessions"]
if "expected_5hz_steps" in session
}
def audit(
feature_root: Path,
*,
raw_manifest: Path | None,
) -> dict[str, Any]:
feature_root = feature_root.resolve()
raw_steps = load_raw_steps(raw_manifest)
sessions: list[dict[str, Any]] = []
issues: list[str] = []
all_targets: list[np.ndarray] = []
for session in ALL_SESSIONS:
arrays: dict[str, np.ndarray] = {}
shapes: dict[str, list[int]] = {}
dtypes: dict[str, str] = {}
for modality, expected_dim in FEATURE_SPECS.items():
path = feature_root / modality / f"{session}.npy"
if not path.is_file():
issues.append(f"{session}: missing {path}")
continue
array = np.load(path, mmap_mode="r")
arrays[modality] = array
shapes[modality] = list(array.shape)
dtypes[modality] = str(array.dtype)
if array.ndim != 2 or array.shape[1] != expected_dim:
issues.append(
f"{session}: {modality} expected [T,{expected_dim}], got {array.shape}"
)
if len(arrays) != len(FEATURE_SPECS):
continue
lengths = {name: int(array.shape[0]) for name, array in arrays.items()}
if len(set(lengths.values())) != 1:
issues.append(f"{session}: unequal modality lengths {lengths}")
target = np.asarray(arrays["target_perframe"])
if not np.all((target == 0) | (target == 1)):
issues.append(f"{session}: target contains values outside {{0,1}}")
allowed_rows = {(1.0, 0.0, 0.0), (0.0, 1.0, 0.0), (0.0, 0.0, 1.0), (0.0, 1.0, 1.0)}
observed_rows = {tuple(row) for row in np.unique(target, axis=0)}
unexpected = observed_rows - allowed_rows
if unexpected:
issues.append(f"{session}: unexpected target rows {sorted(unexpected)}")
all_targets.append(target)
num_steps = lengths["target_perframe"]
record: dict[str, Any] = {
"session": session,
"split": "train" if session in TRAIN_SESSIONS else "test",
"num_steps": num_steps,
"duration_seconds_at_5hz": num_steps / TASK_FPS,
"shapes": shapes,
"dtypes": dtypes,
"class_positive_steps": target.sum(axis=0).astype(int).tolist(),
"overlap_wearer_other_steps": int(
np.sum((target[:, 1] == 1) & (target[:, 2] == 1))
),
}
if session in raw_steps:
record["raw_manifest_expected_steps"] = raw_steps[session]
record["feature_minus_raw_steps"] = num_steps - raw_steps[session]
if abs(record["feature_minus_raw_steps"]) > 1:
issues.append(
f"{session}: feature/raw mismatch "
f"{record['feature_minus_raw_steps']:+d} steps"
)
sessions.append(record)
combined = np.concatenate(all_targets, axis=0)
unique_rows, row_counts = np.unique(combined, axis=0, return_counts=True)
split_steps = {
split: sum(item["num_steps"] for item in sessions if item["split"] == split)
for split in ("train", "test")
}
return {
"feature_root": str(feature_root),
"raw_manifest": str(raw_manifest.resolve()) if raw_manifest is not None else None,
"task_fps": TASK_FPS,
"step_seconds": 1.0 / TASK_FPS,
"sessions": sessions,
"total_steps": int(combined.shape[0]),
"total_duration_hours": float(combined.shape[0] / TASK_FPS / 3600),
"split_steps": split_steps,
"class_positive_steps": combined.sum(axis=0).astype(int).tolist(),
"target_row_counts": {
str(row.astype(int).tolist()): int(count)
for row, count in zip(unique_rows, row_counts)
},
"issues": issues,
}
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--feature-root", type=Path, required=True)
parser.add_argument("--raw-manifest", type=Path)
parser.add_argument("--output", type=Path)
return parser.parse_args()
def main() -> None:
args = parse_args()
report = audit(args.feature_root, raw_manifest=args.raw_manifest)
if args.output is not None:
args.output.parent.mkdir(parents=True, exist_ok=True)
args.output.write_text(json.dumps(report, indent=2) + "\n")
print(f"Saved report to {args.output}")
print(
f"{report['total_steps']} steps at {report['task_fps']} Hz = "
f"{report['total_duration_hours']:.6f} hours"
)
print(f"split steps: {report['split_steps']}")
print(f"class-positive steps [background, wearer, other]: {report['class_positive_steps']}")
print(f"target rows: {report['target_row_counts']}")
if report["issues"]:
print(f"Found {len(report['issues'])} issue(s):")
for issue in report["issues"]:
print(f" - {issue}")
else:
if report["raw_manifest"] is not None:
print("All modality lengths, dimensions, labels, and raw steps match.")
else:
print("All feature-modality lengths, dimensions, and labels are internally consistent.")
if __name__ == "__main__":
main()

321
easycom_manifest.py Normal file
View File

@@ -0,0 +1,321 @@
#!/usr/bin/env python3
"""Build and validate an EgoSpeak-compatible EasyCom raw-media manifest.
EasyCom stores each session as approximately one-minute MP4 segments named by
their session-relative start time (``mm-ss-SSS.mp4``). This script preserves
that global time origin and optionally checks EgoSpeak ``target_perframe``
arrays at 5 Hz.
Example:
python easycom_manifest.py \
--dataset-root /data1/yfzx/wangx/Datasets/easycom/extracted_easycom \
--target-root /path/to/EasyCom/target_perframe \
--output easycom_manifest.json
"""
from __future__ import annotations
import argparse
import json
import re
import subprocess
from fractions import Fraction
from pathlib import Path
from typing import Any
TASK_FPS = 5
TRAIN_SESSIONS = (
"Session_7",
"Session_11",
"Session_10",
"Session_6",
"Session_9",
"Session_4",
"Session_8",
"Session_5",
"Session_12",
)
TEST_SESSIONS = ("Session_3", "Session_2", "Session_1")
ALL_SESSIONS = TRAIN_SESSIONS + TEST_SESSIONS
TIMECODE_RE = re.compile(r"^(\d{2})-(\d{2})-(\d{3})$")
def parse_timecode(stem: str) -> float:
match = TIMECODE_RE.fullmatch(stem)
if match is None:
raise ValueError(f"Expected mm-ss-SSS filename, got {stem!r}")
minutes, seconds, milliseconds = (int(value) for value in match.groups())
if seconds >= 60:
raise ValueError(f"Invalid seconds in timecode {stem!r}")
return minutes * 60.0 + seconds + milliseconds / 1000.0
def probe_video(path: Path) -> dict[str, Any]:
completed = subprocess.run(
[
"ffprobe",
"-v",
"error",
"-select_streams",
"v:0",
"-show_entries",
"stream=width,height,avg_frame_rate,nb_frames,duration:format=duration",
"-of",
"json",
str(path),
],
check=True,
capture_output=True,
text=True,
)
payload = json.loads(completed.stdout)
streams = payload.get("streams", [])
if not streams:
raise RuntimeError(f"No video stream found in {path}")
stream = streams[0]
rate_text = stream.get("avg_frame_rate", "0/0")
if rate_text in ("0/0", "N/A"):
raise RuntimeError(f"No usable video frame rate in {path}")
fps = float(Fraction(rate_text))
frame_text = stream.get("nb_frames")
if frame_text not in (None, "N/A"):
num_frames = int(frame_text)
else:
duration_text = stream.get("duration") or payload.get("format", {}).get("duration")
if duration_text is None:
raise RuntimeError(f"No frame count or duration in {path}")
num_frames = round(float(duration_text) * fps)
# EgoSpeak processes EasyCom at 20 video frames/s and groups four frames
# per 5 Hz task step. The released MP4s are nominally already 20 FPS.
duration = num_frames / fps
frames_at_20hz = round(duration * 20.0)
width = int(stream["width"])
height = int(stream["height"])
return {
"width": width,
"height": height,
"fps": fps,
"num_frames": num_frames,
"video_duration_seconds": duration,
"frames_at_20hz": frames_at_20hz,
# EasyCom videos may contain an ID banner to the left of the 1920px
# video. EgoSpeak crops it and retains the rightmost video region.
"egospeak_crop": {
"x": max(0, width - 1920),
"y": 0,
"width": min(1920, width),
"height": height,
},
}
def find_target(target_root: Path, session: str) -> Path | None:
direct = target_root / f"{session}.npy"
if direct.is_file():
return direct
matches = sorted(target_root.rglob(f"{session}.npy"))
if len(matches) > 1:
raise RuntimeError(
f"Multiple targets found for {session}: "
+ ", ".join(str(path) for path in matches)
)
return matches[0] if matches else None
def target_metadata(path: Path) -> dict[str, Any]:
try:
import numpy as np
except ImportError as exc:
raise RuntimeError("Install numpy to inspect target_perframe arrays") from exc
target = np.load(path, mmap_mode="r")
if target.ndim not in (1, 2):
raise ValueError(f"Unexpected target shape in {path}: {target.shape}")
if target.ndim == 2 and target.shape[1] != 3:
raise ValueError(f"Expected three classes in {path}, got {target.shape}")
return {
"path": str(path.resolve()),
"shape": list(target.shape),
"num_steps": int(target.shape[0]),
"duration_seconds_at_5hz": float(target.shape[0] / TASK_FPS),
}
def build_manifest(
dataset_root: Path,
*,
target_root: Path | None,
probe_media: bool,
) -> dict[str, Any]:
dataset_root = dataset_root.resolve()
video_root = dataset_root / "Main" / "Video_Compressed"
if not video_root.is_dir():
raise FileNotFoundError(f"EasyCom video directory not found: {video_root}")
sessions: list[dict[str, Any]] = []
all_issues: list[str] = []
all_discontinuities: list[str] = []
for session in ALL_SESSIONS:
session_dir = video_root / session
if not session_dir.is_dir():
all_issues.append(f"missing session directory: {session_dir}")
continue
clip_paths = sorted(
session_dir.glob("*.mp4"), key=lambda path: parse_timecode(path.stem)
)
if not clip_paths:
all_issues.append(f"no MP4 clips in {session_dir}")
continue
clips: list[dict[str, Any]] = []
session_issues: list[str] = []
session_discontinuities: list[str] = []
previous_source_end: float | None = None
concatenated_20hz_frames = 0
for path in clip_paths:
source_start = parse_timecode(path.stem)
media = probe_video(path) if probe_media else None
duration = media["video_duration_seconds"] if media is not None else None
source_end = source_start + duration if duration is not None else None
if previous_source_end is not None:
delta = source_start - previous_source_end
if abs(delta) > 0.05:
kind = "gap" if delta > 0 else "overlap"
session_discontinuities.append(
f"{kind} of {abs(delta):.3f}s before {path.name}"
)
if source_end is not None:
previous_source_end = source_end
record: dict[str, Any] = {
"path": str(path.resolve()),
"timecode": path.stem,
"source_session_start_seconds": source_start,
}
if media is not None and source_end is not None:
concat_start_frame = concatenated_20hz_frames
concatenated_20hz_frames += int(media["frames_at_20hz"])
record.update(
{
**media,
"source_session_end_seconds": source_end,
"egospeak_concat_start_seconds": concat_start_frame / 20.0,
"egospeak_concat_end_seconds": concatenated_20hz_frames / 20.0,
"target_start_index": concat_start_frame // 4,
"target_end_index_exclusive": concatenated_20hz_frames // 4,
}
)
clips.append(record)
source_session_end = (
max(float(clip["source_session_end_seconds"]) for clip in clips)
if probe_media
else None
)
split = "train" if session in TRAIN_SESSIONS else "test"
session_record: dict[str, Any] = {
"session": session,
"split": split,
"clips": clips,
"issues": session_issues,
"source_timeline_discontinuities": session_discontinuities,
}
if source_session_end is not None:
session_record["source_session_end_seconds"] = source_session_end
session_record["egospeak_concat_duration_seconds"] = (
concatenated_20hz_frames / 20.0
)
session_record["expected_5hz_steps"] = concatenated_20hz_frames // 4
if target_root is not None:
target_path = find_target(target_root.resolve(), session)
if target_path is None:
session_issues.append(f"target_perframe missing for {session}")
else:
target = target_metadata(target_path)
session_record["target"] = target
if source_session_end is not None:
difference = target["num_steps"] - (concatenated_20hz_frames // 4)
session_record["target_minus_media_steps"] = difference
# Small encoder/container rounding differences are normal;
# anything beyond one task step needs manual inspection.
if abs(difference) > 1:
session_issues.append(
f"target/media length mismatch: {difference:+d} steps"
)
all_issues.extend(f"{session}: {issue}" for issue in session_issues)
all_discontinuities.extend(
f"{session}: {item}" for item in session_discontinuities
)
sessions.append(session_record)
return {
"dataset": "EasyCom",
"dataset_root": str(dataset_root),
"video_root": str(video_root),
"task_fps": TASK_FPS,
"prediction_step_seconds": 1.0 / TASK_FPS,
"class_names": [
"background",
"camera wearer speaking",
"other people speaking",
],
"future_target_offsets": list(range(1, 11)),
"train_sessions": list(TRAIN_SESSIONS),
"test_sessions": list(TEST_SESSIONS),
"sessions": sessions,
"source_timeline_discontinuities": all_discontinuities,
"issues": all_issues,
}
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--dataset-root", type=Path, required=True)
parser.add_argument("--target-root", type=Path)
parser.add_argument("--output", type=Path, required=True)
parser.add_argument(
"--no-probe",
action="store_true",
help="Skip ffprobe duration and continuity checks.",
)
return parser.parse_args()
def main() -> None:
args = parse_args()
manifest = build_manifest(
args.dataset_root,
target_root=args.target_root,
probe_media=not args.no_probe,
)
args.output.parent.mkdir(parents=True, exist_ok=True)
args.output.write_text(json.dumps(manifest, indent=2) + "\n")
num_clips = sum(len(session["clips"]) for session in manifest["sessions"])
print(
f"Saved {len(manifest['sessions'])} sessions / {num_clips} clips "
f"to {args.output}"
)
if manifest["issues"]:
print(f"Found {len(manifest['issues'])} issue(s):")
for issue in manifest["issues"]:
print(f" - {issue}")
else:
print("No missing-session or target-length issues found.")
discontinuities = manifest["source_timeline_discontinuities"]
if discontinuities:
print(
f"Observed {len(discontinuities)} source-timeline discontinuity(ies); "
"EgoSpeak removes these when concatenating retained clips."
)
if __name__ == "__main__":
main()

357
prepare_easycom_session.py Normal file
View File

@@ -0,0 +1,357 @@
#!/usr/bin/env python3
"""Prepare one EasyCom session for causal Qwen2.5-Omni extraction.
The released EasyCom session consists of retained, approximately one-minute
MP4 clips. EgoSpeak removes source-timeline gaps by concatenating those clips
and crops the left participant-ID banner. This script reproduces that dense
timeline and writes fixed-global-rate 2 FPS and/or 5 FPS AV files.
It intentionally never fills gaps inferred from clip filenames. The number
of 5 Hz task steps is derived from ``target_perframe`` (preferred) or from the
sum of the released 20 FPS video frames.
"""
from __future__ import annotations
import argparse
import json
import math
import shlex
import subprocess
from dataclasses import asdict, dataclass
from fractions import Fraction
from pathlib import Path
from easycom_manifest import parse_timecode
SOURCE_FPS = 20
TASK_FPS = 5
KEEP_WIDTH = 1920
KEEP_HEIGHT = 1080
@dataclass(frozen=True)
class ClipInfo:
path: str
timecode_seconds: float
width: int
height: int
fps: float
frames: int
def _run_json(command: list[str]) -> dict:
completed = subprocess.run(
command, check=True, capture_output=True, text=True
)
return json.loads(completed.stdout)
def probe_clip(path: Path) -> ClipInfo:
payload = _run_json(
[
"ffprobe",
"-v",
"error",
"-select_streams",
"v:0",
"-show_entries",
"stream=width,height,avg_frame_rate,nb_frames",
"-of",
"json",
str(path),
]
)
streams = payload.get("streams", [])
if not streams:
raise RuntimeError(f"No video stream in {path}")
stream = streams[0]
if stream.get("nb_frames") in (None, "N/A"):
raise RuntimeError(f"Exact nb_frames is required for {path}")
fps = float(Fraction(stream["avg_frame_rate"]))
return ClipInfo(
path=str(path.resolve()),
timecode_seconds=parse_timecode(path.stem),
width=int(stream["width"]),
height=int(stream["height"]),
fps=fps,
frames=int(stream["nb_frames"]),
)
def load_target_steps(path: Path) -> int:
try:
import numpy as np
except ImportError as exc:
raise RuntimeError("numpy is required when --target is used") from exc
target = np.load(path, mmap_mode="r")
if target.ndim != 2 or target.shape[1] != 3:
raise ValueError(f"Expected target shape [T,3], got {target.shape}")
return int(target.shape[0])
def ffconcat_quote(path: str) -> str:
# ffconcat uses single-quoted paths; escape an embedded quote as '\''.
return "'" + path.replace("'", "'\\''") + "'"
def validate_inventory(clips: list[ClipInfo], target_steps: int | None) -> int:
if not clips:
raise ValueError("No clips found")
for clip in clips:
if not abs(clip.fps - SOURCE_FPS) < 1e-6:
raise ValueError(f"Expected 20 FPS, got {clip.fps} in {clip.path}")
if clip.width < KEEP_WIDTH or clip.height != KEEP_HEIGHT:
raise ValueError(
f"Cannot take rightmost {KEEP_WIDTH}x{KEEP_HEIGHT} from "
f"{clip.width}x{clip.height}: {clip.path}"
)
total_frames = sum(clip.frames for clip in clips)
if total_frames % (SOURCE_FPS // TASK_FPS):
raise ValueError(
f"Total source frames {total_frames} are not divisible into 5 Hz steps"
)
media_steps = total_frames // (SOURCE_FPS // TASK_FPS)
if target_steps is not None and media_steps != target_steps:
raise ValueError(
f"Dense media gives {media_steps} steps but target has {target_steps}"
)
return media_steps
def build_ffmpeg_command(
clips: list[ClipInfo],
output: Path,
*,
fps: float,
dense_duration: float,
) -> list[str]:
# Trim every decoded stream independently before concatenation. Released
# AAC clips can contain roughly one codec-padding frame beyond their video
# duration; concatenating them first would accumulate AV drift.
command = [
"ffmpeg",
"-hide_banner",
"-loglevel",
"warning",
"-y",
]
for clip in clips:
command.extend(["-i", clip.path])
filters: list[str] = []
concat_inputs: list[str] = []
for index, clip in enumerate(clips):
duration = clip.frames / SOURCE_FPS
filters.append(
f"[{index}:v:0]trim=start=0:end={duration:.9f},"
f"setpts=PTS-STARTPTS,"
f"crop={KEEP_WIDTH}:{KEEP_HEIGHT}:iw-{KEEP_WIDTH}:0[v{index}]"
)
filters.append(
f"[{index}:a:0]atrim=start=0:end={duration:.9f},"
f"asetpts=PTS-STARTPTS[a{index}]"
)
concat_inputs.extend([f"[v{index}]", f"[a{index}]"])
filters.append(
"".join(concat_inputs)
+ f"concat=n={len(clips)}:v=1:a=1[vcat][acat]"
)
# fps=start_time=0 establishes one global visual lattice after the clips
# have become a dense timeline. The final atrim is a defensive guard
# against decoder rounding at the last sample.
filters.append(
f"[vcat]fps=fps={fps}:start_time=0:round=down,format=yuv420p[vout]"
)
filters.append(
f"[acat]atrim=start=0:end={dense_duration:.9f},"
"asetpts=PTS-STARTPTS[aout]"
)
command.extend(
[
"-filter_complex",
";".join(filters),
"-map",
"[vout]",
"-map",
"[aout]",
"-c:v",
"libx264",
"-preset",
"medium",
"-crf",
"18",
"-c:a",
"aac",
"-b:a",
"128k",
"-movflags",
"+faststart",
str(output),
]
)
return command
def validate_output(path: Path, *, fps: float, dense_duration: float) -> dict:
payload = _run_json(
[
"ffprobe",
"-v",
"error",
"-show_entries",
"stream=index,codec_type,avg_frame_rate,nb_frames,duration,sample_rate,channels",
"-of",
"json",
str(path),
]
)
video = next(
(stream for stream in payload.get("streams", []) if stream["codec_type"] == "video"),
None,
)
audio = next(
(stream for stream in payload.get("streams", []) if stream["codec_type"] == "audio"),
None,
)
if video is None or audio is None:
raise RuntimeError(f"Prepared file must contain video and audio: {path}")
actual_fps = float(Fraction(video["avg_frame_rate"]))
actual_frames = int(video["nb_frames"])
# FFmpeg may either retain or discard the final partial visual interval,
# depending on timestamps delivered by the concat filter. Both are valid
# fixed grids; for 1388.4 s at 2 FPS this permits 2776 or 2777 frames.
frame_product = dense_duration * fps
expected_frames = {
int(math.floor(frame_product + 1e-7)),
int(math.ceil(frame_product - 1e-7)),
}
if not math.isclose(actual_fps, fps, rel_tol=0.0, abs_tol=1e-6):
raise RuntimeError(f"Prepared FPS is {actual_fps}, expected {fps}: {path}")
if actual_frames not in expected_frames:
raise RuntimeError(
f"Prepared frame count is {actual_frames}, expected one of "
f"{sorted(expected_frames)}: {path}"
)
audio_duration = float(audio["duration"])
# AAC frame padding can shift the container duration by a few milliseconds.
if abs(audio_duration - dense_duration) > 0.05:
raise RuntimeError(
f"Prepared audio duration {audio_duration:.6f}s differs from dense "
f"timeline {dense_duration:.6f}s: {path}"
)
return {
"video_fps": actual_fps,
"video_frames": actual_frames,
"allowed_video_frames": sorted(expected_frames),
"video_duration_seconds": float(video["duration"]),
"audio_duration_seconds": audio_duration,
"audio_sample_rate": int(audio["sample_rate"]),
"audio_channels": int(audio["channels"]),
}
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--session-dir", type=Path, required=True)
parser.add_argument("--output-dir", type=Path, required=True)
parser.add_argument("--target", type=Path)
parser.add_argument(
"--fps",
type=float,
nargs="+",
default=[2.0, 5.0],
help="Fixed global visual rates to generate (default: 2 5).",
)
parser.add_argument("--dry-run", action="store_true")
parser.add_argument(
"--reuse-existing",
action="store_true",
help="Validate and reuse an existing output MP4 instead of re-encoding it.",
)
return parser.parse_args()
def main() -> None:
args = parse_args()
clip_paths = sorted(
args.session_dir.glob("*.mp4"), key=lambda path: parse_timecode(path.stem)
)
clips = [probe_clip(path) for path in clip_paths]
target_steps = load_target_steps(args.target) if args.target else None
num_steps = validate_inventory(clips, target_steps)
session = args.session_dir.name
duration = num_steps / TASK_FPS
args.output_dir.mkdir(parents=True, exist_ok=True)
concat_path = args.output_dir / f"{session}.ffconcat"
concat_text = "ffconcat version 1.0\n" + "".join(
f"file {ffconcat_quote(clip.path)}\n" for clip in clips
)
concat_path.write_text(concat_text)
outputs: dict[str, str] = {}
output_media: dict[str, dict] = {}
commands: list[list[str]] = []
for fps in args.fps:
if fps <= 0:
raise ValueError("Every --fps value must be positive")
fps_label = f"{fps:g}".replace(".", "p")
output = args.output_dir / f"{session}_{fps_label}fps.mp4"
outputs[f"{fps:g}"] = str(output.resolve())
command = build_ffmpeg_command(
clips,
output,
fps=fps,
dense_duration=duration,
)
commands.append(command)
if args.dry_run:
print(shlex.join(command))
elif args.reuse_existing and output.is_file():
print(f"Reusing {output} ...", flush=True)
output_media[f"{fps:g}"] = validate_output(
output, fps=fps, dense_duration=duration
)
else:
print(f"Preparing {output} ...", flush=True)
subprocess.run(command, check=True)
output_media[f"{fps:g}"] = validate_output(
output, fps=fps, dense_duration=duration
)
metadata = {
"session": session,
"source_session_dir": str(args.session_dir.resolve()),
"source_clip_order": [asdict(clip) for clip in clips],
"source_num_clips": len(clips),
"source_frames_20hz": sum(clip.frames for clip in clips),
"task_fps": TASK_FPS,
"num_steps": num_steps,
"dense_duration_seconds": duration,
"crop": {
"policy": "rightmost_1920x1080",
"width": KEEP_WIDTH,
"height": KEEP_HEIGHT,
},
"visual_sampling": "fixed_session_global_grid",
"outputs": outputs,
"output_media": output_media,
"target": str(args.target.resolve()) if args.target else None,
"dry_run": bool(args.dry_run),
"reuse_existing": bool(args.reuse_existing),
"commands": [shlex.join(command) for command in commands],
}
metadata_path = args.output_dir / f"{session}.json"
metadata_path.write_text(json.dumps(metadata, indent=2) + "\n")
print(
f"{session}: {len(clips)} clips, {num_steps} steps, "
f"{duration:.3f}s dense timeline"
)
print(f"Saved metadata to {metadata_path}")
if __name__ == "__main__":
main()

270
qwen_omni_features.py Normal file
View File

@@ -0,0 +1,270 @@
"""Load Qwen2.5-Omni-3B and extract causal audio-video features.
This research example does not generate text or speech: only the Thinker
(multimodal understanding model) is loaded.
Example:
python qwen_omni_features.py --video clip.mp4 \
--use-audio-in-video --output clip_features.pt
"""
from __future__ import annotations
import argparse
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Literal, Mapping
import torch
from qwen_omni_utils import process_mm_info
from transformers import (
Qwen2_5OmniProcessor,
Qwen2_5OmniThinkerForConditionalGeneration,
)
MODEL_ID = "Qwen/Qwen2.5-Omni-3B"
@dataclass
class OmniFeatures:
"""Features returned for one or more audio-video examples."""
token_features: torch.Tensor # [B, T, D]
attention_mask: torch.Tensor # [B, T]
last_token: torch.Tensor # [B, D]
tail_mean: torch.Tensor # [B, D]
hidden_states: tuple[torch.Tensor, ...] | None = None
def choose_dtype(dtype: Literal["auto", "bfloat16", "float16", "float32"]):
if dtype == "auto":
return "auto"
return getattr(torch, dtype)
def load_qwen_omni(
model_id: str = MODEL_ID,
*,
dtype: Literal["auto", "bfloat16", "float16", "float32"] = "auto",
device_map: str | Mapping[str, str | int] = "auto",
use_flash_attention_2: bool = False,
) -> tuple[Qwen2_5OmniThinkerForConditionalGeneration, Qwen2_5OmniProcessor]:
"""Load the multimodal Thinker and processor without the Talker."""
model_kwargs: dict[str, Any] = {
"torch_dtype": choose_dtype(dtype),
"device_map": device_map,
}
if use_flash_attention_2:
model_kwargs["attn_implementation"] = "flash_attention_2"
model = Qwen2_5OmniThinkerForConditionalGeneration.from_pretrained(
model_id, **model_kwargs
)
processor = Qwen2_5OmniProcessor.from_pretrained(model_id)
model.eval()
return model, processor
def build_conversation(
*,
video_path: str | Path | None = None,
audio_path: str | Path | None = None,
prompt: str = (
"Observe the conversation up to the current time and represent the "
"multimodal evidence useful for predicting the next speaking turn."
),
) -> list[dict[str, Any]]:
"""Build one Qwen chat containing video, optional audio, and instruction."""
if video_path is None and audio_path is None:
raise ValueError("At least one of video_path and audio_path is required")
content: list[dict[str, str]] = []
if video_path is not None:
content.append({"type": "video", "video": str(Path(video_path).resolve())})
if audio_path is not None:
content.append({"type": "audio", "audio": str(Path(audio_path).resolve())})
content.append({"type": "text", "text": prompt})
return [
{
"role": "system",
"content": [{
"type": "text",
"text": (
"You are a causal multimodal conversation observer. "
"Use only information available up to the current time."
),
}],
},
{"role": "user", "content": content},
]
def prepare_inputs(
processor: Qwen2_5OmniProcessor,
conversations: list[list[dict[str, Any]]],
*,
use_audio_in_video: bool,
model_device: torch.device,
model_dtype: torch.dtype,
):
"""Convert a batch of conversations into model tensors."""
texts = processor.apply_chat_template(
conversations, add_generation_prompt=True, tokenize=False
)
audios, images, videos = process_mm_info(
conversations, use_audio_in_video=use_audio_in_video
)
inputs = processor(
text=texts,
audio=audios,
images=images,
videos=videos,
padding=True,
return_tensors="pt",
use_audio_in_video=use_audio_in_video,
)
# BatchFeature.to(dtype) casts floating tensors only; IDs/masks stay integer.
return inputs.to(model_device).to(model_dtype)
def _last_valid_token(states: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
positions = torch.arange(states.shape[1], device=states.device).unsqueeze(0)
indices = positions.masked_fill(~mask.bool(), -1).max(dim=1).values.clamp_min(0)
batch_indices = torch.arange(states.shape[0], device=states.device)
return states[batch_indices, indices]
def _tail_mean(
states: torch.Tensor, mask: torch.Tensor, tail_tokens: int
) -> torch.Tensor:
if tail_tokens <= 0:
raise ValueError("tail_tokens must be positive")
# Ranking valid tokens from the right supports both left and right padding.
reverse_ranks = mask.long().flip(1).cumsum(dim=1).flip(1)
tail_mask = mask.bool() & (reverse_ranks <= tail_tokens)
weights = tail_mask.unsqueeze(-1).to(states.dtype)
return (states * weights).sum(dim=1) / weights.sum(dim=1).clamp_min(1)
@torch.inference_mode()
def extract_features(
model: Qwen2_5OmniThinkerForConditionalGeneration,
processor: Qwen2_5OmniProcessor,
conversations: list[list[dict[str, Any]]],
*,
use_audio_in_video: bool = False,
layer: int = -1,
tail_tokens: int = 8,
return_all_layers: bool = False,
output_device: str | torch.device = "cpu",
) -> OmniFeatures:
"""Extract fused Thinker states without generation.
``last_token`` is the recommended one-vector-per-window baseline. The
Thinker is causal, so it summarizes all earlier prompt and AV tokens.
"""
inputs = prepare_inputs(
processor,
conversations,
use_audio_in_video=use_audio_in_video,
model_device=model.device,
model_dtype=model.dtype,
)
# A direct forward ends at the observation boundary; generate() would append
# answer tokens and change which representation is being extracted.
outputs = model(
**inputs,
use_audio_in_video=use_audio_in_video,
output_hidden_states=True,
use_cache=False,
return_dict=True,
)
if outputs.hidden_states is None:
raise RuntimeError("The model did not return hidden states")
states = outputs.hidden_states[layer]
mask = inputs["attention_mask"]
return OmniFeatures(
token_features=states.to(output_device),
attention_mask=mask.to(output_device),
last_token=_last_valid_token(states, mask).to(output_device),
tail_mean=_tail_mean(states, mask, tail_tokens).to(output_device),
hidden_states=(
tuple(x.to(output_device) for x in outputs.hidden_states)
if return_all_layers else None
),
)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--video", type=Path)
parser.add_argument("--audio", type=Path)
parser.add_argument("--output", type=Path, required=True)
parser.add_argument("--model-id", default=MODEL_ID)
parser.add_argument("--prompt", default=None)
parser.add_argument("--use-audio-in-video", action="store_true")
parser.add_argument("--tail-tokens", type=int, default=8)
parser.add_argument("--layer", type=int, default=-1)
parser.add_argument(
"--dtype",
choices=("auto", "bfloat16", "float16", "float32"),
default="auto",
)
parser.add_argument("--flash-attention-2", action="store_true")
parser.add_argument("--device", default="cuda:0")
return parser.parse_args()
def main() -> None:
args = parse_args()
if args.video is None and args.audio is None:
raise SystemExit("Pass --video, --audio, or both")
if args.use_audio_in_video and args.video is None:
raise SystemExit("--use-audio-in-video requires --video")
model, processor = load_qwen_omni(
args.model_id,
dtype=args.dtype,
device_map={"": args.device},
use_flash_attention_2=args.flash_attention_2,
)
conversation_kwargs: dict[str, Any] = {
"video_path": args.video,
"audio_path": args.audio,
}
if args.prompt is not None:
conversation_kwargs["prompt"] = args.prompt
conversation = build_conversation(**conversation_kwargs)
features = extract_features(
model,
processor,
[conversation],
use_audio_in_video=args.use_audio_in_video,
layer=args.layer,
tail_tokens=args.tail_tokens,
)
args.output.parent.mkdir(parents=True, exist_ok=True)
torch.save({
"model_id": args.model_id,
"layer": args.layer,
"token_features": features.token_features,
"attention_mask": features.attention_mask,
"last_token": features.last_token,
"tail_mean": features.tail_mean,
}, args.output)
print(f"Saved features to {args.output}")
print(f"token_features: {tuple(features.token_features.shape)}")
print(f"last_token: {tuple(features.last_token.shape)}")
print(f"tail_mean: {tuple(features.tail_mean.shape)}")
if __name__ == "__main__":
main()

View File

@@ -0,0 +1,458 @@
{
"feature_root": "/Users/wangx/.mounty/MX500B/Code/EgoSpeak/dataset/EasyCom",
"raw_manifest": null,
"task_fps": 5,
"step_seconds": 0.2,
"sessions": [
{
"session": "Session_7",
"split": "train",
"num_steps": 7270,
"duration_seconds_at_5hz": 1454.0,
"shapes": {
"rgb_kinetics_resnet50": [
7270,
2048
],
"flow_kinetics_resnet50": [
7270,
2048
],
"audio_wav2vec2": [
7270,
5120
],
"target_perframe": [
7270,
3
]
},
"dtypes": {
"rgb_kinetics_resnet50": "float32",
"flow_kinetics_resnet50": "float32",
"audio_wav2vec2": "float32",
"target_perframe": "float64"
},
"class_positive_steps": [
1529,
1994,
4198
],
"overlap_wearer_other_steps": 451
},
{
"session": "Session_11",
"split": "train",
"num_steps": 6837,
"duration_seconds_at_5hz": 1367.4,
"shapes": {
"rgb_kinetics_resnet50": [
6837,
2048
],
"flow_kinetics_resnet50": [
6837,
2048
],
"audio_wav2vec2": [
6837,
5120
],
"target_perframe": [
6837,
3
]
},
"dtypes": {
"rgb_kinetics_resnet50": "float32",
"flow_kinetics_resnet50": "float32",
"audio_wav2vec2": "float32",
"target_perframe": "float64"
},
"class_positive_steps": [
1662,
965,
4676
],
"overlap_wearer_other_steps": 466
},
{
"session": "Session_10",
"split": "train",
"num_steps": 7199,
"duration_seconds_at_5hz": 1439.8,
"shapes": {
"rgb_kinetics_resnet50": [
7199,
2048
],
"flow_kinetics_resnet50": [
7199,
2048
],
"audio_wav2vec2": [
7199,
5120
],
"target_perframe": [
7199,
3
]
},
"dtypes": {
"rgb_kinetics_resnet50": "float32",
"flow_kinetics_resnet50": "float32",
"audio_wav2vec2": "float32",
"target_perframe": "float64"
},
"class_positive_steps": [
1920,
2014,
3757
],
"overlap_wearer_other_steps": 492
},
{
"session": "Session_6",
"split": "train",
"num_steps": 8819,
"duration_seconds_at_5hz": 1763.8,
"shapes": {
"rgb_kinetics_resnet50": [
8819,
2048
],
"flow_kinetics_resnet50": [
8819,
2048
],
"audio_wav2vec2": [
8819,
5120
],
"target_perframe": [
8819,
3
]
},
"dtypes": {
"rgb_kinetics_resnet50": "float32",
"flow_kinetics_resnet50": "float32",
"audio_wav2vec2": "float32",
"target_perframe": "float64"
},
"class_positive_steps": [
1520,
3232,
4967
],
"overlap_wearer_other_steps": 900
},
{
"session": "Session_9",
"split": "train",
"num_steps": 8528,
"duration_seconds_at_5hz": 1705.6,
"shapes": {
"rgb_kinetics_resnet50": [
8528,
2048
],
"flow_kinetics_resnet50": [
8528,
2048
],
"audio_wav2vec2": [
8528,
5120
],
"target_perframe": [
8528,
3
]
},
"dtypes": {
"rgb_kinetics_resnet50": "float32",
"flow_kinetics_resnet50": "float32",
"audio_wav2vec2": "float32",
"target_perframe": "float64"
},
"class_positive_steps": [
2957,
884,
4808
],
"overlap_wearer_other_steps": 121
},
{
"session": "Session_4",
"split": "train",
"num_steps": 8658,
"duration_seconds_at_5hz": 1731.6,
"shapes": {
"rgb_kinetics_resnet50": [
8658,
2048
],
"flow_kinetics_resnet50": [
8658,
2048
],
"audio_wav2vec2": [
8658,
5120
],
"target_perframe": [
8658,
3
]
},
"dtypes": {
"rgb_kinetics_resnet50": "float32",
"flow_kinetics_resnet50": "float32",
"audio_wav2vec2": "float32",
"target_perframe": "float64"
},
"class_positive_steps": [
1386,
1013,
6681
],
"overlap_wearer_other_steps": 422
},
{
"session": "Session_8",
"split": "train",
"num_steps": 7642,
"duration_seconds_at_5hz": 1528.4,
"shapes": {
"rgb_kinetics_resnet50": [
7642,
2048
],
"flow_kinetics_resnet50": [
7642,
2048
],
"audio_wav2vec2": [
7642,
5120
],
"target_perframe": [
7642,
3
]
},
"dtypes": {
"rgb_kinetics_resnet50": "float32",
"flow_kinetics_resnet50": "float32",
"audio_wav2vec2": "float32",
"target_perframe": "float64"
},
"class_positive_steps": [
1859,
1778,
4533
],
"overlap_wearer_other_steps": 528
},
{
"session": "Session_5",
"split": "train",
"num_steps": 8558,
"duration_seconds_at_5hz": 1711.6,
"shapes": {
"rgb_kinetics_resnet50": [
8558,
2048
],
"flow_kinetics_resnet50": [
8558,
2048
],
"audio_wav2vec2": [
8558,
5120
],
"target_perframe": [
8558,
3
]
},
"dtypes": {
"rgb_kinetics_resnet50": "float32",
"flow_kinetics_resnet50": "float32",
"audio_wav2vec2": "float32",
"target_perframe": "float64"
},
"class_positive_steps": [
2256,
2814,
4083
],
"overlap_wearer_other_steps": 595
},
{
"session": "Session_12",
"split": "train",
"num_steps": 8894,
"duration_seconds_at_5hz": 1778.8,
"shapes": {
"rgb_kinetics_resnet50": [
8894,
2048
],
"flow_kinetics_resnet50": [
8894,
2048
],
"audio_wav2vec2": [
8894,
5120
],
"target_perframe": [
8894,
3
]
},
"dtypes": {
"rgb_kinetics_resnet50": "float32",
"flow_kinetics_resnet50": "float32",
"audio_wav2vec2": "float32",
"target_perframe": "float64"
},
"class_positive_steps": [
1252,
3353,
6099
],
"overlap_wearer_other_steps": 1810
},
{
"session": "Session_3",
"split": "test",
"num_steps": 8400,
"duration_seconds_at_5hz": 1680.0,
"shapes": {
"rgb_kinetics_resnet50": [
8400,
2048
],
"flow_kinetics_resnet50": [
8400,
2048
],
"audio_wav2vec2": [
8400,
5120
],
"target_perframe": [
8400,
3
]
},
"dtypes": {
"rgb_kinetics_resnet50": "float32",
"flow_kinetics_resnet50": "float32",
"audio_wav2vec2": "float32",
"target_perframe": "float64"
},
"class_positive_steps": [
2024,
2547,
4521
],
"overlap_wearer_other_steps": 692
},
{
"session": "Session_2",
"split": "test",
"num_steps": 7706,
"duration_seconds_at_5hz": 1541.2,
"shapes": {
"rgb_kinetics_resnet50": [
7706,
2048
],
"flow_kinetics_resnet50": [
7706,
2048
],
"audio_wav2vec2": [
7706,
5120
],
"target_perframe": [
7706,
3
]
},
"dtypes": {
"rgb_kinetics_resnet50": "float32",
"flow_kinetics_resnet50": "float32",
"audio_wav2vec2": "float32",
"target_perframe": "float64"
},
"class_positive_steps": [
1718,
1350,
5152
],
"overlap_wearer_other_steps": 514
},
{
"session": "Session_1",
"split": "test",
"num_steps": 6942,
"duration_seconds_at_5hz": 1388.4,
"shapes": {
"rgb_kinetics_resnet50": [
6942,
2048
],
"flow_kinetics_resnet50": [
6942,
2048
],
"audio_wav2vec2": [
6942,
5120
],
"target_perframe": [
6942,
3
]
},
"dtypes": {
"rgb_kinetics_resnet50": "float32",
"flow_kinetics_resnet50": "float32",
"audio_wav2vec2": "float32",
"target_perframe": "float64"
},
"class_positive_steps": [
1538,
2266,
3858
],
"overlap_wearer_other_steps": 720
}
],
"total_steps": 95453,
"total_duration_hours": 5.302944444444444,
"split_steps": {
"train": 72405,
"test": 23048
},
"class_positive_steps": [
21621,
24210,
57333
],
"target_row_counts": {
"[0, 0, 1]": 49622,
"[0, 1, 0]": 16499,
"[0, 1, 1]": 7711,
"[1, 0, 0]": 21621
},
"issues": []
}

9
requirements.txt Normal file
View File

@@ -0,0 +1,9 @@
torch
numpy
transformers==4.52.3
accelerate
qwen-omni-utils[decord]
librosa
Pillow
# Optional on compatible NVIDIA GPUs: flash-attn

147
smoke_test.py Normal file
View File

@@ -0,0 +1,147 @@
#!/usr/bin/env python3
"""Minimal single-GPU smoke test for a local Qwen2.5-Omni checkpoint.
Examples:
python smoke_test.py --model /data1/yfzx/wangx/Models/Qwen2.5-Omni-3B
python smoke_test.py --model /data1/yfzx/wangx/Models/Qwen2.5-Omni-3B \
--video /data1/example.mp4
Qwen2.5-Omni generation currently is more reliable on a single GPU for this
use case. In particular, do not replace the explicit device map below with
``device_map=\"auto\"``: it can place cache tensors on different GPUs.
"""
import argparse
import time
import torch
from qwen_omni_utils import process_mm_info
from transformers import Qwen2_5OmniForConditionalGeneration, Qwen2_5OmniProcessor
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--model", required=True, help="Local model directory or HF model ID")
parser.add_argument("--video", help="Optional local .mp4 file for an audio-video test")
parser.add_argument("--device", default="cuda:0", help="One CUDA device, e.g. cuda:0")
parser.add_argument("--max-new-tokens", type=int, default=64)
parser.add_argument(
"--video-fps",
type=float,
default=1.0,
help="Frames per second sampled by the Omni video processor (default: 1.0)",
)
parser.add_argument(
"--video-max-pixels",
type=int,
default=256 * 28 * 28,
help="Maximum pixels per sampled video frame (default: 200704)",
)
parser.add_argument(
"--flash-attn2",
action="store_true",
help="Use FlashAttention-2; requires the flash_attn package.",
)
return parser.parse_args()
def move_tensors_to_device(batch: dict, device: str) -> dict:
"""Move processor tensors without accidentally converting integer IDs."""
return {
name: value.to(device) if torch.is_tensor(value) else value
for name, value in batch.items()
}
def main() -> None:
args = parse_args()
if not torch.cuda.is_available():
raise RuntimeError("A CUDA GPU is required for this smoke test.")
print(f"GPU: {torch.cuda.get_device_name(args.device)}")
print(f"CUDA: {torch.version.cuda}")
print(f"PyTorch: {torch.__version__}")
print(f"Loading: {args.model}")
started = time.perf_counter()
# A 3B model fits comfortably on one L20 (48 GB). Keep all weights and
# generation cache on the same device; device_map='auto' may split them.
model_load_kwargs = dict(
torch_dtype=torch.bfloat16,
device_map={"": args.device},
)
if args.flash_attn2:
model_load_kwargs["attn_implementation"] = "flash_attention_2"
model = Qwen2_5OmniForConditionalGeneration.from_pretrained(
args.model,
**model_load_kwargs,
)
model.disable_talker() # We need Thinker text states only; saves VRAM.
model.eval()
processor = Qwen2_5OmniProcessor.from_pretrained(args.model)
print(f"Model loaded in {time.perf_counter() - started:.1f}s")
if args.video:
content = [
{
"type": "video",
"video": args.video,
# These are consumed by qwen_omni_utils while sampling the
# video. They prevent quadratic vision attention from
# allocating tens of GB for a short, high-resolution clip.
"fps": args.video_fps,
"max_pixels": args.video_max_pixels,
},
{
"type": "text",
"text": "Briefly describe the people, speech, and interaction in this video.",
},
]
print(
f"Video: {args.video} "
f"(sample_fps={args.video_fps}, max_pixels={args.video_max_pixels})"
)
else:
content = [
{"type": "text", "text": "Reply with exactly: Qwen Omni smoke test passed."},
]
messages = [{"role": "user", "content": content}]
prompt = processor.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
)
audios, images, videos = process_mm_info(
messages, use_audio_in_video=bool(args.video)
)
inputs = processor(
text=prompt,
audio=audios,
images=images,
videos=videos,
return_tensors="pt",
padding=True,
use_audio_in_video=bool(args.video),
)
inputs = move_tensors_to_device(inputs, args.device)
started = time.perf_counter()
with torch.inference_mode():
output_ids = model.generate(
**inputs,
max_new_tokens=args.max_new_tokens,
do_sample=False,
return_audio=False,
)
# Remove the prompt tokens before decoding the generated completion.
generated_ids = output_ids[:, inputs["input_ids"].shape[1] :]
answer = processor.batch_decode(
generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
)[0]
print(f"Generation completed in {time.perf_counter() - started:.1f}s")
print(f"\nModel response:\n{answer.strip()}")
if __name__ == "__main__":
main()

1866
streaming_turn_model.py Normal file

File diff suppressed because it is too large Load Diff