Add no-box full-frame Qwen AV baselines

2026-09-17 02:51:26 +08:00
parent a9f748d6a6
commit 62134a0fa0
10 changed files with 1164 additions and 17 deletions
--- a/easycom_head.py
+++ b/easycom_head.py
@@ -67,6 +67,13 @@ CURRENT_FRAME_PERSON_SET_FUSION_MODES = (
    "audio-person-frame-set",
    "gated-person-state-adapter",
    "highrate-person-av-attention",
+    "fullframe-spatial-attention",
+)
+FULLFRAME_SPATIAL_FUSION_MODES = ("fullframe-spatial-attention",)
+FULLFRAME_TEMPORAL_FUSION_MODES = ("fullframe-spatiotemporal-attention",)
+FULLFRAME_FUSION_MODES = (
+    *FULLFRAME_SPATIAL_FUSION_MODES,
+    *FULLFRAME_TEMPORAL_FUSION_MODES,
 )
 HIGHRATE_PERSON_AV_FUSION_MODES = ("highrate-person-av-attention",)
 TRACK_PERSON_FUSION_MODES = (
@@ -74,12 +81,17 @@ TRACK_PERSON_FUSION_MODES = (
    "causal-person-adapter-v2",
    "causal-person-adapter-egospeak-gru",
 )
+HISTORY_VISUAL_SET_FUSION_MODES = (
+    *TRACK_PERSON_FUSION_MODES,
+    *FULLFRAME_TEMPORAL_FUSION_MODES,
+)
 PERSON_ADAPTER_V2_MODES = (
    "causal-person-adapter-v2",
    "causal-person-adapter-egospeak-gru",
 )
 PERSON_SET_FUSION_MODES = (
    *CURRENT_FRAME_PERSON_SET_FUSION_MODES,
+    *FULLFRAME_TEMPORAL_FUSION_MODES,
    *TRACK_PERSON_FUSION_MODES,
 )
 LOCAL_FUSION_MODES = ("pooled", *PERSON_SET_FUSION_MODES)
@@ -92,6 +104,12 @@ LOCAL_FUSION_SCOPES = {
    "highrate-person-av-attention": (
        "global-omni-traditional-audio-base-plus-zero-init-25hz-person-residual"
    ),
+    "fullframe-spatial-attention": (
+        "f2-global-omni-qwen-audio-latest-fullframe-spatial-tokens"
+    ),
+    "fullframe-spatiotemporal-attention": (
+        "f3-global-omni-qwen-audio-causal-fullframe-token-history"
+    ),
    "track-gru-residual": "causal-anonymous-track-gru-pooled-baseline-residual",
    "causal-person-adapter-v2": (
        "causal-anonymous-track-gru-dynamic-gated-logit-residual"
@@ -206,7 +224,10 @@ class TurnPredictionHead(nn.Module):
        fused_size = hidden_size
        projected_modalities = (
            (("audio", local_audio_dim),)
-            if local_fusion in CURRENT_FRAME_PERSON_SET_FUSION_MODES
+            if local_fusion in (
+                *CURRENT_FRAME_PERSON_SET_FUSION_MODES,
+                *FULLFRAME_TEMPORAL_FUSION_MODES,
+            )
            else (("visual", local_visual_dim), ("audio", local_audio_dim))
        )
        for name, dimension in projected_modalities:
@@ -291,6 +312,13 @@ class TurnPredictionHead(nn.Module):
                    batch_first=True,
                )
                self.audio_person_gate = nn.Parameter(torch.zeros(()))
+            if local_fusion in FULLFRAME_TEMPORAL_FUSION_MODES:
+                self.person_track_gru = nn.GRU(
+                    person_attention_dim,
+                    person_attention_dim,
+                    num_layers=1,
+                    batch_first=True,
+                )
            if local_fusion == "gated-person-state-adapter":
                # Screening version of F6 for precomputed frozen states.  It
                # injects selected person evidence after the Thinker readout,
@@ -303,6 +331,8 @@ class TurnPredictionHead(nn.Module):
                    nn.Linear(person_attention_dim, hidden_size),
                )
                self.person_state_adapter_gate = nn.Parameter(torch.zeros(()))
+            elif local_fusion in FULLFRAME_TEMPORAL_FUSION_MODES:
+                fused_size += person_attention_dim + 1
            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
@@ -419,7 +449,10 @@ class TurnPredictionHead(nn.Module):

    @property
    def local_input_size(self) -> int:
-        if self.local_fusion in CURRENT_FRAME_PERSON_SET_FUSION_MODES:
+        if self.local_fusion in (
+            *CURRENT_FRAME_PERSON_SET_FUSION_MODES,
+            *FULLFRAME_TEMPORAL_FUSION_MODES,
+        ):
            return self.local_audio_dim + 1
        return sum(d + 1 for d in (self.local_visual_dim, self.local_audio_dim) if d)

@@ -567,6 +600,65 @@ class TurnPredictionHead(nn.Module):
            audio_valid,
        )

+    def _fullframe_temporal_context(
+        self,
+        states: torch.Tensor,
+        local: torch.Tensor,
+        token_histories: torch.Tensor | None,
+        history_mask: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Encode same-grid-position visual histories, then query the grid."""
+
+        if states.ndim != 2:
+            raise ValueError("F3 expects batched [B,D] global states")
+        if local.shape != (len(states), self.local_audio_dim + 1):
+            raise ValueError("F3 requires aligned [Qwen audio,valid] features")
+        if token_histories is None or token_histories.ndim != 4:
+            raise ValueError("F3 requires token histories [B,K,L,visual_dim]")
+        if (
+            token_histories.shape[0] != len(states)
+            or token_histories.shape[2] != self.person_history_frames
+            or token_histories.shape[3] != self.local_visual_dim
+        ):
+            raise ValueError(
+                "F3 history shape does not match its configured length/dimension"
+            )
+        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
+        token_states, token_valid = self._encode_person_track_histories(
+            token_histories,
+            history_mask,
+        )
+        audio_values = local[:, : self.local_audio_dim]
+        audio_valid = local[:, self.local_audio_dim :]
+        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(
+            (self.null_person.expand(len(states), -1, -1), token_states),
+            dim=1,
+        )
+        memory_valid = torch.cat(
+            (
+                torch.ones(
+                    (len(states), 1), dtype=torch.bool, device=states.device
+                ),
+                token_valid,
+            ),
+            dim=1,
+        )
+        context, _ = self.person_attention(
+            query[:, None, :],
+            memory,
+            memory,
+            key_padding_mask=~memory_valid,
+            need_weights=False,
+        )
+        visual_valid = token_valid.any(dim=1, keepdim=True).to(states.dtype)
+        return audio_projected, audio_valid, context[:, 0, :], visual_valid
+
    def _encode_person_track_histories(
        self,
        person_histories: torch.Tensor,
@@ -685,6 +777,21 @@ class TurnPredictionHead(nn.Module):
                    (states, audio, audio_valid, person_context, person_valid),
                    dim=-1,
                )
+        elif self.local_fusion in FULLFRAME_TEMPORAL_FUSION_MODES:
+            if local is None:
+                raise ValueError("F3 requires local Qwen audio features")
+            audio, audio_valid, visual_context, visual_valid = (
+                self._fullframe_temporal_context(
+                    states,
+                    local,
+                    person_features,
+                    person_mask,
+                )
+            )
+            states = torch.cat(
+                (states, audio, audio_valid, visual_context, visual_valid),
+                dim=-1,
+            )
        elif self.local_input_size:
            if local is None or local.shape != (*states.shape[:-1], self.local_input_size):
                raise ValueError(
@@ -1554,7 +1661,26 @@ def load_session_data(
            name,
            dataset=dataset,
        )
-        if local_fusion in TRACK_PERSON_FUSION_MODES:
+        if local_fusion in FULLFRAME_FUSION_MODES:
+            if dataset != "easycom" or local_modalities != "av":
+                raise ValueError(
+                    "Full-frame F2/F3 require EasyCom and --local-modalities av"
+                )
+            from fullframe_qwen_av import load_fullframe_av_features
+
+            local, local_people, contract = load_fullframe_av_features(
+                local_path,
+                payload,
+                sample_id=name,
+                dataset=dataset,
+                expected_contract=local_contract,
+                history_max_gap_seconds=(
+                    person_track_max_gap_seconds
+                    if local_fusion in FULLFRAME_TEMPORAL_FUSION_MODES
+                    else None
+                ),
+            )
+        elif local_fusion in TRACK_PERSON_FUSION_MODES:
            if local_modalities != "av":
                raise ValueError(
                    "Track person fusion requires --local-modalities av"
@@ -1610,7 +1736,7 @@ def load_session_data(
    elif local_contract is not None:
        raise ValueError("--local-feature-root is required for this checkpoint")
    elif local_fusion != "pooled":
-        raise ValueError("Person-set fusion requires --local-feature-root")
+        raise ValueError("Set/token fusion requires --local-feature-root")
    return SessionData(
        name,
        states,